Hierarchical Reasoning Model — part 2/4
Cleaned-text code points [14504, 30231). Bibliography omitted here; source.md includes it.
The one-step gradient approximation is theoretically grounded in the mathematics of Deep Equilibrium Models (DEQ) 36 which employs the Implicit Function Theorem (IFT) to bypass BPTT, as detailed next. Consider an idealized HRM behavior where, during high-level cycle k k , the L-module repeatedly updates until its state z L z_{L} converges to a local fixed point z L ⋆ z_{L}^{\star} . This fixed point, given the current high-level state z H k − 1 z_{H}^{k-1} , can be expressed as z L ⋆ = f L ( z L ⋆ , z H k − 1 , x ~ , θ L ) . z_{L}^{\star}=f_{L}(z_{L}^{\star},z_{H}^{k-1},\tilde{x};\theta_{L})\;. The H-module then performs a single update using this converged L-state: z H k = f H ( z H k − 1 , z L ⋆ , θ H ) . z_{H}^{k}=f_{H}(z_{H}^{k-1},z_{L}^{\star};\theta_{H})\;. With a proper mapping ℱ \mathcal{F} , the updates to the high-level state can be written in a more compact form as z H k = ℱ ( z H k − 1 , x ~ , θ ) z_{H}^{k}=\mathcal{F}(z_{H}^{k-1};\tilde{x},\theta) , where θ = ( θ I , θ L ) \theta=(\theta_{I},\theta_{L}) , and the fixed-point can be written as z H ⋆ = ℱ ( z H ⋆ , x ~ , θ ) z_{H}^{\star}=\mathcal{F}(z_{H}^{\star};\tilde{x},\theta) . Let J ℱ = ∂ ℱ ∂ z H J_{\mathcal{F}}=\frac{\partial\mathcal{F}}{\partial z_{H}} be the Jacobian of ℱ \mathcal{F} , and assume that the matrix I − J ℱ I-J_{\mathcal{F}} is invertible at z H ⋆ z_{H}^{\star} and that the mapping ℱ \mathcal{F} is continuously differentiable. The Implicit Function Theorem then allows us to calculate the exact gradient of fixed point z H ⋆ z_{H}^{\star} with respect to the parameters θ \theta without explicit backpropagation: ∂ z H ⋆ ∂ θ = ( I − J ℱ | z H ⋆ ) − 1 ∂ ℱ ∂ θ | z H ⋆ . \frac{\partial z_{H}^{\star}}{\partial\theta}=\left(I-J_{\mathcal{F}}\big|_{z_{H}^{\star}}\right)^{-1}\frac{\partial\mathcal{F}}{\partial\theta}\bigg|_{z_{H}^{\star}}\;. (1)
Calculating the above gradient requires evaluating and inverting matrix ( I − J ℱ ) (I-J_{\mathcal{F}}) that can be computationally expensive. Given the Neumann series expansion, ( I − J ℱ ) − 1 = I + J ℱ + J ℱ 2 + J ℱ 3 + … , (I-J_{\mathcal{F}})^{-1}=I+J_{\mathcal{F}}+J_{\mathcal{F}}^{2}+J_{\mathcal{F}}^{3}+\dots\,, the so-called 1-step gradient 37 approximates the series by considering only its first term, i.e. ( I − J ℱ ) − 1 ≈ I (I-J_{\mathcal{F}})^{-1}\approx I , and leads to the following approximation of Equation 1 : ∂ z H ∗ ∂ θ H ≈ ∂ f H ∂ θ H , ∂ z H ∗ ∂ θ L ≈ ∂ f H ∂ z L ∗ ⋅ ∂ z L ∗ ∂ θ L , ∂ z H ∗ ∂ θ I ≈ ∂ f H ∂ z L ∗ ⋅ ∂ z L ∗ ∂ θ I . \frac{\partial z_{H}^{*}}{\partial\theta_{H}}\approx\frac{\partial f_{H}}{\partial\theta_{H}},\quad\frac{\partial z_{H}^{*}}{\partial\theta_{L}}\approx\frac{\partial f_{H}}{\partial z^{*}_{L}}\cdot\frac{\partial z^{*}_{L}}{\partial\theta_{L}},\quad\frac{\partial z_{H}^{*}}{\partial\theta_{I}}\approx\frac{\partial f_{H}}{\partial z^{*}_{L}}\cdot\frac{\partial z^{*}_{L}}{\partial\theta_{I}}\;. (2) The gradients of the low-level fixed point, ∂ z L ∗ ∂ θ L \frac{\partial z^{*}_{L}}{\partial\theta_{L}} and ∂ z L ∗ ∂ θ I \frac{\partial z^{*}_{L}}{\partial\theta_{I}} , can also be approximated using another application of the 1-step gradient: ∂ z L ∗ ∂ θ L ≈ ∂ f L ∂ θ L , ∂ z L ∗ ∂ θ I ≈ ∂ f L ∂ θ I . \frac{\partial z^{*}_{L}}{\partial\theta_{L}}\approx\frac{\partial f_{L}}{\partial\theta_{L}},\quad\frac{\partial z^{*}_{L}}{\partial\theta_{I}}\approx\frac{\partial f_{L}}{\partial\theta_{I}}\;. (3) By substituting Equation 3 back into Equation 2 , we arrive at the final simplified gradients.
Before defining our loss function, we must first introduce two key elements of our proposed method: deep supervision and adaptive computational time .
#### Deep supervision
Inspired by the principle that periodic neural oscillations regulate when learning occurs in the brain 38 , we incorporate a deep supervision mechanism into HRM, as detailed next.
Given a data sample ( x , y ) (x,y) , we run multiple forward passes of the HRM model, each of which we refer to as a segment . Let M M denote the total number of segments executed before termination. For each segment m ∈ { 1 , … , M } m\in\{1,\dots,M\} , let z m = ( z H m N T , z L m N T ) z^{m}=(z^{mNT}_{H},z^{mNT}_{L}) represent the hidden state at the conclusion of segment m m , encompassing both high-level and low-level state components.
At each segment m m , we apply a deep supervision step as follows: 1. Given the state z m − 1 z^{m-1} from the previous segment, compute the next state z m z^{m} and its associated output y ^ m \hat{y}^{m} through a forward pass in the HRM model: ( z m , y ^ m ) ← HRM ( z m − 1 , x , θ ) (z^{m},\hat{y}^{m})\leftarrow\textsc{HRM}(z^{m-1},x;\theta)
2. Compute the loss for the current segment: L m ← Loss ( y ^ m , y ) L^{m}\leftarrow\textsc{Loss}(\hat{y}^{m},y)
3. Update parameters: θ ← OptimizerStep ( θ , ∇ θ L m ) \theta\leftarrow\textsc{OptimizerStep}(\theta,\nabla_{\theta}L^{m})
The crucial aspect of this procedure is that the hidden state z m z^{m} is “detached” from the computation graph before being used as the input state for the next segment. Consequently, gradients from segment m + 1 m+1 do not propagate back through segment m m , effectively creating a 1-step approximation of the gradient of the recursive deep supervision process 39 , 40 . This approach provides more frequent feedback to the H-module and serves as a regularization mechanism, demonstrating superior empirical performance and enhanced stability in deep equilibrium models when compared to more complex, Jacobian-based regularization techniques 39 , 41 . Figure 4 shows pseudocode of deep supervision training.
#### Adaptive computational time (ACT)
The brain dynamically alternates between automatic thinking (“System 1”) and deliberate reasoning (“System 2”) 42 . Neuroscientific evidence shows that these cognitive modes share overlapping neural circuits, particularly within regions such as the prefrontal cortex and the default mode network 43 , 44 . This indicates that the brain dynamically modulates the “runtime” of these circuits according to task complexity and potential rewards 45 , 46 .
Inspired by the above mechanism, we incorporate an adaptive halting strategy into HRM that enables “thinking, fast and slow”. This integration leverages deep supervision and uses the Q-learning algorithm 47 to adaptively determine the number of segments. A Q-head uses the final state of the H-module to predict the Q-values Q ^ m = ( Q ^ halt m , Q ^ continue m ) \hat{Q}^{m}=(\hat{Q}^{m}_{\text{halt}},\hat{Q}^{m}_{\text{continue}}) of the “halt” and “continue” actions: Q ^ m = σ ( θ Q ⊤ z H m N T ) , \hat{Q}^{m}=\sigma(\theta_{Q}^{\top}z_{H}^{mNT})\,, where σ \sigma denotes the sigmoid function applied element-wise. The halt or continue action is chosen using a randomized strategy as detailed next. Let M max M_{\max} denote the maximum number of segments (a fixed hyperparameter) and M min M_{\min} denote the minimum number of segments (a random variable). The value of M min M_{\min} is determined stochastically: with probability ε \varepsilon , it is sampled uniformly from the set { 2 , ⋯ , M max } \{2,\cdots,M_{\max}\} (to encourage longer thinking), and with probability 1 − ε 1-\varepsilon , it is set to 1. The halt action is selected under two conditions: when the segment count surpasses the maximum threshold M max M_{\max} , or when the estimated halt value Q ^ halt \hat{Q}_{\text{halt}} exceeds the estimated continue value Q ^ continue \hat{Q}_{\text{continue}} and the segment count has reached at least the minimum threshold M min M_{\min} .
The Q-head is updated through a Q-learning algorithm, which is defined on the following episodic Markov Decision Process (MDP). The state of the MDP at segment m m is z m z^{m} , and the action space is { halt , continue } \{\text{halt},\text{continue}\} . Choosing the action “halt” terminates the episode and returns a binary reward indicating prediction correctness, i.e., 𝟏 { y ^ m = y } \mathbf{1}\{\hat{y}^{m}=y\} . Choosing “continue” yields a reward of 0 and the state transitions to z m + 1 z^{m+1} . Thus, the Q-learning targets for the two actions G ^ m = ( G ^ halt m , G ^ continue m ) \hat{G}^{m}=(\hat{G}^{m}_{\text{halt}},\hat{G}^{m}_{\text{continue}}) are given by G ^ halt m \displaystyle\hat{G}_{\text{halt}}^{m} = 𝟏 { y ^ m = y } , \displaystyle=\mathbf{1}\{\hat{y}^{m}=y\}\,, G ^ continue m \displaystyle\hat{G}_{\text{continue}}^{m} = { Q ^ halt m + 1 , if m ≥ N max , max ( Q ^ halt m + 1 , Q ^ continue m + 1 ) , otherwise . \displaystyle=\begin{cases}\hat{Q}_{\text{halt}}^{m+1},&\text{if $m\geq N_{\max}$}\,,\\[6.0pt] \max(\hat{Q}_{\text{halt}}^{m+1},\hat{Q}_{\text{continue}}^{m+1})\,,&\text{otherwise}\;.\end{cases} We can now define the loss function of our learning procedure. The overall loss for each supervision segment combines both the Q-head loss and the sequence-to-sequence loss: L ACT m = Loss ( y ^ m , y ) + BinaryCrossEntropy ( Q ^ m , G ^ m ) . L^{m}_{\text{ACT}}=\textsc{Loss}(\hat{y}^{m},y)+\textsc{BinaryCrossEntropy}(\hat{Q}^{m},\hat{G}^{m})\;. Minimizing the above loss enables both accurate predictions and nearly optimal stopping decisions.
Selecting the “halt” action ends the supervision loop. In practice, sequences are processed in batches, which can be easily handled by substituting any halted sample in the batch with a fresh sample from the dataloader.
Figure 5 presents a performance comparison between two HRM variants: one incorporating ACT and another employing a fixed computational step count equivalent to ACT’s M max M_{\max} parameter. It shows that ACT effectively adapts its computational resources based on task complexity, achieving significant computational savings with minimal impact on performance.
#### Inference-time scaling
An effective neural model should exploit additional computational resources during inference to enhance performance. As illustrated in Figure 5 -(c), HRM seamlessly achieves inference-time scaling by simply increasing the computational limit parameter, M max M_{\max} without requiring further training or architectural modifications.
Additional compute is especially effective for tasks that demand deeper reasoning. On Sudoku—a problem that often requires long-term planning—HRM exhibits strong inference-time scaling. On the other hand, we find that extra computational resources yield minimal gains in ARC-AGI challenge, as solutions generally require only a few transformations.
#### Stability of Q-learning in ACT
The deep Q-learning that underpins our ACT mechanism is known to be prone to instability, often requiring stabilization techniques such as replay buffers and target networks 48 , which are absent in our design. Our approach, however, achieves stability through the intrinsic properties of our model and training procedure. Recent theoretical work by Gallici et al. Super shows that Q-learning can achieve convergence if network parameters are bounded, weight decay is incorporated during training, and post-normalization layers are implemented. Our model satisfies these conditions through its Post-Norm architecture that employs RMSNorm (a layer normalization variant) and the AdamW optimizer. AdamW has been shown to solve an L ∞ L_{\infty} -constrained optimization problem, ensuring that model parameters remain bounded by 1 / λ 1/\lambda 50 .
#### Architectural details
We employ a sequence-to-sequence architecture for HRM. Both input and output are represented as token sequences: x = ( x 1 , … , x l ) x=(x_{1},\ldots,x_{l}) and y = ( y 1 , … , y l ′ ) y=(y_{1},\ldots,y_{l^{\prime}}) respectively. The model includes an embedding layer f I f_{I} that converts discrete tokens into vector representations, and an output head f O ( z , θ O ) = softmax ( θ O z ) f_{O}(z;\theta_{O})=\text{softmax}(\theta_{O}z) that transforms hidden states into token probability distributions y ^ \hat{y} . For small-sample experiments, we replace softmax with stablemax 51 to improve generalization performance. The sequence-to-sequence loss is averaged over all tokens, Loss ( y ^ , y ) = 1 l ′ ∑ i = 1 l ′ log p ( y i ) \textsc{Loss}(\hat{y},y)=\frac{1}{l^{\prime}}\sum_{i=1}^{l^{\prime}}\log p(y_{i}) , where p ( y i ) p(y_{i}) is the probability that distribution y ^ i \hat{y}_{i} assigns to token y i y_{i} . The initial hidden states z 0 z^{0} are initialized by sampling from a truncated normal distribution with standard deviation of 1, truncation of 2, and kept fixed throughout training.
Both the low-level and high-level recurrent modules f L f_{L} and f H f_{H} are implemented using encoder-only Transformer 52 blocks with identical architectures and dimensions. These modules take multiple inputs, and we use straightforward element-wise addition to combine them, though more sophisticated merging techniques such as gating mechanisms could potentially improve performance and is left for future work. For all Transformer blocks in this work—including those in the baseline models—we incorporate the enhancements found in modern LLMs (based on Llama 53 architectures). These improvements include Rotary Positional Encoding 54 , Gated Linear Units 55 , RMSNorm 56 , and the removal of bias terms from linear layers.
Furthermore, both HRM and recurrent Transformer models implement a Post-Norm architecture with weights initialized via truncated LeCun Normal initialization 57 , 58 , 59 , while the scale and bias parameters are excluded from RMSNorm. All parameters are optimized using the Adam-atan2 optimizer 60 , a scale-invariant variant of Adam 61 , combined with a constant learning rate that includes linear warm-up.
## 3 Results
This section begins by describing the ARC-AGI, Sudoku, and Maze benchmarks, followed by an overview of the baseline models and their results. Figure 6 -(a,b,c) presents a visual representation of the three benchmark tasks, which are selected to evaluate various reasoning abilities in AI models.
### 3.1 Benchmarks
ARC-AGI Challenge The ARC-AGI benchmark evaluates general fluid intelligence through IQ-test-like puzzles that require inductive reasoning 27 . The initial version, ARC-AGI-1, presents challenges as input-label grid pairs that force AI systems to extract and generalize abstract rules from just a few examples. Each task provides a few input–output demonstration pairs (usually 2–3) and a test input. An AI model has two attempts to produce the correct output grid. Although some believe that mastering ARC-AGI would signal true artificial general intelligence, its primary purpose is to expose the current roadblocks in AGI progress. In fact, both conventional deep learning methods and CoT techniques have faced significant challenges with ARC-AGI-1, primarily because it requires the ability to generalize to entirely new tasks 28 .
Addressing the limitations identified in ARC-AGI-1, ARC-AGI-2 significantly expands the benchmark by providing a more comprehensive and carefully refined collection of tasks. These new tasks emphasize deeper compositional reasoning, multi-step logic, contextual rule application, and symbolic abstraction. Human calibration studies show these tasks are challenging but doable for people, while being much harder for current AI systems, offering a clearer measure of general reasoning abilities 29 .
Sudoku-Extreme Sudoku is a 9 × \times 9 logic puzzle, requiring each row, column, and 3 × \times 3 block to contain the digits 1–9 exactly once. A prediction is considered correct if it exactly matches the puzzle’s unique solution. Sudoku’s complex logical structure makes it a popular benchmark for evaluating logical reasoning in machine learning 62 , 63 , 64 .