기술 요약: 확률적 초소형 재귀 모델 (PTRM)

문제 정의

초소형 재귀 모델 (TRM) 은 복잡한 추론 작업을 위한 자기회귀형 대규모 언어 모델 (LLM) 에 대한 유망한 아키텍처 대안입니다. 토큰 시퀀스를 생성하는 대신 잠재 상태를 반복적으로 정제함으로써 TRM 은 수두레드 (Sudoku-Extreme) 와 ARC-AGI 와 같은 퍼즐에서 수십억 개의 파라미터 대비 수 orders of magnitude 적은 파라미터 (예: 7M 대 수십억) 로 높은 성능을 달성합니다. 그러나 표준 TRM 추론은 결정론적입니다. 해당 논문은 결정론적 재귀가 종종 잠재 공간 내의 "나쁜 분지 (bad basins)"에 갇힌다는 치명적인 한계를 지적합니다. 이는 잘못된 답변으로 해독되는 영역으로, 모델이 이를 탈출할 수 없습니다.

TRM 은 이러한 문제를 해결할 잠재적 능력을 갖추고 있지만, 표준 추론 절차는 해결 공간을 충분히 탐색하지 못합니다. 여러 확률적 추론 궤적 (예: Chain-of-Thought 변형을 통한) 을 샘플링하고 투표나 검증을 통해 최선의 것을 선택함으로써 테스트 시간 컴퓨팅을 활용할 수 있는 LLM 과 달리, TRM 은 재학습이나 작업별 입력 교란 없이 다양한 궤적을 생성할 수 있는 메커니즘이 부족합니다.

방법론: 확률적 TRM (PTRM)

저자들은 기존 "깊이" 축 (재귀 단계 수) 을 보완하는 새로운 "너비" 축을 도입하여 테스트 시간 컴퓨팅을 확장하는 작업 무관 (task-agnostic) 프레임워크인 **확률적 TRM (PTRM)**을 제안합니다. PTRM 은 학습 절차나 모델 아키텍처를 수정하지 않고 작동합니다.

핵심 메커니즘

PTRM 은 다음 단계를 통해 결정론적 TRM 추론을 확률적 과정으로 변환합니다:

1. 확률적 롤아웃 (Stochastic Rollouts): 단일 궤적 대신 모델이 주어진 입력에 대해 $K$개의 병렬 롤아웃을 실행합니다.
2. 노이즈 주입: 모든 깊은 재귀 단계 (감독 단계) 에서 가우시안 노이즈 ($\epsilon \sim \mathcal{N}(0, \sigma^2 I)$) 를 잠재 상태 $z$에 주입합니다. 이 교란은 서로 다른 롤아웃이 고유한 궤적을 탐색하도록 강제하고 나쁜 분지에서 탈출할 가능성을 높입니다.
3. Q-Head 를 통한 선택: 기존 TRM 아키텍처에는 네트워크와 함께 학습되어 정확도 분류기로 작용하는 Q-Head($f_Q$) 가 포함되어 있습니다 (학습 중 적응형 계산 시간에 사용됨). PTRM 에서 이 Q-Head 는 추론 시 $K$개의 후보 답변 각각을 점수 매기 위해 재사용됩니다. 가장 높은 Q-값을 가진 롤아웃이 최종 출력으로 선택됩니다.

주요 차이점

• 재학습 불필요: PTRM 은 사전 학습된 모든 TRM 체크포인트에 적용 가능합니다.
• 작업 무관성: 작업별 데이터 증강이나 입력 교란이 필요하지 않습니다.
• 병렬화: 너비 확장 ($K$개의 롤아웃) 은 병렬화 가능하여 순차적인 깊이 증가보다 더 실용적인 확장 축을 제공합니다.

주요 기여

1. 잠재 분지 함정 식별: Pencil Puzzle Bench (PPBench) 에 대한 TRM 궤적 분석을 통해 저자들은 실패가 종종 나쁜 분지에서 궤적이 진동함으로써 발생함을 입증했습니다. 그들은 Q-Head 값이 셀 정확도와 밀접하게 연동되어 궤적이 "좋은 분지"에 도달했는지 여부를 나타내는 신뢰할 수 있는 지표임을 보였습니다.
2. 확률적 탈출 메커니즘: 논문은 모든 재귀 단계에서 가우시안 노이즈를 주입함으로써 소수의 롤아웃이 나쁜 분지에서 탈출하여 결정론적 모델이 놓치는 올바른 해법에 수렴할 수 있음을 증명합니다.
3. 테스트 시간 확장 프레임워크: PTRM 은 깊이뿐만 아니라 너비 ($K$) 를 통해 컴퓨팅을 확장하는 방법을 정립하며, 기존 Q-Head 를 검증기로 활용하여 여러 궤적 중 최선의 것을 선택합니다.

실험 결과

저자들은 PTRM 을 결정론적 TRM 베이스라인, 직접 예측 모델, 그리고 최첨단 LLM 과 비교하여 여러 벤치마크에서 평가했습니다.

Pencil Puzzle Bench (PPBench)

• 성능: PTRM($K=100, D=48, \sigma=0.2$) 은 결정론적 TRM 의 집계 정확도를 **62.6%**에서 **91.2%**로 향상시켰습니다.
• 구체적 향상: 결정론적 베이스라인이 어려움을 겪던 퍼즐 유형에서 상당한 개선이 관찰되었습니다. 예를 들어 수두레드 (46.7% $\to$ 97.8%) 와 타파 (40.0% $\to$ 80.0%) 가 있습니다.
• LLM 과의 비교: PTRM 은 단일 최강 최첨단 LLM(직접 전략) 보다 56.5 포인트 (91.2% 대 34.7%) 더 높은 성능을 보였으며, 7 개의 최강 LLM 앙상블 (완벽한 검증자 가정) 보다 36 포인트 (91.2% 대 55.1%) 더 높은 성능을 보였습니다.
• 비용 효율성: PTRM 은 LLM 앙상블 비용의 0.0001 배 미만으로 이러한 결과를 달성했습니다 (올바른 퍼즐당 $0.001 대$38.51).

Sudoku-Extreme

• PTRM 은 결정론적 베이스라인 87.4% 를 상회하는 최첨단 98.75% 정확도 (best-Q@K) 를 달성했습니다.

Maze-Hard 및 ARC-AGI-2

• Maze-Hard: 정확도가 83.8% 에서 86.73%(best-Q@K) 로 향상되었으며, pass@K 는 95.63% 에 도달했습니다.
• ARC-AGI-2: 검증이 더 어려운 작업의 경우 기존 Q-Head 가 올바른 해를 구별하는 데 덜 효과적일 수 있음을 시사하는 modest 한 향상 (pass@1 7.36% 에서 8.47% 로) 이 관찰되었습니다.

중요성 및 주장

이 논문은 PTRM 이 작은 재귀 모델의 잠재적 추론 능력을 해제하는 단순하고 다용도이며 매우 효율적인 방법을 제공한다고 주장합니다. 확률성을 도입하고 모델 자체의 Q-Head 를 선택에 활용함으로써 PTRM 은 더 큰 모델을 학습하는 계산 오버헤드나 작업별 엔지니어링의 복잡성 없이 상당한 정확도 향상을 가능하게 합니다.

저자들은 PTRM 이 고유한 해와 검증 가능한 정확성을 가진 추론 퍼즐에서 뛰어나지만, 그 효과는 검증자로서의 Q-Head 의 품질에 달려 있음을 강조합니다. Q-Head 가 올바른 궤적과 잘못된 궤적을 구별하는 데 어려움을 겪는 작업 (예: ARC-AGI-2) 의 경우 성능 향상은 더 제한적입니다. 이 연구는 향후 개선이 롤아웃 너비 확장뿐만 아니라 더 강력한 검증자 개발에서 나올 수 있음을 시사합니다.

인정된 한계:

• 실험은 계산 제약으로 인해 작은 그리드 크기를 가진 추론 퍼즐로 제한되었습니다.
• 방법의 효과는 작업에 따라 달라지며, 수두레드에 비해 ARC-AGI-2 와 Heyawake 에서는 향상이 더 작습니다.
• Q-Head 가 모든 문제 유형에 대해 완벽한 오라클로 작용할 수 있다는 보장은 없습니다.