diff options
Diffstat (limited to 'papers/txt/trm2025_tiny_recursive.txt')
| -rw-r--r-- | papers/txt/trm2025_tiny_recursive.txt | 796 |
1 files changed, 796 insertions, 0 deletions
diff --git a/papers/txt/trm2025_tiny_recursive.txt b/papers/txt/trm2025_tiny_recursive.txt new file mode 100644 index 0000000..55cf994 --- /dev/null +++ b/papers/txt/trm2025_tiny_recursive.txt @@ -0,0 +1,796 @@ + Less is More: Recursive Reasoning with Tiny Networks + + + + Alexia Jolicoeur-Martineau + Samsung SAIL Montréal + alexia.j@samsung.com + + + Abstract +arXiv:2510.04871v1 [cs.LG] 6 Oct 2025 + + + + + Hierarchical Reasoning Model (HRM) is a + novel approach using two small neural net- + works recursing at different frequencies. This + biologically inspired method beats Large Lan- + guage models (LLMs) on hard puzzle tasks + such as Sudoku, Maze, and ARC-AGI while + trained with small models (27M parameters) + on small data (∼ 1000 examples). HRM holds + great promise for solving hard problems with + small networks, but it is not yet well un- + derstood and may be suboptimal. We pro- + pose Tiny Recursive Model (TRM), a much + simpler recursive reasoning approach that + achieves significantly higher generalization + than HRM, while using a single tiny network + with only 2 layers. With only 7M parameters, + TRM obtains 45% test-accuracy on ARC-AGI- + 1 and 8% on ARC-AGI-2, higher than most + LLMs (e.g., Deepseek R1, o3-mini, Gemini 2.5 + Pro) with less than 0.01% of the parameters. + + + + 1. Introduction + While powerful, Large Language models (LLMs) can + struggle on hard question-answer problems. Given + that they generate their answer auto-regressively, there + is a high risk of error since a single incorrect token can + render an answer invalid. To improve their reliabil- Figure 1. Tiny Recursion Model (TRM) recursively improves + ity, LLMs rely on Chain-of-thoughts (CoT) (Wei et al., its predicted answer y with a tiny network. It starts with the + 2022) and Test-Time Compute (TTC) (Snell et al., 2024). embedded input question x and initial embedded answer + y, and latent z. For up to Nsup = 16 improvements steps, + CoTs seek to emulate human reasoning by having the + it tries to improve its answer y. It does so by i) recursively + LLM to sample step-by-step reasoning traces prior to + updating n times its latent z given the question x, current + giving their answer. Doing so can improve accuracy, answer y, and current latent z (recursive reasoning), and + but CoT is expensive, requires high-quality reasoning then ii) updating its answer y given the current answer y + data (which may not be available), and can be brittle and current latent z. This recursive process allows the model + since the generated reasoning may be wrong. To fur- to progressively improve its answer (potentially address- + ther improve reliability, test-time compute can be used ing any errors from its previous answer) in an extremely + by reporting the most common answer out of K or the parameter-efficient manner while minimizing overfitting. + highest-reward answer (Snell et al., 2024). + + 1 + Recursive Reasoning with Tiny Networks + +However, this may not be enough. LLMs with CoT ers that achieves significantly higher generalization +and TTC are not enough to beat every problem. While than HRM on a variety of problems. In doing so, we +LLMs have made significant progress on ARC-AGI improve the state-of-the-art test accuracy on Sudoku- +(Chollet, 2019) since 2019, human-level accuracy still Extreme from 55% to 87%, Maze-Hard from 75% to +has not been reached (6 years later, as of writing of 85%, ARC-AGI-1 from 40% to 45%, and ARC-AGI-2 +this paper). Furthermore, LLMs struggle on the newer from 5% to 8%. +ARC-AGI-2 (e.g., Gemini 2.5 Pro only obtains 4.9% test +accuracy with a high amount of TTC) (Chollet et al., 2. Background +2025; ARC Prize Foundation, 2025b). + HRM is described in Algorithm 2. We discuss the +An alternative direction has recently been proposed by + details of the algorithm further below. +Wang et al. (2025). They propose a new way forward +through their novel Hierarchical Reasoning Model +(HRM), which obtains high accuracy on puzzle tasks 2.1. Structure and goal +where LLMs struggle to make a dent (e.g., Sudoku The focus of HRM is supervised learning. Given an +solving, Maze pathfinding, and ARC-AGI). HRM is a input, produce an output. Both input and output are +supervised learning model with two main novelties: 1) assumed to have shape [ B, L] (when the shape differs, +recursive hierarchical reasoning, and 2) deep supervision. padding tokens can be added), where B is the batch- +Recursive hierarchical reasoning consists of recurs- size and L is the context-length. +ing multiple times through two small networks ( f L at HRM contains four learnable components: the in- +high frequency and f H at low frequency) to predict the put embedding f I (·; θ I ), low-level recurrent network +answer. Each network generates a different latent fea- f L (·; θ L ), high-level recurrent network f H (·; θ H ), and +ture: f L outputs z H and f H outputs z L . Both features the output head fO (·; θO ). Once the input is embedded, +(z L , z H ) are used as input to the two networks. The the shape becomes [ B, L, D ] where D is the embedding +authors provide some biological arguments in favor of size. Each network is a 4-layer Transformers architec- +recursing at different hierarchies based on the different ture (Vaswani et al., 2017), with RMSNorm (Zhang +temporal frequencies at which the brains operate and & Sennrich, 2019), no bias (Chowdhery et al., 2023), +hierarchical processing of sensory inputs. rotary embeddings (Su et al., 2024), and SwiGLU acti- +Deep supervision consists of improving the answer vation function (Hendrycks & Gimpel, 2016; Shazeer, +through multiple supervision steps while carrying the 2020). +two latent features as initialization for the improve- +ment steps (after detaching them from the computa- 2.2. Recursion at two different frequencies +tional graph so that their gradients do not propagate). Given the hyperparameters used by Wang et al. (2025) +This provide residual connections, which emulates (n = 2 f L steps, 1 f H steps; done T = 2 times), a +very deep neural networks that are too memory ex- forward pass of HRM is done as follows: +pensive to apply in one forward pass. +An independent analysis on the ARC-AGI benchmark x ← f I ( x̃ ) +showed that deep supervision seems to be the primary z L ← f L (z L + z H + x ) # without gradients +driver of the performance gains (ARC Prize Founda- z L ← f L (z L + z H + x ) # without gradients +tion, 2025a). Using deep supervision doubled accuracy + z H ← f H (z L + z H ) # without gradients +over single-step supervision (going from 19% to 39% +accuracy), while recursive hierarchical reasoning only z L ← f L (z L + z H + x ) # without gradients +slightly improved accuracy over a regular model with z L ← z L .detach() +a single forward pass (going from 35.7% to 39.0% ac- z H ← z H .detach() +curacy). This suggests that reasoning across different + z L ← f L (z L + z H + x ) # with gradients +supervision steps is worth it, but the recursion done +in each supervision step is not particularly important. z H ← f H (z L + z H ) # with gradients + ŷ ← argmax( fO (z H )) +In this work, we show that the benefit from recursive +reasoning can be massively improved, making it much +more than incremental. We propose Tiny Recursive where ŷ is the predicted output answer, z L and z H are +Model (TRM), an improved and simplified approach either initialized embeddings or the embeddings of +using a much smaller tiny network with only 2 lay- the previous deep supervision step (after detaching + them from the computational graph). As can be seen, + + 2 + Recursive Reasoning with Tiny Networks + +def hrm(z, x, n=2, T=2): # hierarchical reasoning 2.4. Deep supervision + zH, zL = z + with torch.no grad(): To improve effective depth, deep supervision is used. + for i in range(nT − 2): + zL = L net(zL, zH, x) This consists of reusing the previous latent features + if (i + 1) % T == 0: + zH = H net(zH, zL) + (z H and z L ) as initialization for the next forward pass. + # 1−step grad This allows the model to reason over many iterations + zL = L net(zL, zH, x) + zH = H net(zH, zL) and improve its latent features (z L and z H ) until it + return (zH, zL), output head(zH), Q head(zH) (hopefully) converges to the correct solution. At most +def ACT halt(q, y hat, y true): Nsup = 16 supervision steps are used. + target halt = (y hat == y true) + loss = 0.5∗binary cross entropy(q[0], target halt) + return loss 2.5. Adaptive computational time (ACT) +def ACT continue(q, last step): + if last step: With deep supervision, each mini-batch of data sam- + target continue = sigmoid(q[0]) ples must be used for Nsup = 16 supervision steps + else: + target continue = sigmoid(max(q[0], q[1]))) before moving to the next mini-batch. This is expen- + loss = 0.5∗binary cross entropy(q[1], target continue) + return loss + sive, and there is a balance to be reached between + optimizing a few data examples for many supervision +# Deep Supervision +for x input, y true in train dataloader: steps versus optimizing many data examples with less + z = z init supervision steps. To reach a better balance, a halting + for step in range(N sup): # deep supervision + x = input embedding(x input) mechanism is incorporated to determine whether the + z, y pred, q = hrm(z, x) + loss = softmax cross entropy(y pred, y true) + model should terminate early. It is learned through + # Adaptive computational time (ACT) using Q−learning a Q-learning objective that requires passing the z H + loss += ACT halt(q, y pred, y true) + , , q next = hrm(z, x) # extra forward pass through an additional head and running an additional + loss += ACT continue(q next, step == N sup − 1) forward pass (to determine if halting now rather than + z = z.detach() + loss.backward() later would have been preferable). They call this + opt.step() + opt.zero grad() + method Adaptive computational time (ACT). It is only + if q[0] > q[1]: # early−stopping used during training, while the full Nsup = 16 super- + break + vision steps are done at test time to maximize down- + stream performance. ACT greatly diminishes the time +Figure 2. Pseudocode of Hierarchical Reasoning Models spent per example (on average spending less than 2 +(HRMs). steps on the Sudoku-Extreme dataset rather than the + full Nsup = 16 steps), allowing more coverage of the +a forward pass of HRM consists of applying 6 function dataset given a fixed number of training iterations. +evaluations, where the first 4 function evaluations are +detached from the computational graph and are not 2.6. Deep supervision and 1-step gradient +back-propagated through. The authors uses n = 2 approximations replaces BPTT +with T = 2 in all experiments, but HRM can be gener- +alized by allowing for an arbitrary number of L steps Deep supervision and the 1-step gradient approxima- +(n) and recursions (T) as shown in Algorithm 2. tion provide a more biologically plausible and less + computationally-expansive alternative to Backpropa- +2.3. Fixed-point recursion with 1-step gradient gation Through Time (BPTT) (Werbos, 1974; Rumel- + approximation hart et al., 1985; LeCun, 1985) for solving the temporal + credit assignment (TCA) (Rumelhart et al., 1985; Wer- +Assuming that (z L , z H ) reaches a fixed-point (z∗L , z∗H ) bos, 1988; Elman, 1990) problem (Lillicrap & Santoro, +through recursing from both f L and f H , 2019). The implication is that HRM can learn what + would normally require an extremely large network + z∗L ≈ f L (z∗L + z H + x ) without having to back-propagate through its entire + z∗H ≈ f H (z L + z∗H ) , depth. Given the hyperparameters used by Jang et al. + (2023) in all their experiments, HRM effectively rea- +the Implicit Function Theorem (Krantz & Parks, 2002) + sons over nlayers (n + 1) TNsup = 4 ∗ (2 + 1) ∗ 2 ∗ 16 = +with the 1-step gradient approximation (Bai et al., + 384 layers of effective depth. +2019) is used to approximate the gradient by back- +propagating only the last f L and f H steps. This theo- +rem is used to justify only tracking the gradients of +the last two steps (out of 6), which greatly reduces +memory demands. + + 3 + Recursive Reasoning with Tiny Networks + +2.7. Summary of HRM different from the much smaller n = 2 and T = 2 used + in every experiment of their paper, we observe the +HRM leverages recursion from two networks at dif- + following: +ferent frequencies (high frequency versus low fre- +quency) and deep supervision to learn to improve +its answer over multiple supervision steps (with ACT 1. the residual for z H is clearly well above 0 at every +to reduce time spent per data example). This enables step +the model to imitate extremely large depth without +requiring backpropagation through all layers. This +approach obtains significantly higher performance on 2. the residual for z L only becomes closer to 0 after +hard question-answer tasks that regular supervised many cycles, but it remains significantly above 0 +models struggle with. However, this method is quite +complicated, relying a bit too heavily on uncertain +biological arguments and fixed-point theorems that 3. z L is very far from converged after one f L evalu- +are not guaranteed to be applicable. In the next sec- ation at T cycles, which is when the fixed-point +tion, we discuss those issues and potential targets for is assumed to be reached and the 1-step gradient +improvements in HRM. approximation is used + +3. Target for improvements in Hierarchical Thus, while the application of the IFT theorem and + Reasoning Models 1-step gradient approximation to HRM has some basis + since the residuals do generally reduce over time, a +In this section, we identify key targets for improve- + fixed point is unlikely to be reached when the theorem +ments in HRM, which will be addressed by our pro- + is actually applied. +posed method, Tiny Recursion Models (TRMs). + In the next section, we show that we can bypass the +3.1. Implicit Function Theorem (IFT) with 1-step need for the IFT theorem and 1-step gradient approxi- + gradient approximation mation, thus bypassing the issue entirely. + +HRM only back-propagates through the last 2 of the 6 3.2. Twice the forward passes with Adaptive +recursions. The authors justify this by leveraging the computational time (ACT) +Implicit Function Theorem (IFT) and one-step approx- +imation, which states that when a recurrent function HRM uses Adaptive computational time (ACT) during +converges to a fixed point, backpropagation can be training to optimize the time spent of each data sam- +applied in a single step at that equilibrium point. ple. Without ACT, Nsup = 16 supervision steps would + be spent on the same data sample, which is highly in- +There are concerns about applying this theorem to + efficient. They implement ACT through an additional +HRM. Most importantly, there is no guarantee that + Q-learning objective, which decides when to halt and +a fixed-point is reached. Deep equilibrium models + move to a new data sample rather than keep iterating +normally do fixed-point iteration to solve for the fixed + on the same data. This allows much more efficient +pointz∗ = f (z∗ ) (Bai et al., 2019). However, in the case + use of time especially since the average number of su- +of HRM, it is not iterating to the fixed-point but simply + pervision steps during training is quite low with ACT +doing forward passes of f L and f H . To make matters + (less than 2 steps on the Sudoku-Extreme dataset as +worse, HRM is only doing 4 recursions before stopping + per their reported numbers). +to apply the one-step approximation. After its first +loop of two f L and 1 f H evaluations, it only apply a However, ACT comes at a cost. This cost is not directly +single f L evaluation before assuming that a fixed-point shown in the HRM’s paper, but it is shown in their of- +is reached for both z L and z H (z∗L = f L (z∗L + z H + x ) ficial code. The Q-learning objective relies on a halting +and z∗H = f H (z∗L + z∗H )). Then, the one-step gradient loss and a continue loss. The continue loss requires an +approximation is applied to both latent variables in extra forward pass through HRM (with all 6 function +succession. evaluations). This means that while ACT optimizes + time more efficiently per sample, it requires 2 forward +The authors justify that a fixed-point is reached by + passes per optimization step. The exact formulation is +depicting an example with n = 7 and T = 7 where + shown in Algorithm 2. +the forward residuals is reduced over time (Figure 3 +in Wang et al. (2025)). Even in this setting, which is In the next section, we show that we can bypass the + need for two forward passes in ACT. + + 4 + Recursive Reasoning with Tiny Networks + +3.3. Hierarchical interpretation based on complex of 2 passes). Our approach is described in Algorithm 3 + biological arguments and illustrated in Figure 1. We also provide an ablation + in Table 1 on the Sudoku-Extreme dataset (a dataset +The HRM’s authors justify the two latent variables + of difficult Sudokus with only 1K training examples, +and two networks operating at different hierarchies + but 423K test examples). Below, we explain the key +based on biological arguments, which are very far + components of TRMs. +from artificial neural networks. They even try to match +HRM to actual brain experiments on mice. While in- +teresting, this sort of explanation makes it incredibly Table 1. Ablation of TRM on Sudoku-Extreme comparing % +hard to parse out why HRM is designed the way it Test accuracy, effective depth per supervision step ( T (n + +is. Given the lack of ablation table in their paper, the 1)nlayers ), number of Forward Passes (NFP) per optimization +over-reliance on biological arguments and fixed-point step, and number of parameters +theorems (that are not perfectly applicable), it is hard Method Acc (%) Depth NFP # Params +to determine what parts of HRM is helping what and HRM 55.0 24 2 27M +why. Furthermore, it is not clear why they use two TRM (T = 3, n = 6) 87.4 42 1 5M +latent features rather than other combinations of fea- w/ ACT 86.1 42 2 5M +tures. w/ separate f H , f L 82.4 42 1 10M + no EMA 79.9 42 1 5M +In the next section, we show that the recursive process w/ 4-layers, n = 3 79.5 48 1 10M +can be greatly simplified and understood in a much w/ self-attention 74.7 42 1 7M +simpler manner that does not require any biological w/ T = 2, n = 2 73.7 12 1 5M +argument, fixed-point theorem, hierarchical interpre- w/ 1-step gradient 56.5 42 1 5M +tation, nor using two networks. It also explains why 2 +is the optimal number of features (z L and z H ). + 4.1. No fixed-point theorem required + +def latent recursion(x, y, z, n=6): HRM assumes that the recursions converge to a fixed- + for i in range(n): # latent reasoning point for both z L and z H in order to leverage the 1-step + z = net(x, y, z) + y = net(y, z) # refine output answer gradient approximation (Bai et al., 2019). This allows + return y, z + the authors to justify only back-propagating through +def deep recursion(x, y, z, n=6, T=3): the last two function evaluations (1 f L and 1 f H ). To + # recursing T−1 times to improve y and z (no gradients needed) + with torch.no grad(): bypass this theoretical requirement, we define a full + for j in range(T−1): recursion process as containing n evaluations of f L + y, z = latent recursion(x, y, z, n) + # recursing once to improve y and z and 1 evaluation of f H : + y, z = latent recursion(x, y, z, n) + return (y.detach(), z.detach()), output head(y), Q head(y) z L ← f L (z L + z H + x ) +# Deep Supervision ... +for x input, y true in train dataloader: + y, z = y init, z init z L ← f L (z L + z H + x ) + for step in range(N supervision): + x = input embedding(x input) z H ← f H (z L + z H ) . + (y, z), y hat, q hat = deep recursion(x, y, z) + loss = softmax cross entropy(y hat, y true) + loss += binary cross entropy(q hat, (y hat == y true)) Then, we simply back-propagate through the full re- + loss.backward() cursion process. + opt.step() + opt.zero grad() + if q hat > 0: # early−stopping Through deep supervision, the models learns to take + break any (z L , z H ) and improve it through a full recursion + process, hopefully making z H closer to the solution. +Figure 3. Pseudocode of Tiny Recursion Models (TRMs). This means that by the design of the deep supervi- + sion goal, running a few full recursion processes (even + without gradients) is expected to bring us closer to the +4. Tiny Recursion Models solution. We propose to run T − 1 recursion processes + without gradient to improve (z L , z H ) before running +In this section, we present Tiny Recursion Models + one recursion process with backpropagation. +(TRMs). Contrary to HRM, TRM requires no com- +plex mathematical theorem, hierarchy, nor biological Thus, instead of using the 1-step gradient approxi- +arguments. It generalizes better while requiring only mation, we apply a full recursion process containing +a single tiny network (instead of two medium-size net- n evaluations of f L and 1 evaluation of f H . This re- +works) and a single forward pass for the ACT (instead moves entirely the need to assume that a fixed-point + + 5 + Recursive Reasoning with Tiny Networks + +is reached and the use of the IFT theorem with 1-step While this is intuitive, we wanted to verify whether +gradient approximation. Yet, we can still leverage using more or less features could be helpful. Results +multiple backpropagation-free recursion processes to are shown in Table 2. +improve (z L , z H ). With this approach, we obtain a + More features (> 2): We tested splitting z into dif- +massive boost in generalization on Sudoku-Extreme + ferent features by treating each of the n recursions as +(improving TRM from 56.5% to 87.4%; see Table 1). + producing a different zi for i = 1, ..., n. Then, each + zi is carried across supervision steps. The approach +4.2. Simpler reinterpretation of z H and z L is described in Algorithm 5. In doing so, we found +HRM is interpreted as doing hierarchical reasoning performance to drop. This is expected because, as dis- +over two latent features of different hierarchies due to cussed, there is no apparent need for splitting z into +arguments from biology. However, one might wonder multiple parts. It does not have to be hierarchical. +why use two latent features instead of 1, 3, or more? Single feature: Similarly, we tested the idea of taking +And do we really need to justify these so-called ”hier- a single feature by only carrying z H across supervi- +archical” features based on biology to make sense of sion steps. The approach is described in Algorithm 4. +them? We propose a simple non-biological explana- In doing so, we found performance to drop. This is +tion, which is more natural, and directly answers the expected because, as discussed, it forces the model to +question of why there are 2 features. store the solution y within z. +The fact of the matter is: z H is simply the current Thus, we explored using more or less latent variables +(embedded) solution. The embedding is reversed by on Sudoku-Extreme, but found that having only y and +applying the output head and rounding to the nearest z lead to better test accuracy in addition to being the +token using the argmax operation. On the other hand, simplest more natural approach. +z L is a latent feature that does not directly correspond +to a solution, but it can be transformed into a solution +by applying z H ← f H ( x, z L , z H ). We show an example +on Sudoku-Extreme in Figure 6 to highlight the fact Table 2. TRM on Sudoku-Extreme comparing % Test accu- +that z H does correspond to the solution, but z L does racy when using more or less latent features +not. Method # of features Acc (%) + TRM y, z (Ours) 2 87.4 +Once this is understood, hierarchy is not needed; there TRM multi-scale z n+1 = 7 77.6 +is simply an input x, a proposed solution y (previously TRM single z 1 71.9 +called z H ), and a latent reasoning feature z (previously +called z L ). Given the input question x, current solution +y, and current latent reasoning z, the model recursively +improves its latent z. Then, given the current latent z 4.3. Single network +and the previous solution y, the model proposes a new HRM uses two networks, one applied frequently as a +solution y (or stay at the current solution if its already low-level module f H and one applied rarely as an high- +good). level module ( f H ). This requires twice the number of +Although this has no direct influence on the algorithm, parameters compared to regular supervised learning +this re-interpretation is much simpler and natural. It with a single network. +answers the question about why two features: remem- As mentioned previously, while f L iterates on the la- +bering in context the question x, previous reasoning tent reasoning feature z (z L in HRM), the goal of f H +z, and previous answer y helps the model iterate on is to update the solution y (z H in HRM) given the la- +the next reasoning z and then the next answer y. If tent reasoning and current solution. Importantly, since +we were not passing the previous reasoning z, the z ← f L ( x + y + z) contains x but y ← f H (y + z) does +model would forget how it got to the previous solu- not contains x, the task to achieve (iterating on z versus +tion y (since z acts similarly as a chain-of-thought). If using z to update y) is directly specified by the inclu- +we were not passing the previous solution y, then the sion or lack of x in the inputs. Thus, we considered +model would forget what solution it had and would the possibility that both networks could be replaced +be forced to store the solution y within z instead of by a single network doing both tasks. In doing so, we +using it for latent reasoning. Thus, we need both y and obtain better generalization on Sudoku-Extreme (im- +z separately, and there is no apparent reason why one proving TRM from 82.4% to 87.4%; see Table 1) while +would need to split z into multiple features. reducing the number of parameters by half. It turns + out that a single network is enough. + + 6 + Recursive Reasoning with Tiny Networks + +4.4. Less is more 4.6. No additional forward pass needed with ACT +We attempted to increase capacity by increasing the As previously mentioned, the implementation of ACT +number of layers in order to scale the model. Sur- in HRM through Q-learning requires two forward +prisingly, we found that adding layers decreased gen- passes, which slows down training. We propose a +eralization due to overfitting. In doing the oppo- simple solution, which is to get rid of the continue loss +site, decreasing the number of layers while scaling (from the Q-learning) and only learn a halting proba- +the number of recursions (n) proportionally (to keep bility through a Binary-Cross-Entropy loss of having +the amount of compute and emulated depth approxi- reached the correct solution. By removing the continue +mately the same), we found that using 2 layers (instead loss, we remove the need for the expensive second for- +of 4 layers) maximized generalization. In doing so, we ward pass, while still being able to determine when to +obtain better generalization on Sudoku-Extreme (im- halt with relatively good accuracy. We found no sig- +proving TRM from 79.5% to 87.4%; see Table 1) while nificant difference in generalization from this change +reducing the number of parameters by half (again). (going from 86.1% to 87.4%; see Table 1). +It is quite surprising that smaller networks are bet- +ter, but 2 layers seems to be the optimal choice. Bai 4.7. Exponential Moving Average (EMA) +& Melas-Kyriazi (2024) also observed optimal perfor- On small data (such as Sudoku-Extreme and Maze- +mance for 2-layers in the context of deep equilibrium Hard), HRM tends to overfit quickly and then diverge. +diffusion models; however, they had similar perfor- To reduce this problem and improves stability, we +mance to the bigger networks, while we instead ob- integrate Exponential Moving Average (EMA) of the +serve better performance with 2 layers. This may ap- weights, a common technique in GANs and diffusion +pear unusual, as with modern neural networks, gener- models to improve stability (Brock et al., 2018; Song & +alization tends to directly correlate with model sizes. Ermon, 2020). We find that it prevents sharp collapse +However, when data is too scarce and model size is and leads to higher generalization (going from 79.9% +large, there can be an overfitting penalty (Kaplan et al., to 87.4%; see Table 1). +2020). This is likely an indication that there is too little +data. Thus, using tiny networks with deep recursion 4.8. Optimal the number of recursions +and deep supervision appears to allow us to bypass a +lot of the overfitting. We experimented with different number of recursions + by varying T and n and found that T = 3, n = 3 +4.5. attention-free architecture for tasks with small (equivalent to 48 recursions) in HRM and T = 3, n = 6 + fixed context length in TRM (equivalent to 42 recursions) to lead to optimal + generalization on Sudoku-Extreme. More recursions +Self-attention is particularly good for long-context could be helpful for harder problems (we have not +lengths when L ≫ D since it only requires a matrix of tested it, given our limited resources); however, in- +[ D, 3D ] parameters, even though it can account for the creasing either T or n incurs massive slowdowns. We +whole sequence. However, when focusing on tasks show results at different n and T for HRM and TRM +where L ≤ D, a linear layer is cheap, requiring only a in Table 3. Note that TRM requires backpropagation +matrix of [ L, L] parameters. Taking inspiration from through a full recursion process, thus increasing n too +the MLP-Mixer (Tolstikhin et al., 2021), we can replace much leads to Out Of Memory (OOM) errors. How- +the self-attention layer with a multilayer perceptron ever, this memory cost is well worth its price in gold. +(MLP) applied on the sequence length. Using an MLP +instead of self-attention, we obtain better generaliza- In the following section, we show our main results on +tion on Sudoku-Extreme (improving from 74.7% to multiple datasets comparing HRM, TRM, and LLMs. +87.4%; see Table 1). This worked well on Sudoku 9x9 +grids, given the small and fixed context length; how- 5. Results +ever, we found this architecture to be suboptimal for +tasks with large context length, such as Maze-Hard Following Wang et al. (2025), we test our approach +and ARC-AGI (both using 30x30 grids). We show on the following datasets: Sudoku-Extreme (Wang +results with and without self-attention for all experi- et al., 2025), Maze-Hard (Wang et al., 2025), ARC-AGI- +ments. 1 (Chollet, 2019) and, ARC-AGI-2 (Chollet et al., 2025). + Results are presented in Tables 4 and 5. Hyperparame- + ters are detailed in Section 6. Datasets are discussed + below. + + + 7 + Recursive Reasoning with Tiny Networks + + ity of the MLP on large 30x30 grids). TRM with self- +Table 3. % Test accuracy on Sudoku-Extreme dataset. HRM + attention obtains 85.3% accuracy on Maze-Hard, 44.6% +versus TRM matched at a similar effective depth per super- +vision step ( T (n + 1)nlayers ) + accuracy on ARC-AGI-1, and 7.8% accuracy on ARC- + AGI-2 with 7M parameters. This is significantly higher + HRM TRM + than the 74.5%, 40.3%, and 5.0% obtained by HRM us- + n = k, 4 layers n = 2k, 2 layers + ing 4 times the number of parameters (27M). + k T Depth Acc (%) Depth Acc (%) + 1 1 9 46.4 7 63.2 + 2 2 24 55.0 20 81.9 Table 4. % Test accuracy on Puzzle Benchmarks (Sudoku- + 3 3 48 61.6 42 87.4 Extreme and Maze-Hard) + 4 4 80 59.5 72 84.2 Method # Params Sudoku Maze + 6 3 84 62.3 78 OOM Chain-of-thought, pretrained + 3 6 96 58.8 84 85.8 Deepseek R1 671B 0.0 0.0 + 6 6 168 57.5 156 OOM Claude 3.7 8K ? 0.0 0.0 + O3-mini-high ? 0.0 0.0 + Direct prediction, small-sample training +Sudoku-Extreme consists of extremely difficult Su- Direct pred 27M 0.0 0.0 +doku puzzles (Dillion, 2025; Palm et al., 2018; Park, HRM 27M 55.0 74.5 +2018) (9x9 grid), for which only 1K training samples TRM-Att (Ours) 7M 74.7 85.3 +are used to test small-sample learning. Testing is done TRM-MLP (Ours) 5M/19M 1 87.4 0.0 +on 423K samples. Maze-Hard consists of 30x30 mazes +generated by the procedure by Lehnert et al. (2024) +whose shortest path is of length above 110; both the +training set and test set include 1000 mazes. Table 5. % Test accuracy on ARC-AGI Benchmarks (2 tries) + Method # Params ARC-1 ARC-2 +ARC-AGI-1 and ARC-AGI-2 are geometric puzzles in- Chain-of-thought, pretrained +volving monetary prizes. Each puzzle is designed to Deepseek R1 671B 15.8 1.3 +be easy for a human, yet hard for current AI models. Claude 3.7 16K ? 28.6 0.7 +Each puzzle task consists of 2-3 input–output demon- o3-mini-high ? 34.5 3.0 +stration pairs and 1-2 test inputs to be solved. The final Gemini 2.5 Pro 32K ? 37.0 4.9 +score is computed as the accuracy over all test inputs Grok-4-thinking 1.7T 66.7 16.0 +from two attempts to produce the correct output grid. Bespoke (Grok-4) 1.7T 79.6 29.4 +The maximum grid size is 30x30. ARC-AGI-1 con- Direct prediction, small-sample training +tains 800 tasks, while ARC-AGI-2 contains 1120 tasks. + Direct pred 27M 21.0 0.0 +We also augment our data with the 160 tasks from + HRM 27M 40.3 5.0 +the closely related ConceptARC dataset (Moskvichev +et al., 2023). We provide results on the public evalua- TRM-Att (Ours) 7M 44.6 7.8 +tion set for both ARC-AGI-1 and ARC-AGI-2. TRM-MLP (Ours) 19M 29.6 2.4 + +While these datasets are small, heavy data- +augmentation is used in order to improve gen- 6. Conclusion +eralization. Sudoku-Extreme uses 1000 shuffling +(done without breaking the Sudoku rules) augmenta- We propose Tiny Recursion Models (TRM), a simple +tions per data example. Maze-Hard uses 8 dihedral recursive reasoning approach that achieves strong gen- +transformations per data example. ARC-AGI uses eralization on hard tasks using a single tiny network +1000 data augmentations (color permutation, dihedral- recursing on its latent reasoning feature and progres- +group, and translations transformations) per data sively improving its final answer. Contrary to the +example. The dihedral-group transformations consist Hierarchical Reasoning Model (HRM), TRM requires +of random 90-degree rotations, horizontal/vertical no fixed-point theorem, no complex biological justi- +flips, and reflections. fications, and no hierarchy. It significantly reduces + the number of parameters by halving the number of +From the results, we see that TRM without self- layers and replacing the two networks with a single +attention obtains the best generalization on Sudoku- tiny network. It also simplifies the halting process, +Extreme (87.4% test accuracy). Meanwhile, TRM with removing the need for the extra forward pass. Over- +self-attention generalizes better on the other tasks +(probably due to inductive biases and the overcapac- 1 5M on Sudoku and 19M on Maze + + + + 8 + Recursive Reasoning with Tiny Networks + +all, TRM is much simpler than HRM, while achieving gan training for high fidelity natural image synthe- +better generalization. sis. arXiv preprint arXiv:1809.11096, 2018. +While our approach led to better generalization on 4 Chollet, F. On the measure of intelligence. arXiv +benchmarks, every choice made is not guaranteed to preprint arXiv:1911.01547, 2019. +be optimal on every dataset. For example, we found +that replacing the self-attention with an MLP worked Chollet, F., Knoop, M., Kamradt, G., Landers, B., +extremely well on Sudoku-Extreme (improving test ac- and Pinkard, H. Arc-agi-2: A new challenge +curacy by 10%), but poorly on other datasets. Different for frontier ai reasoning systems. arXiv preprint +problem settings may require different architectures arXiv:2505.11831, 2025. +or number of parameters. Scaling laws are needed +to parametrize these networks optimally. Although Chowdhery, A., Narang, S., Devlin, J., Bosma, M., +we simplified and improved on deep recursion, the Mishra, G., Roberts, A., Barham, P., Chung, H. W., +question of why recursion helps so much compared Sutton, C., Gehrmann, S., et al. Palm: Scaling lan- +to using a larger and deeper network remains to be guage modeling with pathways. Journal of Machine +explained; we suspect it has to do with overfitting, but Learning Research, 24(240):1–113, 2023. +we have no theory to back this explaination. Not all +our ideas made the cut; we briefly discuss some of the Dillion, T. Tdoku: A fast sudoku solver and gener- +failed ideas that we tried but did not work in Section 6. ator. https://t-dillon.github.io/tdoku/, +Currently, recursive reasoning models such as HRM 2025. +and TRM are supervised learning methods rather than + Elman, J. L. Finding structure in time. Cognitive science, +generative models. This means that given an input + 14(2):179–211, 1990. +question, they can only provide a single deterministic +answer. In many settings, multiple answers exist for a Fedus, W., Zoph, B., and Shazeer, N. Switch transform- +question. Thus, it would be interesting to extend TRM ers: Scaling to trillion parameter models with simple +to generative tasks. and efficient sparsity. Journal of Machine Learning Re- + search, 23(120):1–39, 2022. +Acknowledgements + Geng, Z. and Kolter, J. Z. Torchdeq: A library for deep +Thank you Emy Gervais for your invaluable support equilibrium models. arXiv preprint arXiv:2310.18605, +and extra push. This research was enabled in part 2023. +by computing resources, software, and technical as- +sistance provided by Mila and the Digital Research Hendrycks, D. and Gimpel, K. Gaussian error linear +Alliance of Canada. units (gelus). arXiv preprint arXiv:1606.08415, 2016. + + Jang, Y., Kim, D., and Ahn, S. Hierarchical graph +References generation with k2-trees. In ICML 2023 Workshop on +ARC Prize Foundation. The Hidden Drivers of HRM’s Structured Probabilistic Inference Generative Modeling, + Performance on ARC-AGI. https://arcprize. 2023. + org/blog/hrm-analysis, 2025a. [Online; ac- + Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., + cessed 2025-09-15]. + Chess, B., Child, R., Gray, S., Radford, A., Wu, J., +ARC Prize Foundation. ARC-AGI Leaderboard. and Amodei, D. Scaling laws for neural language + https://arcprize.org/leaderboard, 2025b. models. arXiv preprint arXiv:2001.08361, 2020. + [Online; accessed 2025-09-24]. + Kingma, D. P. and Ba, J. Adam: A method for stochas- +Bai, S., Kolter, J. Z., and Koltun, V. Deep equilibrium tic optimization. arXiv preprint arXiv:1412.6980, + models. Advances in neural information processing 2014. + systems, 32, 2019. + Krantz, S. G. and Parks, H. R. The implicit function +Bai, X. and Melas-Kyriazi, L. Fixed point diffusion theorem: history, theory, and applications. Springer + models. In Proceedings of the IEEE/CVF Conference Science & Business Media, 2002. + on Computer Vision and Pattern Recognition, pp. 9430– + 9440, 2024. LeCun, Y. Une procedure d’apprentissage ponr reseau + a seuil asymetrique. Proceedings of cognitiva 85, pp. +Brock, A., Donahue, J., and Simonyan, K. Large scale 599–604, 1985. + + 9 + Recursive Reasoning with Tiny Networks + +Lehnert, L., Sukhbaatar, S., Su, D., Zheng, Q., Mcvay, P., all-mlp architecture for vision. Advances in neural + Rabbat, M., and Tian, Y. Beyond a*: Better planning information processing systems, 34:24261–24272, 2021. + with transformers via search dynamics bootstrap- + ping. arXiv preprint arXiv:2402.14083, 2024. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., + Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, +Lillicrap, T. P. and Santoro, A. Backpropagation I. Attention is all you need. Advances in neural + through time and the brain. Current opinion in neuro- information processing systems, 30, 2017. + biology, 55:82–89, 2019. + Wang, G., Li, J., Sun, Y., Chen, X., Liu, C., Wu, Y., +Loshchilov, I. and Hutter, F. Decoupled weight decay Lu, M., Song, S., and Yadkori, Y. A. Hierarchical + regularization. arXiv preprint arXiv:1711.05101, 2017. reasoning model. arXiv preprint arXiv:2506.21734, + 2025. +Moskvichev, A., Odouard, V. V., and Mitchell, M. The + conceptarc benchmark: Evaluating understanding Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., + and generalization in the arc domain. arXiv preprint Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought + arXiv:2305.07141, 2023. prompting elicits reasoning in large language mod- + els. Advances in neural information processing systems, +Palm, R., Paquet, U., and Winther, O. Recurrent re- + 35:24824–24837, 2022. + lational networks. Advances in neural information + processing systems, 31, 2018. Werbos, P. Beyond regression: New tools for predic- + tion and analysis in the behavioral sciences. PhD +Park, K. Can convolutional neural networks + thesis, Committee on Applied Mathematics, Harvard + crack sudoku puzzles? https://github.com/ + University, Cambridge, MA, 1974. + Kyubyong/sudoku, 2018. + Werbos, P. J. Generalization of backpropagation with +Prieto, L., Barsbey, M., Mediano, P. A., and Birdal, T. + application to a recurrent gas market model. Neural + Grokking at the edge of numerical stability. arXiv + networks, 1(4):339–356, 1988. + preprint arXiv:2501.04697, 2025. + Zhang, B. and Sennrich, R. Root mean square layer +Rumelhart, D. E., Hinton, G. E., and Williams, R. J. + normalization. Advances in Neural Information Pro- + Learning internal representations by error propaga- + cessing Systems, 32, 2019. + tion. Technical report, 1985. +Shazeer, N. Glu variants improve transformer. arXiv + preprint arXiv:2002.05202, 2020. +Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, + Q., Hinton, G., and Dean, J. Outrageously large neu- + ral networks: The sparsely-gated mixture-of-experts + layer. arXiv preprint arXiv:1701.06538, 2017. +Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling + llm test-time compute optimally can be more effec- + tive than scaling model parameters. arXiv preprint + arXiv:2408.03314, 2024. +Song, Y. and Ermon, S. Improved techniques for train- + ing score-based generative models. Advances in + neural information processing systems, 33:12438–12448, + 2020. +Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, + Y. Roformer: Enhanced transformer with rotary + position embedding. Neurocomputing, 568:127063, + 2024. +Tolstikhin, I. O., Houlsby, N., Kolesnikov, A., Beyer, + L., Zhai, X., Unterthiner, T., Yung, J., Steiner, A., + Keysers, D., Uszkoreit, J., et al. Mlp-mixer: An + + 10 + Recursive Reasoning with Tiny Networks + +Hyper-parameters and setup propagating through the whole n + 1 recursions makes + the most sense and works best. +All models are trained with the AdamW opti- +mizer(Loshchilov & Hutter, 2017; Kingma & Ba, 2014) We tried removing ACT with the option of stopping +with β 1 = 0.9, β 2 = 0.95, small learning rate warm- when the solution is reached, but we found that gen- +up (2K iterations), batch-size 768, hidden-size of 512, eralization dropped significantly. This can probably +Nsup = 16 max supervision steps, and stable-max loss be attributed to the fact that the model is spending +(Prieto et al., 2025) for improved stability. TRM uses an too much time on the same data samples rather than +Exponential Moving Average (EMA) of 0.999. HRM focusing on learning on a wide range of data samples. +uses n = 2, T = 2 with two 4-layers networks, while We tried weight tying the input embedding and out- +we use n = 6, T = 3 with one 2-layer network. put head, but this was too constraining and led to a +For Sudoku-Extreme and Maze-Hard, we train for 60k massive generalization drop. +epochs with learning rate 1e-4 and weight decay 1.0. We tried using TorchDEQ (Geng & Kolter, 2023) to +For ARC-AGI, we train for 100K epochs with learning replace the recursion steps by fixed-point iteration as +rate 1e-4 (with 1e-2 learning rate for the embeddings) done by Deep Equilibrium Models (Bai et al., 2019). +and weight decay 0.1. The numbers for Deepseek R1, This would provide a better justification for the 1-step +Claude 3.7 8K, O3-mini-high, Direct prediction, and gradient approximation. However, this slowed down +HRM from the Table 4 and 5 are taken from Wang et al. training due to the fixed-point iteration and led to +(2025). Both HRM and TRM add an embedding of worse generalization. This highlights the fact that +shape [0, 1, D ] on Sudoku-Extreme and Maze-Hard to converging to a fixed-point is not essential. +the input. For ARC-AGI, each puzzle (containing 2-3 +training examples and 1-2 test examples) at each data- +augmentation is given a specific embedding of shape +[0, 1, D ] and, at test-time, the most common answer +out of the 1000 data augmentations is given as answer. +Experiments on Sudoku-Extreme were ran with 1 L40S +with 40Gb of RAM for generally less than 36 hours. +Experiments on Maze-Hard were ran with 4 L40S with +40Gb of RAM for less than 24 hours. Experiments on +ARC-AGI were ran for around 3 days with 4 H100 +with 80Gb of RAM. + +Ideas that failed +In this section, we quickly mention a few ideas that +did not work to prevent others from making the same +mistake. +We tried replacing the SwiGLU MLPs by SwiGLU +Mixture-of-Experts (MoEs) (Shazeer et al., 2017; Fedus +et al., 2022), but we found generalization to decrease +massively. MoEs clearly add too much unnecessary +capacity, just like increasing the number of layers does. +Instead of back-propagating through the whole n + 1 +recursions, we tried a compromise between HRM 1- +step gradient approximation, which back-propagates +through the last 2 recursions. We did so by decou- +pling n from the number of last recursions k that we +back-propagate through. For example, while n = 6 +requires 7 steps with gradients in TRM, we can use +gradients for only the k = 4 last steps. However, we +found that this did not help generalization in any way, +and it made the approach more complicated. Back- + + + 11 + Recursive Reasoning with Tiny Networks + +Algorithms with different number of latent Example on Sudoku-Extreme +features + 8 3 1 + 9 6 8 7 +def latent recursion(x, z, n=6): 3 5 + for i in range(n+1): # latent recursion + z = net(x, z) 6 8 + return z 6 2 +def deep recursion(x, z, n=6, T=3): + 7 4 3 + # recursing T−1 times to improve z (no gradients needed) 9 4 + with torch.no grad(): + for j in range(T−1): + 2 4 1 + z = latent recursion(x, z, n) 6 2 5 7 + # recursing once to improve z + z = latent recursion(x, z, n) Input x + return z.detach(), output head(y), Q head(y) + 5 2 6 7 9 4 8 3 1 +# Deep Supervision +for x input, y true in train dataloader: 3 9 1 2 6 8 4 7 5 + z = z init 4 8 7 3 1 5 2 9 6 + for step in range(N supervision): + x = input embedding(x input) + 1 6 8 5 3 2 7 4 9 + z, y hat, q hat = deep recursion(x, z) 9 3 5 4 7 6 1 8 2 + loss = softmax cross entropy(y hat, y true) 7 4 2 9 8 1 5 6 3 + loss += binary cross entropy(q hat, (y hat == y true)) + z = z.detach() 8 7 3 1 5 9 6 2 4 + loss.backward() 2 5 9 6 4 7 3 1 8 + opt.step() + opt.zero grad() 6 1 4 8 5 3 9 5 7 + if q[0] > 0: # early−stopping + break Output y + 5 2 6 7 9 4 8 3 1 +Figure 4. Pseudocode of TRM using a single-z with deep 3 9 1 2 6 8 4 7 5 +supervision training in PyTorch. 4 8 7 3 1 5 2 9 6 + 1 6 8 5 3 2 7 4 9 + 9 3 5 4 7 6 1 8 2 + 7 4 2 9 8 1 5 6 3 + 8 7 3 1 5 9 6 2 4 +def latent recursion(x, y, z, n=6): + for i in range(n): # latent recursion 2 5 9 6 4 7 3 1 8 + z[i] = net(x, y, z[0], ... , z[n−1]) 6 1 4 8 5 3 9 5 7 + y = net(y, z[0], ... , z[n−1]) # refine output answer + return y, z Tokenized z H (denoted y in TRM) +def deep recursion(x, y, z, n=6, T=3): 5 5 4 9 4 6 3 + # recursing T−1 times to improve y and z (no gradients needed) + with torch.no grad(): 4 3 1 4 6 5 + for j in range(T−1): 4 8 4 3 6 6 4 + y, z = latent recursion(x, y, z, n) + # recursing once to improve y and z 9 6 5 3 5 4 + y, z = latent recursion(x, y, z, n) 3 5 4 3 5 4 4 + return (y.detach(), z.detach()), output head(y), Q head(y) + 6 3 3 3 5 8 8 +# Deep Supervision 3 3 3 6 5 6 6 4 +for x input, y true in train dataloader: 7 5 6 3 3 6 6 + y, z = y init, z init + for step in range(N supervision): 4 3 4 8 3 6 6 4 + x = input embedding(x input) + (y, z), y hat, q hat = deep recursion(x, y, z) Tokenized z L (denoted z in TRM) + loss = softmax cross entropy(y hat, y true) + loss += binary cross entropy(q hat, (y hat == y true)) + loss.backward() + opt.step() + Figure 6. This Sudoku-Extreme example shows an input, ex- + opt.zero grad() pected output, and the tokenized z H and z L (after reversing + if q[0] > 0: # early−stopping the embedding and using argmax) for a pretrained model. + break + This highlights the fact that z H corresponds to the predicted + response, while z L is a latent feature that cannot be decoded +Figure 5. Pseudocode of TRM using multi-scale z with deep to a sensible output unless transformed into z H by f H . +supervision training in PyTorch. + + + + + 12 +
\ No newline at end of file |