Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

기술 요약: 심층 재귀적 트랜스포머(Depth-Recurrent Transformers)에서의 토큰별 고정점 수렴

문제 정의
표준 고정 깊이 트랜스포머는 모델의 품질을 단일 연산 운영 지점에 결속시킵니다. 여러 품질/비용 티어를 서비스하려면 별도의 모델을 학습하고 저장해야 하며, 추론 시 입력 난이도에 따라 연산을 조절하려는 시도는 프롬프트 캐시를 포기하게 만듭니다. 심층 재귀적 트랜스포머(가중치가 공유된 코어를 가변 횟수만큼 적용하는 방식)는 하나의 파라미터 세트로 다양한 깊이에서 작동할 수 있게 함으로써 이론적인 해결책을 제시합니다. 그러나 기존 연구들은 두 가지 구별되는 속성을 혼동하곤 합니다: 깊이 강건성(학습되지 않은 깊이에서도 유용한 출력을 생성하는 능력)과 깊이 할당(특정 토큰에 얼마나 많은 깊이가 필요한지 결정하는 것)입니다. 기존의 토큰별 할당 방식은 상당한 한계를 가집니다. 비공유 스택(non-shared stacks)에서의 조기 종료(early-exit) 방식은 키-값(KV) 캐시 불일치로 인해 실패하며, 혼합 재귀(mixture-of-recursions) 방식은 사전 학습 단계에서 가중치 내에 연산 예산을 고정시켜 사후 재조정(post-hoc retargeting)을 방해합니다. 또한, 공동 학습된 중단 정책(halting policies)은 퇴화된 깊이로 수렴하거나 정교한 튜닝을 요구합니다. 결정적으로, 학습된 정책보다 단순한 학습 불필요 정책(training-free policy)이 더 나은 성능을 낼 수 있는지 확인하기 위해, 훈련된 탄력적 모델의 토큰별 연산량을 직접 측정한 선행 연구는 존재하지 않습니다.

방법론
저자들은 깊이 강건성과 깊이 할당을 분리하는 2단계 접근 방식을 제안합니다:

1. 깊이 강건성 (사전 학습): 최소한의 레시피를 사용하여 무작위 깊이 백본을 학습시킵니다. 각 옵티마이저 단계마다 로그-노멀 분포에서 재귀 횟수 $r$을 샘플링하여 전체 단계에 적용합니다. 모델은 중단 유닛(halting units), 폰더 손실(ponder loss), 또는 FLOP 페널티 없이 표준 교차 엔트로피 손실을 사용합니다. 아키텍처는 고유한 서문(prelude), 가중치가 공유된 코어(상태 붕괴를 방지하기 위한 입력 재주입 및 샌드위치 정규화 포함), 그리고 고유한 결말(coda)로 구성됩니다.
2. 깊이 할당 (사후 정책): 백본이 동결된 후, 토큰별 깊이를 결정하기 위해 동일한 모델에 대해 세 가지 정책을 평가합니다:
   • A0 (균등 깊이): 모든 토큰에 고정된 루프 횟수 $r$을 적용합니다 (베이스라인).
   • A1 (수렴 종료): 연속적인 출력 분포 간의 KL 발산($KL(p_i || p_{i-1})$)이 임계값 $\epsilon$ 미만으로 떨어질 때 토큰을 중단시키는 학습 불필요 규칙입니다.
   • A2 (오라클 레이블 라우터): 모델의 은닉 상태를 기반으로 해당 토큰의 argmax 예측이 향후 루프에서 변경될지 여부를 예측하도록 학습된 선형 프로브입니다.

주요 기여

1. 직접적인 토큰별 측정: 본 연구는 심층 재귀적 언어 모델에서의 토큰별 수렴을 최초로 규명합니다. 연구 결과, 재귀 상태는 토큰별로 고정점으로 수렴함을 보여줍니다. 중앙값 토큰은 6번째 루프에서 수렴하지만, 약 10%의 꼬리 부분 토큰은 훈련 평균 깊이인 8회까지 계속 업데이트됩니다. 수렴 깊이는 토큰 유형에 따라 순서가 정해집니다 (공백은 가장 얕게 수렴하고, 내용어는 가장 깊게 수렴함).
2. 읽기 vs 학습의 수렴: 본 논문은 모델의 상태 역학을 학습하는 것보다 직접 읽는 것이 더 우수함을 입증합니다. 학습 불필요 수렴 종료 방식(A1)은 평균 4.94 루프에서 균등 깊이-8의 품질을 달est합니다 (평균 깊이 38% 감소). 반면, 동일한 모델에서 추출된 수렴 레이블로 학습된 선형 라우터(A2)는 거의 전체 깊이(7.99 루프)를 필요로 하며 아무런 감소 효과를 내지 못합니다.
3. 훈련 평균 포화 가설: 내부 데이터와 외부 일관성 검사를 통해, 저자들은 재귀 모델의 유용한 깊이가 훈련 시의 평균 루프 분포에서 포화된다는 가설을 제시합니다. 평균 4로 학습된 모델은 깊이 4에서 포화되며, 평균 8로 학습된 모델은 깊이 8에서 포화됩니다.

결과

• 탄력성: 무작위 깊이 백본은 깊이 1에서 8까지 단조로운 품질 향상을 보이며, 그 이후에는 32 루프까지도 검증 손실(3.20 nats)이 안정화됩니다. 이는 모델이 특정 깊이에 대한 정교한 튜닝 없이도 깊이 강건함을 확인시켜 줍니다.
• 수렴 역학: 평균 연속 출력 KL 발산은 루프 2에서의 $3.9 \times 10^{-1}$에서 루프 16에서의 $8.5 \times 10^{-6}$으로 떨어집니다. 상태 변화가 단계별로 감쇠하는 것은 시스템이 무의미하게 회전하는 것이 아니라 고정점에 도달함을 확인해 줍니다.
• 할당 효율성: 학습 불료 A1 정책은 평균 4.94 루프에서 균등 깊이-8 모델의 품질과 일치합니다. 이는 얕은 깊이에서는 균등 베이스라인보다 우수하며, 평균 깊이 범위 전반에서 베이스라인과 일치합니다. 학습된 라우터(A2)는 절감 효과를 제공하지 못하며, 베이스라인과 일치하기 위해 거의 전체 깊이를 필요로 합니다.
• 토큰 수준 동작: 구조적 토큰(공백)은 가장 빨리 종료되지만, 내용어는 더 오래 계산합니다. 숫자와 단어 사이의 상대적 순위는 데이터셋(TinyStories vs. FineWeb-Edu) 간에 뒤바뀌는데, 이는 아키텍처의 실패가 아닌 분포 변화에 기인합니다.
• 효율성 경계: 본 연구는 평균 깊이의 감소(FLOP 대리 지표)를 보고하지만, 실제 벽시계 시간(wall-clock) 가속도는 구현되지 않은 깊이별 연속 배치(depth-wise continuous batching)에 달려 있음을 명시합니다. 최선의 경우 벽시계 가속도는 1.48~1.51배로 추정되나, 현재 하네스는 작업을 건너뛰는 대신 마스킹을 수행하므로 실질적인 가속도는 0%입니다.

의의 및 주장
본 논문은 심층 재귀적 트랜스포머의 "탄력성"이 단순한 학습 불필요 정책을 통해 활용 가능한 안정적인 속성이라고 주장합니다. 핵심 발견은 토큰별 수렴 깊이의 변동은 직접 읽고 활용 가능하며, 이로 인해 복잡한 학습 기반 중단 정책이 불필요하고 오히려 덜 효과적이라는 점입니다. 저자들은 이 연구를 근본적인 탄력적 백본에 대한 기여(이는 Huginn이나 LoopFormer와 같은 선행 연구의 공로로 돌림)라기보다는 측정 및 할당 연구로 위치시킵니다.

본 연구의 범위는 다음과 같이 제한적입니다:

• 결과는 단일 규모(135M 클래스 파라미터)와 단일 시드(주 모델에 대해)에서 확립되었습니다.
• FLOP이 일치하는 패리티(parity)에 대한 주장은 없으며, 150배 더 많은 토큰 수를 가진 참조 모델(SmolLM2-135M)과 비교되었습니다.
• 실질적인 가속도는 최종 지표로 보고되지 않고, FLOP 감소량과 최선의 배치 잠재력 사이의 구간으로 보고되었습니다.
• 훈련 평균 깊이가 포화 깊이라는 가설은 제한된 데이터 포인트에 의해 뒷받침되는 열린 가설로 제시되었습니다.

전체 연구는 단일 RTX 4090에서 약 100 GPU-시간 동안 재현 가능했습니다.