테스트 시점 스케일링 효율 높이는 적응형 루프 트랜스포머

핵심 요약

• 루프 트랜스포머는 계층을 재사용해 매개변수 효율성을 보이지만, 테스트 시 출력이 길어질수록 성능 향상을 제대로 비교한 연구가부족했다.
• 기존 루프 트랜스포머는 비루프 모델 대비 정확도-계산 기울기(테스트 시 디코딩 FLOPs 두 배당 정확도 이득)가 가파르지만, 같은 계산량에서는 비루프 모델에 미치지 못한다.
• 고정 깊이 루핑은 모든 토큰에 추가 반복을 쓰는데, 실제로는 많은 토큰이 추가 반복의 이득을 받지 못한다는 것을 발견했다.
• 이를 해결하기 위해 TaH2를 제안하는데, 모델이 루핑 이득을 받는 토큰에만 추가 반복을 집중하게 한다.
• TaH2는 백본과 반복 결정기를 선행깊이 감독을 통해 함께 학습하는데, 온라인 레이블로 추가 반복이 예측을 개선하는지 판단한다.
• 어려운 AIME 벤치마크에서 TaH2는 기울기를 53% 개선하고(2.74 vs. 1.79), 같은 테스트 시 계산량에서 비루프 기준선을 약 3.4포인트 초과한다.
• 최대 반복 깊이가 증가할 때 기존 루프 모델은 성능이 정체되지만, TaH2는 깊이 2에서 +2.8포인트에서 깊이 8에서 +3.9포인트로 계속 향상된다.
• 코드는 GitHub에서 공개되었다.