Universal Reasoning Model — part 2/2

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

The dominant effect arises at position (f), after the MLP expansion, indicating that short-range mixing is most beneficial when applied within an already nonlinear subspace. This supports a functional interpretation in which the MLP—not attention—constitutes the model’s primary source of expressive nonlinearity; augmenting it with W dwconv W_{\mathrm{dwconv}} substantially enhances the model’s nonlinear representational capacity.

As shown in Fig. 4 , the incorporation of short convolution into the MLP significantly enhances channel mixing. While the standard Universal Transformer exhibits relatively sparse and homogeneous attention patterns, the model with ConvSwiGLU produces attention matrices with more diverse and structured distributions. This suggests that short convolution facilitates more effective inter-channel information flow, thereby improving the expressiveness of the attention mechanism.

### 4.5 Truncated Backpropagation Through Loops

As shown in Table 5 , when the total number of inner loops is fixed to 8, truncating gradients for the first two loops—i.e., running the initial two inner-loop iterations in forward-only mode—achieves the best performance. Both pass@1 and pass@1000 peak at this truncation setting, while shorter or longer truncation horizons result in inferior outcomes.

This trend closely resembles truncated backpropagation through time (TBPTT) in recurrent neural networks, where the underlying motivation is largely the same. In full backpropagation through time, gradients are propagated through the entire sequence, which incurs high computational and memory costs and often yields ineffective long-range gradients due to vanishing or exploding behaviors. As a result, practical implementations typically restrict gradient propagation to a fixed recent window, e.g., by backpropagating errors only through the last L L time steps and updating the network parameters accordingly [ 14 , 17 ] .

Similarly, in universal transformers, propagating gradients across all inner-loop iterations can lead to unstable optimization, while overly aggressive truncation limits the model’s ability to coordinate multi-step refinement. Moderately truncating gradient propagation therefore provides a favorable balance between optimization stability and effective long-horizon learning.

We note that all results in this experiment are obtained using a two-layer URM without the short convolution module, which differs from the full URM model reported earlier.

### 4.6 Nonlinearity of Transformers

As shown in Table 4 , the performance on ARC-AGI 1 decreases monotonically as nonlinear components are progressively removed from the model. Among these components, the activation function in the MLP plays a particularly critical role: replacing SwiGLU with simpler nonlinearities such as SiLU or ReLU leads to substantial degradation, while completely removing the attention softmax results in a dramatic collapse in performance. This clear monotonic trend highlights the importance of strong nonlinear transformations for solving complex abstract reasoning tasks.

These results suggest that the expressive power required for ARC-AGI primarily arises from rich nonlinear mappings. Weakening the nonlinearity may systematically limits the model’s ability to represent complex reasoning skills.

We note that the model still retains certain forms of nonlinearity that are not ablated in this study, such as the RMSNorm applied after each layer and the dot-product interaction between queries and keys in attention. However, these components are either difficult to remove without causing training instability or represent relatively weak nonlinear effects compared to explicit activation functions. As ablating them typically leads to training failure, they fall outside the scope of the present analysis.

### 4.7 Muon Optimizer

To evaluate the training efficiency of the Universal Reasoning Model (URM), we compare the Muon (Momentum Updated Orthogonal Newton) optimizer [ 12 ] with a standard adaptive baseline, Adamatan2 [ 6 ] . Muon approximates second-order curvature to apply orthogonal updates to better handle the complex loss landscapes [ 8 ] induced by deep recurrent structures. Both models are trained from scratch under identical experimental settings, including batch size, learning rate schedules, and data augmentation, ensuring that any observed differences arise solely from the choice of optimizer.

Across the ARC-AGI 1 and ARC-AGI 2 benchmarks, Muon demonstrates substantially faster convergence. On ARC-AGI 2, the Muon-optimized model reaches a pass@1 accuracy of 11.5% in approximately 600,000 training steps, whereas the Adamatan2 baseline requires over 1,300,000 steps to achieve the same performance, corresponding to nearly a twofold speedup in optimization. Despite this advantage in early training, both methods converge to similar final accuracies (approximately 53.8% on ARC-AGI 1 and 16.0% on ARC-AGI 2), indicating comparable asymptotic performance.

These results suggest a separation between optimization efficiency and architectural capacity in the URM. While Muon preconditions the challenging spectral properties of recurrent weight matrices [ 13 ] and reduces training cost, it does not lead to improved final generalization.

## 5 Related Work

### 5.1 ARC-AGI

Prior work on the ARC-AGI benchmark [ 3 , 2 ] spans vision-based formulations, large language model (LLM) adaptation, and recurrent reasoning architectures. Vision-centric approaches such as Vision ARC [ 10 ] reformulate ARC as an image-to-image transformation problem and show that standard visual inductive biases can achieve competitive performance, particularly with ensembling and test-time scaling. LLM-based methods explore fine-tuning and test-time training, demonstrating that transient parameter updates outperform static in-context learning on ARC-like tasks. Beyond language and vision models, recurrent architectures emphasize iterative computation as a core mechanism for abstraction. The Hierarchical Reasoning Model (HRM) [ 20 , 7 ] introduces multi-timescale recurrence and achieves strong ARC-AGI results, while subsequent analyses suggest that its gains may largely stem from recurrence rather than explicit hierarchy. The Tiny Recursive Model (TRM) [ 11 ] further simplifies this paradigm, showing that a single lightweight network applied recursively can match or exceed more complex hierarchical designs.

### 5.2 Universal Transformers (Looped Transformers)

The Universal Transformer (UT), also known as the Looped Transformer, was introduced by Dehghani et al. [ 5 ] as an extension of the standard Transformer with recurrent computation and adaptive computation time. Subsequent work has shown that UTs exhibit significantly stronger multi-step reasoning abilities than vanilla Transformers, as the recurrent refinement mechanism helps overcome architectural limitations in multi-hop reasoning tasks [ 4 , 19 ] . In addition, UTs demonstrate improved algorithmic learning capabilities, enabling more effective modeling of iterative and rule-based computations [ 21 ] . By reusing parameters across refinement steps, UTs also achieve higher parameter efficiency, allowing more expressive computation without increasing model size [ 15 ] .

## 6 Conclusion

We systematically investigate the sources of performance gains in Universal Transformer models on complex reasoning tasks. Extensive ablation studies reveal that these gains stem primarily from the recurrent inductive bias and strong nonlinear components of Transformer, rather than from overly complex architectural designs. Motivated by this insight, we propose the Universal Reasoning Model (URM), which enhances nonlinear depth-wise computation via short convolutional gating and improves optimization stability through truncated backpropagation through loops. URM achieves state-of-the-art performance on ARC-AGI 1 and 2.

## 7 Acknowledgement

We thank Benhao Huang for pointing out the typo in the previous version, and we also thank Zhengmao Ye from the Ubiquant AI team for providing infrastructure support.