PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking

Looped language models can run more iterations at test time, but additional iterations need not help them solve longer problems. For sequential state tracking, final-answer supervision does not specify how far each iteration should advance, and repeated updates can change the scale of the recurrent state. We introduce Progress-Aligned Recurrence (PAR), which trains the prediction at iteration k to match the state after the first k operations and rescales recurrent states to their first-iteration norms. On swap-S_5, a Qwen2.5-0.5B model trained on at most 32 moves achieves 95.5\% mean full-state accuracy at 512 moves, or 16times the training length, compared with 55.7\% reported for RLT and 0.8\% for explicit chain of thought on the same backbone. Controlled S_5 ablations show that removing norm anchoring reduces the extrapolation factor from 16times to below 2.2times, and replacing progress-aligned targets with final-answer supervision reduces it to at most 8times. A learned completion readout selects when to answer without a supplied stopping depth. When intermediate labels are unavailable, prefix self-distillation induces the correspondence from shorter problem prefixes using only final answers. The gains apply across sequential state-tracking tasks, establishing the task structure required for depth extrapolation. These findings support training intermediate predictions and state scale together to make additional recurrent depth useful for longer problems.

한 줄 요약

PAR은 반복적 상태 추적에서 길이 일반화를 위해 반복 단계와 작업 단계를 정렬하고 상태 규모를 제어하는 방법이다.

핵심 기여도

PAR은 반복 k에서 첫 k개 작업 후 상태를 예측하도록 훈련하며, RMS 노름을 사용해 상태 규모를 제어한다.
swap-S₅ 데이터셋에서 32단계 훈련 모델이 512단계에서 95.5%의 정확도를 달성해 RLT(55.7%)와 CoT(0.8%)를 초과한다.
norm anchoring 제거 시 길이 일반화 비율이 16배에서 2.2배 미만으로 감소하며, 최종 정답 감독 대체 시 8배로 제한된다.
prefix self-distillation을 통해 중간 라벨 없이도 정답만으로 상태 추적을 학습한다.

핵심 아이디어

기존 반복 언어 모델은 추가 반복이 긴 문제 해결에 도움이 되지 않는다는 문제를 해결하기 위해, PAR은 반복 k에서 첫 k개 작업 후 상태를 예측하도록 훈련한다. 이는 반복과 작업 단계를 정렬하는 **depth–length correspondence**를 구축한다. 또한, 반복 간 상태의 규모가 일관되도록 RMS 노름을 사용해 **norm anchoring**을 적용한다. 이는 상태 드리프트를 방지하고, 반복이 추가될수록 정확도가 유지되도록 한다. 예를 들어, 규칙 테이블 체인 작업에서 반복 k는 k번의 조회 후 값을 예측하도록 학습된다.

기술적 접근법

**Progress-Aligned Recurrence (PAR)**: 반복 k에서 첫 k개 작업 후 상태를 예측하도록 훈련.
**RMS norm anchoring**: 반복 간 상태 벡터를 첫 반복의 RMS 노름으로 재조정.
**Learned halting rule**: 기존 출력 레이어를 활용해 실행 완료 여부를 예측.
**Prefix self-distillation**: 중간 라벨 없이 짧은 문제의 예측을 긴 문제의 반복 타겟으로 사용.
**Qwen2.5-0.5B 모델**: 최대 32단계 훈련 후 512단계에서 95.5% 정확도 달성.

주요 결과

**swap-S₅ 데이터셋**: Qwen2.5-0.5B 모델이 512단계에서 95.5%의 mean full-state 정확도 달성.
**기존 방법 대비 성능**: RLT(55.7%) 대비 +39.8%, CoT(0.8%) 대비 +94.7%.
**norm anchoring 제거 시**: 길이 일반화 비율 16배 → 2.2배 미만.
**최종 정답 감독 사용 시**: 최대 8배로 제한됨.

의의 및 한계

PAR은 반복 단계와 작업 단계를 정렬함으로써 긴 문제 해결 능력을 향상시키며, 상태 규모를 일관되게 유지함으로써 정확도를 보장한다. 이는 순차적 상태 추적 작업에서 길이 일반화를 위한 구조적 접근법을 제시한다. 그러나, PAR은 특정 작업 구조(예: 규칙 테이블, 객체 추적)에 의존하며, 일반적인 언어 모델 아키텍처에 적용 시 추가 조정이 필요할 수 있다. 또한, 학습 데이터가 짧은 문제에만 제한될 경우, 긴 문제에서의 성능은 훈련 방법에 크게 의존한다.

실용적 활용

PAR은 긴 시퀀스를 처리해야 하는 자동화 시스템, 규칙 기반 추론 엔진, 다단계 객체 추적 등에서 활용 가능하다. 특히, 중간 라벨이 부족한 상황에서 prefix self-distillation을 통해 효율적으로 학습할 수 있어, 실제 산업 데이터셋에 유용하다.