학습된 확률적 중단 기법을 통한 루프 트랜스포머의 안정적인 외삽

Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

루프 트랜스포머는 공유 트랜스포머 블록을 반복적으로 적용하며, 가변 길이 알고리즘 작업에 구조적으로 적합합니다. 이들은 훈련 시퀀스의 길이를 초과하는 일반화 능력을 보이지만, 이러한 특성은 불안정하여, 좋은 성능을 보이는 분포 내 솔루션에서도 높은 이상치(OOD) 변동성을 나타냅니다. 우리는 이러한 변동성이 간단한 알고리즘 작업에서 시퀀스 길이와 루프 횟수 간의 인과 관계에서 비롯된다는 것을 확인했습니다. 학습 과정에서 루프 횟수에 확률적 요소를 도입하면 OOD 변동성이 크게 감소하고 추론 시간 동안의 루프 횟수에 따른 예측 안정성이 향상됩니다. 휴리스틱한 무작위화 방식보다 더 나은 성능을 얻기 위해, 강화학습 기반 중단(RL-Halting)을 학습된 확률적 스케줄로 분석했으며, 일반적으로 정확성과 안정성 간의 균형을 개선하는 것을 확인했습니다. 이진 덧셈, Dyck-1, Unique Set 및 Copy 작업에서 학습된 확률적 중단은 종종 이러한 균형을 개선하지만, 최적이 아닌 계산을 안정화시키기도 합니다. 본 연구는 '언제 중단할 것인가'를 추론 시간의 단순한 계산 할당 규칙이 아닌, 훈련 시간을 고려한 설계 선택으로 간주해야 함을 시사합니다.

Looped Transformers, which repeatedly apply a shared transformer block, are an architecturally natural fit for variable-length algorithmic tasks. Although they can exhibit strong length generalization beyond the length of training sequences, this behavior is brittle, yielding high out-of-distribution (OOD) variance, even across well-performing in-distribution solutions. We trace this variance to the spurious correlation in simple algorithmic tasks between sequence length and number of loops. Introducing stochasticity into the number of loops during training sharply reduces OOD variance and stabilizes predictions across inference-time loop counts. To improve upon heuristic randomization schemes, we further analyze RL-Halting as a learned stochastic schedule and find that it generally improves the accuracy-stability trade-off. Across binary addition, Dyck-1, Unique Set, and Copy, learned stochastic stopping often improves this trade-off but can also stabilize a suboptimal computation. Our work suggests that "when to stop" should be treated as a training-time design choice, not merely an inference-time computation-allocation rule.