Universal Reasoning Model — part 1/2

Cleaned-text code points [0, 15770). Bibliography omitted here; source.md includes it.

# Universal Reasoning Model

###### Abstract

Universal transformers (UTs) have been widely used for complex reasoning tasks such as ARC-AGI and Sudoku, yet the specific sources of their performance gains remain underexplored. In this work, we systematically analyze UTs variants and show that improvements on ARC-AGI primarily arise from the recurrent inductive bias and strong nonlinear components of Transformer, rather than from elaborate architectural designs. Motivated by this finding, we propose the Universal Reasoning Model (URM), which enhances the UT with short convolution and truncated backpropagation. Our approach substantially improves reasoning performance, achieving state-of-the-art ∗ 53.8% pass@1 on ARC-AGI 1 and 16.0% pass@1 on ARC-AGI 2. 0 0 footnotetext: ∗ This comparison focuses on pass@1 score of single small models trained from scratch under the same data setting as HRM and TRM, excluding test-time scaling, ensembling, and visual methods such as VARC [ 10 ] . Our code is avaliable at https://github.com/UbiquantAI/URM .

## 1 Introduction

Recent advances in recurrent models [ 20 , 7 , 11 ] have demonstrated the effectiveness of Universal Transformers (UTs) [ 5 ] in addressing complex reasoning tasks, such as ARC-AGI and Sudoku [ 3 , 2 ] . UT-based small models, despite being trained from scratch on these tasks without internet-scale pre-training, consistently outperform most standard Transformer-based Large Language models (LLMs) by a significant margin [ 20 ] .

While this contrast highlights the potential of UTs for depth-intensive iterative reasoning, the function and impact of gating mechanisms remain insufficiently explored beyond their initial intuition

Prior studies often attribute improvements to high-level architectural innovations [ 20 , 7 , 11 ] , yet our analysis reveals that the core performance gain actually arises from the often-overlooked recurrent inductive bias intrinsic to the Universal Transformer. In particular, nonlinear depth-wise computation plays a much larger role than previously acknowledged, suggesting that architectural modifications that enhance recurrent processing can yield substantial downstream improvements. Motivated by this insight, we further investigate and strengthen this inductive bias via a simplified yet effective enhancement to the UT framework, enabling stronger abstraction capabilities while preserving parameter efficiency.

Our main contributions are as follows: • Through extensive ablation studies, we show that the performance of models on ARC-AGI–style complex reasoning tasks primarily stems from their nonlinearity. Moreover, we reveal that the true source of reasoning capability beyond standard Transformers comes from the recurrent mechanism of Universal Transformers rather than overly elaborate design in prior work.

• By introducing short convolutions and truncated backpropagation into the Universal Transformer, we achieve a state-of-the-art 53.8% pass@1 accuracy on ARC-AGI 1 and 16.0% on ARC-AGI 2.

## 2 Preliminaries

### 2.1 Standard Transformer

Let 𝒱 \mathcal{V} denote the vocabulary of size V V , and let 𝐱 = ( x 1 , … , x N ) ∈ 𝒱 N \mathbf{x}=(x_{1},\dots,x_{N})\in\mathcal{V}^{N} be an input sequence of length N N . We define the token embedding function as ϕ : 𝒱 N → ℝ N × d \phi:\mathcal{V}^{N}\to\mathbb{R}^{N\times d} , mapping discrete tokens to a d d -dimensional continuous representation. Conversely, the unembedding function (or language modeling head) is denoted by ψ : ℝ N × d → ℝ N × V \psi:\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times V} , which projects hidden states back to the vocabulary logit space.

A single Transformer layer, parameterized by θ \theta , is defined as a function 𝒯 θ : ℝ N × d → ℝ N × d \mathcal{T}_{\theta}:\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times d} . This function typically composes a Multi-Head Self-Attention (MHSA) module and a Position-wise Feed-Forward Network (FFN), each wrapped with residual connections and layer normalization:

𝒯 θ ​ ( H ) = FFN ​ ( LN ​ ( H ′ + H ) ) , where ​ H ′ = MHSA ​ ( LN ​ ( H ) ) \begin{split}\mathcal{T}_{\theta}(H)&=\text{FFN}(\text{LN}(H^{\prime}+H)),\\ \text{where }H^{\prime}&=\text{MHSA}(\text{LN}(H))\end{split}

A standard, non-recursive Transformer model ℳ std \mathcal{M}_{\text{std}} of depth L L is constructed by stacking L L layers with distinct parameters Θ = { θ 1 , … , θ L } \Theta=\{\theta_{1},\dots,\theta_{L}\} . The forward pass is the composition of these layers:

ℳ std ​ ( 𝐱 ) = ψ ∘ 𝒯 θ L ∘ ⋯ ∘ 𝒯 θ 1 ∘ ϕ ⁡ ( 𝐱 ) \mathcal{M}_{\text{std}}(\mathbf{x})=\psi\circ\mathcal{T}_{\theta_{L}}\circ\dots\circ\mathcal{T}_{\theta_{1}}\circ\phi(\mathbf{x})

Here, the operator ∘ \circ denotes function composition. The computational cost and parameter count both scale linearly with L L , creating a rigid coupling between model capacity and inference compute.

### 2.2 Universal Transformer

The Universal Transformer (UT) [ 5 ] extends the standard Transformer [ 18 ] by introducing recurrent computation over depth . Instead of stacking L L distinct layers, the UT applies a single transition block repeatedly to refine token representations. For an input sequence 𝐱 \mathbf{x} with embedding matrix 𝐇 0 ∈ ℝ n × d \mathbf{H}^{0}\in\mathbb{R}^{n\times d} , the UT updates states as

𝐇 t + 1 = LayerNorm ⁡ ( 𝐇 t + MHA ⁡ ( 𝐇 t ) ) , \mathbf{H}^{t+1}=\mathrm{LayerNorm}\!\left(\mathbf{H}^{t}+\mathrm{MHA}\!\left(\mathbf{H}^{t}\right)\right),

followed by a shared position-wise transition function

𝐇 t + 1 ← LayerNorm ( 𝐇 t + 1 + Transition ( 𝐇 t + 1 ) ) , t = 0 , … , T − 1 , \mathbf{H}^{t+1}\leftarrow\mathrm{LayerNorm}\!\left(\mathbf{H}^{t+1}+\mathrm{Transition}\!\left(\mathbf{H}^{t+1}\right)\right),\qquad t=0,\dots,T-1,

where Transition \mathrm{Transition} is either a feed-forward network or separable convolution. To encode both position and refinement depth, UT adds 2-D sinusoidal embeddings at each step.

#### 2.2.1 Parameter Sharing

A key design of UT is weight tying across depth. The attention and transition parameters

Θ UT = { 𝐖 h Q , 𝐖 h K , 𝐖 h V , 𝐖 O , Θ Transition } \Theta_{\mathrm{UT}}=\{\mathbf{W}_{h}^{Q},\mathbf{W}_{h}^{K},\mathbf{W}_{h}^{V},\mathbf{W}^{O},\Theta_{\mathrm{Transition}}\}

are reused for all t t . Thus, the model performs iterative representation refinement with a flexible number of steps T T , enabling (i) depth adaptation at inference and (ii) higher theoretical expressivity than fixed-depth Transformers.

#### 2.2.2 Adaptive Computation Time (ACT)

With ACT [ 9 ] , different tokens may halt at different recurrent steps. At step t t , each position predicts a halting probability

p t , i = σ ⁡ ( 𝐰 ⊤ ​ 𝐡 t , i + b ) , p_{t,i}=\sigma(\mathbf{w}^{\top}\mathbf{h}_{t,i}+b),

accumulated until reaching threshold 1 − ϵ 1-\epsilon . The final token representation is a weighted mixture

𝐡 i final = ∑ t Δ t , i ​ 𝐡 t , i , \mathbf{h}^{\mathrm{final}}_{i}=\sum_{t}\Delta_{t,i}\,\mathbf{h}_{t,i},

where Δ t , i \Delta_{t,i} is the truncated allocation. ACT allows UT to allocate more computation to complex tokens and less to simpler ones.

## 3 Universal Reasoning Model

The base architecture of our Universal Reasoning Model (URM) closely follows that of the Universal Transformer [ 5 ] , with the difference being its decoder-only design. This aspect is consistent with previous works such as HRM [ 20 ] and TRM [ 11 ] . Our work differs from previous models [ 20 , 11 ] by introducing the following ConvSwiGLU module and a Truncated Backpropagation Through Loops mechanism.

### 3.1 ConvSwiGLU

To strengthen the non-linearity of Universal Transformer, we introduce a ConvSwiGLU (motivation see Section 4.6 ), which augments the standard SwiGLU feed-forward block with a depthwise short convolution. Unlike the conventional point-wise SwiGLU [ 16 ] , which treats each token independently, our design explicitly injects local contextual interactions into the gating mechanism, introducing lightweight channel mixing in token space without increasing sequence-level complexity [ 1 , 22 ] .

Given an input sequence X ∈ ℝ T × d X\in\mathbb{R}^{T\times d} , we first project it into an expanded intermediate representation:

[ 𝐆 , 𝐔 ] = X ​ W up ∈ ℝ T × 2 ​ m . [\mathbf{G},\mathbf{U}]=XW_{\text{up}}\in\mathbb{R}^{T\times 2m}.

The SwiGLU activation produces a gated representation:

𝐇 ffn = SiLU ​ ( 𝐆 ) ⊙ 𝐔 . \mathbf{H}_{\text{ffn}}=\text{SiLU}(\mathbf{G})\odot\mathbf{U}.

To integrate short-range token interactions, we apply a depthwise 1D convolution over the gated features:

𝐇 conv = σ ⁡ ( 𝐖 dwconv ∗ 𝐇 ffn ) , \mathbf{H}_{\text{conv}}=\sigma\bigl(\mathbf{W}_{\text{dwconv}}*\mathbf{H}_{\text{ffn}}\bigr),

where 𝐖 dwconv ∈ ℝ m × 1 × k \mathbf{W}_{\text{dwconv}}\in\mathbb{R}^{m\times 1\times k} is a depthwise convolution kernel of size k = 2 k=2 .

Finally, the output is projected back to the hidden dimension:

𝐘 = [ σ ⁡ ( 𝐖 dwconv ∗ ( SiLU ​ ( 𝐆 ) ⊙ 𝐔 ) ) ] ​ W down . \boxed{\mathbf{Y}=\bigl[\sigma(\mathbf{W}_{\text{dwconv}}*(\text{SiLU}(\mathbf{G})\odot\mathbf{U}))\bigr]W_{\text{down}}.}

### 3.2 Truncated Backpropagation Through Loops

When the number of recurrent reasoning loops becomes large, the gradients propagated from early loops may hinder optimization due to noise accumulation and instability (see empirical evidence in Section 4.5 ). To alleviate this issue, we employ Truncated Backpropagation Through Loops (TBPTL) and only compute gradients for the later loops.

Consider a D D -layer Universal Reasoning Model unrolled for M M iterative loops during training. Let 𝐡 t ( d ) \mathbf{h}_{t}^{(d)} denote the hidden representation of layer d ∈ { 1 , … , D } d\in\{1,\ldots,D\} at iteration t ∈ { 1 , … , M } t\in\{1,\ldots,M\} . The recurrent transition is defined as:

𝐡 t ( d ) = F θ ( d ) ​ ( 𝐡 t ( d − 1 ) , 𝐡 t − 1 ( d ) ) , \mathbf{h}_{t}^{(d)}=F_{\theta}^{(d)}\big(\mathbf{h}_{t}^{(d-1)},\mathbf{h}_{t-1}^{(d)}\big),

where F θ ( d ) F_{\theta}^{(d)} denotes the parameterized transformation at layer d d with trainable parameters θ \theta .

Instead of backpropagating through all M M loops, we partition the rollout into forward-only and trainable segments. Specifically, for a truncation index N < M N<M : { 1 , 2 , … , N } ⏟ no backward pass , { N + 1 , … , M } ⏟ forward + backward . \underbrace{\{1,2,\ldots,N\}}_{\text{no backward pass}},\qquad\underbrace{\{N+1,\ldots,M\}}_{\text{forward + backward}}. During training, we compute gradients only on the loss accumulated in the latter ( M − N ) (M-N) loops: ℒ TBPTL ​ ( θ ) = ∑ t = N + 1 M ℒ ⁡ ( 𝐡 t ( D ) , y ) , \mathcal{L}_{\text{TBPTL}}(\theta)=\sum_{t=N+1}^{M}\mathcal{L}\big(\mathbf{h}_{t}^{(D)},y\big), where ℒ ⁡ ( ⋅ ) \mathcal{L}(\cdot) is cross-entropy loss function. The gradients with respect to θ \theta are thus: ∇ θ ℒ TBPTL = ∑ t = N + 1 M ∂ ℒ ∂ 𝐡 t ( D ) ​ ∂ 𝐡 t ( D ) ∂ θ . \nabla_{\theta}\mathcal{L}_{\text{TBPTL}}=\sum_{t=N+1}^{M}\frac{\partial\mathcal{L}}{\partial\mathbf{h}_{t}^{(D)}}\frac{\partial\mathbf{h}_{t}^{(D)}}{\partial\theta}.

## 4 Experiment

### 4.1 Experiment Settings

Our experimental setup largely follows HRM and TRM [ 20 , 11 ] . We use the same datasets and augmented data as in prior work, and apply an exponential moving average (EMA) to model parameters to improve training stability, following [ 11 ] . All models are trained with the AdamAtan2 optimizer [ 6 ] . For ARC-AGI 1 and ARC-AGI 2, the main model learning rates are set to 1 × 10 − 4 1\times 10^{-4} and 3 × 10 − 4 3\times 10^{-4} , respectively, while the puzzle embedding uses a learning rate of 1 × 10 − 2 1\times 10^{-2} ; for Sudoku, the puzzle embedding learning rate is 1 × 10 − 4 1\times 10^{-4} . Weight decay is set to 0.1 for both the main model and puzzle embedding on ARC-AGI 1 and ARC-AGI 2, and to 1.0 for Sudoku, consistent with prior work. The model has 4 layers with hidden size 512 and 8 attention heads. The inner loop runs for 8 steps, with the first two steps being forward-only, while the outer loop employs Adaptive Computation Time (ACT) [ 9 ] with a maximum of 16 steps.

### 4.2 Main Results

As shown in Table 1, the Universal Reasoning Model (URM) achieves substantial improvements over prior UT-based approaches across all benchmarks. On ARC-AGI 1, URM reaches 53.8% pass@1, outperforming TRM (40.0%) and HRM (34.4%) by large margins. On ARC-AGI 2, URM obtains 16.0% pass@1, nearly tripling HRM and more than doubling TRM. A similar advantage appears on Sudoku, where URM achieves 77.6% accuracy, surpassing both TRM and HRM.

Notably, URM’s gains further widen under larger sampling budgets (e.g., pass@1000), indicating that iterative refinement enables richer candidate generation rather than brittle one-step predictions.

### 4.3 Why Universal Transformer?

Table 2 demonstrates that the performance gains of Universal Transformers (UTs) on ARC-AGI 1 arise from substantially higher parameter efficiency rather than increased model scale or computational budget. With only 4× parameters, a UT achieves a pass@1 score of 40.0, dramatically outperforming vanilla Transformers that employ up to 32× more parameters yet remain markedly weaker. Simply scaling depth or width in vanilla Transformers yields diminishing returns and can even lead to performance degradation, highlighting a fundamental inefficiency in how parameters are used to support multi-step reasoning.

Crucially, this advantage persists even when computation is held constant. At 32× FLOPs, reallocating computation from deep, non-shared layers to recurrent refinement improves pass@1 from 23.75 for vanilla Transformers to 40.0 for UTs. This behavior is consistent with analyses of previous works [ 15 ] , which argue that many reasoning tasks benefit more from iterative computation than from increasing the number of independent layers. In standard Transformers, additional FLOPs are often spent on redundant refinement in higher layers, whereas recurrent computation converts the same budget into increased effective depth [ 23 , 15 ] .

This superior efficiency is driven by the recurrent inductive bias introduced by parameter sharing across depth. Through repeated application of a shared transformation, UTs realize iterative refinement that better aligns with the structure of algorithmic reasoning, while avoiding any increase in parameter count. Consequently, under both fixed parameter and fixed FLOPs budgets, UTs consistently outperform vanilla Transformers on reasoning tasks, making them particularly well suited for reasoning-intensive settings such as ARC-AGI, where multi-step abstraction is more critical than sheer scale.

### 4.4 Short Convolution

To strengthen the nonlinear inductive bias of the Universal Transformer, we introduce a depthwise short convolution module parameterized by W dwconv W_{\mathrm{dwconv}} (see Section 3.1 for details), which provides token-local mixing while preserving the per-step computational budget. Since ARC-AGI performance correlates strongly with nonlinear capacity (Section 4.6 ), we evaluate how inserting this module at different locations affects the recurrent transition.

We examine six insertion points: (a) after the SDPA output; (b) after the value projection; (c) after the key projection; (d) after the query projection; (e) between multi-head concatenation and the output projection; and (f) after the MLP expansion.

As shown in Figure 3 , inserting the W dwconv W_{\mathrm{dwconv}} module inside the attention pathway, positions (a)–(d), does not yield improvements and often degrades performance, suggesting that local perturbations interfere with the geometric structure of attention’s linear projections. A mild gain appears at position (e), where the perturbation acts only on aggregated multi-head features.