1. 문제 정의 (Problem)

최근 'Test-time Scaling'은 추론 시 추가적인 계산 자원을 할당하여 모델 성능을 향상시키는 유망한 방향으로 부상했습니다. 특히 Transformer 기반의 순환 (Recurrent) 또는 반복 (Iterative) 모델 (예: Universal Transformer, PonderLM 등) 은 추론 시 숨은 상태 (hidden states) 를 반복적으로 정제하여 성능을 높입니다.

그러나 기존 접근법에는 다음과 같은 한계가 존재합니다:

• 고정된 반복 횟수: 대부분의 사전 학습된 순환 모델은 추론 시 고정된 횟수만큼만 반복을 수행합니다.
• 비효율적인 자원 할당: 이는 쉬운 토큰 (trivial tokens) 에는 불필요한 계산 자원을 낭비하고, 복잡한 토큰 (complex tokens) 에는 충분한 계산 시간을 제공하지 못하는 결과를 초래합니다.
• 수동 튜닝의 필요성: 기존 적응적 방법들은 종종 수동으로 조정된 가지치기 비율 (pruning ratios) 이나 외부 감독 신호 (hard labels) 에 의존합니다.

2. 방법론 (Methodology)

AdaPonderLM은 PonderLM 아키텍처를 기반으로 하며, 사전 학습 단계에서 토큰별 조기 종료 (early exiting) 를 학습하는 엔드 - 투 - 엔드 자기지도 프레임워크를 도입합니다.

2.1. MLP 기반 게이트 및 모델 프레임워크

• 반복적 정제 (Recurrent Refinement): 동일한 Transformer 백본을 여러 번 반복 실행하여 토큰 표현을 정제합니다.
• 반복별 게이트 (Iteration-specific MLP Gates): 각 반복 단계 $i$마다 토큰별 게이트 로짓 (logit) 을 생성하는 전용 2 층 MLP 를 사용합니다. (공유 MLP 는 불안정하여 각 반복마다 독립적인 MLP 를 사용함).
• 단조 가지치기 (Monotonic Pruning): 게이트 확률 $s_i$가 임계값 $\tau$ (예: $10^{-4}$) 미만일 때 해당 토큰을 가지치기 (정지) 합니다. 일단 정지된 토큰은 이후 반복에서 다시 활성화되지 않도록 **지속적 마스크 (persistent mask)**를 유지합니다.
• 확률 가중 업데이트: PonderLM 방식처럼 다음 토큰 분포를 기반으로 임베딩을 업데이트하고, 게이트된 잔차 연결을 통해 토큰 표현을 갱신합니다.

2.2. KV 재사용 메커니즘 (KV Reuse Mechanism)

• 훈련 - 추론 일관성 보장: 토큰이 가지치기되면 이후 반복 단계에서 새로운 Key/Value 상태를 생성하지 않고, 이전 반복 단계의 캐시된 KV 상태를 재사용합니다.
• 실용적 가속화: 이 메커니즘은 Attention 연산 자체를 변경하지 않으므로 Flash Attention 과 같은 최적화된 커널과 호환되며, 불규칙한 제어 흐름이나 불필요한 재계산을 방지하여 실제 추론 속도를 높입니다.

2.3. 사전 학습 알고리즘

• 2 단계 학습 전략:
  1. Warm-up: pondering 정규화 없이 일반적인 언어 모델 손실 ($L_{CE}$) 만으로 학습.
  2. Regularization: pondering 정규화 항 ($L_{ponder}$) 을 추가. 이는 모든 토큰과 반복 단계의 게이트 값 중 하위 $k$% (bottom-$K$) 에 패널티를 부과하여 모델이 불필요한 반복을 줄이도록 유도합니다.
• Warmup Schedule: $k$ (가지치기 목표 비율) 를 학습 단계에 따라 점진적으로 증가시켜 학습 안정성을 확보합니다.

3. 주요 기여 (Key Contributions)

1. 자기지도 학습 기반 적응적 깊이: 추가적인 감독 신호 없이 사전 학습 단계에서 토큰별 정지 행동을 학습하는 엔드 - 투 - 엔드 프레임워크 제안.
2. 자동 자원 할당: 토큰별 또는 레이어별 가지치기 비율을 수동으로 조정할 필요 없이, 모델이 토큰의 난이도에 따라 pondering 예산을 자동으로 분배.
3. KV 재사용 메커니즘: 순환 디코딩 시 가지치기된 토큰의 KV 캐시를 재사용하여 훈련과 추론의 일관성을 유지하면서 실제 가속화를 실현.
4. 광범위한 실험: Pythia 아키텍처 (70M, 410M, 1.4B, 2.8B 파라미터) 를 대상으로 한 사전 학습 및 계속 학습 (continued pretraining) 검증.

4. 실험 결과 (Results)

• 성능 vs 효율성 트레이드오프:
  • 70M/410M 모델: PonderLM 및 다른 순환 모델 (Loop Transformer, Pause Token 등) 과 비교하여 추론 FLOPs 를 약 8~10% 감소시키면서도 언어 모델링 퍼플렉시티 (PPL) 와 검증 손실을 동등하거나 더 나은 수준으로 유지했습니다.
  • 1.4B/2.8B 모델: 계속 학습 (continued pretraining) 을 통해 확장된 모델에서도 PonderLM 의 성능을 유지하면서 추론 오버헤드를 4 배에서 3.7~3.8 배 수준으로 낮췄습니다.
• 하류 작업 (Downstream Tasks): LAMBADA, PIQA, WinoGrande 등 다양한 벤치마크에서 AdaPonderLM 은 고정된 계산량을 가진 PonderLM 과 비교하여 약 10% 의 FLOPs 절감을 달성하면서도 경쟁력 있는 정확도를 유지했습니다 (2.8B 기준 0-shot 평균 정확도 +2.2%, 5-shot +3.5% 향상).
• 적응적 행동 분석:
  • 학습된 게이트는 높은 NLL(난이도) 을 가진 토큰에 더 많은 계산 단계를 할당하고, 쉬운 토큰은 초기 단계에서 정지시킵니다.
  • Iso-FLOPs 조건 비교: 학습된 적응적 정책 (Adaptive Policy) 은 균일 분포 (Uniform) 나 기하급수적 분포 (Geometric) 와 같은 고정된 가지치기 전략보다 낮은 평가 손실을 보여주어, 토큰별 동적 깊이 할당의 우수성을 입증했습니다.

5. 의의 및 결론 (Significance)

AdaPonderLM 은 고정된 반복 횟수의 비효율성을 해결하고, 토큰별 난이도에 기반한 동적 계산 할당을 자기지도 학습으로 가능하게 한 획기적인 접근법입니다.

• 효율성: 불필요한 반복 계산을 제거하여 추론 비용을 절감하면서도 모델의 표현 능력을 유지합니다.
• 실용성: KV 재사용 메커니즘을 통해 이론적인 효율성이 실제 하드웨어 가속화로 이어지도록 설계되었습니다.
• 미래 방향: 이 연구는 테스트 시간 확장 (Test-time Scaling) 을 위한 자기지도 적응적 순환 메커니즘이 대규모 언어 모델의 효율성을 높이는 유망한 방향임을 시사합니다.

요약하자면, AdaPonderLM 은 "어떤 토큰이 얼마나 생각 (pondering) 해야 하는지"를 모델 스스로 학습하게 함으로써, 계산 자원을 가장 필요한 곳에 집중시키는 지능적인 언어 모델을 구현했습니다.