Temporal Recurrence Favors Fewer Layers

Introduction

Across many deep learning tasks, models benefit from progressively refining their internal representations through a sequence of latent transformations [25, 5, 59, 12]. This refinement is commonly implemented by stacking layers, allowing later layers to build on intermediate states produced by earlier ones. Depth therefore controls one important source of sequential computation available to a model.

In natural streaming tasks, sequential computation can also accumulate across time. The model is repeatedly applied as the task unfolds, for example, across tokens in language modeling, frames in video, or observation steps in reinforcement learning. A recurrent model can carry a latent state across these steps, allowing each update to build on computations performed at earlier steps. Recurrence can therefore do more than preserve information from the past: it can also continue refining representations or plans for future predictions and actions. This temporal recurrence is different from depth recurrence, as used in looped Transformers, where shared blocks are repeatedly applied on the same input before producing an output [12, 11, 32, 17]. Here, computation continues across successive token or environment steps as new inputs arrive.

Whether sequential computation can accumulate across external steps depends on the structure of the task. It can do so when several consecutive steps contribute to the same broader problem, allowing computation performed at one step to remain useful later. One favorable setting is when relatively easy predictions or actions come before a harder one: an agent may refine a plan while walking through a corridor, and a language model may continue refining its representation of a passage while generating locally predictable tokens. In contrast, if each step presents a new independent problem that must be solved immediately, earlier computation has no value for later steps.

Prior work demonstrates that models can learn useful computation across external steps. State-tracking studies train models to predict the current state after each input operation [41]. On difficult permutation-composition tasks, single-layer recurrent models succeed across the tested sequence lengths, whereas Transformers require increasing depth to maintain accuracy. Evidence also comes from reinforcement learning, where recurrent agents learn planning-like behavior

Our code is available at fewer-layers-with-rec. in fully observed Sokoban environments ([22]). Subsequent analysis shows that their internal plans develop across environment steps and causally influence their actions ([6]).

This raises the central question of the paper: How much within-step depth is still needed once temporal recurrence is available? Prior work has shown that recurrence can make shallow models competitive ([41], [22], [14], [46]). We instead formulate this question as a compute-allocation problem, comparing similar recurrent and non-recurrent models. Let $S$ denote the number of sequential stages within that step, and let $P$ denote the parallel work performed at each stage. The total work per step is then approximately $W \approx S \cdot P$. At fixed $W$, allocations with fewer sequential stages spend more work on parallel computation at each stage. We ask whether temporal recurrence changes the best allocation: do recurrent models require a smaller $S$ than non-recurrent models to achieve comparable performance?

To realize this comparison, we use recurrent and non-recurrent models with the same three allocation axes, illustrated in Figure 1. The within-step depth $L$ controls the number of successive stages (layers) $S$. Within each stage, multiple experts are evaluated in parallel and their outputs are combined, allowing $P$ to be increased in two ways: by adding more experts $E$, or by increasing the hidden dimension $d$ of each expert. The recurrent variant carries a latent state across external steps, while the non-recurrent variant does not. Varying $L$, $E$, and $d$ under matched work lets us compare how much within-step depth remains useful with and without temporal recurrence.

Our experiments support two main conclusions:

• Temporal recurrence shifts the best compute allocation toward fewer layers. Across Sokoban and streamed FineWeb, the best observed recurrent allocations use substantially fewer layers per step than the best non-recurrent allocations, while achieving comparable or better performance under approximately matched per-step work.

• Recurrence and parallel capacity play different roles. Increasing per-layer expert count benefits both recurrent and non-recurrent models, while gains from additional layers saturate earlier with recurrence. Recurrence therefore reduces the performance cost of using fewer layers, while expert count and width provide useful ways to spend the reallocated compute.

More broadly, we hypothesize that temporal recurrence will shift the best compute allocation toward shallower within-step updates in other streaming settings where successive steps contribute to a shared problem and carried state can support continued refinement.

Why Fewer Layers Matter

Why Fewer Layers Matter

The compute-allocation tradeoff between sequential stages $S$ and parallel work $P$ matters also for latency. Under a fixed work budget, shifting computation from $S$ into $P$ reduces the serial part of the update, potentially decreasing latency if enough hardware parallelism is available. The actual speedups also depend on implementation details, memory bandwidth, and communication overhead.

This potential reduction in per-update latency is particularly important for general nonlinear RNNs, whose state updates must generally be unrolled sequentially across tokens during language-model pretraining. Because the latency of each update accumulates at every token in the sequence, deeper updates lengthen both the forward unroll and the backward pass through the sequence. If comparable performance can be preserved while reallocating work from sequential stages $S$ to parallel work $P$, reducing the latency of each update could improve training throughput and make large-scale pretraining of general nonlinear RNNs more practical.

Other sequence architectures address this scalability problem by changing how computation is organized across tokens. Transformers process sequence positions in parallel during training, while structured recurrent models, such as Mamba, constrain their state transitions so that they can be evaluated with parallel scans across the sequence ([57, 21, 20]). These gains in temporal parallelism come with an expressivity trade-off. Under standard precision assumptions and complexity-theoretic conjectures, theory places fixed-depth Transformers and parallelizable linear recurrent models in computational classes believed to be weaker than attainable by general nonlinear RNNs ([41, 42]). This trade-off motivates preserving general nonlinear recurrence and studying whether work can be reallocated from $S$ to $P$ without sacrificing performance, as a potential route to improving training throughput.

There are also pieces of evidence supporting the possibility of parallelization in general. Recent transformer studies suggest that some computations normally implemented as strict depth may be more exchangeable than the architecture implies. [53] find that middle layers of pre-trained transformers share substantial representational structure, and that some tasks tolerate skipping, reordering, or parallelizing layers with limited degradation. [16] train vision transformers to tolerate randomized attention-layer execution order, showing that robustness to arbitrary layer order can be learned, though with an accuracy cost.

Preliminaries

We conduct our study in two settings that support updating a latent state over long external sequences: deep reinforcement learning and autoregressive language modeling. In both cases, the model repeatedly receives new input, updates an internal state when recurrence is used, and produces an action or prediction.

Deep Reinforcement Learning. In reinforcement learning, an agent interacts with an environment over timesteps $t = 1, 2, \ldots$ ([54, 43, 50]). At each step it receives an observation $o_t$, updates an internal latent state $s_t$, selects an action $\pi(a_t \mid s_t)$, and receives a reward $r_t$. The goal is to maximize the expected discounted return, $J(\pi) = \mathbb{E}_{\pi} \left[\sum_{t=1}^{\infty} \gamma^{t-1}r_t\right]$, where $\gamma\in[0,1)$ is a discount factor that determines the relative importance of immediate and future rewards. The policy $\pi(a_t \mid s_t)$ maps the current latent state to a distribution over actions, and the expectation is taken over trajectories induced by the policy and environment dynamics. Our experiments use actor-critic agents ([37, 43]).

Recurrent architectures are commonly used to maintain belief-like latent states that summarize trajectory history in partially observable environments and support long-horizon decision making ([24, 26, 58, 27]). In this work, we emphasize a complementary use of recurrent computation in reinforcement learning: recurrence can support reasoning-like latent refinement, not merely memory-like tracking.

Autoregressive Sequence Modeling. Autoregressive sequence modeling aims to predict each token from its preceding context. Given a token sequence $x_{1:T}$, the objective is to model $p(x_{1:T}) = \prod_{t=1}^{T} p(x_t \mid x_{<t})$, or equivalently to predict the next-token distribution at each position ({cite:p}`bengio2003_a_neural et al., 2003; [57]. This requires representations that accumulate and refine contextual information over long token streams.

Classical recurrent sequence models maintain a latent state across token updates ([28]) and update that state sequentially during training. Modern long-sequence architectures often designed to improve training efficiency by enabling parallel scanning computation over sequence positions ([57]; [10]; [21]; [20]). Our version of the recurrent Transformer is not designed to support parallel scanning; instead, it emphasizes expressive latent-state updates with feedback connections through the carried state. This is intended to allow computation to compose across successive token steps rather than being concentrated within a single deep feed-forward pass.

Experimental Framework

We introduce the Parallel Experts framework for studying how compute should be allocated between sequential and parallel computation. It is designed to expose a clean tradeoff between within-step depth and parallel capacity while allowing latent state to be carried across external steps. At each external step, several expert modules read the current input and a shared latent state, compute candidate updates in parallel, and merge their outputs into the shared state. In recurrent models, the final state is carried to the next external step. This construction exposes three allocation axes: within-step depth $L$, expert count $E$ per depth level, and expert width $d$.

The depth $L$ counts the number of sequential depth level within one step. The expert count $E$ counts experts evaluated in parallel at each depth level, and $d$ controls per-expert capacity by adjusting expert’s hidden dimension (width). If $C(d)$ is the cost of one expert transformation, the approximated work is

We use this proxy to construct matched-budget allocation sweeps that are intended to compare where computation should be placed to maximize the performance for recurrent and non-recurrent architectures.

At each environment or sequence step, the model receives an external input $x_t$. Let $z_t = e(x_t)$ denote the encoded representation of this external input. At step $t$, the model also receives an incoming shared latent state $s_t^0$, expert-local activations $h_t^{0,i}$, and architecture-specific memories $m_{t-1}^{\ell,i}$, where $\ell\in\{1,\ldots,L\}$ indexes depth levels and $i \in\{1,\ldots,E\}$ indexes per depth level parallel experts. Each depth level applies all experts in parallel:

This normalized sum aggregates expert proposals while keeping the variance approximately the same as the number of contributing experts changes. To maximize the sequential computation available on the carried state, the recurrent variant uses top-down feedback, feeding the final state into the first depth level of the next step ([61]; [22]; [14]):

Here $y_t$ denotes action logits and a value estimate in reinforcement learning, or next-token logits in language modeling.

More specifically, in the LSTM instantiation $m_t^{\ell,i}$ corresponds to the cell state $c_t^{\ell,i}$, and in the transformer — the key–value (KV) cache. Full pseudocode and architecture-specific instantiations are given in Appendix A and alternative aggregation rules are studied in Appendix D.

The Parallel Experts framework is intended as a probe framework rather than a claim of architectural optimality. We use parallel experts because they expose a clean tradeoff between within-step depth and parallel capacity while optionally allowing a carried latent state. Other MoE mechanisms, including routed MLP experts ([38]), Soft MoE variants ([49]), or alternative recurrent communication schemes, may shift the quantitative frontier, but we expect the qualitative conclusion to stay the same.

Experiments

Our experiments test different compute allocations. The Parallel Experts framework exposes three ways to spend compute: within-step depth $L$, per-layer expert count $E$, or width $d$ of each expert. The central question is whether, once recurrent state is carried across the external steps, the best allocations are still require deep sequential computation within each step.

For our main experiments we evaluate different allocations in two settings: Sokoban, using the Sokoban-v0 environment from Jumanji ([4]), and streamed FineWeb language modeling ([48]). We report episodic return for reinforcement learning and validation cross-entropy loss for language modeling. Unless stated otherwise, in RL each configuration is run with at least three random seeds. Additional implementation details are reported in Appendix G.

Sokoban

Sokoban is a fully observable puzzle environment in which an agent must push four boxes onto four target locations on a $10 \times10$ grid. Levels are procedurally generated, and mistakes can be irreversible: for example, pushing a box into a corner without a target makes the level unsolvable. Planning is therefore important even though the agent observes the entire board at every step. Prior work showed that DRC agents learn planning-like behavior in Sokoban ([22]), and later analysis found that their internal plans develop progressively through recurrence and causally affect their actions ([6]). These properties make Sokoban a natural setting for studying whether temporal recurrence reduces the within-step depth needed for planning.

We use a DRC-inspired architecture with ConvLSTM experts ([22]), organized according to our Parallel Experts framework. The original DRC setup combines temporal recurrence with depth-recurrence, repeated passes through the recurrent stack on the same observation. To focus on temporal recurrence, we instead make a single pass through the $L$ depth levels per environment step, without additional internal looping. We compare models trained with state carried across environment steps against models trained without carried state, keeping the remaining architecture and training setup the same. For both variants, we vary within-step depth $L$, expert count $E$, and expert width $d$ under approximately matched per-step compute budgets. This comparison tests whether temporal recurrence changes how compute is best allocated between within-step depth and parallel capacity.

Figure 2 shows the allocation sweeps for both variants under the small compute budget. With recurrence, the performance saturates at $L = 4$, while the non-recurrent version saturates at depth 8 with much worse performance. Temporal recurrence changes not only the achieved performance but also the preferred compute allocation, favoring fewer sequential stages and more parallel capacity.

We repeat the allocation sweeps at medium and large compute budgets. Figure 3 summarizes all three budgets, reporting the best observed mean return over expert counts, widths, and depths up to $L$. Across the tested budgets, recurrent performance saturates by depth 2-4, whereas non-recurrent models benefit from greater within-step depth, for example, still improving at a within-step depth of 16 under the large budget. At the same time, recurrent model performance is better across all settings, and larger budgets continue to improve both architectures at every tested within-step depth, showing that additional parallel compute remains a useful scaling axis across different configurations. The full medium- and large-budget sweeps are reported in Appendix J.

For a more direct architectural comparison, we also test ConvGLU models, whose experts apply a $3 \times3$ convolution followed by GLU activation. These models omit the LSTM cell-state machinery retained in the non-recurrent ConvLSTM baseline. Here, the only architectural difference is one extra projection in the recurrent variant, feeding the previous step’s final residual stream into the first depth level. We observe the same earlier depth saturation with recurrence under this single architectural intervention (see Appendix K).

To examine how planning develops at different within-step depths, we apply the spatial probes of [6] to the final residual representations of two large budget recurrent models with $L = 2$ and $L = 4$. The probes predict the direction of the agent’s movements in each spatial location. Figure 4 compares recall across environment steps, with thresholds calibrated to approximately 93% precision. The $L = 4$ model initially has higher recall, but the $L = 2$ model catches up over the next few environment steps, after which both maintain similar recall. Both models also have almost identical performance. This pattern is consistent with the shallower model compensating for less planning within each update by continuing the computation across environment steps, catching up early enough to preserve task performance. Full probing details and additional depth comparisons are provided in Appendix F.

Language Modeling

We next test whether the compute-allocation trends observed in Sokoban extend to language modeling. We train recurrent and non-recurrent models for next-token prediction on approximately one billion tokens from FineWeb and compare validation loss under approximately matched per-token compute budgets. We use Transformer blocks within our Parallel Experts framework, with the non-recurrent model reducing to a standard Transformer when $E = 1$. The recurrent variant feeds the final residual stream into the first layer at the next token step, with additional adjustments for stability and memory efficiency described in Appendix A.

During training, recurrent models process continuous token streams, carrying their state and sliding-window KV caches across training sequences and resetting them only at shard boundaries. Non-recurrent models are trained on independent training sequences, allowing parallel computation across token positions. We use fixed KV cache lengths of 128 and 1024 for recurrent models against non-recurrent training sequence lengths of 256 and 2048 tokens, respectively, approximately matching the average attention-context length during training. Both variants use the same total training-token budget and number of tokens per optimizer update. At evaluation, both variants process the validation stream continuously using sliding-window KV caches of the same length (128 or 1024). Further streaming implementation details are provided in Appendix B.

Results. Figure 5 shows the allocation sweeps under the small compute budget with a 128-token attention window. The recurrent model reaches a validation loss of approximately 3.51 with only two depth levels ($L = 2$), comparable to the best non-recurrent allocation at $L = 16$. As in Sokoban, the preferred allocation shifts toward fewer sequential stages and greater parallel capacity when recurrence is available.

We extend the comparison to a medium compute budget and to the models trained and evaluated with an average 1024-token attention window. Figure 3 summarizes these settings, reporting the lowest validation loss among allocations with depth at most $L$. Across the tested settings, recurrent models reach their best observed loss at $L = 2$ or 4, whereas non-recurrent models continue to improve up to $L = 16$ or 32. Although the preferred depth depends on the setting, recurrent models consistently achieve comparable validation loss with substantially less within-step depth.

What Makes the Earlier Saturation Work?

The early saturation appears to reflect two approximately separable effects. In Figure 6, increasing per-layer expert count improves both recurrent and non-recurrent models, suggesting that parallel capacity is broadly useful regardless of recurrence. In contrast, the depth-response curve depends strongly on recurrence: non-recurrent performance is much more sensitive to reducing $L$, whereas recurrent models saturate at substantially shallower within-step depth. Thus, recurrence primarily lowers the cost of reallocating compute away from depth, while experts/width provide a useful parallel place for that compute. Finally, parallel experts must communicate through shared state. In the communication-radius ablation (Appendix E), Full shared-state access performs best while the narrower communication radii lower performance monotonically. Interestingly, Sokoban performance continues to improve beyond the eight-expert regime explored in prior deep RL Soft MoE work ([44]; [52]).

Related Work

Temporal Recurrence and Model Depth. Several prior works provide evidence that temporal recurrence across external steps, such as token steps in streamed language modeling or environment steps in reinforcement learning, can make shallower recurrent neural networks more competitive with standard feed-forward architectures without temporal recurrence. Feedback Transformer [14] expose top-level representations from previous tokens to the current token’s computation and show that such feedback makes shallow models with fewer parameters more competitive with standard Transformers that have more layers.

In reinforcement learning, DRC [22] carries recurrent state across steps of the Sokoban environment and outperforms deep feed-forward baselines. A closely related classical work, RHN [61] studies how to scale within-step recurrent depth under a fixed parameter budget by trading off depth against width, with depth saturating at approximately 8–9 levels. Importantly, however, compute still grows substantially with depth in the RHN setup, even though the parameter count is held fixed. Similarly, recent recurrent Transformer work [46] finds favorable depth–width tradeoffs under a fixed parameter budget. In contrast, our work keeps compute approximately fixed, rather than parameter count, and studies the tradeoff between within-step depth, expert width, and per-layer expert count for recurrent and non-recurrent models. In this sense, our work complements and extends prior studies.

Depth Recurrence, Mixture-of-Experts and Compute Allocation. A long line of work studies neural computation as an iterative refinement process in which representations are progressively updated through recurrent/looping transformations on the same input, rather than being computed in a single feed-forward pass. Universal Transformers [12] reformulated Transformer layers as recurrent updates applied across depth, while Adaptive Computation Time [19] and PonderNet [1] explored dynamically allocating additional sequential computation to difficult inputs.

A few recent works have also combined depth recurrence with Mixture-of-Experts approaches, motivated by the lack of parameter capacity in these looping architectures; examples include recurrent modular architectures such as Recurrent Independent Mechanisms (RIMs) [18] and recurrent MoE systems such as MoEUT [9]. Mixture-of-experts (MoE) architectures scale model capacity through conditional or parallel computation [31, 34]. Modern sparse MoE systems demonstrate that parameter count can scale substantially while keeping per-token computation nearly constant [51, 38, 15, 13]. Most prior work integrates experts inside fixed-depth feed-forward architectures, where the dominant scaling axis remains serial layer composition. Other works study improved routing and expert utilization, including differentiable aggregation mechanisms such as Soft MoE [49, 44, 60]. In our work, we focus mainly on temporal recurrence rather than depth recurrence, and MoE appears in our framework as a way to study a sequential-to-parallel computation allocation problem. This perspective is also related to broader work on depth–width tradeoffs and compute-optimal scaling ([39]; [7]; [29]).

Discussion and Limitations

Our results suggest that temporal recurrence does not eliminate sequential computation, but changes where it occurs. Non-recurrent models must complete the computation required for each output within the current step, whereas recurrent models can preserve intermediate representations and continue refining them as new inputs arrive. In the settings studied here, this shifts the best observed compute allocation away from within-step depth and toward parallel capacity. Our Sokoban probing results are consistent with this interpretation: the shallower recurrent agent initially forms less complete plans than the deeper agent, but catches up over subsequent environment steps while achieving comparable task performance.

This tradeoff should depend on the structure of the task. Computation can accumulate across time only when consecutive inputs contribute to a shared problem and the carried state preserves useful intermediate computation. We therefore expect recurrence to reduce the need for within-step depth most clearly in settings with persistent context, progressively evolving representations, or opportunities to refine decisions over multiple steps. The same benefit may not arise when each input presents an independent problem, when every decision must be solved immediately, or when the recurrent state is too limited to retain useful intermediate information.

Our empirical study covers two domains, a finite set of architectures, and a restricted range of compute budgets, training durations, and task difficulties. The reported allocations are therefore the best among those evaluated, rather than globally compute-optimal configurations. Isolating the effect of recurrence also requires compromises in architectural parity. ConvLSTM contains state machinery that may favor the recurrent variant, while ConvGLU provides a more direct comparison in which recurrence is introduced through feedback from the previous step’s final representation. Similarly, recurrent Transformers require stabilization and memory-management choices that are absent from the standard non-recurrent baseline. Changes in architecture, scale, or architecture-specific tuning may shift the preferred allocation for either model class.

Optimization may also contribute to the earlier depth saturation observed in recurrent models. We train recurrent models using truncated backpropagation through time, which may make credit assignment through deep recurrent updates more difficult ([2]; [47]). Non-recurrent models constructed from similar blocks continue to benefit from greater depth, arguing against a general inability to optimize deep networks, but this does not rule out optimization difficulties specific to the recurrent setting. Finally, our work proxy captures approximately matched computation rather than realized hardware efficiency. Although reallocating computation from depth to parallel capacity shortens the serial path through each update, actual latency depends on hardware utilization, memory bandwidth, kernel efficiency, and communication overhead. We provide preliminary latency and throughput measurements in Appendix H, but leave a comprehensive systems evaluation to future work.

Conclusion

Across Sokoban and streamed FineWeb, temporal recurrence shifts the best observed compute allocation toward substantially fewer layers per step, with comparable or better performance than non-recurrent models under approximately matched per-step work. For example, in the small-budget language modeling FineWeb comparison, a recurrent model with $L = 2$ reaches a validation loss comparable to the best non-recurrent allocation at $L = 16$. Recurrence and parallel capacity play distinct roles in this shift: recurrence reduces the performance cost of using fewer layers, while per-layer expert count and width provide useful ways to spend the reallocated work. Increasing parallel capacity remains useful both with and without recurrence.

Our Sokoban decoding analysis is consistent with computation being redistributed across time, with the shallower model catching up in planning while achieving comparable task performance. Overall, our results support the broader hypothesis that temporal recurrence can shift the best compute allocation toward fewer layers per step in other streaming tasks where useful computation can accumulate across time. This motivates exploring shallower, more parallel updates as a potential route to reducing latency and improving training throughput, making general nonlinear recurrent models more practical at scale.

Acknowledgment

We acknowledge the support from the Canada CIFAR AI Chair Program and from the Canada Excellence Research Chairs Program. The research was enabled in part by computational resources provided by the Digital Research Alliance of Canada and Mila Quebec AI Institute. IA thanks Andrei Mircea Romascanu for helpful discussions and comments. We thank the Frontier Intelligence Group at Sakana AI for a helpful discussion and suggestions on the paper title, especially Ciaran Regan, whose suggestions helped shape the final title. We would also like to thank the Python community ([56, 45]) for developing tools that enabled this work, including NumPy ([23]), Matplotlib ([30]), Jupyter ([36]), and Pandas ([40]).

Architectures

All architecture variants follow the same pattern: an external input is encoded into $z$, a shared latent state $s$ is read by several experts in parallel, expert outputs are aggregated, and the updated shared state is carried forward through the external sequence in the recurrent version, as illustrated in Figure 7. The reinforcement-learning models use LSTM experts, while the FineWeb next-token model uses transformer experts. Figs. 8 and 9 then specialize the expert block for reinforcement learning and FineWeb next-token prediction, and Algorithms 1–2 give the high-level update rules.

LSTM

Figure 8 illustrates one depth level with parallel LSTM experts used in reinforcement learning experiments. Each expert receives the encoded observation $z$, the shared latent state $s$, and its own local LSTM tensors $(h_i, c_i)$. The experts are evaluated in parallel. At each depth level, their hidden outputs are aggregated into the shared state $s$, which is passed to the next depth level; the policy and value heads read $s$ after the final depth level. In the non-recurrent LSTM version, all $c_i$ are always zero, and $h_i$ and $s$ inputs for the experts in the first depth level are omitted.

Algorithm 1 Parallel experts LSTM

1. Encode external input $z \leftarrow e(x_t)$

2. Receive or initialize shared latent state $s$, rolling hidden states $h_{1:E}$, and cells $c_{1:E}^{1:L}$

3. for within-step depth layer $\ell= 1,\ldots,L$ do

4. &nbsp;&nbsp;&nbsp;&nbsp;for all experts $i = 1,\ldots,E$ in parallel do

5. &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$(h_i,c_i^\ell) \leftarrow f_{\theta_i^\ell}(z,s,h_i,c_i^\ell)$

6. &nbsp;&nbsp;&nbsp;&nbsp;end for

7. &nbsp;&nbsp;&nbsp;&nbsp;$s \leftarrow A(h_1,\ldots,h_E)$

8. end for

9. $(\pi_t,V_t) \leftarrow\operatorname{head}(s)$

10. Return $\pi_t$, $V_t$, updated shared latent state $s$, and LSTM states $h_{1:E}$, $c_{1:E}^{1:L}$

Algorithm 1 describes the same computation across $L$ sequential depth levels. The cell $c_i^\ell$ is local to expert $i$ at depth level $\ell$.

The Sokoban actor–critic use a spatial recurrence. The observation encoder produces a spatial feature map, and the shared latent state keeps the same board-like spatial structure. The encoder is RMS-normalized over the spatial dimensions. At each environment step, recurrent convolutional LSTM (ConvLSTM) experts read the encoded observation, the shared spatial state, and the incoming hidden state and their previous cell states. The ConvLSTM gates are produced by a $3 \times3$ convolution followed by RMS normalization of the gate preactivations. Expert hidden outputs are aggregated into the shared state, which is carried to the next environment step in case of recurrence or to the next depth level. The Sokoban head aggregates the spatial expert output with attention pooling and global average pooling, applies a dense GLU, and then uses separate linear projections for action logits and the scalar value estimate.

Transformer

The FineWeb architecture uses standard NanoGPT baseline ([57], [35], [33]) components: learned token embeddings, grouped-query attention with RoPE, RMS-normalized queries and keys, a squared-ReLU feed-forward block, final RMSNorm, an LM head, and logit soft-capping. To make a parallel expert version of this non-recurrent Transformer we simply execute $E$ transformer blocks in parallel and then merge the outputs using the aggregation in (2).

Recurrent Transformer

To make the recurrent version of it we were needed to introduce two changes in addition to recurrence to stabilize the training, and address memory constrains. We have tried to keep these modifications minimal. First, to stabilize the training each block at all depth levels recieves embeddings of the current token in addition to residual stream. Plus, the output of each expert is normilized with parameter-free RMSNorm. The second addition is that we share the KV cache between experts at each depth level. In our experiments, sharing the KV cache increases validation loss by approximately 0.02, but helps greatly to reduce memory pressure and increases throughput.

In detail, in the parallel expert Recurrent Transformer each expert reads the token embedding $z$ and shared latent state $s$, forms an expert input, and prepares $k_i$, $v_i$, $q_i$. The per-expert queries remain separate, while the key–value proposals are aggregated across experts into one shared key–value entry. This entry is pushed into a rolling KV cache $(K,V)$ local to that depth level, so all experts at that level attend to the same shared KV cache. The rolling KV cache is a fixed-length FIFO buffer: each new aggregated KV entry is appended, and the oldest entry is discarded once the buffer is full. The expert attention outputs pass through the transformer feed-forward path, and RMS normalization, producing expert outputs $h_i$ that are aggregated into the next shared latent state. After the final depth level, the shared state is read by the head. The KV caches and shared state are then carried forward by the streamed-document language-modeling pipeline.

Algorithm 2 Parallel experts Recurrent Transformer

1. Encode token $z \leftarrow e(x_t)$

2. Receive or initialize shared latent state $s$ and FIFO KV caches $(K^{1:L}, V^{1:L})$

3. for within-step depth level $\ell= 1, \ldots, L$ do

4. for all experts $i = 1, \ldots, E$ in parallel do

5. $(k_i, v_i, q_i) \leftarrow p_{\theta_i}(z, s)$

6. end for

7. $k \leftarrow A(k_1, \ldots, k_E)$

8. $v \leftarrow A(v_1, \ldots, v_E)$

9. $(K^\ell, V^\ell) \leftarrow\operatorname{push}_{\mathrm{FIFO}}((K^\ell, V^\ell), (k, v))$

10. for all experts $i = 1, \ldots, E$ in parallel do

11. $h_i \leftarrow g_{\theta_i}(q_i, K^\ell, V^\ell)$

12. end for

13. $s \leftarrow A(h_1, \ldots, h_E)$

14. end for

15. $\operatorname{logits}_t \leftarrow\operatorname{head}(s)$

16. Return next-token logits $\operatorname{logits}_t$, shared latent state $s$, and KV caches $(K^{1:L}, V^{1:L})$

FineWeb Streaming Implementation

For a shard with $S$ tokens and global batch size $B$, the loader partitions the shard into $B$ contiguous streams of length $\lfloor S/B \rfloor$. Batch slot $b$ reads only from stream $b$. At stream position $q$, slot $b$ reads a window of length $T + 1$, where $T$ is the training sequence length. The first $T$ tokens are inputs and the next $T$ shifted tokens are labels.

The recurrent state is also organized by batch slot. It is reset at the first window of a shard, carried across later adjacent windows in the same slot, and reset again when a new shard begins.

Algorithm 3 FineWeb shard streaming

1. Let $B$ be the global batch size and $T$ the training sequence length

2. for all training shards with $S$ tokens do

3. Split the shard into $B$ contiguous streams of length $\lfloor S/B \rfloor$

4. Initialize recurrent state for all batch slots

5. for $q = 0, T, 2T, \ldots$ while all streams have $T + 1$ tokens available do

6. for all batch slots $b = 1, \ldots, B$ do

7. $w_b \leftarrow$ tokens from stream $b$ at positions $q : q + T + 1$

8. $x_b \leftarrow w_b[0:T],\quad y_b \leftarrow w_b[1:T+1]$

9. Reset slot state only if $q = 0$

10. end for

11. Compute next-token loss and updated recurrent state on $(x_{1:B}, y_{1:B})$

12. Update parameters; carry the detached recurrent state to the next window

13. end for

14. end for

Sokoban Solve Rate Results

We evaluate Sokoban policies trained on 50M environment steps on the 1,000-level unfiltered-test split with ConvLSTM architecture. For each fixed-compute budget and within-step depth $L \in\{1, 2, 4, 8, 16\}$, we select the configuration with the highest mean training return. All selected policies act greedily for at most 120 steps per test level. Table 1 reports the results.

Recurrent Budget	Recurrent $L$	Recurrent $E$	Recurrent $d$	Recurrent Unfiltered Test (%)	Feedforward Budget	Feedforward $L$	Feedforward $E$	Feedforward $d$	Feedforward Unfiltered Test (%)
small	1	1	128	86.03 ± 3.71	small	1	4	64	3.20 ± 0.00
small	2	2	64	91.93 ± 1.08	small	2	2	64	30.50 ± 0.96
small	4	1	64	92.70 ± 1.04	small	4	1	64	59.63 ± 3.83
small	8	2	32	86.13 ± 2.40	small	8	2	32	74.63 ± 3.91
small	16	1	32	79.00 ± 0.96	small	16	1	32	73.23 ± 2.97
medium	1	1	256	93.23 ± 1.06	medium	1	4	128	3.93 ± 0.12
medium	2	2	128	96.20 ± 0.70	medium	2	2	128	39.03 ± 0.78
medium	4	1	128	96.67 ± 0.67	medium	4	1	128	68.93 ± 1.95
medium	8	2	64	93.40 ± 3.84	medium	8	2	64	86.73 ± 1.34
medium	16	1	64	88.67 ± 1.72	medium	16	1	64	84.10 ± 4.32
large	1	4	512	96.23 ± 0.91	large	1	58	128	4.57 ± 0.50
large	2	8	256	98.47 ± 0.15	large	2	2	512	52.93 ± 0.68
large	4	1	512	98.70 ± 0.17	large	4	4	256	78.47 ± 0.55
large	8	7	128	95.67 ± 1.44	large	8	7	128	92.00 ± 0.26
large	16	13	64	88.57 ± 3.80	large	16	13	64	93.00 ± 0.40

For external context, [8] report Thinker and a reproduction of DRC ([22]) on the same unfiltered Sokoban test set after 50 million real-environment frames. [55] report DRC(3,3) and ResNet point estimates after two billion environment steps on the same split with the same 120-step episode limit.

Method	Training steps	Unfiltered Test
Thinker	50M	94.5 $\pm0.8$
DRC(3,3) (Thinker reproduction)	50M	91.3 $\pm1.2$
DRC(3,3)	2B	99.3
ResNet	2B	97.9

Aggregation Ablations

Aggregation rules. Let $h_i$ denote the output of expert $i$. For routed variants, $\alpha_i$ denotes the learned softmax routing coefficient for expert $i$. The four ablations choices are mean pooling, sqrt-normalized sum pooling, softmax-gated weighted averaging, and L2-normalized softmax-gated aggregation. We label these as Mean, Sqrt-sum, Softmax mean, and Softmax L2-sum in Figure 10:

The sqrt-normalized variants are variance-preserving normalizations: if independent expert outputs have comparable scale, the aggregate scale stays roughly stable as the effective number of contributors changes. The routed mean is the usual softmax-gated weighted average. In the implementation, denominators in the routed formulas are protected with a small $\epsilon$ for numerical stability.

Routing coefficients. The coefficients are produced by a scaled dot-product router. For dense experts, the router context is the concatenation of the encoded observation or token representation and the current shared latent state, $c = [x_{\mathrm{enc}}; s]$. For convolutional recurrent experts, both tensors are first spatially average-pooled, then concatenated. The router maps this context to a query and compares it to learned expert keys:

Here $e_i$ is a learned embedding for expert $i$, $d$ is the router embedding dimension, and $\ell_i$ are the router logits.

Default choice. We use Sqrt-sum as the default aggregation rule used throughout the paper, except in experiments that explicitly ablate aggregation. It is parameter-free, keeps the aggregate scale stable as the number of contributing experts changes, and is consistently competitive across domains. For the FineWeb shared-global transformer, we use it for both expert outputs and the merged key-value proposals.

Communication Radius

We vary how many other experts each expert can read at each depth level in a Sokoban model with $E = 8$ ConvLSTM experts per layer, width $d = 32$, and within-step depth $L = 2$. Experts are arranged in a ring and read their own state plus those of a chosen number of preceding neighbours. The visible states are summed and divided by the square root of their count, keeping the input width and expert computation fixed.

Figure 11 shows that full communication achieves the highest mean return. These results indicate that sharing information across experts contributes to the model’s performance. We only use full communication for other experiments in the paper.

Sokoban Planning Analysis

We follow the spatial probing method of [6], who decoded iteratively refined plans in DRC agents. We apply their agent-move probes to our models’ final residual representations and compare how decoded plans develop across environment steps at different depth allocations.

We extend Figure 4 to four large-budget recurrent ConvLSTMs with fixed width $d = 256$ and $(L, E) \in(1, 15), (2, 8), (4, 4), (8, 2)$. For each allocation, we select the best of three training seeds by solve rate on 1,000 Unfiltered Test levels, using greedy actions with a 120-step limit.

Following [6] for each frozen policy, we train convolution $3 \times3$ probes on 3,000 Unfiltered Train trajectories. Each probe reads the final residual representation at each step. At each level location, the target is the direction of the agent in this location in the recorded trajectory, or NEVER if agent does not visit it.

We evaluate decoding on the 906 Test levels solved by all four policies. Recall is the fraction of non-NEVER targets correctly predicted above threshold, and precision is the fraction of emitted directions that are correct. We choose one fixed threshold per model to achieve approximately 93% precision for all prediction over first 20 steps.

The $L = 4$ model initially has higher recall than the $L = 2$ model, indicating a more complete decoded plan. The $L = 2$ model catches up over the next few environment steps, after which both maintain similar recall. Together with their nearly identical solve rates, this pattern is consistent with the shallower model compensating for less planning within each update by continuing the computation across environment steps, catching up early enough to preserve task performance. A similar pattern appears for the $L = 1$ and $L = 8$ models, which also achieve similar test solve rates: $L = 8$ initially has higher recall, while $L = 1$ narrows the gap over subsequent steps, although it does not fully close it.

Visualization. To make the analysis self-contained, we also include a visualization of how plans develop in the two layer model (Figure 13). The main DRC model analyzed by [6] combines within-step looping with temporal recurrence, and our models do not use within-step looping.

Implementation Details and Hyperparameters Used in Experiments

This section reports the training hyperparameters used across the main experiments.

Parameter	Value
Sokoban
Encoder hidden dimension	32
Discount rate $\gamma$	0.97
Lambda for general advantage estimation	0.97
Entropy coefficient	0.01
Value function coefficient	1.0
Normalize advantages	No
Unroll length	20
Number of environments	32
Total environment steps	50,000,000
Learning rate	$4 \times10^{-4}$
Anneal learning rate	Yes
Optimizer	Adam
Adam epsilon	$10^{-6}$
Maximum gradient norm for clipping	1.0

Parameter	Value
FineWeb next-token prediction
Vocabulary size	50,304
Attention heads	8 query, 4 key-value default; scaled with width proportionally
Token batch size	131,072
Training steps	10,000
Optimizer	Muon core; AdamW embed/readout
Gate parameters	Muon
Optimizer betas	(0.8, 0.95)
Gradient clipping	1.0
Learning rate schedule	Warmup-cosine decay
Peak learning rate	$2 \times10^{-2}$
Minimum learning rate	$2 \times10^{-4}$
Warmup steps	100
Embedding learning rate	$0.2(d/768)^{-1/2}$
Readout learning rate	$0.004(d/768)^{-1/2}$
Weight decay	0
Cautious weight decay	0.01
Numerical precision	bfloat16
Transformer
Training sequence length	256 or 2048
Recurrent Transformer
Shared KV cache length	128 or 1024
Training sequence length	64

Compute Resources

The reinforcement-learning experiments were run on NVIDIA RTX 8000 and L40S GPUs. Most Sokoban runs use one or two GPUs and take a few hours for small configurations, while the largest fixed-compute and within-step-depth-heavy configurations can take up to roughly one day. The FineWeb next-token prediction experiments were run on NVIDIA H100 GPUs. All FineWeb runs use four GPUs and take several hours for standard configurations, while the largest fixed-compute and within-step-depth-heavy configurations can take up to roughly one day.

Latency and Training Throughput

These measurements are preliminary. Neither implementation was extensively optimized for inference, so the latency results should be interpreted with caution. Both implementations use JAX; Recurrent Transformer experts at the same depth level are evaluated with jax.vmap whenever they can execute in parallel. Measurements use four 80GB NVIDIA H100 GPUs, hidden dimension 768, eight query heads, four key-value heads, and budget $L \cdot E = 16$. For the Recurrent Transformer, $L$ is within-step depth and $E$ is the number of parallel experts per depth level. In latency measurements, $B$ is the global inference batch. In training throughput measurements, $B$ is the global batch size and $C$ is the training sequence length, giving $CB$ tokens per optimizer update.

Inference Latency

Table 5 reports latency after pre-filling the KV caches of either model to length $C$.

Model	$C = 128$, $B = 16$	$C = 2048$, $B = 16$
Transformer $L = 16$, $E = 1$	0.565	0.670
Recurrent $L = 1$, $E = 16$	0.108	0.131
Recurrent $L = 2$, $E = 8$	0.431	0.554
Recurrent $L = 4$, $E = 4$	0.552	0.637
Recurrent $L = 8$, $E = 2$	0.815	0.935
Recurrent $L = 16$, $E = 1$	1.402	1.577

Training Throughput

Table 6 measures the complete FineWeb training update, including the normal dataloader, forward and backward passes, and optimizer update.

131k effective tokens per optimizer update
Model	$C = 32$ $B = 4{,}096$	$C = 64$ $B = 2{,}048$	$C = 128$ $B = 1{,}024$	$C = 256$ $B = 512$	$C = 1024$ $B = 128$	$C = 2048$ $B = 64$
Transformer $L = 16$, $E = 1$	967.7k	1,027.4k	1,025.3k	1,054.3k	1,016.9k	968.1k
Recurrent $L = 1$, $E = 16$	505.4k	445.7k	349.5k	229.4k	51.1k	OOM
Recurrent $L = 2$, $E = 8$	265.8k	209.5k	145.2k	85.8k	OOM	OOM
Recurrent $L = 4$, $E = 4$	189.7k	148.7k	99.1k	OOM	OOM	OOM
Recurrent $L = 8$, $E = 2$	118.3k	88.4k	OOM	OOM	OOM	OOM
262k effective tokens per optimizer update
Model	$C = 32$ $B = 8{,}192$	$C = 64$ $B = 4{,}096$	$C = 128$ $B = 2{,}048$	$C = 256$ $B = 1{,}024$	$C = 1024$ $B = 256$	$C = 2048$ $B = 128$
Transformer $L = 16$, $E = 1$	886.9k	1,031.5k	1,095.8k	1,132.0k	1,095.8k	1,071.5k
Recurrent $L = 1$, $E = 16$	556.2k	526.4k	447.7k	313.1k	OOM	OOM
Recurrent $L = 2$, $E = 8$	306.0k	264.6k	190.2k	OOM	OOM	OOM
Recurrent $L = 4$, $E = 4$	199.9k	170.8k	OOM	OOM	OOM	OOM
Recurrent $L = 8$, $E = 2$	OOM	OOM	OOM	OOM	OOM	OOM

FineWeb Fixed-Compute Sweeps

Figure 14 show FineWeb fixed-compute allocations at small budget with long context and Figure 15 shows FineWeb fixed-compute allocations the medium budget, comparing Recurrent Transformer and Transformer architectures.

Sokoban Fixed-Compute Sweeps

Figure 2, Figure 16 and Figure 17 show the Sokoban the ConvLSTM allocations at the small, medium, and large budgets under recurrent and feedforward settings. The large-budget snapshot includes only completed training runs.

Recurrent	Budget	$L$	$E$	$d$	Mean	Variance	$N$
Recurrent	Small	1	1	128	7.750	0.0076	3
Recurrent	Small	1	4	64	7.664	0.1200	3
Recurrent	Small	1	12	32	5.670	1.0279	3
Recurrent	Small	2	2	64	8.947	0.0685	3
Recurrent	Small	2	6	32	7.336	0.0454	3
Recurrent	Small	4	1	64	9.212	0.0005	3
Recurrent	Small	4	3	32	7.952	0.4108	3
Recurrent	Small	8	2	32	7.929	0.1292	3
Recurrent	Small	16	1	32	6.297	0.0277	3
Recurrent	Medium	1	1	256	8.795	0.0179	3
Recurrent	Medium	1	4	128	8.566	0.1426	3
Recurrent	Medium	1	14	64	8.382	0.1108	3
Recurrent	Medium	1	45	32	7.224	0.3117	3
Recurrent	Medium	2	2	128	9.558	0.0101	3
Recurrent	Medium	2	7	64	9.241	0.0038	3
Recurrent	Medium	2	23	32	8.928	0.0031	3
Recurrent	Medium	4	1	128	9.778	0.0197	3
Recurrent	Medium	4	3	64	9.677	0.0443	3
Recurrent	Medium	4	11	32	9.200	0.0221	3
Recurrent	Medium	8	2	64	9.067	0.9128	3
Recurrent	Medium	8	6	32	8.919	0.1909	3
Recurrent	Medium	16	1	64	8.068	0.0537	3
Recurrent	Medium	16	3	32	7.378	0.1338	3
Recurrent	Large	1	1	1024	9.479	0.0288	3
Recurrent	Large	1	4	512	9.503	0.0033	3
Recurrent	Large	1	15	256	9.394	0.0132	3
Recurrent	Large	1	58	128	9.278	0.0469	3
Recurrent	Large	1	208	64	9.001	0.0354	3
Recurrent	Large	2	2	512	10.092	0.0085	3
Recurrent	Large	2	8	256	10.219	0.0113	3
Recurrent	Large	2	29	128	10.023	0.0074	3
Recurrent	Large	2	104	64	9.963	0.0267	3
Recurrent	Large	4	1	512	10.189	0.0077	3
Recurrent	Large	4	4	256	10.163	0.0138	3
Recurrent	Large	4	14	128	10.164	0.0074	3
Recurrent	Large	4	52	64	9.944	0.0239	3
Recurrent	Large	8	2	256	9.429	0.4706	3
Recurrent	Large	8	7	128	9.551	0.1403	3
Recurrent	Large	8	26	64	8.575	0.2624	3
Recurrent	Large	16	1	256	7.735	0.1215	3
Recurrent	Large	16	4	128	7.667	0.2837	3
Recurrent	Large	16	13	64	8.241	0.7027	3
Feedforward	Small	1	1	128	-9.080	0.0007	3
Feedforward	Small	1	4	64	-8.940	0.0185	3
Feedforward	Small	1	12	32	-9.174	0.0033	3
Feedforward	Small	2	2	64	-2.050	0.2530	3
Feedforward	Small	2	6	32	-2.255	0.2295	3
Feedforward	Small	4	1	64	3.697	0.1589	3
Feedforward	Small	4	3	32	2.938	0.0237	3
Feedforward	Small	8	2	32	5.902	0.5632	3
Feedforward	Small	16	1	32	5.491	0.0821	3
Feedforward	Medium	1	1	256	-8.889	0.0042	3
Feedforward	Medium	1	4	128	-8.829	0.0086	3
Feedforward	Medium	1	14	64	-8.874	0.0006	3
Feedforward	Medium	1	45	32	-8.917	0.0004	3
Feedforward	Medium	2	2	128	0.204	0.1498	3
Feedforward	Medium	2	7	64	-0.357	0.2066	3
Feedforward	Medium	2	23	32	-0.297	0.1262	3
Feedforward	Medium	4	1	128	5.571	0.1257	3
Feedforward	Medium	4	3	64	5.039	0.1118	3
Feedforward	Medium	4	11	32	4.753	0.0918	3
Feedforward	Medium	8	2	64	8.367	0.0050	3
Feedforward	Medium	8	6	32	7.324	0.0180	3
Feedforward	Medium	16	1	64	7.674	0.1781	3
Feedforward	Medium	16	3	32	7.336	0.0311	3
Feedforward	Large	1	1	1024	-8.810	0.0143	3
Feedforward	Large	1	4	512	-8.801	0.0122	3
Feedforward	Large	1	15	256	-8.885	0.0265	3
Feedforward	Large	1	58	128	-8.760	0.0083	3
Feedforward	Large	1	208	64	-8.765	0.0041	3
Feedforward	Large	2	2	512	2.356	0.0064	3
Feedforward	Large	2	8	256	2.216	0.0352	3
Feedforward	Large	2	29	128	1.797	0.0835	3
Feedforward	Large	2	104	64	1.497	0.0736	3
Feedforward	Large	4	1	512	6.807	0.0178	3
Feedforward	Large	4	4	256	7.295	0.0715	3
Feedforward	Large	4	14	128	6.938	0.0739	3
Feedforward	Large	4	52	64	6.226	0.2677	3
Feedforward	Large	8	2	256	9.090	0.0501	3
Feedforward	Large	8	7	128	9.215	0.0544	3
Feedforward	Large	8	26	64	8.875	0.1905	3
Feedforward	Large	16	1	256	8.650	0.1802	3
Feedforward	Large	16	4	128	9.271	0.0152	3
Feedforward	Large	16	13	64	9.417	0.0358	3

ConvGLU Sokoban

Figure 18 shows the similar saturation as scaling effects but for the ConvGLU architecture in Sokoban. Figure 19 and Figure 20 show the recurrent and feedforward allocations for the ConvGLU sweeps.