유형: 논문 · 저자: Alexia Jolicoeur-Martineau · 출처: Samsung SAIT AI Lab · 별점: ⭐️⭐️⭐️

전제

문제의식

• LLM은 hard QA에서는 좋은 성과를 내지 못 함
• CoT나 test-time compute로 보완해도 ARC-AGI류에서 한계가 있음
• HRM은 작은 모델과 작은 데이터로 Sudoku, Maze, ARC-AGI에서 LLM 대비 강한 성능을 보였지만 구조가 복잡하여 일반화 및 최적화가 어려움

LLM이 hard QA에서 불리한 이유

• 오토리그레시브 생성은 토큰 1개 오류가 전체 답의 무효로 이어질 수 있어 정답 형식이 엄격한 퍼즐에 취약
• 대안: CoT나 test-time compute
• 대안의 한계
  • CoT는 비용이 크고 고품질 추론 데이터가 필요하며, 생성된 추론 자체가 틀릴 수 있어 취약
  • TTC는 K번 샘플링 후 다수결/최고보상 선택으로 보완하지만, ARC-AGI-2 같은 과제에서 여전히 낮은 정확도를 보임
    • 많이 찍어본다고 해서 본질적으로 못 푸는 문제가 해결되지는 않는다

Hierarchical Reasoning Model (HRM)

• 두 개의 작은 네트워크를 서로 다른 빈도로 재귀적(recursive)으로 구성함으로써 어려운 문제에 접근

  입력 x
     ↓
  중간 상태 z₁ → z₂ → z₃ → ...
     ↓
  출력 y
  

  • 두 개의 네트워크를 쓰는 이유
    • zL (low-level): 빠르게 바뀌어야 하는 정보
    • zH (high-level): 천천히 바뀌어야 하는 정보
    • 하나의 신경망을 쓰면 위 두 개가 구분이 안 된다. ⇒ zL과 zH를 번갈아 업데이트하면서 z 전체가 점점 안정된 상태로 수렴하도록 만든다

입력 x
   ↓
zL ↔ zL ↔ zL ↔ ...
   ↑        ↓
  zH       zH

• 고정점 fixed point 와 IFT
  • 어떤 함수를 반복 적용했을 때 더 이상 값이 변하지 않는 상태 ⇒ $f(z) = z$
  • “z 업데이트를 계속하면 언젠가는 z가 더 이상 변하지 않는 안정된 상태에 도달한다”고 가정
  • 그 고정점을 기준으로 학습
• ACT(Adaptive Computation Time)
  • “z를 몇 번이나 업데이트해야 충분한가?”
  • 모델이 스스로 언제 멈출지 정하게 하자는 발상
• 한계: 복잡하고 생물학적 근거 및 고정점 이론에 의존하고 있어 성능 이해 및 확장이 어려움

TRM(Tiny Recursion Model)

TRM의 가정

1. 작은 네트워크를 반복 적용함으로써 깊은 추론 능력을 확보할 수 있다.
2. 재귀적 추론(recursive reasoning) 구조만으로 충분히 높은 일반화 성능을 얻을 수 있다.
3. 극단적으로 작은 파라미터 수를 유지하면서도 의미 있는 성능을 얻을 수 있다

TRM 특징

1. 단일 신경망 구조
   • 단일 신경망을 재귀적으로 사용 (2-layer, 7M 파라미터)
   • 계산이 단순하고 메모리 부담이 적다
2. 재귀적(reasoning) 업데이트 과정
   • 입력 질문, 현재 답변, 잠재 변수(latent) z 를 같이 재귀적으로 업데이트
   • 입력과 현재 답에 대해 여러 번 같은 신경망을 거치며 답을 점진적으로 개선
3. deep supervision
   • “최종 답만 맞히는 학습”은 중간 과정이 엉망이어도 최종 우연으로 맞을 수 있음
   • deep supervision은 여러 단계에서 답이 점점 좋아지는 방향으로 학습 압력을 줌
4. 고정점에 의존하지 않고 학습
   • fixed point와 1-step gradient approximation같은 정당화가 없어도 된다를 주장

Tiny Recursion Model

• HRM에서 제거한 것들
  1. z를 zL/zH로 나누는 구조와 두 개 네트워크를 하나로 통합
  2. 사고 과정을 명시적 계층 구조로 표현하는 방식
  3. 고정점 이론과 IFT 기반의 이론적 정당화
  4. ACT 기반의 동적 멈춤 결정
• HRM에서 유지/강화한 것들
  1. 재귀적 반복 업데이트 계산 방식
  2. 중간 단계마다 답이 개선되도록 강제하는 deep supervision 학습
  3. 추론 깊이를 반복 횟수로 제어한다는 관점

TRM 플로우

1. 입력 질문(x) 을 숫자 형태로 변환
2. 초기 답(y) 을 무작위로 설정
3. 현재 상태(z)를 초기화
4. TRM 신경망에 현재 질문(x), 답(y), 상태(z)를 넣음
5. z를 먼저 업데이트
6. 그 결과로 나온 z를 기반으로 답(y) 도 다시 계산
7. 이 과정을 여러 번 반복

성능

모델	파라미터 개수	ARC-AGI-1 정확도	ARC-AGI-2 정확도
TRM (7M)	약 7M	45%	8%
LLM (Deepseek R1, o3-mini 등)	수십억 이상	더 낮은 경우 많음	더 낮은 경우 많음

어떤 과제에서는 self-attention을 MLP로 바꾸는 식의 단순화가 유리할 수 있음

• self-attention
  • 입력 시퀀스 전체에서 어떤 토큰이 어떤 토큰과 관련 있는지를 동적으로 학습
  • 계산량: 입력 길이를 N이라고 하면 O(N²)
  • 문제점
    • 입력이 조금만 바뀌어도 attention 분포가 크게 바뀔 수 있다
    • 특히 작은 모델, 작은 데이터 환경에서는 attention이 쉽게 과적합되거나 붕괴된다
    • attention head마다 다른 분포를 학습하면서 학습 안정성이 낮아질 수 있다
    • 반복 적용 시 attention 분포가 점점 치우치거나 특정 토큰에만 집중하는 현상이 발생하기 쉽다
  • 관계 구조가 불규칙하고 가변적인 문제에 최적화된 연산
    • Hard QA는 어떤 토큰이 중요한가보다 “전체 상태가 규칙을 만족하는가”가 핵심
• MLP
  • 계산량: O(N) 또는 O(N·D)
  • 고정된 가중치로 입력 벡터 전체를 일관된 방식으로 변환
  • 입력 변화에 대해 출력 변화가 상대적으로 부드럽다
  • 작은 모델, 작은 데이터, 반복 구조에서는 MLP가 수치적으로 더 안정적일 수 있다
  • 반복 적용 시 출력 분포가 비교적 예측 가능하고 폭주/붕괴 위험이 낮다