[논문 리뷰] Universal Reasoning Model

https://arxiv.org/abs/2512.14693

Universal Reasoning Model

Background

• 최근 연구들에서는 Universal Transformer(UT)가 복잡한 추론 문제(ex. ARC-AGI & Sudoku)에서 강력한 성능을 보인다.
• 작은 UT 모델조차 인터넷 규모의 사전 학습 없이 훈련했음에도 불구, 대부분의 표준 Transformer 기반 LLM을 성능 면에서 능가한다.
• 기존 연구는 UT 성능 향상을 주로 architectural innovation 때문이라고 설명했으나, 실제 UT의 recurrent inductive bias가 핵심 원인

Universal Transformer(UT)

1. 구조적 차이
   • Standard Transformer와 달리 UT는 하나의 transition block을 반복적으로 적용하여 토큰 표현을 점진적으로 정제함
   • 상태 업데이트: 입력 시퀀스 $x$의 임베딩 행렬 $H_0\in\mathbb{R}^{N\times d}$에 대해 반복적으로 계산
     1. Self-Attention 업데이트: $H_{t+1}=\mathrm{LayerNorm}(H_t+\mathrm{MHA}(H_t))$
     2. Transition block 적용: $H_{t+1}\leftarrow \mathrm{LayerNorm}(H_{t+1}+\mathrm{Transition}(H_{t+1}))$
   • 여기서 Transition은 FFN 또는 separable convolution 사용 가능
   • 2D sinusoidal embedding을 추가하여 토큰 위치와 반복 단계 정보를 동시에 인코딩
2. 파라미터 공유
   • UT의 핵심 설계: 깊이에 걸쳐 가중치 재사용
     $$\Theta_\mathrm{UT}=\{W_h^Q,W_h^K,W_h^V,W^O,\Theta_\mathrm{Transition}\}$$
   • 모든 반복 단계 t에서 동일 파라미터 사용 → 모델 크기 증가 없이 반복적 표현 정제 가능
   • 장점:
     1. 추론 시 유연한 반복 단계 수 T 조절 가능 → depth adaption
     2. 고정 깊이 Transformer보다 이론적 표현력 증가
3. Adaptive Computation Time
   • 각 토큰별로 반복 단계에서 연산을 조기 종료 가능
   • 단계 $t$에서 각 위치 $i$가 종료할 확률:
     $$p_{t,i}=\sigma(w^Th_{t,i}+b)$$
   • 누적 확률이 $1-\epsilon$에 도달하면 해당 토큰의 계산 종료
   • 최종 표현은 가중합 형태:
     $$h_i^\mathrm{final}=\sum_t\Delta_{t,i}h_{t,i}$$
   • 장점: 복잡한 토큰은 더 많은 연산, 단순한 토큰은 적은 연산 배정 가능

Universal Reasoning Model(URM)

• Universal Transformer와 유사한 depth-wise 반복 구조
• UT와의 차별점
  • Decoder-only 구조 채택
  • 새로운 ConvSwiGLU 모듈
  • Truncated Backpropagation Through Loops(TBPTL) 도입

Figure 2. Illustration of Universal Reasoning Model (URM) architecture.

ConvSwiGLU: 비선형성 + 국소 문맥 결합

기존 SwiGLU의 한계

• Point-wise FFN 구조 → 각 토큰을 독립적으로 처리
• UT의 반복 추론에서 토큰 간 국소 상호작용 부족

ConvSwiGLU의 핵심 아이디어

• 게이팅 이후 depthwise short convolution을 적용
• 토큰 공간에서 local 문맥 혼합을 유도
• 시퀀스 전체 복잡도는 증가시키지 않음

연산 흐름

1. 차원 확장
   $$[G,U]=XW_\mathrm{up}\in\mathbb R^{T\times 2m}$$
2. SwiGLU 게이팅
   $$H_{ffn}=\mathrm{SiLU}(G)\odot U$$
3. Depthwise 1D Convolution (k=2)
   $$H_\mathrm{conv}=\sigma(W_\mathrm{dwconv}*H_\mathrm{ffn})$$
4. 차원 축소
   $$Y=H_\mathrm{conv}W_\mathrm{down}$$

효과

• 비선형성 강화
• 국소 패턴 인식 능력 증가
• 파라미터 효율성과 반복 추론 특성 유지

Truncated Backpropagation Through Loops (TBPTL)

문제 인식

• 반복 추론 루프 수 $M$이 커질수록:
  • 초기 루프에서 전달되는 gradient가 노이즈 누적
  • 학습 불안정 및 성능 저하 발생

해결 전략

• 초기 반복은 forward-only
• 후반 반복만 gradient 계산

수식적 정의

• D-layer URM을 M번 반복 전개(unroll)
  $$h_t^{(d)}=F_\theta^{(d)}(h_t^{(d-1)},h_{t-1}^{(d)})$$
• Truncation index $N < M$ 설정
  • $t=1~N$: forward only
  • $t=N+1 ~ M$: forward + backward
• 손실 함수
  $$L_\mathrm{TBPTL}(\theta)=\sum_{t=N+1}^ML(h_t^{(D)},y)$$
• gradient는 후반 루프에 대해서만 계산됨

Experiment

Main Results

• pass@10, pass@100, pass@1000으로 갈수록 URM의 격차가 더 커짐
  → URM은 단발 예측이 아니라 반복 추론을 통해 더 "다양하고 질 좋은 후보 해답 공간"을 생성
• Ablation 결과
  • Short Convolution 제거 → 성능 큰 폭 하락
  • Truncated Backpropagation 제거 → 성능 더 큰 폭 하락

Table 1. The performance of URM, TRM, and HRM on three complex reasoning tasks.

Short Convolution 실험

depthwise convolution이 UT의 추론 성능을 높이는 이유는 무엇이며, 어디에 넣어야 가장 효과적인가?

1. 삽입 위치 실험
   • (a) SDPA 출력 직후
   • (b) Value projection 이후
   • (c) Key projection 이후
   • (d) Query projection 이후
   • (e) Multi-head concat 이후, output projection 이전
   • (f) MLP 확장(FFN expansion) 이후
2. 커널 크기 실험
   • ConvSwiGLU에서 사용하는 kernel size 변화에 따른 ARC-AGI pass@1 비교(k=2일때가 성능이 제일 높음)

Figure 3. ARC-AGI pass@1 results for inserting the short convolution module at different positions within the UT transition (left), and varying the kernel size of the ConvSwiGLU module applied after the MLP expansion (right).

Truncated Backpropagation Through Loops

UT의 inner loop에서 모든 반복에 대해 gradient를 전파하는 것이 정말 최선인가? 에 대한 실험결과

• 모든 loop에 gradient를 흘리는 것은 오히려 해롭다.
• 적당한 truncation이 성능과 안정성을 동시에 개선
• UT의 반복 추론은 RNN과 동일하게 TBPTL이 필수적인 학습 기법
• 이는 URM의 성능 향상이 단순 구조 변경이 아니라 학습 동역학까지 고려한 설계임을 보여준다.

Table 3. Effect of Truncated Backpropagation Through Loops (TBPTL) across inner loops oon ARC-AGI 1.

Muon Optimizer

• URM의 높은 성능은 더 좋은 optimizer 덕분인가, 아니면 구조적 귀납 바이어스 덕분인가? 에 대한 실험결과
• Muon은 곡률 정보를 활용해 직교성을 유지하는 방향으로 업데이트하여 gradient noise 감소, 초기 학습 안정화, 빠른 수렴이 이루어진다. 하지만 빨리 도달하는지를 개선했으나 성능을 높이는 것은 바꾸지 못함.
• 즉, 추론 능력의 상한선은 옵티마이저가 아니라 아키텍처가 결정

figure 4. ARC-AGI pass@1 and pass@1000 performance of Adam and Muon optimizers on ARC-AGI 1 and ARC-AGI 2 benchmarks.

Conclusion

1. Short Convolutional Gating (ConvSwiGLU)
   • 비선형 서브스페이스에서 국소 토큰 혼합
   • 반복 추론의 표현력 극대화
2. Truncated Backpropagation Through Loops
   • 반복 깊이에서의 gradient 불안정 해결
   • 학습 안정성과 성능 동시 개선

→ UT는 반복되는 비선형 추론이며, URM은 이를 가장 단순하고 효과적인 형태로 구현한 모델이다.