Recurrent-Depth Transformers: An Open Theoretical Implementation for Long-Horizon Reasoning

Traditional Transformer architectures typically improve representation capacity by stacking more layers. Increasing depth enables more complex feature transformations, but it also increases parameter count, training cost, and inference latency. More importantly, a fixed depth is not necessarily suitable for every task. Simple tasks may require only limited computation, while mathematical proofs, long-horizon planning, and multi-stage logical analysis may benefit from substantially longer reasoning processes.

Recurrent-depth Transformers offer an alternative approach: use a relatively small set of Transformer layers and execute the same computation repeatedly in a continuous hidden space. The model does not need to produce readable text at every step or convert intermediate reasoning into additional tokens. Instead, it directly updates its hidden state. This decouples computational depth from parameter scale and provides a theoretical foundation for building stronger multi-step reasoning systems.

This type of implementation should be understood as a reproduction and validation of a publicly studied research direction, not as an implementation of any particular commercial model. Its value lies in providing an architectural hypothesis that can be experimentally evaluated, analyzed, and iterated.

1. From Fixed Depth to Recurrent Depth

The computation performed by a standard Transformer can be abstracted as follows:

h₀ = Embedding(x)
h₁ = Layer₁(h₀)
h₂ = Layer₂(h₁)
……
hₙ = Layerₙ(hₙ₋₁)
y = Head(hₙ)

Each layer usually has its own parameters. The model depth is determined before training and inference, and the number of layers executed typically does not change significantly across inputs.

A recurrent-depth model instead uses a computation unit with shared parameters:

h₀ = Embedding(x)
for t in range(T):
    hₜ = RecurrentBlock(hₜ₋₁, t)
y = Head(hₜ)

Here, T represents the number of recurrent iterations. It can also be interpreted as the number of reasoning steps performed in the hidden space. With fewer iterations, the system can complete simple tasks quickly. For more complex tasks, the number of iterations can be increased, allowing the same parameters to continually refine and deepen the current state.

This design has three important characteristics.

First, parameter reuse. Repeated execution does not increase the parameter count linearly. As a result, the model can obtain greater effective computational depth at a similar parameter scale.

Second, expandable inference. Training can use a fixed or limited number of recurrent iterations, while inference can attempt to use more iterations to test whether the model has some ability to extrapolate to greater reasoning depth.

Third, continuous intermediate states. Each computation step operates on hidden vectors rather than generated text. This can reduce the overhead associated with formatting, tokenization, and repeatedly concatenating textual context.

2. Why Recurrent Computation Is Difficult to Train

Recurrent computation is not simply a matter of repeatedly calling the same layer. If an update function is applied to a hidden state over many iterations, small numerical deviations may be amplified or rapidly attenuated. Over time, this can result in exploding gradients, vanishing gradients, or representation collapse.

The recurrent update can be approximated as:

hₜ₊₁ = F(hₜ, x)

Within a local region, if the Jacobian of F has a large spectral radius, state perturbations may grow as the number of iterations increases. If the spectral radius is too small, state changes may quickly disappear, preventing the model from accumulating useful information.

Therefore, training a recurrent Transformer requires more than optimizing the loss function. It also requires controlling the stability of the underlying recurrent dynamical system.

A common approach is to impose contraction constraints on state updates, keeping the spectral radius of key transformations close to, but below, 1. Intuitively, each iteration should neither amplify the state without limit nor erase existing information entirely. Practical techniques may include:

1. Scaling residual branches to limit the magnitude of each update;
2. Applying spectral normalization or norm constraints to selected linear mappings;
3. Using stable normalization placement to reduce numerical drift across iterations;
4. Monitoring hidden-state norms, gradient norms, and attention distributions;
5. Starting training with fewer recurrent iterations and gradually increasing computational depth.

Spectral-radius constraints alone cannot solve every training problem. Attention, feed-forward networks, normalization, and residual connections jointly determine the overall dynamics of a Transformer block. Stability analysis should therefore cover the complete recurrent block rather than focusing on a single matrix.

3. Recurrent Indices and Positional Information

If every iteration uses exactly the same computational form, the model may have difficulty distinguishing between the first processing step and the tenth. Recurrent-depth models therefore generally need to inject information about the current computation step.

One approach is to add a recurrent-index embedding:

hₜ₊₁ = F(hₜ + eₜ, x)

Here, eₜ is the embedding associated with iteration t. It may be a learned vector, a periodic function, or another deterministic encoding.

The purpose of the recurrent index is not to represent a token’s position in the input sequence. Instead, it allows the model to recognize its current reasoning stage. For example, early iterations may focus on extracting local features, middle iterations may combine relationships, and later iterations may correct or refine the result.

To prevent the model from relying too heavily on a fixed number of iterations, training can randomize the recurrent length or use different maximum depths for different samples.

It is also important to distinguish among three types of position information:

• Input sequence position: the location of a token in the original sequence;
• Recurrent computation position: the current iteration number;
• Intra-layer structural position: the computation stage within attention or feed-forward modules.

Treating these concepts as interchangeable can create positional-encoding conflicts and harm generalization.

4. Two Directions for Attention Modules

A recurrent model can use standard multi-head attention or a more cache-efficient form of compressed attention.

Standard multi-head attention directly stores keys and values. In long-context or multi-iteration scenarios, the key-value cache can become a bottleneck for memory capacity and bandwidth. Compressed attention first maps keys and values into lower-dimensional latent representations, which are then used for computation or reconstructed as needed.

The basic idea is:

k_compressed = Compress(k)
v_compressed = Compress(v)
attention = Query(h) · k_compressed
output = Aggregate(attention, v_compressed)

This approach can reduce cache size, but it introduces compression error and additional projection computation. The design must therefore balance:

• Compression dimension and information fidelity;
• Training throughput and inference cache usage;
• Attention accuracy and memory-access efficiency;
• Whether errors are continuously amplified as the number of recurrent iterations increases.

For a research prototype, it is generally preferable to begin with a clear grouped-query attention baseline and then gradually replace it with compressed attention. This separates issues related to the recurrent mechanism from issues related to attention compression and reduces experimental complexity.

5. Sparse Mixture-of-Experts and Computational Efficiency

As the overall model scale increases, a mixture-of-experts architecture can be introduced. The basic idea is to prepare multiple feed-forward experts and use a router to select only a small number of them according to the current hidden state.

routing = Router(h)
experts = TopK(routing)
h_next = Combine(experts(h))

Sparse expert architectures can substantially increase the total parameter count while limiting the number of parameters activated for each token. To improve stability, fine-grained experts can be combined with both shared and sparse experts.

Shared experts process general patterns common to all inputs, such as basic syntax, common semantics, and fundamental transformations. Sparse experts can handle domain-specific features, complex reasoning, or specialized structural patterns. This reduces the risk of forcing the router to distribute basic capabilities across separate experts.

The routing mechanism should focus on three issues:

1. Load balancing: preventing a small number of experts from receiving too many tokens;
2. Routing jitter: preventing frequent expert switching across adjacent recurrent iterations;
3. Expert degradation: preventing some experts from receiving too little training.

These issues can be managed through load-balancing losses, capacity limits, routing-temperature adjustment, and expert-utilization monitoring. In recurrent settings, the router should consider not only the current input but also the hidden states from previous iterations. Otherwise, the model may make unnecessary expert switches at every step.

6. Adaptive Computation Time

A fixed number of recurrent iterations is simple to implement, but it cannot adapt to tasks with different levels of difficulty. Adaptive computation allows the model to stop early once the hidden state has become sufficiently stable.

One option is to measure the difference between adjacent hidden states:

delta = Norm(hₜ₊₁ - hₜ)
if delta < threshold:
    stop

Another option is to add a halting prediction head that estimates the benefit of continuing computation. When the state change becomes sufficiently small or the output distribution stabilizes, the model can terminate the recurrent process.

Early stopping must still enforce both a minimum and a maximum number of iterations. The minimum ensures that the model performs enough computation for basic tasks, while the maximum prevents abnormal inputs from entering an unbounded loop. Training should also apply moderate regularization to halting behavior so that the model does not stop prematurely simply to reduce computation.

Adaptive computation should not be evaluated only by average speed. The analysis should also cover:

• The actual iteration distribution across tasks of different difficulty;
• Whether early stopping harms long-chain reasoning accuracy;
• Whether halting decisions genuinely correlate with state convergence;
• Whether extreme inputs trigger the maximum iteration limit.

7. Training Process and Experimental Design

A relatively robust training process can be divided into four stages.

Stage 1: Short-cycle pretraining. Use a small number of recurrent iterations so that the model can first learn basic language and structural representations. At this stage, monitor loss reduction, gradient ranges, and hidden-state stability.

Stage 2: Recurrent-depth expansion. Gradually increase the maximum number of iterations and introduce recurrent-index embeddings so that the model can adapt to deeper computation in the hidden space.

Stage 3: Task-specific training. Add mathematical, planning, code-analysis, and multi-step logic data. Evaluate whether greater recurrent depth produces actual reasoning gains rather than merely increasing repetitive computation.

Stage 4: Adaptive-computation training. Introduce halting behavior and a computational-cost penalty, allowing the model to balance accuracy against the amount of computation used.

Experiments should not compare only parameter counts. They should also control for active parameter count, the number of training tokens, inference computation, and context length. At minimum, the following baselines and ablations are recommended:

• A fixed-depth Transformer with a similar parameter count;
• A recurrent model with similar active computation;
• The same model evaluated with different numbers of recurrent iterations;
• An ablation with and without stability constraints;
• An ablation with and without recurrent-index embeddings;
• A comparison with and without adaptive halting.

Only under these controls is it possible to determine whether an improvement comes from the recurrent structure, increased computation, or another training factor.

8. Key Engineering Considerations

In practical implementations, recurrent models are often more difficult to debug. Each iteration’s hidden state, gradient norm, expert load, attention entropy, and halting probability should be recorded in a monitoring system to make anomalies easier to locate.

Memory management is another major concern. Although recurrent execution reuses parameters, automatic differentiation may still retain intermediate states from every iteration. Depending on the training objective, gradient checkpointing, truncated backpropagation, or segmented backward passes can be used to balance memory consumption and training accuracy.

During inference, operator fusion, cache layout, and batching strategy also require attention. Because different samples may stop at different iterations, dynamic batching becomes more complex. An implementation can regroup samples according to their current iteration or define a small number of discrete computation tiers to reduce scheduling overhead.

Conclusion

The central value of recurrent-depth Transformers is not simply repeating a network layer. It is treating model depth as an adjustable computational resource. Through hidden-space recurrence, stability constraints, recurrent indices, sparse experts, compressed attention, and adaptive halting, this architecture can create a more flexible balance among parameter scale, reasoning depth, and computational cost.

Significant challenges remain. These include the reliability of depth extrapolation, the accumulation of recurrent errors, the possibility of providing systematic stability guarantees, and whether additional recurrent computation consistently translates into gains on real tasks.

For developers and researchers, a practical path is to first build a small-scale prototype that is observable, reproducible, and easy to ablate. Sparse routing, compressed caching, and dynamic computation can then be introduced incrementally. Only through carefully controlled experiments can we determine whether recurrence genuinely improves long-horizon reasoning.