Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach — part 4/5

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

From another angle, these models can be described as neural networks learning a fixed-point iteration, as studied in deep equilibrium models ( Bai et al., 2019 ; Bai et al., 2022 ) . They are further related to diffusion models ( Song and Ermon, 2019 ) , especially latent diffusion models ( Rombach et al., 2022 ) , but we note that language diffusion models are usually run with a per-sequence, instead of a per-token, iteration count ( Lee et al., 2018 ) . A key difference of our approach to both equilibrium models and diffusion models is in the training objective, where equilibrium methods solve the “direct” problem ( Geiping and Moeller, 2019 ) , diffusion models solve a surrogate training objective, and our work suggests that truncated unrolling is a scalable alternative.

More generally, all architectures that recur in depth can also be understood as directly learning the analog to the gradient of a latent energy-based model ( LeCun and Huang, 2005 ; LeCun, 2022 ) , to an implicitly defined intermediate optimization layer ( Amos and Kolter, 2017 ) , or to a Kuramoto layer ( Miyato et al., 2024 ) . Analogies to gradient descent at inference time also show the connection to test time adaptation ( Sun et al., 2020 ) , especially test-time adaptation of output states ( Boudiaf et al., 2022 ) .

Aside from full recurrent-depth architectures, there also exist a number of proposals for hybrid architectures, such as models with latent sub-networks ( Li et al., 2020a ) , LoRA adapters on top of weight-shared layers ( Bae et al., 2024 ) , or (dynamic) weight-tying of trained models ( Hay and Wolf, 2023 ; Liu et al., 2024b ) .

As mentioned in Section 6 , while we consider the proposed recurrent depth approach to be a very natural way to learn to reason in continuous latent space from the ground up, the works of Hao et al. (2024) ; Cheng and Durme (2024) and Liu et al. (2024a) discuss how to finetune existing fixed-depth transformers with this capability. These works have a similar aim to ours, enabling reasoning in latent space, but approach this goal from separate directions.

For additional discussions related to the idea of constructing a prior that incentivizes reasoning and algorithm learning at the expense of memorization of simple patterns, we also refer to Chollet (2019) , Schwarzschild (2023) , Li et al. (2020b) and Moulton (2023) .

## 9 Future Work

Aside from work extending and analyzing the scaling behaviors of recurrent depth models, there are many questions that remain unanswered. For example, to us, there are potentially a large number of novel post-training schemes that further enhance the capabilities of these models, such as fine-tuning to compress the recurrence or reinforcement learning with data with different hardness levels ( Zelikman et al., 2024 ) , or to internalize reasoning from CoT data into the recurrence ( Deng et al., 2024 ) .

Another aspect not covered in this work is the relationship to other modern architecture improvements. Efficient sequence mixing operations, especially those that are linear in sequence dimension, such as linear attention ( Katharopoulos et al., 2020 ; Yang et al., 2024b ) , are limited in the number of comparisons that can be made. However, with recurrent depth, blocks containing linear operators can repeat until all necessary comparisons between sequence elements are computed ( Suzgun et al., 2019 ) . For simplicity, we also focus on a single recurrence, where prior work has considered multiple successive recurrent stages ( Takase and Kiyono, 2023 ; Csordás et al., 2024 ) .

Finally, the proposed architecture is set up to be compute-heavy , with more “materialized” parameters than there are actual parameters. This naturally mirrors mixture-of-expert models (MoE), which are parameter-heavy , using fewer active parameters per forward pass than exist within the model ( Shazeer et al., 2017 ; Fedus et al., 2022 ) . We posit that where the recurrent-depth setup excels at learning reasoning patterns, the MoE excels at effectively storing and retrieving complex information. Their complementarity supports the hypothesis that a future architecture would contain both modifications. While in a standard MoE model, each expert can only be activated once per forward pass, or skipped entirely, a recurrent MoE model could also refine its latent state over multiple iterations, potentially routing to the same expert multiple times, before switching to a different one ( Tan et al., 2023 ; Csordás et al., 2024 ) . While MoE models are the currently leading solution to implement this type of “memory” in dense transformers, these considerations also hold for other memory mechanisms suggested for LLMs ( Sukhbaatar et al., 2019 ; Fan et al., 2021 ; Wu et al., 2022 ; He et al., 2024 ) .

## 10 Conclusions

The models described in this paper are ultimately still a proof-of-concept. We describe how to train a latent recurrent-depth architecture, what parameters we chose, and then trained a single model at scale. Future training runs are likely to train with more optimized learning rate schedules, data mixes and accelerators. Still we observe a number of interesting behaviors emerging naturally from recurrent training. The most important of these is the ability to use latent reasoning to dramatically improve performance on reasoning tasks by expending test-time computation. In addition, we also observe context-dependent convergence speed, path independence, and various zero-shot abilities. This leads us to believe that latent reasoning is a promising research direction to complement existing approaches for test-time compute scaling. The model we realize is surprisingly powerful given its size and amount of training data, and we are excited about the potential impact of imbuing generative models with the ability to reason in continuous latent space without the need for specialized data at train time or verbalization at inference time.

## Acknowledgements

This project was made possible by the INCITE program: An award for computer time was provided by the U.S. Department of Energy’s (DOE) Innovative and Novel Computational Impact on Theory and Experiment (INCITE) Program. This research used resources of the Oak Ridge Leadership Computing Facility at the Oak Ridge National Laboratory, which is supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC05-00OR22725. Work on the LLNL side was prepared by LLNL under Contract DE-AC52-07NA27344 and supported by the LLNL-LDRD Program under Project No. 24-ERD-010 and 24-ERD-058 (LLNL-CONF-872390). This manuscript has been authored by Lawrence Livermore National Security, LLC under Contract No. DE-AC52-07NA27344 with the U.S. Department of Energy. The United States Government retains a non-exclusive, paid-up, irrevocable, world-wide license to publish or reproduce the published form of this manuscript, or allow others to do so, for United States Government purposes.

JG further acknowledges the support of the Hector II foundation. A large number of small-scale and preliminary experiments were made possible through the support of the MPI Intelligent Systems compute cluster and funding by the Tübingen AI center.

UMD researchers were further supported by the ONR MURI program, DARPA TIAMAT, the National Science Foundation (IIS-2212182), and the NSF TRAILS Institute (2229885). Commercial support was provided by Capital One Bank, the Amazon Research Award program, and Open Philanthropy. Finally, we thank Avi Schwarzschild for helpful comments on the initial draft.

## Appendix A Additional Information

## Potential Implications of This Work

This work describes a novel architecture and training objective for language modeling with promising performance, especially on tasks that require the model to reason. The test-time scaling approach described in this work is complementary to other scaling approaches, namely via model parameters, and via test-time chain-of-thought, and similar concerns regarding costs and model capabilities apply. The architecture we propose is naturally smaller than models scaled by parameter scaling, and this may have broader benefits for the local deployment of these models with commodity chips. Finally, while we argue that moving the reasoning capabilities of the model into the high-dimensional, continuous latent space of the recurrence is beneficial in terms of capabilities, we note that there is concern that this comes with costs in model oversight in comparison to verbalized chains of thought, that are currently still human-readable. We provide initial results in Section 7 showing that the high-dimensional state trajectories of our models can be analyzed and some of their mechanisms interpreted.

### A.1 Classical Reasoning Problems

We include a small study of the classical problem of multi-operand arithmetic in Figure 14 .

### A.2 Implementation Details

#### Device Speed Details

Nominally, each MI250X ( AMD, 2021 ) achieves 383 TFLOP in bfloat16, i.e. 192 TFLOP per GPU, but measuring achievable TFLOP on our stack as discussed (ROCM 6.2.0, PyTorch 2.6 pre-release 11/02) for arbitrary matrix multiplication shapes (i.e. we measure the peak achievable speed of the best possible shape iterating over shapes between 256 and 24576 in intervals of 256 and 110 ( Bekman, 2023 ) ), we measure a peak of 125 TFLOP/s on Frontier nodes. Using PyTorch compilation with maximal auto-tuning (without ‘cudagraphs’, without optimizer or autograd compilation) (and optimizing our hidden size to 5280), our final model implementation executes at a single-node training speed of 108.75 TFLOP/s, i.e. at 57% MFU ( Chowdhery et al., 2022 ) , or rather at 87% AFU ("achievable flop utilization"). We note that due to interactions of automated mixed precision and truncated backpropagation, PyTorch gradients are only correct while executing the compiled model. We further circumvent issues with the flash attention implementation shipped with PyTorch sdpa using the AMD fork of the original flash attention repository 6 6 6 https://github.com/Dao-AILab/flash-attention/ , which can be found at https://github.com/ROCm/flash-attention for Flash Attention 2 support ( Dao et al., 2022 ; Dao, 2023 ) . We experiment with fused head and loss implementations 7 7 7 https://github.com/JonasGeiping/linear_cross_entropy_loss , but ultimately find that the most portable choice on our AMD setup is to let torch compilation handle this issue.

#### Parallelization Strategy

As mentioned in the main body, because our depth-recurrent model is compute-heavy, it is optimal to run the model using only distributed data parallel training across nodes and zero-1 optimizer sharding within nodes ( Rajbhandari et al., 2020 ) , if we make use of gradient checkpointing at every step of the recurrent iteration. This allows us to eschew more communication-heavy parallelization strategies that would be required for models with the same FLOP footprint, but more parameters, which require substantial planning on this system ( Singh et al., 2024 ; Singh and Bhatele, 2022 ) . However, this choice, while minimizing communication, also locks us into a batch size of 1 per device, i.e. 4096 in total, and 16M tokens per step.

#### RCCL Interconnect Handling

Due to scheduling reasons, we settled on targeting 512 node allocation segments on Frontier, i.e. 4096 GPUs. However, this posed a substantial network interconnect issue. The connection speed between frontier nodes is only acceptable, if RCCL (AMD GPU communication collectives) commands are routed through open fabrics interface calls, which happens via a particular plugin 8 8 8 https://github.com/ROCm/aws-ofi-rccl . To achieve sufficient bus bandwidth above 100GB/s requires NCCL_NET_GDR_LEVEL=PHB , a setting that, on NVIDIA systems, allows packages to go through the CPU, and only uses direct interconnect if GPU and NIC are on the same (NUMA) node ( Wu and Stock, 2024 ) . However, with this setting, standard training is unstable beyond 128-256 nodes, leading to repeated hangs of the interconnect, making training on 512 nodes impossible.

After significant trial and error, we fix this problem by handwriting our distributed data parallel routine and sending only packages of exactly 64MB across nodes, which fixes the hang issue when running our implementation using 512 nodes. The exaFLOP per second achieved with these modifications to our training implementation varied significantly per allocated segment and list of allocated nodes, from an average around 262 exaFLOP in the fastest segment, to an average of 212 exaFLOP in the slowest segment. This is a range of 52-64 TFLOP/s per GPU, i.e. 41%-51% AFU, or 1-1.2M tokens per second.

#### Pretraining Metrics.

During the pretraining run, we run a careful tracking of optimizer and model health metrics, tracking effective Adam learning rates per layer, optimizer RMS ( Wortsman et al., 2023a ) , L 2 L^{2} and L 1 L^{1} parameter and gradient norms, recurrence statistics such as ‖ s k − s k − 1 ‖ ‖ s k ‖ \frac{||s_{k}-s_{k-1}||}{||s_{k}||} , ‖ s k ‖ ||s_{k}|| , ‖ s 0 − s k ‖ ||s_{0}-s_{k}|| . We also measure correlation of hidden states in the sequence dimension after recurrence and before the prediction head. We hold out a fixed validation set and measure perplexity when recurring the model for [ 1 , 4 , 8 , 16 , 32 , 64 ] [1,4,8,16,32,64] steps throughout training.

## Appendix B Latent Space Visualizations

On the next pages, we print a number of latent space visualizations in more details than was possible in Section 7 . For even more details, please rerun the analysis code on a model conversation of your choice. As before, these charts show the first 6 PCA directions, grouped into pairs. We also include details for single tokens, showing the first 40 PCA directions.

## Appendix C Pretraining Data