기술 요약: 완전 루핑된 트랜스포머를 통한 루프 단순 안정화

문제 제기

대형 언어 모델 (LLM) 이 뛰어난 일반화 능력을 입증해 왔음에도 불구하고, 모델 크기를 늘리는 전통적인 확장 패러다임은 고품질 공개 텍스트 데이터의 제한된 공급으로 인해 점점 더 제약받고 있습니다. 확장된 체인 오브 씽킹 (chain-of-thought) 추론을 통한 테스트 시간 확장 (test-time scaling) 과 같은 대안적 접근 방식은 종종 컨텍스트 길이와 추론 비용의 상당한 증가를 수반합니다.

**루핑된 트랜스포머 (Looped Transformer, LT)**는 동일한 트랜스포머 백본 블록을 반복적으로 재사용하여 파라미터 효율적인 대안을 제시하며, 파라미터 수를 증가시키지 않고 추가 계산을 통해 성능을 향상시킵니다. 그러나 LT 는 루프 반복 횟수가 증가함에 따라 심각한 학습 불안정성을 겪습니다. 저자들은 단순히 루프 반복 횟수를 늘리는 것이 종종 학습 붕괴를 초래하여 모델이 더 깊은 계산 예산을 효과적으로 활용하지 못하게 만든다고 규명했습니다.

불안정성 진단

초기 학습 역학 (최초 2,000 단계의 옵티마이저 단계) 에 대한 분석을 통해 저자들은 LT 의 불안정성을 유발하는 두 가지 구체적인 현상을 진단했습니다:

1. 그래디언트 진동 (Gradient Oscillation): 초기 학습 동안 그래디언트는 강한 진동을 보입니다. 이는 공유된 루핑된 블록을 통한 그래디언트 누적 때문으로, 시간 역전 (Backpropagation Through Time, BPTT) 으로 학습된 순환 신경망 (RNN) 이 직면한 어려움과 유사합니다.
2. 잔여 폭발 (Residual Explosion): 루프 반복 횟수가 증가함에 따라 잔여 상태의 노름이 급격히 증가합니다. 모델이 고정된 깊이의 백본을 훨씬 더 깊은 유효 계산 그래프로 풀어나가기 때문에, 각 단계에서의 경미한 증폭조차 누적되어 최적화를 방해하는 폭발적인 잔여 노름을 초래합니다.

방법론: 완전 루핑된 트랜스포머 (FLT)

이러한 문제들을 추가적인 학습 가능한 파라미터나 계산 오버헤드를 도입하지 않고 해결하기 위해, 저자들은 **완전 루핑된 트랜스포머 (Fully Looped Transformer, FLT)**를 제안합니다. 이 아키텍처는 두 가지 파라미터 없는 수정 사항을 포함합니다:

1. 완전 루핑된 아키텍처 (Fully Looped Architecture, FLA)

기존의 LT(바닐라 LT) 에서는 이전 루프 반복의 출력이 다음 반복의 첫 번째 레이어로만 전달됩니다. 이는 더 깊은 레이어들이 긴 변환 체인을 통해 순환 상태에 접근하도록 강제합니다.

• 메커니즘: FLA 는 이전 루프 반복 ($h^{(t-1)}_L$) 의 순환 신호를 현재 반복 ($t$) 의 모든 레이어에 분배합니다.
• 수식: 각 레이어 $l$에 대해 업데이트는 $h^{(t)}_l = f^{(l)}_\theta(h^{(t)}_{l-1}, h^{(t-1)}_L)$로 정의됩니다.
• 목표: 이는 깊은 잔여 네트워크에서 발견되는 단축 연결 (shortcut connections) 을 모방하여, 유효 순환 깊이와 관련된 최적화 어려움을 완화하고 잔여 폭발을 줄입니다.

2. 어텐션 주입 (Attention Injection, AI)

그래디언트 진동을 억제하기 위해 저자들은 LSTM 과 GRU 의 게이트 메커니즘에서 영감을 받아 순환 신호 흐름을 제어하는 메커니즘을 도입합니다.

• 메커니즘: 이전 은닉 상태를 잔여 흐름에 직접 추가하는 대신, AI 는 기존 셀프 어텐션 블록을 크로스 어텐션 (Cross-Attention) 블록으로 재사용합니다.
• 구현:
  • 첫 번째 루프 ($t=1$) 에서는 표준 셀프 어텐션을 사용합니다.
  • 후속 루프 ($t>1$) 에서는 이전 반복의 은닉 상태 ($h^{(t-1)}_L$) 를 **쿼리 (Query, Q)**로 사용합니다.
  • 현재 레이어의 출력 ($z^{(t)}_l$) 은 **키 (Key, K)**와 **값 (Value, V)**로 모두 사용됩니다.
  • 수식: $a^{(t)}_l = \text{Attention}(Q = W_Q h^{(t-1)}_L, K = W_K z^{(t)}_l, V = W_V z^{(t)}_l)$.
• 목표: 순환 상태를 쿼리를 통해 라우팅함으로써, 주입된 신호는 이전 잔여 상태의 직접적인 증폭이 아니라 소프트맥스에 의해 제어되는 현재 값 벡터들의 정규화된 혼합물이 됩니다. 이는 그래디언트 진동을 억제하는 제어된 혼합 메커니즘을 제공합니다.

주요 기여

1. 진단: 루핑된 트랜스포머의 학습 불안정성을 유발하는 주된 원인으로 그래디언트 진동과 잔여 폭발을 규명했습니다.
2. 아키텍처: FLA 와 AI 를 통해 학습을 안정화시키는 파라미터 없는 개선안인 완전 루핑된 트랜스포머를 제안했습니다.
3. 실증적 검증: FLT 가 기준선 루핑된 트랜스포머가 붕괴되는 영역인 12 회 루프 반복까지 안정적으로 학습될 수 있음을 입증했습니다.

실험 결과

저자들은 사전 학습 (FineWeb-Edu) 과 하류 작업 (Wikitext-2, 상식 추론 벤치마크) 에서 FLT 를 기준선 (원래 LT, 입력 주입이 적용된 LT, 첫 번째 레이어 어텐션 주입이 적용된 LT) 과 비교 평가했습니다.

• 안정성: 기준선 LT 모델은 9 회 또는 12 회 루프 반복에서 붕괴되었으나 (높은 손실 플래토와 폭발적인 잔여 상태 발생), FLT 는 더 매끄러운 손실 감소와 유계된 잔여 노름을 보이는 안정된 학습 역학을 유지했습니다.
• 성능:
  • 기준선이 붕괴하지 않은 "더 온화한" 설정에서 FLT 는 평균 하류 작업 성능을 최대 **13.2%**까지 향상시켰습니다 (예: 베이스 크기 모델에서 6 루프 시).
  • 베이스 크기 (318M 파라미터) 에서 FLT 는 9 루프 시 가장 높은 평균 점수 (41.72) 와 가장 낮은 Wiki2 퍼플렉시티 (38.44) 를 달성한 반면, 원래 LT 는 붕괴했습니다.
• 애블레이션 연구:
  • FLA 만으로는 잔여 폭발을 줄였지만 학습 붕괴를 완전히 방지하지는 못했습니다.
  • 어텐션 주입만 (첫 번째 레이어에서) 은 입력 주입과 유사하게 작용했으나 고루프 횟수에는 불충분했습니다.
  • FLA 와 AI 의 결합이 고루프 횟수 (예: 12 회) 에서 안정된 학습에 필수적이었습니다.
• 테스트 시간 적응: 고정된 루프 반복 횟수 (예: 9 또는 12) 로 학습된 모델들은 추론 루프 횟수가 학습 범위 내에 있을 때 성능과 추론 루프 횟수 간에 양의 상관관계를 보였습니다. 그러나 추론 루프가 학습 횟수를 초과하면 성능이 예측 불가능해져서, 견고한 적응을 위해서는 가변적인 루프 횟수에 대한 명시적 학습이 필요할 수 있음을 시사합니다.

중요성과 주장

본 논문은 완전 루핑된 트랜스포머가 루핑된 모델의 학습 불안정성에 대한 간단하고 파라미터 없는 해결책을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 깊은 루핑 가능: 파라미터 수나 컨텍스트 길이를 증가시키지 않고 모델이 훨씬 더 많은 테스트 시간 계산 (최대 12 회 반복) 을 활용할 수 있도록 합니다.
• 효율성: 데이터가 제한된 환경에서 계산과 성능 향상을 교환할 수 있는 메커니즘을 제공합니다.
• 적응성: 추론 예산이 학습 분포 내에 유지된다는 전제하에, 추론 시 루프 반복 횟수를 조정하여 성능과 계산 예산을 균형 있게 맞출 수 있다는 예비 증거를 제공합니다.

저자들은 그들의 발견이 루핑된 트랜스포머 학습 역학에 대한 이해를 심화시키고, 안정적이고 파라미터 효율적이며 적응적인 아키텍처로 나아가는 길을 제시한다고 결론지었습니다.