https://arxiv.org/abs/2506.21734
Hierarchical Reasoning Model
Reasoning, the process of devising and executing complex goal-oriented action sequences, remains a critical challenge in AI. Current large language models (LLMs) primarily employ Chain-of-Thought (CoT) techniques, which suffer from brittle task decompositi
arxiv.org
현재 LLM 기반 reasoning은 CoT(Chain-of-Thought) 방식에 의존하지만 다음과 같은 한계들이 있다.
- 취약한 작업 분해 (Token-level reasoning)
- 대규모 데이터 요구 (reasoning 학습에 필요한 데이터)
- 높은 지연 시간
이를 해결하기 위해 본 논문에서는 "인간 뇌의 계층적/다중 시간척도 처리"에서 영향을 받은 HRM(Hierachical Reasoning Model)을 제안한다.
HRM은 2개의 상호작용하는 순환 모듈로 구성된 새로운 순환 아키텍처이다.
- High-level module: 느리지만 추상적인 계획 담당
- Low-level module: 빠르고 세밀한 계산 담당
이러한 순환 아키텍처를 통해 CoT 없이도, 단일 forward pass로 복잡한 reasoning을 수행할 수 있다.
HRM(Hierachical Reasoning Model)
1. 계층적 처리
- 뇌는 여러 계층의 피질 영역을 통해 정보를 처리한다.
- 상위 영역: 추상적이고 장기적인 정보 통합 → High-Level module
- 하위 영역: 구체적이고 빠른 센서/운동 처리 → Low-Level module
2. 시간적 분리
- 뇌의 계층적 처리 수준은 서로 다른 시간 리듬으로 동작한다.
- 상위 영역: 느린 리듬 → High-Level module은 느리게
- 하위 영역: 빠른 리듬 → Low-Level module은 빠르게
- 이렇게 속도를 다르게 가져가면, 느린 상위 처리가 빠른 하위 처리를 안정적으로 조율할 수 있다.
3. 순환 연결
- 뇌는 순환적 연결이 매우 풍부
→ 상향/하향 피드백 루프가 끊임없이 반복되며 정보 표현을 점진적으로 정교하게 다듬음 - 이러한 반복 구조를 통해 더 정밀하고 문맥에 민감한 판단을 가능하게 해준다.
- 하지만, 뇌는 BPTT(Back Propagation Through TIme) 같은 복잡한 비효율적인 학습 방식 없이도 학습
→ BPTT 없이도 학습 가능한 구조와 학습 알고리즘 설계
HRM 계산 과정
- 입력 처리
입력 $x$를 네트워크에 넣기 전, 내부 표현으로 변환:
$$\tilde x=f_I(x;\theta_I)$$ - 시간 구조
총 계산은 N개의 High-level 사이클, 각 사이클은 T개의 Low-level 스텝으로 구성
→ 총 시간 스텝: $i=1,\dots,\mathrm N\times \mathrm T$ - Low-level 상태 업데이트 (모든 스텝마다)
$$z_L^{(i)}=f_L(z_L^{(i-1)},z_H^{(i-1)},\tilde x;\theta_L)$$ - High-level 상태 업데이트 (T번에 1번만)
$$z_H^{(i)}=\begin{cases}f_H(z_H^{(i)},z_L^{(i-1)};\theta_H)\quad \mathrm{if}\ i\equiv 0 \pmod T\\z_H^{(i)}\quad \quad \quad \quad \quad \quad \quad \ \mathrm{otherwise}\end{cases}$$ - 최종 출력 생성
$$\hat y=f_O(z_H^{(NT)};\theta_O)$$
→ 입력 $x$를 내부 표현으로 변환한 뒤, Low-level 모듈이 T번 상태를 업데이트하고, 이후 High-level 모듈이 한 번 상태를 갱신한다. 이 과정을 N번 반복한 후, 최종 High-level 상태를 출력 네트워크에 넣어 예측 결과 $\hat y$를 생성한다.
→ 매 사이클마다 High-level 상태 $z_H$가 새롭게 갱신되기 때문에, Low-level 모듈은 매번 서로 다른 수렴 경로를 따라 연산을 수행할 수 있다.
※ 위 과정이 끝난 뒤, 예측이 충분히 확실하면 정지하고 $\hat y$를 출력하고, 아니라면 다음 forward pass로 이어감
Hierarchical convergence
- HRM: 계층적 수렴 방식으로 High는 느리게 수렴, Low은 반복 수렴(local 수렴)하며 High가 리셋함
- RNN: 빠르게 수렴하며, residual이 급격히 0으로 가며, 깊은 계산이 어려움
- DNN: 앞/뒤에서만 의미 있는 변화 발생, 중간층은 vanishing gradient 문제로 인해 gradient가 거의 없어짐
Approximate gradient
HRM은 고정적 수렴 특성을 활용해, 모든 상태를 저장하지 않고 마지막 상태만을 기반으로 하는 1-step gradient approximation을 사용한다. 이 방법은 다음과 같은 경로로 그래디언트를 계산한다.
Output → H 모듈의 마지막 상태 → L 모듈의 마지막 상태 → 입력 표현
이는 Deep Equilibrium Model(DEQ)의 수학적 이론을 기반으로 하며, Implicit Function Theorem(IFT)을 통해 전체 연산 과정을 단일 단계로 근사한다. 고정점 근처에서 ($I-J_F)^{-1}\approx I$로 가정하면, 복잡한 역전파를 생략하고 다음과 같이 단순한 근사식으로 파라미터별 그래디언트를 계산할 수 있다.
$${\partial z_H^*\over\partial\theta_H}\approx{\partial f_H\over\partial \theta_H}\\ {\partial z_H^*\over \partial\theta_L}\approx{\partial f_H\over\partial z_L^*}\cdot{\partial z_L^*\over\partial\theta_L}\\{\partial z_H^*\over \partial\theta_I}\approx{\partial f_H\over\partial z_L^*}\cdot{\partial z_L^*\over\partial \theta_I}$$
이 방식은 $O(1)$의 메모리로 동작한다.
Deep Supervision
뇌의 주기적 신경 진동이 학습 타이밍을 조절한다는 생물학적 원리에 착안하여, 학습 중간에도 손실을 계산하는 방식으로 입력 $(x, y)$에 대해 모델을 여러 번 forward pass하는데, 이를 segment라고 부른다. 각 segment $m$ 마다 다음 절차를 수행한다.
- 이전 segment의 hidden state $z^{m-1}$를 받아, HRM을 한 번 실행해 현재 상태 $z^m$ 과 예측값 $\hat y^m$을 계산
- 현재 예측에 대한 loss $\mathcal L_m = \mathrm{LOSS}(\hat y^m, y)$ 계산
- 2. 에서 구한 loss에 대해 파라미터 $\theta$를 한 번 업데이트
이전 segment의 hidden state를 detach하기 때문에, gradient는 현재 segment까지만 흐르고 이전으로는 전파되지 않는다. 즉, 1-step gradient approximation을 segment 단위로 적용한 셈이다.
이 방식은 High-level module에 더 잦은 피드백을 주고, Jacobian 기반 정규화보다 안정적이며 학습 성능도 우수하다고 보고되었다.
Adaptive Computational Time (ACT)
각 segment가 끝날 때마다 High-level module의 최종 상태를 기반으로 Q-head가 "계속"할지 "멈출"지 결정하는 Q-value $(\hat Q_\mathrm{halt}, \hat Q_\mathrm{continue})$를 예측한다. 멈출 결정은 다음 두 조건 중 하나를 만족할 때 이루어진다.
멈춤 결정은 다음 두 조건 중 하나를 만족할 때 이루어진다.
1. segment 개수가 최대 허용치 $M_\mathrm{max}$에 도달했을 때
2. $M_\mathrm{min}$ 이상 segment를 진행했고, $\hat Q_{halt}>\hat Q_\mathrm{continue}$인 경우
각 segment에서는 예측 정확도에 따라 보상이 주어지고, 이 과정은 에피소드 기반 Q-learning으로 학습된다. 최종적으로, loss는 예측 손실과 Q-head의 이진 분류 손실을 더해 구성된다.
$$\mathcal L_\mathrm{ACT}^{(m)}=\mathrm{Loss}(\hat y_m, y)+\mathrm{BinaryCrossEntropy}(\hat Q_m, \hat G_m)$$
→ 모델이 과제 난이도에 따라 연산 자원을 유연하게 조절할 수 있게 하며, 고정 계산량을 사용할 때보다 더 효율적인 성능을 보임
Inference-time Scaling
HRM은 학습이나 구조 변경 없이, 단순 최대 세그먼트 수 M_\mathrm{max}를 늘리는 것만으로 자연스럽게 이 기능을 구현한다.
- ex. 스도쿠처럼 깊은 추론이 필요한 문제에서는 계산량을 늘릴수록 성능이 뚜렷하게 향상됨
- ex. ARC-AGI 처럼 간단한 변환만 필요한 작업에서는 추가 계산이 큰 효과를 주지 않음
→ 추론 시간에서의 유연한 성능 향상 가능
Stability of Q-learning in ACT
Q-learning이 수렴할 조건
- 모델 파라미터가 유한한 범위 내에 있을 것
- 학습 시 weight decay 사용
- 정규화 레이어 적용
HRM
- RMSNorm: post-norm 구조로 작동하는 레이어 정규화
- AdamW Optimizer: weight decay 포함, 모든 파라미터를 {1\over \lambda}이하로 제한
→ HRM은 Q-learning이 수렴할 조건을 모두 만족
→ HRM은 추가적인 안정화 기법 없이도 Q-learning을 안정적으로 구현할 수 있음
Results
1. Benchmarks
- ARC-AGI: AI가 유도 추론을 통해 보지 못한 문제를 일반화할 수 있는지를 평가하기 위해 설계된 벤치마크
- Sudoku-Extreme: 고난이도 스도쿠 퍼즐로 구성된 벤치마크 데이터셋
- Maze-Hard: 30x30 미로에서 최단 경로를 찾는 과제
2. Evaluation Details
- HRM은 Pretrained 나 CoT 없이도 일반적인 Transformer 모델을 크게 능가
- 특히 reasoning이 중요한 벤치마크(ARC, Sudoku, Maze)에서 소량 데이터만으로도 우수한 성능
- ACT와 계층적 반복 구조가 모델의 핵심 경쟁력
3. Visualization of intermediate timesteps
- Maze-Hard: 여러 경로 동시 탐색 → 필터링 → 구조 형성 → 반복 개선
- Sudoku: 깊이 우선 탐색(DFS) 처럼 작동 + 후보 해법 탐색하다가 막다른 길이면 백트래킹
- ARC-AGI: 정답에 가까워지도록 단계별로 개선
→ HRM은 단일한 추론 방식에 고정되지 않고, 과제에 따라 유연하게 전략을 바꿈
Brain Correspondence
Participation Ratio(PR): 뉴런 활동의 유효 차원 수를 측정하는 지표
→ 정보 표현이 얼마나 고차원 공간에서 다양하게 펼쳐져 있는지를 나타냄
생쥐 대뇌피질의 계층적 PR 분포
- (a) - 생쥐 뇌의 기능 영역 시각화
- (b) - 각 영역의 PR 계층 수준에 따라 비교
- Spearman 상관계수 $\rho$=0.79 → 계층이 높을수록 고차원 표현
학습된 HRM의 PR 분석
- (c) - 다양한 태스크에서 HRM의 고수준 상태의 PR은 태스크 수가 많아질수록 증가
- HRM은 더 다양한 문제를 처리하기 위해 표현 공간을 확장함
- 저수준 상태의 PR은 안정적 → 역할이 보다 구체적, 제한적
- (d) - 100개 Sudoku 풀이 과정에서 측정한 PR
- 고수준 모듈이 훨씬 고차원 표현을 학습
→ 학습을 통해 HRM은 명확한 계층적 표현 구조를 형성
학습되지 않은 HRM과의 비교
- (e) - 무작위 가중치를 가진 HRM은 태스크 수가 늘어도 PR 증가 없음
- (f) - 두 모듈 모두 낮고 비슷한 PR 값
→ 고차원적 계층 표현은 문제 해결 능력을 학습하면서 자연스럽게 나타나는 특징
Conclusion
HRM의 성공은 기존 비계층적 AI 모델 패러다임에 도전하며, 뇌의 방식같은 계층적 추론이 CoT 기반 방식 외에 또 하나의 강력한 대안적 추론 프레임워크로서 계층적 추론 모델의 가능성을 제시한다. 이는 Turing-complete universal computation 을 향한 기반 모델의 방향성을 제시할 수 있다.