Latent Reasoning in TRMs is Secretly a Policy Improvement Operator

기술 요약: TRM의 잠재적 추론(Latent Reasoning)은 사실 정책 개선 연산자(Policy Improvement Operator)이다

1. 문제 정의

Tiny Recursive Models (TRMs) 및 Hierarchical Reasoning Models (HRMs)와 같이 잠재 재귀(latent recursion)를 활용하는 최근의 소형 모델들은 ARC-AGI와 같은 복잡한 추론 작업에서 유망한 결과를 보여주었다. 지배적인 이론은 이러한 모델들이 더 적은 파라미터로 더 큰 모델을 모사할 수 있도록 유효 네트워크 깊이를 증가시킴으로써 성능 향상을 달una다는 것이다. 그러나 경험적 관찰에 따르면, 재귀적으로 추가된 레이어는 동일한 피드포워드 깊이를 가진 원패스(one-pass) 모델에 비해 성능이 떨어지는 경우가 종종 발생한다. 이는 루프 구조를 가진 아키텍처에서 모든 재귀 단계가 의미 있게 기여하지 못하고, 결과적으로 "죽은 연산(dead compute)"이 발생함을 나타낸다.

본 연구가 다루는 핵심 질문은 다음과 같다:

• 언제, 그리고 왜 잠재적 추론이 성능을 향상시키는가?
• 언제 그것이 중복된 연산을 초래하는가?
• 이러한 모델의 모든 단계가 효과적으로 작동하도록 하려면 훈련을 어떻게 개선할 수 있는가?

현재 TRM의 훈련 방식은 주로 최종 출력에만 감독(supervision)을 적용한다. 이는 초기 정교화 단계(early refinement steps)에 약한 신용 할당(credit assignment) 문제를 남겨두어, 특히 데이터가 적은 환경에서 최적화를 취약하게 만든다.

2. 방법론

저자들은 잠재적 추론에 대한 이론적 재정의와 그에 상응하는 훈련 방법론을 제안한다.

이론적 정식화: 정책 개선으로서의 잠재적 추론
본 논문은 잠재 재귀 추론을 **정책 개선 연산자(policy improvement operator)**로 정식화한다.

• 참조 정책 ($\hat{\pi}_t$): 잠재 추론 단계 전의 모델 출력 분포 (즉, "첫 번째 추측").
• 개선된 정책 ($\pi^+_t$): 한 번의 잠재 추론 업데이트 후의 출력 분포.
• 어드밴티지 신호 (Advantage Signal): 개선된 정책과 참조 정책 간의 로그 확률 차이는 암시적 어드밴티지 신호($A_t(s, a) = \log \pi^+_t(a|s) - \log \hat{\pi}_t(a|s)$) 역할을 한다.
• 메커니즘: 저자들은 단일 TRM 단계가 최적성 조건(optimality condition)에 따라 참조 정책을 재가중치화하는 정책 업데이트를 구현함을 입증한다. 구체적으로, 잠재 추론 단계는 정답이 개선된다는 사건에 조건화된 정책처럼 동작한다. 이는 잠재 공간 알고리즘을 강화 학습(RL)의 잘 알려진 정책 개선 프레임워크 및 확산 모델(diffusion models)의 분류기 없는 가이드(Classifier-Free Guidance, CFG)와 연결한다.

제안된 훈련 체계: 심층 개선 감독 (Deep Improvement Supervision, DIS)
정책 개선에 대한 통찰을 바탕으로, 저자들은 "죽은 연산" 단계를 제거하기 위해 **심층 개선 감독 (DIS)**을 도입한다.

• 단계별 타겟: 최종 출력만을 감독하는 대신, DIS는 각 재귀 단계 $s$에 대해 중간 타겟 $\{y^\dagger_s\}$를 제공한다. 이 타겟들은 정답 출력 $y^*$에 대한 **단조적 이산 오염 스케줄(monotonic discrete corruption schedule)**을 통해 생성된다.
• 훈련 목적: 모델은 이전 타겟 $y^\dagger_{s-1}$ (참조 정책을 통해)로부터 다음의 덜 오염된 타겟 $y^\dagger_s$ (개선된 정책을 통해)로 매핑되도록 훈련된다.
• 손실 함수: 훈련은 포스트-추론 로짓(post-reasoning logits)과 단계별 타겟 $y^\dagger_s$ 사이의 교차 엔트로피 손실을 최소화한다. 이는 모든 단계에서 정답 토큰에 대한 로그 비율 어드밴티지가 양수가 되도록 강제하여, 각 재귀 단계가 중복된 연산이 아닌 방향성 있는 개선을 수행하도록 보장한다.

아키텍처 단순화
제안된 방법은 다음과 같이 TRM 아키텍처를 단순화한다:

• 학습된 정지 메커니즘(Adaptive Computation Time)의 필요성을 제거한다.
• 감독 단계 및 잠재 추론 단계의 수를 줄인다.
• 오염 스케줄에 의해 결정되는 고정된 재귀 단계 수($N_{sup}$)를 사용한다.

3. 주요 기여

1. 이론적 통찰: 본 논문은 잠재 추론이 수학적으로 정책 개선 연산자와 동일함을 확립한다. 이는 잠재 공간 추론 알고리즘과 RL 프레임워크 사이의 간극을 메우며, 어드밴티지 극대화의 관점에서 모델의 동작을 설명한다.
2. 심층 개선 감독 (DIS): 단계별 중간 타겟을 제공하는 새로운 훈련 체계를 제시한다. 이는 장기적 신용 할당이라는 어려운 문제를, 각 단계가 이전 상태로부터 명시적으로 개선되도록 훈련되는 다루기 쉬운 지도 학습 과제로 변환한다.
3. 효율성 향상: 제안된 방법은 경쟁력 있는 성능을 달성하면서도 계산 오버헤드를 크게 줄인다. 이 방법은 정지 모듈을 위한 학습을 피하고, 정확도를 유지하거나 향상시키면서 표준 TRM 설정에 비해 총 순전파(forward pass) 횟수를 18배 줄인다.

4. 실험 결과

저자들은 N-Queens 문제와 ARC-AGI 벤치마크(ARC-AGI-1 및 ARC-AGI-2)에서 그들의 방법을 평가하였다.

• N-Queens: 컴팩트 모델(0.8M 파라미터)을 사용하여, DIS 방법은 0.69 정확도를 달성하였으며, 이는 정지 헤드와 16단계의 감독 단계를 가진 표준 TRM의 성능과 일치하지만, 훨씬 적은 추론 단계($T=1, n=2$ vs. $T=3, n=6$)를 사용하였다.
• ARC-AGI:
  • 컴팩트 모델 (0.8M 파라미터): DIS는 ARC-AGI-1에서 24.0% 정확도를 달성하여, 기존의 TRM-compact(12.0%) 및 외부 지식 없이 작동하는 대부분의 오픈 소스 LLM을 능가하였다.
  • 미디엄 모델 (7M 파라미터): DIS는 ARC-AGI-1에서 41.3%, ARC-AGI-2에서 **6.0%**를 달성하여, 표준 TRM-medium(40.4% / 3.3%)을 능가하고 기존 TRM(40.4% / 3.3%)과 대등한 성능을 보이면서도 더 적은 잠재 추론 단계를 사용하였다.
• 연산 감소: 이 방법은 베이스라인 TRM 구성에 비해 총 순전파 횟수를 18배 감소시켜, 중간 단계의 명시적 감독이 과도한 잠재 추론 루프의 필요성을 제거함을 입증하였다.

5. 의의 및 주장

본 논문은 잠재적 추론을 정책 개선 과정으로 바라봄으로써, 왜 표준 재귀 모델이 때때로 실패하는지(단계가 어드밴티지 마진을 개선하지 못하는 "죽은 연산" 때문)와 이를 어떻게 해결할 수 있는지를 설명할 수 있다고 주장한다.

주된 의의는 심층 개선 감독(DIS) 프레임워크에 있으며, 이는 다음을 수행한다:

• 모든 재귀 단계가 비사소한(non-trivial) 개선을 수행하도록 보장하는 건설적인 훈련 규칙을 제공한다.
• 복잡한 정지 메커니즘이나 취약한 신용 할당의 필요성을 제거함으로써 재귀 모델의 훈련을 단순화한다.
• 훈련 목적 함수가 정책 개선의 이론적 메커니즘과 일치할 때, 소형 파라미터 효율적 모델이 복잡한 추론 작업(ARC-AGI 등)에서 최첨단(SOTA) 결과를 달성할 수 있음을 입증한다.

저자들은 자신들의 접근 방식이 잠재 재귀의 내재적 알고리즘 특성과 훈련 신호를 일치시킴으로써, 더 효율적이고 효과적인 추론 모델을 향한 경로를 제공한다고 결론짓는다.