탄력적 깊이 루프 트랜스포머 예산 기반 추론을 위한 단축 경로 조절

초록

LoopFormer는 루프형 트랜스포머에 시간 t와 단계 크기 Δt를 조건으로 부여하고, 다양한 길이의 루프 궤적을 동시에 학습함으로써 사용자가 지정한 연산 예산에 맞춰 유연하게 깊이를 조절할 수 있는 모델이다. 단축‑일관성 손실을 통해 짧은 루프에서도 의미 있는 표현을 얻고, 긴 루프에서는 점진적인 정제를 보장한다. 실험 결과, 언어 모델링과 추론 벤치마크에서 낮은 FLOP에서도 경쟁력 있는 퍼플렉시티를 유지하면서 예산이 늘어날수록 성능이 매끄럽게 향상되는 것을 확인하였다.

상세 분석

LoopFormer의 핵심 아이디어는 “궤적 기반 조건부 루프”이다. 기존 루프형 트랜스포머는 고정된 반복 횟수 L을 전제로 학습·추론했지만, 실제 서비스 환경에서는 연산 예산이 가변적이다. 이를 해결하기 위해 저자들은 각 루프 단계 i를 (t_{i‑1}, Δ_i)라는 두 연속 변수에 조건화한다. 여기서 t_{i‑1}는 현재 누적 정규화 시간(0 ~ 1)이고, Δ_i는 현재 단계가 차지하는 시간 비율이다. 두 변수는 사인‑코사인 주파수 임베딩을 거쳐 작은 MLP에 입력돼 RMSNorm 스케일(γ₁,γ₂)과 MHSA·FFN 잔차 게이트(α₁,α₂)를 생성한다. 이 메커니즘은 “시간‑스텝” 정보를 직접 네트워크에 주입해, 동일 파라미터가 서로 다른 궤적(예: 10 step vs 2 step)에서도 일관된 동작을 하게 만든다.

학습 단계에서는 최대 L step 궤적과, 임의로 샘플링한 짧은 궤적 S (1 ≤ S < L) 두 가지를 동시에 사용한다. 손실은 세 부분으로 구성된다. 첫째, 전체 L step에 대한 표준 언어 모델링 손실 L_L; 둘째, 샘플된 S step에 대한 동일 손실 L_S; 셋째, “단축‑일관성” 손실 L_cons으로, 짧은 궤적의 토큰 로짓을 stop‑gradient 처리된 전체 궤적 로짓에 맞추어 정규화한다. 이 일관성 손실은 짧은 루프가 최종 목표 표현을 미리 예측하도록 강제함으로써, 예산이 감소해도 의미 있는 표현을 유지하게 만든다.

구조적으로 LoopFormer는 기존 루프형 디코더와 동일한 Transformer 블록(k = 1 혹은 2)을 공유한다. 차별점은 각 루프마다 위에서 설명한 시간·스텝 임베딩이 RMSNorm과 게이트에 삽입된다는 점이다. 이는 DiT와 같은 확산 모델에서 차용한 adaLN 방식과 유사하지만, 언어 모델링에 맞게 “스텝 크기”까지 확장한 것이 특징이다. 또한, 저자들은 “균일 단계”와 “비균일 단계” 두 종류의 스케줄을 모두 지원하도록 설계했으며, 추론 시 사용자는 원하는 예산 M과 Δ_M을 자유롭게 지정할 수 있다.

실험에서는 1.3B 파라미터 규모의 LoopFormer를 다양한 데이터셋(예: WikiText‑103, C4, GSM‑8K)에 적용했다. 결과는 다음과 같다. (1) 동일 FLOP 대비 비루프(Non‑looped) 모델보다 퍼플렉시티가 5 ~ 10 % 개선되었으며, 특히 0.3 L 정도의 짧은 루프에서도 손실이 크게 증가하지 않았다. (2) 추론 예산을 단계적으로 늘릴 때 성능이 거의 선형에 가깝게 상승했으며, 이는 기존 고정‑L 모델이 짧은 루프에서 급격히 성능이 떨어지는 현상과 대조된다. (3) Zero‑shot 추론 벤치마크(예: ARC‑Easy, BoolQ)에서도 짧은 예산에서 경쟁력 있는 점수를 기록했고, 전체 예산에서는 최첨단 루프형 모델에 근접했다.

추가 분석에서는 CKA, 곡률, 엔트로피, 이방성 등 여러 기하학·정보 이론적 지표를 사용해 루프 단계별 표현 변화를 측정했다. 일반 루프형 모델은 깊이가 늘어날수록 표현이 정체되거나 수축되는 경향이 있었지만, LoopFormer는 단계마다 지속적인 변동성을 유지했고, 특히 단축‑일관성 손실이 적용된 경우 고차원 공간에서의 흐름이 부드럽게 이어졌다. 이는 “잠재적 추론(latent reasoning)” 능력이 예산 변화에 강인함을 의미한다.

전반적으로 LoopFormer는 (1) 시간·스텝 조건화, (2) 단축‑일관성 학습, (3) 예산‑조건 추론이라는 세 축을 통해 루프형 트랜스포머의 유연성을 실현했으며, 향후 대규모 LLM에 적용해 비용‑효율적인 동적 추론을 구현할 수 있는 기반을 제공한다.