PaperlayerView exact revision on arXiv

Source-linked AI summary

Temporal Recurrence Favors Fewer Layers

Ivan Anokhin, Johan Obando-Ceron, Irina Rish, Sebastian Risi

arXiv:2609.12531v1cs.LG

TL;DR

The paper asks how much within-step depth is needed when recurrence can carry computation across streaming steps. It compares recurrent and non-recurrent compute allocations under approximately matched per-step work, finding that recurrence favors substantially fewer layers across Sokoban and streamed FineWeb with comparable or better performance.

  • Problem

    The paper asks how much within-step depth remains useful once temporal recurrence provides computation across steps.

    Source: §1 · p. 2

  • Method

    The study varies within-step depth, parallel expert count, and expert width to compare recurrent and non-recurrent allocations under approximately matched per-step computation.

    Source: §1 · p. 2

  • Results

    Across Sokoban and streamed FineWeb, recurrent allocations use substantially fewer layers per step with comparable or better performance under approximately matched work.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support exploring shallower, more parallel updates as a potential route to reducing latency and improving training throughput.

    Source: §8 · p. 10

    The empirical study covers two domains, finite architectures, and restricted compute budgets, durations, and task difficulties, so reported allocations are best among those evaluated rather than globally compute-optimal.

    Source: §7 · p. 9

Abstract

from arXiv · show

In streaming tasks, recurrent models can carry latent computation across time, allowing each update to build on representations produced earlier. This raises a basic question: once temporal recurrence provides sequential computation across steps, how much depth is still needed within each step? Prior work has shown that recurrence can make shallow models competitive. We instead study this question as a compute-allocation problem, varying within-step depth, expert width, and the number of parallel experts per layer across several compute budgets. For each budget, we compare the best observed recurrent and non-recurrent allocations and the performance they achieve under approximately matched per-step computation. Across Sokoban and autoregressive FineWeb language modeling, we find that temporal recurrence shifts the best observed compute allocation toward substantially fewer layers, with comparable or better performance.

1 INTRODUCTION

The introduction asks how much within-step depth remains useful when recurrence can carry computation across external steps. It frames this as a matched-work allocation comparison and reports that recurrence favors fewer layers while parallel capacity remains useful.

1 INTRODUCTION

  • Carried state can support continued refinement across steps when successive inputs contribute to a shared problem, but offers no such benefit when each step is independent.The introduction gives locally predictable tokens and corridor navigation before a harder prediction or action as favorable settings.Source: §1 · p. 1
  • Prior work reports useful computation across external steps in state tracking and planning, including single-layer recurrent success on tested permutation-composition sequence lengths.The cited Sokoban analysis reports that internal plans develop across environment steps and causally influence actions.Source: §1 · pp. 1–2
  • The central question is whether recurrence lets models achieve comparable performance with fewer sequential stages S by reallocating fixed per-step work W ≈ S · P toward parallel work P.The comparison varies within-step depth and parallel computation in recurrent and non-recurrent models.Source: §1 · p. 2
  • The study varies within-step depth, parallel expert count, and expert width under matched work, distinguishing temporal recurrence from computation repeated on the same input.Figure 1 illustrates sequential, multi-level parallel, and single-level parallel allocations under the same approximate work budget.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Across Sokoban and streamed FineWeb, recurrent allocations use substantially fewer layers than non-recurrent ones with comparable or better performance under approximately matched per-step work.Recurrence reduces the performance cost of fewer layers, while expert count and width provide ways to spend reallocated compute.Source: §1 · p. 2

2 WHY FEWER LAYERS MATTER

The section motivates fewer sequential stages as a potential way to shorten update latency while preserving the expressivity of general nonlinear recurrence. It notes that realized speedups depend on hardware and implementation.

2 WHY FEWER LAYERS MATTER

  • Reallocating work from sequential stages S to parallel work P can shorten each update’s serial path and potentially reduce latency if hardware parallelism is available.Actual speedups also depend on hardware utilization, memory bandwidth, and communication overhead.Source: §2 · p. 3
  • Because general nonlinear RNN updates are unrolled sequentially across tokens, deeper updates lengthen both the forward unroll and the backward pass, potentially reducing training throughput.Preserving comparable performance while shifting work from S to P could make large-scale pretraining more practical.Source: §2 · p. 3
  • Parallelizable linear recurrent models and fixed-depth Transformers offer temporal parallelism, but the cited theory places them in computational classes believed weaker than general nonlinear RNNs under stated assumptions.The section motivates reallocating work while retaining general nonlinear recurrence.Source: §2 · p. 3
  • Some Transformer studies report that middle layers share representational structure or tolerate reordered and parallel execution, with randomized order carrying an accuracy cost in one study.These findings concern potential exchangeability of computations normally implemented as strict depth.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries introduce reinforcement learning and autoregressive language modeling as streaming settings where inputs arrive over time. They distinguish state-tracking recurrence from the paper’s focus on recurrent updates that can refine latent representations.

Deep Reinforcement Learning

  • In reinforcement learning, an agent receives observations, updates a latent state, chooses actions, and aims to maximize expected discounted return.The experiments use actor–critic agents.Source: §3 · p. 3
  • Recurrent models can maintain belief-like states for history and long-horizon decisions; this work emphasizes recurrence as reasoning-like latent refinement rather than memory-like tracking.The distinction is framed as complementary to the common memory-like use of recurrence in partially observable environments.Source: §3 · p. 3

Autoregressive Sequence Modeling

  • Autoregressive sequence modeling predicts each token from preceding context, requiring representations that accumulate and refine information over long token streams.The objective is to model the sequence probability through next-token predictions.Source: §3 · pp. 3–4
  • Unlike architectures designed for parallel scans across sequence positions, this recurrent Transformer emphasizes expressive latent-state updates with feedback through the carried state.Its design aims to compose computation across successive token steps rather than concentrate it in one deep feed-forward pass.Source: §3 · p. 4

4 EXPERIMENTAL FRAMEWORK

The experimental framework compares recurrent and non-recurrent models by varying sequential depth, parallel expert count, and expert width under an approximate work budget. In recurrent variants, the final latent state is carried into the next external step.

4 EXPERIMENTAL FRAMEWORK

  • The Parallel Experts framework varies within-step depth L, parallel expert count E per level, and expert width d while allowing recurrent models to carry latent state across steps.Experts read the current input and shared state, compute candidate updates in parallel, and merge them into the shared state.Source: §4 · p. 4
  • Approximate work is W ≈ L · E · C(d), used to construct matched-budget allocation sweeps for recurrent and non-recurrent models.C(d) denotes the cost of one expert transformation.Source: §4 · p. 4
  • At each external step, experts receive the encoded input and incoming shared state, while expert-local activations track depth levels and parallel experts.The supplied passages identify the encoded input as z_t and the indices for depth level and expert.Source: §4 · p. 4
  • The recurrent variant feeds the final state into the first depth level of the next step, and the model output represents action/value predictions or next-token logits.The aggregation of expert proposals is normalized to keep variance approximately stable as the number of experts changes.Source: §4 · p. 4
  • The framework is a probe of allocation tradeoffs, not a claim of architectural optimality; alternative expert mechanisms or recurrent communication schemes may shift the quantitative frontier.The authors expect the qualitative conclusion to remain the same.Source: §4 · p. 4

5 EXPERIMENTS

Across Sokoban and FineWeb, recurrent models reach their best observed performance with substantially fewer within-step layers than non-recurrent models. The experiments vary depth and parallel expert capacity under approximately matched compute, and find that expert scaling remains useful while recurrent depth gains saturate earlier.

5 EXPERIMENTS; 5.1 SOKOBAN; 5.2 LANGUAGE MODELING

  • The experiments vary within-step depth, expert count, and expert width under matched compute, evaluating episodic return in Sokoban and validation cross-entropy in streamed FineWeb.Sokoban uses a DRC-inspired ConvLSTM with one pass through the depth levels per environment step; FineWeb uses Transformer blocks for next-token prediction.Source: §5 · p. 5, §5.1 · p. 5, §5.2 · pp. 6–7

5.1 SOKOBAN; 5.2 LANGUAGE MODELING

  • At the small budgets, the best Sokoban recurrent allocation uses four layers and beats the eight-layer non-recurrent allocation, while FineWeb reaches approximately 3.51 loss at L = 2, comparable to the L = 16 non-recurrent allocation.Both comparisons use approximately matched per-step or per-token work.Source: Figure 2 · p. 5, §5.1 · p. 5, Figure 5 · p. 7
  • Across tested budgets, recurrent Sokoban models saturate at depths 2–4 while non-recurrent models benefit from greater depth; FineWeb recurrent models reach best observed loss at L = 2 or 4 versus L = 16 or 32.Recurrent Sokoban performance is better across the tested settings, while recurrent FineWeb achieves comparable validation loss with substantially less depth.Source: Figure 3 · p. 6, §5.1 · p. 6, §5.2 · p. 7

5.1 SOKOBAN

  • The earlier depth saturation also appears in ConvGLU models, where the recurrent variant differs by feeding the previous step’s final residual stream into the first depth level.The non-recurrent ConvGLU baseline omits LSTM cell-state machinery.Source: §5.1 · p. 6
  • In large-budget Sokoban probes, the four-layer model initially has higher plan recall, but the two-layer model catches up over the next few environment steps and both maintain similar recall and performance.Probe thresholds were calibrated to approximately 93% precision.Source: Figure 4 · p. 6, §5.1 · p. 6

5.2 LANGUAGE MODELING

  • FineWeb training carries recurrent state and sliding-window KV caches across continuous token streams, resetting state only at shard boundaries; non-recurrent models train on independent sequences.The variants use the same total token budget and tokens per optimizer update, with context lengths approximately matched during training.Source: §5.2 · p. 7

5.3 WHAT MAKES THE EARLIER SATURATION WORK?

  • Across both model types, increasing per-layer expert count helps, while additional depth gains saturate earlier with recurrence; shared-state access also outperforms narrower communication radii.Sokoban performance continues improving beyond the eight-expert regime explored in prior deep RL Soft MoE work (Obando Ceron et al., 2024; Sokar et al., 2025).Source: §5.3 · p. 7, Figure 6 · p. 8

6 RELATED WORK

Prior work studies recurrence across external steps, recurrence across depth, and expert-based parallel computation. This paper distinguishes its focus by examining temporal recurrence and depth–width allocation under approximately fixed computation.

6 RELATED WORK

  • Feedback Transformer uses previous-token top-level representations to make shallow models with fewer parameters more competitive with deeper standard Transformers.DRC carries state across Sokoban steps and outperforms deep feed-forward baselines.Source: §6 · p. 8
  • RHN studies depth–width tradeoffs under fixed parameter count, with depth saturating near 8–9 levels even as compute grows substantially with depth.Recent recurrent Transformer work also studies favorable depth–width tradeoffs at fixed parameter count; this paper instead keeps compute approximately fixed.Source: §6 · p. 8
  • Universal Transformers apply recurrent updates across depth, while Adaptive Computation Time and PonderNet explore dynamically allocating sequential computation to difficult inputs.These approaches refine representations through repeated transformations of the same input.Source: §6 · p. 8
  • Prior recurrent MoE work includes RIMs and MoEUT, while most prior MoE systems place experts inside fixed-depth feed-forward architectures.Here, MoE serves to study sequential-to-parallel allocation, with the main focus on temporal rather than depth recurrence.Source: §6 · p. 8
  • The paper’s compute-allocation perspective also relates to broader work on depth–width tradeoffs and compute-optimal scaling.The cited works are Lu et al. (2017), Cheng et al. (2017), and Hoffmann et al. (2022).Source: §6 · p. 9

7 DISCUSSION AND LIMITATIONS

The results suggest temporal recurrence shifts computation away from within-step depth, with shallower Sokoban agents catching up in plan completeness over later steps. The findings are bounded by task structure, the evaluated architectures and budgets, possible recurrent-optimization difficulties, and the distinction between matched computation and hardware efficiency.

7 DISCUSSION AND LIMITATIONS

  • Temporal recurrence shifts the best observed allocation from within-step depth toward parallel capacity, while shallower Sokoban agents catch up in plan completeness over later steps.The shallower agent achieves comparable task performance despite initially forming less complete plans.Source: §7 · p. 9
  • Recurrence may reduce the need for within-step depth when consecutive inputs contribute to a shared problem and the state preserves useful intermediate computation.The benefit may not arise for independent inputs, decisions that must be solved immediately, or a state too limited to retain useful information.Source: §7 · p. 9
  • The reported allocations are best among those evaluated, not globally compute-optimal, because the study covers two domains and restricted architectures, budgets, training durations, and task difficulties.Architectural parity also involves compromises, including state machinery in ConvLSTM and stabilization and memory-management choices for recurrent Transformers.Source: §7 · p. 9
  • Truncated backpropagation through time may hinder credit assignment through deep recurrent updates, and approximately matched computation does not establish hardware latency efficiency.Realized latency also depends on utilization, memory bandwidth, kernel efficiency, and communication overhead; comprehensive systems evaluation remains future work.Source: §7 · p. 9

8 CONCLUSION

Across Sokoban and streamed FineWeb, recurrence shifts the best observed compute allocation toward fewer layers per step while achieving comparable or better performance at approximately matched work. The conclusion also connects this allocation shift to planning across environment steps and the potential for shallower, more parallel updates.

8 CONCLUSION

  • Across Sokoban and streamed FineWeb, recurrent allocations use substantially fewer layers per step with comparable or better performance under approximately matched per-step work.In the small-budget FineWeb comparison, L = 2 reaches validation loss comparable to the best non-recurrent allocation at L = 16.Source: §8 · p. 9
  • Sokoban decoding is consistent with the L = 2 model catching up in planning across environment steps while maintaining comparable task performance.The conclusion presents shallower, more parallel updates as a potential route to reducing latency and improving training throughput.Source: §8 · pp. 9–10

A ARCHITECTURES / A.1 LSTM

  • Each external step encodes a fresh input, updates shared state through within-step experts, and produces an output; recurrence routes the top state to the next step.The LSTM experts aggregate their outputs into shared state at each depth level, which the policy and value heads read after the final level.Source: Figure 7 · p. 16, §A.1 · p. 16

A.2 TRANSFORMER / A.3 RECURRENT TRANSFORMER

  • The FineWeb recurrent Transformer adds token embeddings at each depth and shares a KV cache across experts, trading approximately 0.02 higher validation loss for reduced memory pressure and increased throughput.Experts aggregate key–value proposals into a rolling cache and carry the cache and shared state forward across the streamed document.Source: §A.3 · pp. 17–18

B FINEWEB STREAMING IMPLEMENTATION

  • FineWeb shard streaming partitions tokens into contiguous batch-slot streams, carries recurrent state across adjacent windows, and resets it at shard boundaries.Each window uses T input tokens and the next T shifted tokens as labels.Source: B FINEWEB STREAMING IMPLEMENTATION · p. 18

C SOKOBAN SOLVE RATE RESULTS

  • Sokoban policies are evaluated on 1,000 unfiltered test levels after 50M environment steps, selecting the highest-training-return configuration at each compute budget and depth.The table reports solve rates and standard deviations across three seeds; state is either carried or reset at each environment step.Source: C SOKOBAN SOLVE RATE RESULTS · p. 19, Table 1 · p. 19

D AGGREGATION ABLATIONS

  • Aggregation ablations compare mean pooling, square-root-normalized sum, softmax-gated weighted averaging, and L2-normalized softmax-gated aggregation; square-root normalization is the default.The paper describes square-root normalization as keeping aggregate scale stable as the number of contributing experts changes.Source: D AGGREGATION ABLATIONS · p. 20

E COMMUNICATION RADIUS

  • In Sokoban, full communication among experts achieves the highest mean return, so the paper uses full communication in its other experiments.The comparison varies the number of preceding neighbors whose states each expert can read.Source: E COMMUNICATION RADIUS · p. 21, Figure 11 · p. 21

F SOKOBAN PLANNING ANALYSIS

The Sokoban planning analysis probes decoded agent-move plans from final residual representations across environment steps. Shallower models catch up in decoded-plan recall over subsequent steps, consistent with computation being redistributed across time.

SOKOBAN PLANNING ANALYSIS

  • Convolutional probes read final residual representations and predict the agent’s movement direction at each location, or NEVER if the agent does not visit it.Probes are trained on 3,000 unfiltered training trajectories and evaluated on the 906 test levels solved by all four policies.Source: F SOKOBAN PLANNING ANALYSIS · p. 22
  • The decoding comparison uses recall and precision, with a fixed threshold per model calibrated to approximately 93% precision across predictions over the first 20 steps.Recall is computed over non-NEVER targets; precision is computed over emitted directions.Source: F SOKOBAN PLANNING ANALYSIS · p. 22
  • L = 4 initially has higher decoded-plan recall than L = 2, but L = 2 catches up over the next few environment steps and both maintain similar recall.The models have nearly identical solve rates; L = 1 similarly narrows, but does not fully close, its initial recall gap with L = 8.Source: F SOKOBAN PLANNING ANALYSIS · p. 22, Figure 12 · p. 22

G IMPLEMENTATION DETAILS AND HYPERPARAMETERS USED IN EXPERIMENTS

The implementation section documents training settings, compute resources, and preliminary latency measurements. It cautions that neither inference implementation was extensively optimized, so latency results should be interpreted carefully.

G IMPLEMENTATION DETAILS AND HYPERPARAMETERS USED IN EXPERIMENTS

G.1 COMPUTE RESOURCES

  • Sokoban runs used RTX 8000 and L40S GPUs, while FineWeb runs used four H100 GPUs; the largest configurations in both settings could take roughly one day.Most smaller configurations took a few hours for Sokoban and several hours for standard FineWeb runs.Source: §G.1 · p. 24

H LATENCY AND TRAINING THROUGHPUT

  • Latency measurements are preliminary, and neither implementation was extensively optimized for inference, so the results warrant caution.The measurements used four 80GB H100 GPUs and Recurrent Transformer experts were evaluated with jax.vmap when parallel execution was possible.Source: H LATENCY AND TRAINING THROUGHPUT · p. 24

H.1 INFERENCE LATENCY

  • Table 5 reports median latency in milliseconds per token after pre-filling either model’s KV cache to length C.The supplied table caption identifies the measure and cache-prefill condition but provides no latency values.Source: §H.1 · p. 24, Table 5 · p. 25

H.2 TRAINING THROUGHPUT

The training-throughput subsection specifies what the FineWeb update measurement includes, while the following fixed-compute sweeps compare recurrent and non-recurrent allocations at small and medium budgets.

H.2 TRAINING THROUGHPUT

  • Table 6 measures FineWeb training throughput in tokens per second for a complete update, including the dataloader, forward and backward passes, and optimizer update.Each panel fixes tokens per optimizer update; C is training sequence length and B is global batch size.Source: §H.2 · p. 25, Table 6 · p. 25

I FINEWEB FIXED-COMPUTE SWEEPS

J SOKOBAN FIXED-COMPUTE SWEEPS

The section presents Sokoban fixed-compute allocation sweeps for ConvLSTM and ConvGLU models under recurrent and feedforward settings, alongside ConvGLU scaling analyses.

J SOKOBAN FIXED-COMPUTE SWEEPS

  • ConvLSTM allocations are shown under recurrent and feedforward settings at small, medium, and large budgets; the large-budget snapshot includes only completed runs.Figure 16 and Figure 17 cover the medium and large budgets, respectively.Source: J SOKOBAN FIXED-COMPUTE SWEEPS · p. 26, Figure 16 · p. 26, Figure 17 · p. 27
  • Table 7 summarizes ConvLSTM matched-compute runs with mean and sample variance of final-five episodic returns, plus seed counts.The sample variance uses denominator N −1.Source: Table 7 · p. 27

K CONVGLU SOKOBAN

  • ConvGLU Sokoban sweeps show mean return across expert counts at fixed depths, depth saturation at small and medium budgets, and recurrent-versus-feedforward allocations.The return plots report one standard error across three seeds; Figures 19 and 20 show small- and medium-budget allocations, respectively.Source: K CONVGLU SOKOBAN · p. 27, Figure 18 · p. 28, Figure 19 · p. 28
Loading 2609.12531v1…