Pretraining Recurrent Networks without Recurrence

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Pretraining Recurrent Networks without Recurrence

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Pretraining Recurrent Networks without Recurrence

Akarsh Kumar

Phillip Isola MIT

y0̂

y1̂

̂ yT≤1

yT̂

̂ yt̂ yt+1 … yT̂

Readout

fθ

yT≤1

Readout

Updater

yt yt+1 … yT

Readout

m−

SMT

Readout

arXiv:2606.06479v1 [cs.LG] 4 Jun 2026

BPTT

Decoder

Updater

fθ

…

mT≤1

Updater

fθ

Encoder

xT x0

…

Loss Gradient Path

Updater

fθ

m̂ t+1

mt+1

Encoder

xt+1

…

xt+1

Figure 1: BPTT vs SMT. Left: BPTT trains an RNN by recurrently unrolling the “updater” network in time, and backpropagating gradients through the entire graph. Right: Supervised Memory Training (SMT) trains an RNN with supervised learning on one-step memory transition labels, which are generated by a Transformer encoder-decoder model pair trained to produce predictive states. SMT is fully time-parallel. In SMT, the longest gradient path between tokens is O(1) (compared to O(T ) in BPTT), which stabilizes gradients, making learning long-range dependencies qualitatively easier.

Abstract Training recurrent neural networks (RNNs) requires assigning credit across long sequences of computations. Standard backpropagation through time (BPTT) addresses this problem poorly: it is sequential in time, limiting parallelism, and suffers from vanishing or exploding gradients, making long-range associations difficult to learn. We propose Supervised Memory Training (SMT), a method for training nonlinear RNNs that sidesteps recurrent credit propagation entirely by reducing RNN training to supervised learning on one-step memory transition labels (mt , xt+1 ) → mt+1 . SMT acquires these memory labels by training a Transformerbased encoder on a predictive state objective—retaining only information from the past necessary to predict the future. By decoupling what to remember from how to update memory, SMT enables time-parallel RNN training with a stable O(1) length gradient path between any two tokens—without ever unrolling the RNN. We find that SMT outperforms BPTT when pretraining various RNN architectures on tasks like language modeling and pixel sequence modeling. SMT enables nonlinear RNNs to better capture long-range dependencies and train in parallel, potentially unlocking the scaling of models that build temporal abstractions of past experience. Project Page: akarshkumar.com/smt Source Code: github.com/akarshkumar0101/smt

Preprint.

Introduction

Recurrent neural networks (RNNs) store information about the past that will only become useful in the future. The core training challenge is that the utility of a memory may be delayed: many intermediate computations intervene between writing information and eventually using it. These intervening steps confound learning the correct associations, a problem known as credit assignment [89]. The standard approach, backpropagation through time (BPTT), assigns credit across a sequence by unrolling the RNN in time and propagating gradients backward through the resulting computation graph [102, 131]. Although conceptually well-motivated, BPTT is sequential in time and suffers from unstable high variance gradients that may vanish or explode [97]. The lack of time-parallelism makes BPTT scale poorly, while its gradient instability makes learning long-range associations difficult, as credit must propagate across up to O(T ) steps [11]. Is recurrent credit propagation unavoidable? In this paper, we propose Supervised Memory Training (SMT), a method to train nonlinear RNNs that sidesteps recurrent credit propagation by reducing the problem to supervised learning. Suppose we had access to the optimal memory state at each timestep, m∗t . Then, RNN training reduces to learning the one-step update (m∗t , xt+1 ) → m∗t+1 using standard supervised objectives. The challenge, of course, is how to actually obtain such memory labels. In this paper, we assert that an effective memory is a sufficient statistic of the past for predicting the future, i.e., a predictive state [75]. The past is typically viewed as a sequence, suggesting that memory must be computed sequentially over time. Our key insight is that, by augmenting each observation with its timestamp, the past can instead be losslessly represented as a set of timestamped events, rather than a sequence. Under this reparameterization, the optimal memory becomes a permutation-invariant function of this set, and can therefore be estimated using models that operate in parallel over time. This reframing allows us to train memory representations without recurrently propagating credit through time. In practice, we train a Transformer encoder model to embed the past context into a memory that a separate decoder can use to predict the future. This objective operationalizes the notion of a predictive state: a representation of the past that retains only the information needed to predict the future and nothing more. Once this teacher encoder has learned to construct such memory representations, the RNN can then focus on learning the now much simpler task of updating that memory over time. In essence, SMT decouples learning what to remember (memory representation), which is a nonsequential problem, from learning how to update memory (memory dynamics), which is a sequential process but can be supervised one-step at a time. This decoupling enables time-parallel training of nonlinear RNNs without unrolling, and creates a stable O(1) gradient path for long-range associations. Indeed, Transformers solved time-parallelism and credit assignment in the same way [124], and have since revolutionized sequence modeling [14]. However, Transformers do not possess a compressed memory of the past in the way RNNs or human brains do [38, 64]. Instead, Transformers store the entire history of past token representations and attend to all of them when processing each new token. As a result, their memory size grows with sequence length, leading to prohibitive computational costs for unbounded sequences, such as a human lifetime of experience [121]. Sliding-window transformers mitigate this issue by storing only the most recent tokens, but have the severe drawback that they lose access to information before the context window [21]. In contrast, no known biological intelligence operates in this manner—accessing its entire experiential history for every new decision—but instead constructs a temporally compressed abstraction of past experience, like an RNN [12]. Linear attention RNN models also exhibit time-parallel training and relatively stable credit assignment, while maintaining a fixed memory size [67, 40, 39, 24]. But, because their transition function is linear, the class of functions they can represent is fundamentally constrained [84], which can lead to failure on important sequential tasks such as state tracking [83, 77]. SMT aims to combine the best of all worlds: time-parallel training, stable O(1) long-range credit assignment, fixed-memory inference, and maximal expressivity via nonlinear dynamics. Our results confirm that, on language modeling and pixel sequence modeling tasks, SMT outperforms BPTT in learning long-range dependencies while requiring less sequential computation. SMT should primarily be used for pretraining RNNs, followed by some lightweight post-training to mitigate drift from the teacher memory trajectories and adapt to specific downstream tasks. In fact, post-training is necessary to go beyond the limitations of the teacher encoder [81]. Beyond its role as a training approach for RNNs, SMT can also be seen as a new method for learning representations (mappings from data to latent variables) and for learning world models (transitions from state at time t to state at time t + 1). 2

Methods

2.1

Background

Causal Conditional Sequence Modeling Let x = [x0 , . . . , xT ] and y = [y0 , . . . , yT ] denote input and output sequences. The objective is to learn a model of the conditional distribution p(y | x). We assume each output yt depends only on x0 , . . . , xt . Formally we model this distribution with QT t=0 pθ (yt | x≤t ). Autoregressive sequence modeling is a special case when xt = yt−1 . Recurrent Neural Networks (RNNs) An RNN models this problem using a fixed-size latent state, mt , that summarizes past inputs. At each timestep, this state is updated according to: mt+1 = fθ (mt , xt+1 ) (1) where fθ is the transition function. The predicted output token distribution is then pθ (yt | x≤t ) = softmax(gθ (mt )), where gθ is the readout function. Ideally, mt “remembers” important information from the past and intentionally “forgets” unimportant information, i.e., mt is a memory. Backpropagation Through Time (BPTT) Traditionally, RNNs are trained with BPTT [102, 131]. In the forward pass, fθ is recurrently unrolled over the sequence. The input sequence xt is provided via teacher forcing, while the memory sequence mt is generated by the RNN’s transition and is used to compute the output predictions. Conceptually, the computation graph takes the form: mt = fθ (. . . fθ (fθ (m∅ , x0 ), x1 ), . . . , xt ) with m∅ = 0 Gradients are then computed end-to-end on this unrolled computation graph, propagating from the output prediction losses backward through the trajectory of the nonlinear dynamical system. Thus, this gradient credit assignment signal may have to travel for a path length of up to O(T ) steps. Depending on the singular values of the Jacobian of fθ , gradients may vanish or explode in time. BPTT has two well-known limitations: 1. Equation 1 is usually implemented with a recurrent for-loop, preventing parallelization [96]. 2. BPTT often produces unstable high variance gradients [11]. When gradients vanish, the RNN experiences a recency bias, hindering the learning of long-range associations [100]. When gradients explode, the induced dynamical system is chaotic, causing training instability [97]. 2.2

Supervised Memory Training (SMT)

We propose Supervised Memory Training (SMT) for pretraining nonlinear RNNs without BPTT. The core idea is to decouple the learning of memory representation from memory dynamics. Motivation Consider a hypothetical oracle memory-encoding model Q, that takes as input the sequence of tokens up to timestep t, xctx t = [x0 , x1 , . . . , xt ], and outputs an effective compressed memory for that timestep m∗t = Q(xctx t ). This memory retains all information from the past input that is relevant for predicting the future output, ytfut = [yt , . . . , yT ], while deliberately discarding unimportant details. For example, the oracle would remember the personalities of characters in a story, but discard details of what they were wearing on a specific day, just as humans do. Running Q at different points along the sequence produces a corresponding sequence of memory labels [m∗0 , m∗1 , . . . , m∗T ]. With Q, the RNN’s problem of learning a temporal update collapses to standard supervised learning on oracle memory transitions labels (m∗t , xt+1 ) → m∗t+1 . Our key insight is that Q does not have to be a recurrent function over [x0 , x1 , . . . , xt ], but can instead be represented as a permutation-invariant function over the set {(x0 , 0), (x1 , 1), . . . , (xt , t)} (details in Appendix E). In practice, SMT approximates Q by training a time-parallel model (e.g. a Transformer) to compress the past input into a memory representation that a separate decoder model can use to predict the future output. This future-predicting objective operationalizes the notion of a predictive state [75]. Formulation Formally, we have the RNN fθ , bidirectional encoder Eϕ , and causal decoder Dψ . Eϕ and Dψ are time-parallel Transformer architectures. Given x = [x0 , x1 , . . . , xT ] and y = [y0 , y1 , . . . , yT ], we consider, for each timestep t, a decomposition into the past and future: xctx xfut ytfut = [yt , . . . , yT ] t = [x0 , . . . , xt ] t = [xt+1 , . . . , xT ] The encoder maps each context to a memory state with mt = Eϕ (xctx t ). Then, the decoder predicts the future output distribution using the memory of the past and teacher forced future inputs: T Y fut pϕ,ψ (ytfut | xctx pψ (yτ | mt , xt+1:τ ) = Dψ (mt , xfut t , xt ) = t ) τ =t

The future decoding loss for timestep t is (CE denotes the sequence level cross-entropy loss): fut fut ctx fut Ldec t = CE yt , pϕ,ψ (yt | xt , xt )

(2)

We have the RNN predict the next memory given the current memory and the next input with m̂t+1 = fθ (mt , xt+1 ). This prediction is supervised with the next timestep’s memory: Ldyn = MSE(m̂t+1 , mt+1 ) t

(3)

This dynamics loss has two distinct purposes: 1) to train the RNN and 2) to explicitly shape the encoder memory representations to be Markovian (i.e. mt+1 is predictable solely from (mt , xt+1 )). We add a uniformity loss [130] to prevent the memory space from collapsing: Lunif = log Eta ,tb ∼[0,...,T ] exp(−2∥mta − mtb ∥22 )

(4)

The full objective is a weighted sum of all three losses: h i Lsmt = λdec Et Ldec + λdyn Et Ldyn + λunif Lunif t t

(5)

where the λ terms control the trade-off between memory representation, dynamics, and collapse. Practice Theoretically, it should be enough to train Eϕ and Dψ with only Ldec , and separately train fθ with only Ldyn (proof in Appendix F). However, in practice we find it beneficial to jointly train all models in one stage with Lsmt , since that explicitly optimizes mt to be Markovian, and provides additional temporal credit propagation benefits described in Section 3.6. fut For experiments, we truncate xctx t to a context length Tc and yt to a future length Tf . For computational efficiency, we estimate the expectation in Lsmt by randomly sampling a single timestep t, rather than computing all timesteps in the sequence. This yields SMT a smaller training memory footprint than BPTT: O(M + T ) instead of O(M T ), where M is the memory size.

Properties of SMT In SMT, the encoder model constructs appropriate memory representations of the past, while the RNN is responsible for learning the now much simpler task of updating that memory in one-step, thereby decoupling memory representation from memory dynamics. In contrast, under BPTT training, the RNN must learn both tasks simultaneously. Since the memory labels are acquired with a “teacher” encoder-decoder pair, SMT inherits all of its properties, such as time-parallelism, O(1) credit path for long-range associations, and gradient stability. 2.3

DAgger Memory Training (DMT)

After SMT, the RNN achieves low one-step error SMT RNN in predicting (mt , xt+1 ) → mt+1 when mt comes Trajectory from the encoder. However, at evaluation time, the model is unrolled autoregressively, using its own preEncoder dicted memories rather than the encoder memories DMT Trajectory mt as input. This train–test mismatch causes small preθt Loss diction errors to accumulate over time, leading to a m̂ t growing drift between the RNN-generated memory trajectory [m̂0 , . . . , m̂T ] and the encoder trajectory Figure 2: SMT vs DMT. SMT trains the [m0 , . . . , mT ], even with teacher forced input tokens. RNN with behavior cloning on the encoderThis drift is quantified as δt = MSE(m̂t , mt ). generated memory states (off-policy imitation We introduce DAgger Memory Training (DMT), a learning). DMT unrolls the RNN with its own finetuning phase that corrects this drift via on-policy memory states and then imitates the encoder imitation learning [101]. By exposing the RNN to its trajectory (on-policy imitation learning). Figown induced memory state distribution, DMT trains ure design inspired by Jacobs et al. [59]. the RNN to autocorrect its errors to stay aligned with the encoder trajectory (Figure 2). Concretely, given x, we first compute the encoder trajectory [m0 , . . . , mT ] using only Eϕ and then the RNN trajectory [m̂0 , . . . , m̂T ] using fθ . Instead of training on SMT labels (mt , xt+1 ) → mt+1 , we train on DMT labels (m̂t , xt+1 ) → mt+1 . Equivalently, the training loss is: Ldmt = Et [MSE(m̂t , mt )] 4

(6)

During DMT, we freeze the encoder and decoder and only train the RNN with a small learning rate. Note that DMT unrolls the RNN memories, but still uses teacher forced xt inputs. Although DMT unrolls the RNN and gradients may optionally propagate through time, its objective is fundamentally different than standard BPTT, since long-range credit is already assigned in the encoder memory labels, mt . DMT is not time-parallel. That said, DMT should primarily be viewed as a lightweight fine-tuning phase following SMT. Table 1 shows the resource requirements for the different methods. Table 1: Resource requirements. T is token sequence length. Tc is SMT encoder context length. For RNNs, M is the memory state size. We ignore log terms for simplicity. LA denotes linear attention (in its parallel and recurrent form). Complexity classes are from Merrill et al. [84]. Training (T -Length Sequence) Method Transformer LA (parallel) LA (recurrent) BPTT (RNN) SMT (RNN) DMT (RNN)

Memory

Compute

Sequential Operations

O(T ) O(T ) O(M ) O(M T ) O(M + Tc ) O(M T )

O(T 2 ) O(T 2 ) O(T ) O(T ) O(Tc2 ) O(Tc2 T )

O(1) O(1) O(1) O(T ) O(1) O(T )

Inference (One-Step)

Credit Path Length

Memory

Compute

O(1) O(1) O(1) O(T ) O(1) O(1)

O(T ) O(M ) O(M ) O(M ) O(M ) O(M )

O(T ) O(1) O(1) O(1) O(1) O(1)

Complexity Class

TC0 PNC1 PNC1 L/P L/P L/P

Experiments

We study the properties of SMT and compare against BPTT, the standard RNN training algorithm. We restrict our analysis to nonlinear RNNs, the primarily setting BPTT is applied. Transformers and linear RNNs are excluded as they are qualitatively distinct model classes [84, 38]. “BPTT RNN” denotes the BPTT baseline. “SMT Encoder∗ ” generates memories mt with the SMTtrained encoder and predicts next tokens using the decoder. This method is essentially a Transformer baseline with the same memory bottleneck as our RNNs. Since it serves as the teacher during SMT and DMT, it provides a reference upper bound on RNN performance. “SMT→DMT RNN” denotes the RNN pretrained with SMT and finetuned with DMT, which constitutes our full method. Architectures We use RNN architectures based on a Transformer, MLP, and GRU [17] backbone. Datasets We consider character-level language modeling on TinyStories [27] as a naturalistic task requiring long-range memory [118]. As a more challenging problem, we test our method on raster-scan order pixel sequence modeling of sparse images from MNIST [71] and Sketchy [105]. This is a hard problem for RNNs [70, 123]. Imagine you are an ant traversing an image pixel by pixel, row by row. When you see a new white pixel, in order to recognize the shape and slope of the stroke it belongs to, you must remember the white pixels you saw in the previous rows, which may be hundreds of timesteps ago, buried among black pixels. RNNs must achieve this with finite memory, meaning no direct attention to earlier pixels, and thus forcing long-range memory to emerge. We term this “Attneave’s task”, based on classic work from perceptual psychology [5]. More details on architectures, datasets, and experiments are in Appendix B. 3.1 Synthetic Task Experiments We first evaluate BPTT and SMT on synthetic tasks designed to isolate and probe specific properties of the training algorithm. The RNN architecture with the Transformer backbone is used for these experiments. For these synthetic experiments, we set Tc = Tf = T and train all timesteps in the Lsmt expected value. Our tasks include the following (details of tasks are in Appendix B.2.1): 1. Retrieval to test Gradient Stability (sweep sequence length and noise level). 2. String Copy to test Memory Capacity (sweep sequence length and memory state size). 3. Stack Operations to test State Tracking (sweep sequence length and state complexity). 4. Keys-Values to test Associative Recall (sweep number of and complexity of associations). 5. Modular Arithmetic to test In-Context Learning (sweep difficulty and number of examples). Figure 3 shows that SMT→DMT outperforms BPTT in all settings of all tasks. BPTT struggles to learn as sequences get longer, even when the task is simple, e.g. retrieval. It also struggles to utilize memory capacity fully, do associative recall, and perform in-context learning, all of which require solving long-range credit assignment. In contrast, SMT seems agnostic to the sequence length, and is able to solve all of the harder credit assignment problems except associative recall. We attribute these differences to BPTT’s O(T ) credit path length, compared to SMT’s O(1). Further analysis in Section 3.7 confirms the difference in gradient stability in both methods.

4 8 16 32 64 Sequence Length

4 8 16 32 64 4 8 16 32 64 Sequence Length

Difficulty Difficulty

0 1 2 3 4

1 2 3 4 5 1 2 3 4 5

1 2 3 4 5

2 4 6 8 10

2 4 6 8 10 Number of Key-Values

4 8 16 32 62

Difficulty

4 8 16 32 64

0 1 2 3 4

Association Complexity

4 8 16 32 64

In-Context Learning

Association Complexity

State Complexity

4 8 16 32 64

Associative Recall

Association Complexity

4 8 16 32 64 Sequence Length

256 512 4 102 8 204 6 409

4 8 16 32 64

State Complexity

Memory Size

Memory Size 4 8 16 32 64

256 512 4 102 8 204 6 409 4 8 16 32 64

State Tracking

State Complexity

0.0 2 0.0 5 0.0 0.1 0.3

4 8 16 32 64

256 512 4 102 8 204 6 409

Memory Size

Noise Noise

0.0 2 0.0 5 0.0 0.1 0.3

Noise

SMT Encoder *

0.0 2 0.0 5 0.0 0.1 0.3

Memory Capacity

Test Loss

SMT DMT RNN

BPTT RNN

Gradient Stability

0 1 2 3 4 4 8 16 32 62 In Context Examples

Figure 3: Synthetic Task Experiments. We evaluate BPTT, SMT, and SMT→DMT using five synthetic tasks with various settings to probe different properties of the algorithms. ∗ signifies that the SMT Encoder is the teacher Transformer (not an RNN) and is used only as a reference. Across all tasks and task settings, SMT→DMT outperforms BPTT, signaling that SMT has better gradient properties, memory utilization, state tracking, associative recall, and in-context learning than BPTT. MNIST Dataset Samples

3.2 Attneave’s Pixel Sequence Modeling We now evaluate on Attneave’s tasks. Figure 4 shows the stark difference between MNIST samples generated by RNNs trained with BPTT and SMT→DMT. Figure 5 shows images generated by an SMT→DMT RNN trained on Sketchy.

Samples Generated by BPTT RNN (Transformer Backbone)

Along with the synthetic experiments, these results confirm that SMT doesn’t suffer from a recency bias like BPTT, allowing it to properly attribute credit across long sequences.

Samples Generated by SMT DMT RNN (Transformer Backbone)

3.3 Sequential Compute and Data We now evaluate BPTT and SMT across real domains and various RNN architectures. Each method is allowed N optimization steps on token batches of shape B × T (number of sequences × sequence length). We sweep N , B, and T for each method to profile how much sequential compute and data each method uses to achieve a target performance. Sequential compute, measured in sequential FLOPs, is a metric proportional to the amount of inherently serial steps required to do the computation (∼ time it would take on an infinitely parallel computer). Sequential compute is a useful quantity because modern hardware is highly parallel, making it the primary constraint in large-scale model training [55]. Data is measured by the number of tokens processed by the model during training. We elaborate on how sequential FLOPs and data is calculated for each method in Appendix A.

Samples Generated by BPTT RNN (GRU Backbone)

Figure 4: Attneave’s MNIST Generation. BPTT fails to effectively capture the long-range dependencies required for pixel sequence modeling, even with a GRU. SMT→DMT captures these dependencies with a non-gated RNN architecture. More samples are in Appendix Figure 17. Sketchy Dataset Samples

Samples Generated by SMT DMT RNN (Transformer Backbone)

Figure 5: Attneave’s Sketchy Generation. SMT→DMT captures the stroke structure of human-drawn sketches through only pixel seFigure 6 shows the results. In sequential com- quence modeling on sparse images. More samples pute, SMT Encoder and SMT→DMT RNN are are in Appendix Figure 18. significantly more efficient than BPTT with the Transformer and MLP backbones. In data, SMT Encoder and SMT→DMT RNN has approximately 6

TinyStories RNN Backbone → Transformer

MNIST

MLP

1.6

GRU

1.0

Transformer

MLP

GRU

Test Loss

Sequential Compute

1.4 1.2 1.0 0.8 0.6

0.8

0.6

0.4 107

108

107

108

107

108

107

108

SeqFLOPs 1.6

107

108

1.0

Test Loss

1.4

Test Loss

Data Processed

108

SeqFLOPs

1.2 1.0 0.8 0.6

0.8

0.6

0.4 107

107

109

107

109

107

109

Tokens

107

109

Tokens

BPTT RNN Best BPTT Loss

SMT → DMT RNN Best SMT → DMT Loss

SMT Encoder *

Figure 6: Sequential Compute and Data Efficiency. We sweep training hyperparameters for BPTT, SMT, and SMT→DMT and plot the resulting runs’ performance along sequential compute (SeqFLOPs) used and data processed (Tokens), across different RNN architectures and datasets. Runs are capped at one day on an H200 GPU. ∗ signifies that the SMT Encoder is the teacher Transformer (not an RNN) and is used only as a reference. Generally, SMT and SMT→DMT are more efficient than BPTT in sequential compute, and around the same or better efficiency in data. the same data efficiency as BPTT with the Transformer and MLP backbones on TinyStories. However on MNIST, SMT Encoder and SMT→DMT RNN shows significantly better data efficiency. This result is explained by the short vs long range memory information requirements of natural language [30] vs pixel sequence modeling [120]. SMT→DMT is unable to train GRU RNNs, because the GRU architecture induces memory space collapse during SMT training, degrading RNN rollout. 3.4

Scaling Laws

We evaluate the scaling behavior of SMT→DMT along three axes: context length, memory state size, and model parameter count. For the first two, we logarithmically sweep Tc and the number of memory tokens in the Transformer-based RNN (Figure 15). For model scaling, we vary the width and depth of the RNN, encoder, and decoder. We use the TinyStories domain for these experiments. Figure 7 shows that SMT→DMT exhibits smooth, predictable performance improvements with larger context length and bigger memory state size. Together with the previous experiments, these results reaffirm that SMT effectively leverages long contexts and large memory states. Figure 8 presents the parameter scaling results. The SMT encoder follows a standard power-law-like scaling trend. The SMT→DMT RNN also improves smoothly with scale, albeit with a differently shaped scaling curve. Interestingly, the RNN appears to more closely match the encoder’s performance at larger scales. Context and Memory Scaling Test Loss

0.7 0.6 16

128

Context Length, Tc

256

512

Figure 7: Scaling Context and Memory. SMT→DMT shows smooth performance improvements as you increase the context length and the memory size in TinyStories.

SMT Encoder * Test Loss

Memory Size 256 512 1024 2048 4096

0.8

Model Scaling

1.2

1.0

SMT DMT RNN

0.8 Model Width

0.6

0.4

64 128 256 512

106

107

108

0.4

Model Parameters

64 128 256 512

106

107

108

Figure 8: Scaling Model Size. Sweeping the width and depth of the RNN and teacher shows smooth performance improvements in TinyStories. The RNN imitates the teacher performance better at larger scale. 7

3.5

Compression as a Scaling Axis

Scaling Laws for Compression

4096

1.0 7

0.5

0.90

Neural scaling laws predict the relationship between a resource (e.g. compute, data) and a desired property (e.g. validation loss, benchmark accuracy) [65, 54]. Can the desired property instead be compression [58]? For RNNs, compression can be interpreted as achieving the same performance with a smaller memory state size. Thus, to answer this question, we train a set of SMT models on TinyStories across a sweep of memory state sizes and training compute budgets.

Test Loss

1024

0.63

0.97

512

0.87

0.83 0.80 0.77 0.73

256 1.00×1014

Figure 9 shows the scaling curve, confirming that SMT can achieve more compression when allocated more compute. Since compression is often speculated as being a core property of intelligent systems [113, 69], scaling along this compression axis may be a desired direction forward for future sequence models. Notably, Transformers perform no compression of the past [38], which may explain their training efficiency. 3.6

0.6 0.6

0.93

Memory Size

2048

2.65×1014

FLOPs

7.04×1014

0.70

1.87×1015

0.5

Figure 9: Scaling Laws for Compression. We plot iso-loss contours for SMT-trained encoder models across a range of memory state sizes and training compute budgets. For a fixed target performance, SMT can achieve higher compression (smaller memory size) using additional compute. This result suggests a new property to scale when given more training compute: memory state compression.

Ablations

Predictive State and Detached RNN The impact of the predictive state objective (Equation 2) is evaluated by sweeping the future length Tf , while keeping Tc large enough to see the whole sequence. The impact of the dynamics objective (Equation 3) on memory representation is tested by detaching the model computation graph with stop grads at two locations such that the gradients from Ldyn t flow to the RNN, but not the encoder (detached); the non-detached SMT baseline is referred to as joint. This ablation isolates the contribution of explicitly training mt to be a Markovian representation. Figure 10 shows the results on the needle retrieval task. To solve the task, and thus have proper credit assignment, SMT requires either large enough Tf or joint training. When Tf is large enough, there is a O(1) credit path length between the needle and the answer at all timesteps. Interestingly, when Tf is small, there exists no credit path to learn early timestep memories, yet joint training still learns effectively, even when Tf = 1. Credit must be propagating through the RNN dynamics from mT to mT −1 , and so on, to m0 . But because the RNN is never unrolled, there is no computation graph for credit to propagate directly. The only explanation is that credit is being amortized into gradient optimization steps. Each optimization step sends information from mt to mt−1 through fθ ; T such gradients steps sends information T steps back in the sequence. This implies that solving T sequence length credit assignment task when Tf = 1, requires at least T gradient optimization steps. This credit amortization phenomenon is reminiscent of value bootstrapping in RL [119].

Needle Retrieval T = 64

4 8 16 Future Length, Tf

Figure 10: Joint SMT Ablation. Here, the task requires credit assignment across T timesteps. When the RNN is detached during SMT, Tf must be large enough to capture the task signal (Tf = T ). With joint training, SMT solves the task even when Tf is small.

10 1 10 4 10 7 10 10

BPTT SMT

Start of Sequence

Detached SMT Joint SMT

Failed Solved

Gradient Magnitude of Memory mt, || mLt ||

Gradient Properties at Initialization

102

Exploding Gradients

Vanishing Gradients 32 48 64 80 Timestep in Sequence, t

Stable Gradients 96

112

End of Sequence

SMT DMT RNN Loss

λ Coefficients The values of λdyn and λunif are swept here to check their effects. Figure 16 shows the results. The best RNNs require λdyn = 0.1, and λunif = 0.001. When λunif = 0, although the RNN performance is preserved, the memory space is collapsed, as indicated by Lunif .

128

Figure 11: Gradient Properties of BPTT and SMT. In the needle retrieval task, the loss is applied at the last timestep. BPTT propagates gradients backward through all timesteps, risking vanishing/exploding gradients for each mt , depending on the weight initialization. SMT is non-recurrent and has a O(1) credit path length, making its gradients agnostic to initialization and time-horizon. 8

0.8 0.6 0.4 0.2 0.0

RNN Before DMT RNN After DMT

RNN Loss

R 2 Error of mt, mt

DMT Improves Performance 4 3

RNN Before DMT RNN After DMT

1 0.6

100 150 Timestep in Sequence, t

200

250

0.6

103

Rollout R 2 Error of mt, mt

DMT Reduces Drift

1.0

1 2 3 RNN Loss Before DMT

Rollout Drift is not Fully Predicted RNN Before DMT RNN After DMT

102 101

Correlation=0.50

100 10 1

Correlation=0.45 10 3

10 2

10 1

100

One Step R 2 Error of mt, mt

Figure 12: Impact of DMT across many runs with different SMT λdec and λdyn hyperparameters. Left: Applying DMT reduces the drift of the RNN rollout (measured with 1 − R2 of RNN memory prediction m̂t of encoder ground truth mt ). Middle: DMT significantly improves RNN performance across settings. Right: The one-step drift of the RNN only partially correlates with the rollout drift. Drift and DMT As described in Section 2.3, RNN suffers from drift post-SMT. Figure 12 shows an analysis of drift and DMT’s mitigation of it. From a dynamical systems perspective, DMT seems to discover RNNs which have an initially higher drift, but which plateau at a much lower equilibrium drift. Interestingly, this equilibrium drift value is not fully predicted by the one-step drift, inviting future investigations into predicting and mitigating rollout drift in one-step during SMT. 3.7

Analysis

Gradient Properties of BPTT and SMT The fundamental difference between BPTT and SMT in long-range credit assignment is dictated by their gradients. Figure 11 shows the gradient magnitude ∂L of mt , ∥ ∂m ∥, at different t for both methods with different model weight initializations on the needle t retrieval task. In BPTT, gradients vanish or explode over time, due to BPTT’s gradient propagation through recurrent modules. In SMT, gradient magnitudes are independent of t, because the credit path length between tokens is independent of the sequence length. This result explains why SMT does not suffer a recency bias and is able to do stably perform long-horizon credit assignment.

Sequence Length Generalization

0.4 Test Loss

Benefit of RNNs over Transformers SMT trains an RNN to mimic a Transformer encoder model, raising the question of why an RNN is needed at all, given the Transformer. RNNs are qualitatively more efficient than Transformers at inference, requiring O(1) rather than O(T ) memory and compute per generated token (Table 1). RNNs also constitute a more expressive class of models [84, 77].

Transformer Encoder SMT DMT RNN Training Sequence Length

0.3 0.2 0.1 0.0

128

256

512

Here, we compare an SMT→DMT RNN against a Sequence Length, T Transformer on the synthetic stack state tracking task. For a fair comparison, we use the SMT encoder as Figure 13: Sequence Length Generalizathe Transformer baseline, since it imposes the same tion. An SMT→DMT trained RNN generalmemory-information bottleneck as the RNN. izes better than its Transformer teacher when Figure 13 shows the Transformer outperforms the evaluated on sequence lengths longer than RNN on training sequence lengths, but significantly training. The task is synthetic state tracking. underperforms the RNN on sequence lengths longer than training. Prior work on length generalization reports similar findings [99]. This result reflects the distinct inductive biases of the architectures: Transformers behave like growing lookup tables in context, while RNNs update finite states [38]. The latter is a better inductive bias for generalization. Memory Space To better understand what SMT is learning, we train smaller SMT models that have a 2D memory state and directly visualize their memory space across three synthetic tasks in Figure 14. In the retrieval tasks, SMT learns to collapse many sequence states into only a few effective memory states: an initial state, a state indicating the next token is the needle, and states corresponding to the needle value. Then, the RNN learns finite-state machine behavior to transition between these states. In contrast, string copying requires lossless sequence compression and thus SMT cannot alias distinct memory states together. It learns to create a tree-like memory geometry to store all possible sequences, matching the tree structure of all possible strings. Figure 21 and Figure 22 show memory visualizations for models trained on MNIST. These results indicate SMT memories form effective temporal abstractions of the past depending on what the future requires. 9

Memory Space Axis 1

Memory Space for Retrieval (2 Needles) Needle =1

Init

Memory Space for Retrieval (5 Needles) 4

Needle =2

1 3

Needle Coming

Memory Space for String Copy 2 1 Init 2212 21 11 32

31 13 23

Init Memory Space Axis 2

Encoder Memories Annotated Memory State

Memory Transitions RNN Transition Field (for input token = 2)

Figure 14: Memory Space Visualization. The encoder learns different memory geometries for different tasks. In Retrieval, the encoder collapses many sequence states into a few memory states, creating finite-state machine like behavior. In String Copy, the encoder constructs a tree-like memory geometry to compress all possible sequences. Some geometries induce more complex RNN transition fields.

Related Works

Recurrent Neural Networks (RNNs) RNNs were studied extensively early in AI because their recurrence mechanism resembles biological brains [80] and can be applied to any sequential task [28]. Many different algorithms were proposed for learning, including random guessing [108], evolutionary algorithms [87, 3, 115, 104, 106], hebbian learning [45, 56, 86, 92], real-time recurrent learning [132], and other algorithms [93, 10, 70, 8, 63]. BPTT is the only widely adopted algorithm [131]. However, it has repeatedly been shown that BPTT produces unstable gradients that vanish, explode, or exhibit high variance [11, 53, 97, 7]. Several directions address this issue. One direction focuses on architectural modifications, including residual connections [114, 44] and gating mechanisms [18], culminating in the development of the LSTM [52] and GRU [17]. A parallel direction addressed gradient instability through orthogonal weight parameterizations to prevent exponential growth or decay across time [107, 4, 134, 85, 126, 47]. Others explored external memory [98, 37, 41], hierarchical modeling [48, 19, 127, 61], other unique directions [60, 78, 88]. Recently, there has been renewed interest in RNNs in the form of linear state space models [40, 112, 39], linear attention models [67, 117, 24, 137, 136], and even nonlinear RNN models [9, 15, 90, 94]. Recurrent computation more generally has been reappearing across paradigms including in diffusion [51], looped Transformers [34], and reasoning [42, 33]. Time-Parallel Training Transformers revolutionized sequence modeling [14] largely because they have time-parallel training [124] (unlike prior attention methods [6]), which is crucial for leveraging modern hardware [49, 55] to scale performance [65, 54]. Linear RNNs gained popularity [67, 24, 137, 31] after it was realized they can be parallelized with the associative scan algorithm [13, 79, 129]. A recent line of work attempts to parallelize nonlinear RNNs as well [74, 22]. Rather than computing [m0 , . . . , mT ] with mt+1 = fθ (mt ), they formulate the forward pass as an iterative optimization procedure. Starting with an initial guess [m00 , . . . , m0T ], they construct a system of T equations, {mt+1 − fθ (mt ) = 0}Tt=0 , and solve this system with Newton’s method [95]. Many works have further built on this approach [23, 35]. Although appealing, this approach approximates BPTT and hence will suffer from its O(T ) credit path length and corresponding gradient instability, along with the added convergence worries of Newton’s method [36]. In contrast, SMT uses an encoder to train mt to be a predictive state while satisfying mt+1 ≈ fθ (mt ) and providing an O(1) credit path length. Computation Complexity Class of Models A model’s architecture determines the problems it can theoretically solve [57]. Some tasks are inherently sequential and cannot be efficiently parallelized [2]; the circuit depth of a task is the minimum number of sequential steps required to solve it on an infinitely parallel computer [103, 20]. Every neural network has a corresponding sequential depth— the longest nonlinear computation path from input to output—which bounds the class of problems it can solve [82]. Models with constant or logarithmic sequential depth per layer, such as Transformers and linear RNNs, are provably limited to tasks with equivalently low circuit depth [81, 82, 138]. While such models succeed on tasks amenable to parallelization (e.g. parity tracking via associative scan [76, 72]), they systematically fail on tasks requiring deep sequential computation (e.g. tracking a chess board [83]). Interestingly, the aspect that makes models parallelizable, limits their performance on harder problems [81]. Nonlinear RNNs are one of the few classes of models where its sequential depth grows with the input sequence length [84]. Although these constraints were seen as theoretical, there is growing evidence they affect models in practice as well [77]. 10

In SMT, we train a nonlinear RNN (which is fully expressive), using a time-parallel teacher Transformer (which has limits). We note this limitation but argue that SMT is a pretraining algorithm, which should be used with a lightweight post-training algorithm to solve downstream tasks [32]. Predictive State Representations (PSRs) A PSR is a way of modeling a partially observed dynamical system by representing its state only in terms of predictions about future observations [75, 110], a representation that is sufficient for optimal decision making [111]. Early works interpreted PSRs as a literal vector of probabilities of future events, but have since been generalized [25]. Belief states are a related concept, which also defines a sufficient statistic of the past [62]. PSRs have been previously incorporated into RNNs [26, 46]. Venkatraman et al. [125] introduce an auxiliary objective for RNNs that trains hidden states to predict statistics of future observations using a decoder. However, these works still unroll the RNN and use BPTT, and thus are not time-parallelizable and have a O(T ) credit path. Other Related Work Our work is related to the literature on cross-architecture teacher-student distillation [66, 128, 43, 16, 91], but these works do not address the challenges of training nonlinear RNNs. In concurrent work, Teoh et al. [122] introduced Next-Latent Prediction (NextLat), which trains an RNN with memory state supervision from a Transformer. With a particular setting of the hyperparameters, SMT and NextLat are equivalent. However, Teoh et al. [122] focuses its experiments on the case where the latent representations are still trained with BPTT, which can be optionally truncated to T = 1, whereas we focus primarily on the T = 1 case. Additionally, NextLat’s goal is to regularize the Transformer to learn compact world models, rather than training the RNN without BPTT. The Recurrent Transformer is an RNN architecture that attends to all past hidden states, creating an O(1) gradient path that stabilizes credit assignment [94]. However, because it retains all past hidden states, its memory grows unboundedly during inference—making it more akin to a Transformer than a fixed-memory RNN. Crucially, training still requires sequential unrolling and BPTT. SMT, by contrast, replaces BPTT and supports arbitrary fixed-memory RNN architectures and enables time-parallel training by never unrolling the RNN. Other works similarly combine Transformers with recurrent processing but also train with sequential unrolling and BPTT [29, 15, 135]. A new line of work uses principles from diffusion models to train blocks of a feed-forward network in parallel, avoiding global backpropagation [73, 109].

Discussion

In SMT, the teacher model is time-parallel, and is thus constrained in expressivity [81], implying that SMT-trained RNNs may suffer the same problem. Therefore, BPTT finetuning might be required to achieve expressivity beyond the teacher. Additionally, while SMT is useful for learning how to encode sequences, it is not necessarily to be used for learning reasoning since intermediate steps are not supervised. The same limitation applies to Transformers yet post-training allows them to effectively solve longer-horizon tasks than the training horizon; the same might be true for SMT-trained RNNs. The current SMT variant computes and trains only a single mt within a sequence. We found that training on all memories [m0 , . . . , mT ] offered no improvement in our settings, but this may not hold at larger scales. After SMT, the RNN experiences drift away from the teacher memory trajectory. DMT provides one solution but is not time-parallel; however, it may be parallelized via DEER [74]. RNNs have the promise of solving problems that extend over unbounded horizons, such as the entire lifetime of an agent. However, training methods for RNN have been hindered by the inability of BPTT to assign credit effectively over such a long horizon. Our method circumvents the credit assignment issue with an O(1) connection path. In the regimes we studied, this effectively allows for learning memories that are only useful many steps later, an ability that is crucial for lifelong learning.

Acknowledgments and Disclosure of Funding This work was supported by an NSF GRFP Fellowship to A.K., a Packard Fellowship and Sloan Research Fellowship to P.I., and ONR MURI grant N00014-22-1-2740. This work was also supported under project ID 43 as part of the Swiss AI Initiative, through a grant from the ETH Domain and computational resources provided by the Swiss National Supercomputing Centre (CSCS) under the Alps infrastructure. We thank Alyosha Efros for suggesting the Attneave framing for pixel sequence modeling and recommending the Sketchy dataset. We thank Alexander Huth for initially motivating A.K. to work on memory many years ago. We thank Assaf Ben-Kish for reviewing an earlier draft of this manuscript. We thank Han Guo and Oliver Sieberling for technical advice on algorithmic complexity.

References [1] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022. [2] Gene M Amdahl. Validity of the single processor approach to achieving large scale computing capabilities. In Proceedings of the April 18-20, 1967, spring joint computer conference, pages 483–485, 1967. [3] Peter J Angeline, Gregory M Saunders, and Jordan B Pollack. An evolutionary algorithm that constructs recurrent neural networks. IEEE transactions on Neural Networks, 5(1):54–65, 1994. [4] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference on machine learning, pages 1120–1128. PMLR, 2016. [5] Fred Attneave. Some informational aspects of visual perception. Psychological review, 61(3): 183, 1954. [6] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2016. URL https://arxiv.org/abs/1409.0473. [7] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.01271, 2018. [8] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in neural information processing systems, 32, 2019. [9] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. Advances in Neural Information Processing Systems, 37: 107547–107603, 2024. [10] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks, 2015. URL https://arxiv.org/abs/ 1506.03099. [11] Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157–166, 1994. [12] Max S Bennett. A brief history of intelligence: evolution, AI, and the five breakthroughs that made our brains. HarperCollins, 2023. [13] Guy E Blelloch. Prefix sums and their applications. 1990. [14] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877– 1901, 2020. 12

[15] Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:11079–11091, 2022. [16] Yingfa Chen, Zhen Leng Thai, Zihan Zhou, Zhu Zhang, Xingyu Shen, Shuo Wang, Chaojun Xiao, Xu Han, and Zhiyuan Liu. Hybrid linear attention done right: Efficient distillation and effective architectures for extremely long contexts, 2026. URL https://arxiv.org/abs/ 2601.22156. [17] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder–decoder approaches. In Proceedings of SSST-8, eighth workshop on syntax, semantics and structure in statistical translation, pages 103–111, 2014. [18] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014. [19] Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint arXiv:1609.01704, 2016. [20] Stephen A Cook. A taxonomy of problems with fast parallel algorithms. Information and control, 64(1-3):2–22, 1985. [21] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th annual meeting of the association for computational linguistics, pages 2978–2988, 2019. [22] Federico Danieli, Miguel Sarabia, Xavier Suau Cuadros, Pau Rodriguez, and Luca Zappella. Deeppcr: Parallelizing sequential operations in neural networks. Advances in Neural Information Processing Systems, 36:47598–47625, 2023. [23] Federico Danieli, Pau Rodriguez, Miguel Sarabia, Xavier Suau, and Luca Zappella. Pararnn: Unlocking parallel training of nonlinear rnns for large language models, 2025. URL https: //arxiv.org/abs/2510.21450. [24] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024. [25] Carlton Downey, Ahmed Hefny, and Geoffrey Gordon. Practical learning of predictive state representations. arXiv preprint arXiv:1702.04121, 2017. [26] Carlton Downey, Ahmed Hefny, Boyue Li, Byron Boots, and Geoffrey Gordon. Predictive state recurrent neural networks, 2017. URL https://arxiv.org/abs/1705.09353. [27] Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.07759, 2023. [28] Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):179–211, 1990. [29] Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Addressing some limitations of transformers with feedback memory. arXiv preprint arXiv:2002.09402, 2020. [30] Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling?, 2025. URL https://arxiv.org/abs/2410.23771. [31] Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadeghi. Were rnns all we needed? arXiv preprint arXiv:2410.01201, 2024. [32] Yulu Gan and Phillip Isola. Neural thickets: Diverse task experts are dense around pretrained weights. arXiv preprint arXiv:2603.12228, 2026. 13

[33] Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05171, 2025. [34] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pages 11398–11442. PMLR, 2023. [35] Xavier Gonzalez, Andrew Warrington, Jimmy T Smith, and Scott W Linderman. Towards scalable and stable parallelization of nonlinear rnns. Advances in Neural Information Processing Systems, 37:5817–5849, 2024. [36] Xavier Gonzalez, Leo Kozachkov, David M Zoltowski, Kenneth L Clarkson, and Scott W Linderman. Predictability enables parallelization of nonlinear state space models. arXiv preprint arXiv:2508.16817, 2025. [37] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014. [38] Albert Gu. On the tradeoffs of state space models and transformers, 2025. URL https: //goombalab.github.io/blog/2025/tradeoffs/. [39] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. [40] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. [41] Steven Stenberg Hansen. Long timescale credit assignment in neuralnetworks with external memory, 2017. URL https://arxiv.org/abs/1701.03866. [42] Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models to reason in a continuous latent space, 2025. URL https://arxiv.org/abs/2412.06769. [43] Lukas Hauzenberger, Niklas Schmidinger, Thomas Schmied, Anamaria-Roberta Hartl, David Stap, Pieter-Jan Hoedt, Maximilian Beck, Sebastian Böck, Günter Klambauer, and Sepp Hochreiter. Effective distillation to hybrid xlstm architectures, 2026. URL https://arxiv. org/abs/2603.15590. [44] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [45] Donald Olding Hebb. The organization of behavior: A neuropsychological theory. Psychology press, 1949. [46] Ahmed Hefny, Zita Marinho, Wen Sun, Siddhartha Srinivasa, and Geoffrey Gordon. Recurrent predictive state policy networks, 2018. URL https://arxiv.org/abs/1803.01489. [47] Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference on Machine Learning, pages 1969–1978. PMLR, 2018. [48] Salah Hihi and Yoshua Bengio. Hierarchical recurrent neural networks for long-term dependencies. Advances in neural information processing systems, 8, 1995. [49] W Daniel Hillis and Guy L Steele Jr. Data parallel algorithms. Communications of the ACM, 29(12):1170–1183, 1986. [50] Geoffrey E Hinton and James A Anderson. Parallel models of associative memory: updated edition. Psychology press, 2014. 14

[51] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. [52] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9 (8):1735–1780, 1997. [53] Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001. [54] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, DDL Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 10, 2022. [55] Sara Hooker. The hardware lottery. Communications of the ACM, 64(12):58–65, 2021. [56] John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554–2558, 1982. [57] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359–366, 1989. [58] Marcus Hutter. Universal artificial intelligence: Sequential decisions based on algorithmic probability, volume 300. Springer, 2005. [59] Mozes Jacobs, Thomas Fel, Richard Hakim, Alessandra Brondetta, Demba Ba, and T Andy Keller. Block-recurrent dynamics in vision transformers. arXiv preprint arXiv:2512.19941, 2025. [60] Herbert Jaeger. The “echo state” approach to analysing and training recurrent neural networkswith an erratum note. Bonn, Germany: German national research center for information technology gmd technical report, 148(34):13, 2001. [61] Alexia Jolicoeur-Martineau. Less is more: Recursive reasoning with tiny networks, 2025. URL https://arxiv.org/abs/2510.04871. [62] Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99–134, 1998. [63] Anil Kag and Venkatesh Saligrama. Training recurrent neural networks via forward propagation through time. In International Conference on Machine Learning, pages 5189–5200. PMLR, 2021. [64] Eric R Kandel. Principles of neural science, 2000. [65] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. [66] Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into rnns, 2021. URL https://arxiv.org/abs/2103.13076. [67] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. [68] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint arXiv:2212.04458, 2022. [69] Andrei Nikolaevic Kolmogorov. Three approaches to the quantitative definition of information. International journal of computer mathematics, 2(1-4):157–168, 1968. [70] Alex M Lamb, Anirudh Goyal ALIAS PARTH GOYAL, Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio. Professor forcing: A new algorithm for training recurrent networks. Advances in neural information processing systems, 29, 2016. 15

[71] Yann LeCun and Corinna Cortes. The MNIST database of handwritten digits, 1998. URL http://yann.lecun.com/exdb/mnist/. [72] Belinda Z. Li, Zifan Carl Guo, and Jacob Andreas. (how) do language models track state?, 2025. URL https://arxiv.org/abs/2503.02854. [73] Qinyu Li, Yee Whye Teh, and Razvan Pascanu. Noprop: Training neural networks without full back-propagation or full forward-propagation. arXiv preprint arXiv:2503.24322, 2025. [74] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length, 2024. URL https://arxiv.org/ abs/2309.12252. [75] Michael Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, 14, 2001. [76] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749, 2022. [77] Yuxi Liu, Konpat Preechakul, Kananart Kuwaranancharoen, and Yutong Bai. The serial scaling hypothesis. arXiv preprint arXiv:2507.12549, 2025. [78] Mantas Lukoševičius and Herbert Jaeger. Reservoir computing approaches to recurrent neural network training. Computer science review, 3(3):127–149, 2009. [79] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length, 2018. URL https://arxiv.org/abs/1709.04057. [80] Warren S McCulloch and Walter Pitts. A logical calculus of the ideas immanent in nervous activity. The bulletin of mathematical biophysics, 5(4):115–133, 1943. [81] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531–545, 2023. [82] William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought, 2024. URL https://arxiv.org/abs/2310.07923. [83] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.08819, 2024. [84] William Merrill, Hongjian Jiang, Yanhong Li, and Ashish Sabharwal. Why are linear rnns more parallelizable? arXiv preprint arXiv:2603.03612, 2026. [85] Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey. Efficient orthogonal parametrisation of recurrent neural networks using householder reflections. In International Conference on Machine Learning, pages 2401–2409. PMLR, 2017. [86] Thomas Miconi, Jeff Clune, and Kenneth O. Stanley. Differentiable plasticity: training plastic neural networks with backpropagation, 2018. URL https://arxiv.org/abs/1804.02464. [87] Geoffrey F Miller, Peter M Todd, and Shailesh U Hegde. Designing neural networks using genetic algorithms. In ICGA, volume 89, pages 379–384, 1989. [88] John Miller and Moritz Hardt. Stable recurrent models. arXiv preprint arXiv:1805.10369, 2018. [89] Marvin Minsky. Steps toward artificial intelligence. Proceedings of the IRE, 49(1):8–30, 1961. [90] Mayank Mishra, Shawn Tan, Ion Stoica, Joseph Gonzalez, and Tri Dao. M2 rnn: Nonlinear rnns with matrix-valued states for scalable language modeling. arXiv preprint arXiv:2603.14360, 2026. [91] Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodríguez, Luca Zappella, and Federico Danieli. Attention to mamba: A recipe for cross-architecture distillation. arXiv preprint arXiv:2604.14191, 2026. 16

[92] Elias Najarro and Sebastian Risi. Meta-learning through hebbian plasticity in random networks, 2022. URL https://arxiv.org/abs/2007.02686. [93] Yann Ollivier, Corentin Tallec, and Guillaume Charpiat. Training recurrent networks online without backtracking. arXiv preprint arXiv:1507.07680, 2015. [94] Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi, Alexandru Meterez, Mujin Kwun, and Sham Kakade. The recurrent transformer: Greater effective depth and efficient decoding. arXiv preprint arXiv:2604.21215, 2026. [95] James M Ortega and Werner C Rheinboldt. Iterative solution of nonlinear equations in several variables. SIAM, 2000. [96] Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. How to construct deep recurrent neural networks. arXiv preprint arXiv:1312.6026, 2013. [97] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310–1318. Pmlr, 2013. [98] Leonid Peshkin, Nicolas Meuleau, and Leslie Kaelbling. Learning policies with external memory. arXiv preprint cs/0103003, 2001. [99] Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021. [100] Shauli Ravfogel, Yoav Goldberg, and Tal Linzen. Studying the inductive biases of rnns with synthetic variations of natural languages. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 3532–3542, 2019. [101] Stephane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning, 2011. URL https://arxiv.org/ abs/1011.0686. [102] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, 323(6088):533–536, 1986. [103] Walter L Ruzzo. On uniform circuit complexity. Journal of Computer and System Sciences, 22(3):365–383, 1981. [104] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as a scalable alternative to reinforcement learning. arXiv preprint arXiv:1703.03864, 2017. [105] Patsorn Sangkloy, Nathan Burnell, Cusuh Ham, and James Hays. The sketchy database: learning to retrieve badly drawn bunnies. Acm Transactions on Graphics (TOG), 35(4):1–12, 2016. [106] Bidipta Sarkar, Mattie Fellows, Juan Agustin Duque, Alistair Letcher, Antonio León Villares, Anya Sims, Clarisse Wibault, Dmitry Samsonov, Dylan Cope, Jarek Liesen, et al. Evolution strategies at the hyperscale. arXiv preprint arXiv:2511.16652, 2025. [107] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013. [108] Jürgen Schmidhuber, Sepp Hochreiter, and Yoshua Bengio. Evaluating benchmark problems by random guessing. A Field Guide to Dynamical Recurrent Networks, pages 231–235, 2001. [109] Makoto Shing, Masanori Koyama, and Takuya Akiba. Diffusionblocks: Block-wise neural network training via diffusion interpretation. arXiv preprint arXiv:2506.14202, 2025. [110] Satinder Singh, Michael James, and Matthew Rudary. Predictive state representations: A new theory for modeling dynamical systems. arXiv preprint arXiv:1207.4167, 2012. 17

[111] Satinder P Singh, Michael L Littman, Nicholas K Jong, David Pardoe, and Peter Stone. Learning predictive state representations. In Proceedings of the 20th International Conference on Machine Learning (ICML-03), pages 712–719, 2003. [112] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933, 2022. [113] Ray J Solomonoff. A formal theory of inductive inference. part i. Information and control, 7 (1):1–22, 1964. [114] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint arXiv:1505.00387, 2015. [115] Kenneth O Stanley and Risto Miikkulainen. Evolving neural networks through augmenting topologies. Evolutionary computation, 10(2):99–127, 2002. [116] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. [117] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. [118] Ilya Sutskever, James Martens, and Geoffrey E Hinton. Generating text with recurrent neural networks. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 1017–1024, 2011. [119] Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998. [120] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers, 2020. URL https://arxiv.org/abs/2011.04006. [121] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022. [122] Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Pratyusha Sharma, Riashat Islam, Alex Lamb, and John Langford. Next-latent prediction transformers learn compact world models, 2025. URL https://arxiv.org/abs/2511.05963. [123] Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pages 1747–1756. PMLR, 2016. [124] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [125] Arun Venkatraman, Nicholas Rhinehart, Wen Sun, Lerrel Pinto, Martial Hebert, Byron Boots, Kris M. Kitani, and J. Andrew Bagnell. Predictive-state decoders: Encoding the future into recurrent networks, 2017. URL https://arxiv.org/abs/1709.08520. [126] Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. In International conference on machine learning, pages 3570–3578. PMLR, 2017. [127] Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori. Hierarchical reasoning model, 2025. URL https://arxiv.org/ abs/2506.21734. [128] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models, 2025. URL https://arxiv.org/ abs/2408.15237. 18

[129] Shang Wang, Yifan Bai, and Gennady Pekhimenko. Bppsa: Scaling back-propagation by parallel scan algorithm, 2020. URL https://arxiv.org/abs/1907.10134. [130] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International conference on machine learning, pages 9929–9939. PMLR, 2020. [131] Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):1550–1560, 1990. [132] Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270–280, 1989. [133] David J Willshaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. holographic associative memory. Nature, 222(5197):960–962, 1969.

Non-

[134] Scott Wisdom, Thomas Powers, John Hershey, Jonathan Le Roux, and Les Atlas. Full-capacity unitary recurrent neural networks. Advances in neural information processing systems, 29, 2016. [135] Qingyang Wu, Zhenzhong Lan, Kun Qian, Jing Gu, Alborz Geramifard, and Zhou Yu. Memformer: A memory-augmented transformer for sequence modeling. In Findings of the association for computational linguistics: AACL-IJCNLP 2022, pages 308–318, 2022. [136] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, 2025. URL https://arxiv.org/abs/2412.06464. [137] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length, 2025. URL https://arxiv.org/ abs/2406.06484. [138] Morris Yau, Sharut Gupta, Valerie Engelmayer, Kazuki Irie, Stefanie Jegelka, and Jacob Andreas. Sequential-parallel duality in prefix scannable models, 2026. URL https://arxiv. org/abs/2506.10918.

Definitions

Credit Assignment Path Length For any differentiable computational graph, backpropagation propagates gradients from the scalar loss backward through the graph to each leaf node (typically model weights). We define the credit assignment path length as the maximum distance between any two nodes (e.g. tokens) in the computation graph. Distance is measured as the number of intervening non-identity operations that modify gradients (e.g. matrix multiplications or nonlinearities). The longer this path, the less effective backpropagation is for properly learning associations between distant nodes and assigning credit [11]. Under this definition, BPTT has O(T ) credit assignment path length, whereas Transformers and SMT have O(1) path length between any two tokens. Sequential Computation (measured in SeqFLOPs) Sequential computation is the amount of serial (non-parallelizable) work required to complete a computation. Some computations may require substantial total work but little sequential work (e.g. matrix multiplication). As parallel hardware such as GPUs continues to scale, total work matters less than the amount of inherently sequential work required [55]. Sequential compute is measured by analyzing the computation graph required to execute an algorithm and computing the graph’s critical path: the number of floating point operations that must be executed sequentially on an infinitely parallel computer. We refer to this quantity as sequential FLOPs (SeqFLOPs). For simplicity, we estimate SeqFLOPs by counting the number of sequential atomic deep learning operations (e.g., Linear, LayerNorm) executed over the course of the algorithm. The true SeqFLOPs, measured in floating-point operations, is proportional to this estimate. We compute SeqFLOPs for BPTT, SMT, and DMT. SMT is fully parallelizable in time, incurring O(1) SeqFLOPs per optimization step, independent of the sequence length T . In contrast, BPTT and DMT require unrolling the RNN, which increases SeqFLOPs to O(T ) per optimization step. Data Processed (measured in Tokens) Data Processed is defined as the total number of tokens processed during training, including repeated tokens in multi-epoch settings.

Experiment Details

B.1

Architectures

Our primary architecture used for most experiments is shown in Figure 15. Encoder Architecture We use the same encoder model architecture across all experiments. The model begins with an embedding layer to embed input tokens, and a list of learned memory token registers. The input and register tokens are concatenated and processed by a stack of bidirectional (full attention mask) Transformer blocks. We use bidirectional model because the goal is to create a holistic representation of the entire input sequence. The register tokens are then interpreted as memory tokens at the output. Note that a single memory consists of a list of memory tokens, mt = [m1t , . . . , mM t ]. Within a Transformer block, we use rotary position encodings [116] and RMSNorm instead of LayerNorm. We perform RMSNorm on the output memory tokens for stability. Figure 15 shows the encoder architecture. Decoder Architecture We use the same decoder model architecture across all experiments. The decoder has an embedding layer to embed input tokens, which is weight shared with the encoder model. The memory tokens from the encoder and the embedded future input tokens are concatenated, and then processed by a stack of causally masked Transformer blocks. We use a causal mask because the goal is to learn a generative model of the output sequence. Within the each Transformer block, we use rotary position encodings [116] and RMSNorm instead of LayerNorm. The output predictions are read out at token positions such that ŷt+k is a function of only mt and xt+1 , . . . , xt+k . Figure 15 shows the decoder architecture. Transformer-based RNN Architecture The Transformer-based RNN is our primary RNN architecture and is used for most experiments. It begins with an embedding layer for the current timestep’s input token. The memory tokens are concatenated with the input token and then processed by a stack of bidirectional (full attention mask) Transformer blocks to produce the output memory tokens. We perform RMSNorm on the output memory tokens. 20

… yT̂

Decoder

(Causal Transformer)

mt1

… mtM

Encoder (Bidirectional Transformer)

…

xt+1 … xT

RNN Architecture mt1 mtM xt+1

RNN Readout Architecture

Updater

fθ

(Bidirectional Transformer)

yt̂

m̂ 1t+1

…

̂ yt+1

…

yt̂

Encoder-Decoder Architecture

m̂ M t+1

Readout

(Bidirectional Transformer)

mt1

… mtM

…

Figure 15: Model Architecture for SMT. Left: The encoder reads the input context tokens and a set of learned register tokens, and outputs the memory, mt , which is a set of memory tokens. The decoder takes in this memory and the future input tokens and predicts the future output tokens, using a causal mask. This setup forces information from the context to be compressed into a memory that is useful for predicting the future outputs, given future inputs. Middle: Our RNN maps (mt , xt+1 ) to mt+1 using a Transformer-backbone. Since the memory is a list of tokens and the input is a token, we simply use a full attention Transformer to transform the current memory into the next timestep’s memory. Right: Readout is performed by a full attention Transformer over the memory tokens.

MLP-based RNN Architecture The MLP-based RNN flattens the list of memory tokens into a single vector, concatenates it with the input token embedding, and passes them through an MLP. At the output, the model then unflattens them to be a list of memory tokens again. We perform RMSNorm on the output memory tokens. GRU-based RNN Architecture The GRU-based RNN processes the M memory tokens with a M layer stacked GRU. Layer l reads a single memory token, mlt , and outputs a single memory token for the next timestep, ml+1 t . We do not RMSNorm on the output memory tokens, since that would undermine the GRU’s residual structure. RNN Readout Architecture We use the same readout model architecture for all RNNs. The readout architecture takes in the memory tokens and processes them through a stack of bidirectional (full attention mask) Transformer blocks Figure 15 shows the readout architecture. B.2 B.2.1

Datasets Synthetic Tasks

Retrieval to test Gradient Stability The retrieval task requires the model to remember and reproduce the token immediately following a designated identifier (token 0). For example, given x = [3, 4, 0, 2, 1, 3, 1, 0] the target is y = [∅, ∅, ∅, ∅, ∅, ∅, ∅, 2], where ∅ denotes no prediction target. With probability p, the label is corrupted to a random token. By varying sequence length and noise level, this task probes the algorithm’s capacity for stable gradient credit assignment. String Copy to test Memory Capacity The string copy task requires the model to reproduce a sequence in reverse order after a delimiter (token 0). For example, given x = [3, 4, 1, 2, 0, 0, 0, 0], the target is y = [∅, ∅, ∅, ∅, 2, 1, 4, 3]. By varying the sequence length and the memory state size, this task measures the algorithm’s ability to leverage the RNN’s memory capacity for memorization. Stack Operations to test State Tracking The stack operations task requires the model to track the top element of a stack through a sequence of push and pop operations (denoted by token 0). For example, given x = [1, 0, 2, 3, 0, 1, 0, 0], the target is y = [∅, 1, ∅, ∅, 3, ∅, 1, 2]. By varying the sequence length and state complexity (maximum stack depth), this task evaluates the algorithm’s capacity for state tracking. 21

Keys and Values to test Associative Recall The keys and values task requires the model to store and retrieve associations between keys and values, then recall the value corresponding to a queried key [133, 50]. For example, given x = [b, 1, a, 3, d, 2, 4, a], the target is y = [∅, ∅, ∅, ∅, ∅, ∅, ∅, 3]. By varying the number of associations and association complexity (string length of keys and values), this task evaluates the algorithm’s capacity for associative recall. Modular Arithmetic to test In-Context Learning The modular arithmetic task requires the model to infer a latent linear rule from in-context examples and apply it to novel inputs [1, 68]. For each sequence, parameters a and b are sampled, then the sequence is presented as x = [x0 , y0 , x1 , y1 , x2 , y2 , x3 , y3 , ], where yi = (axi + b) mod V , where V is the vocabulary size. Then, the target is y = [y0 , ∅, y1 , ∅, y2 , ∅, y3 , ∅]. By varying the difficulty (range of values a, b can take on) and the number of in-context examples, this tasks evaluates the algorithm’s ability to induce in-context learning. B.2.2

Natural Tasks

TinyStories TinyStories is a curated dataset of short stories generated by OpenAI’s GPT-4 [27]. We use ASCII character-level tokenization, yielding a vocabulary of 256 tokens. Under this tokenization, the training and test sets contain 1.9B and 19.2M tokens, respectively. MNIST MNIST is a classic image dataset consisting of handwritten digits [71]. Rather than performing classification or 2D image generation, we consider the problem of 1D pixel-sequence modeling. The original 28×28 images are flattened into sequences of length 784 using raster-scan ordering. Each image is represented as a sequence of raw grayscale pixel intensities (0–255), yielding a vocabulary of 256 tokens. The training and test sets contain 47M and 7.8M tokens, respectively. Sketchy Sketchy is an image dataset of human-drawn sketches [105]. Rather than performing classification or 2D image generation, we consider the problem of 1D pixel-sequence modeling. The original images are resized to 64×64 using Lanczos resampling, and the pixels are binarized. Non-overlapping 2×2 patches are tokenized, yielding a vocabulary of 22×2 = 16 tokens. The resulting image is flattened in raster-scan order to form sequences of length 1024. The training and test sets contain 69.5M and 7.7M tokens, respectively. B.3

Algorithms

For all experiments, we use the AdamW optimizer with a weight decay of 0.01 and learning rates tuned separately for each algorithm. For all methods, gradients are clipped to a maximum global norm of 1. Gradient clipping is expected to be particularly beneficial for BPTT. After SMT, we transfer the decoder weights to the RNN readout module. During DMT, this readout head is further finetuned to optimize the next-token prediction loss using RNN-generated memory states. Importantly, this task loss updates only the readout head and not the RNN dynamics function, and therefore does not constitute temporal credit assignment for the RNN itself. Instead, finetuning serves to adapt the readout head to imperfections in the memory states generated by the RNN. Synthetic Experiments We set Tc = Tf = T . To evaluate the expectation in Lsmt , we compute the loss terms at all timesteps t ∈ [0, . . . , T ]. For earlier timesteps, where the available past context is shorter than the required context length, we pad the sequence and modify the attention mask so that padding tokens are ignored. The same procedure is applied to the future context. In these synthetic experiments, the prediction loss is applied only at positions where target output tokens are defined (e.g. the answer token in the needle task). We use batch sizes of 32 sequences during optimization, but this gets expanded to 32 × T input contexts to the encoder. Other Experiments To evaluate the expectation in Lsmt , we compute the loss term at a single timestep sampled uniformly from t ∼ U[0, . . . , T ]. The dataset is represented as one long sequence, meaning padding is not required, as both the past and future contexts extend indefinitely. We compute the Lunif over batches of memories from different sequences. By default for SMT, we set Tc = 256, Tf = 64, λdec = 1.0, λdyn = 0.1, λunif = 0.001 and train for 150000 SGD iterations. Unless otherwise specified, models use a hidden dimension of 256 with 16 22

SMT DMT RNN Loss Tinystories

MNIST

0.592 -3.502

0.571 -3.681

0.605 -3.903

0.617 -3.932

0.618 -3.951

0.64

0.01

0.685 -3.121

0.585 -3.525

0.565 -3.682

0.622 -3.904

0.601 -3.929

0.623 -3.951

0.62

0.1

0.569 -2.763

0.567 -3.294

0.593 -3.684

0.744 -3.893

0.597 -3.926

0.626 -3.951

0.60

1.0

0.588 -2.265

0.597 -2.892

0.592 -3.386

0.630 -3.874

0.812 -3.925

0.993 -3.951

0.58

10.0

0.637 -1.143

0.620 -2.628

0.611 -2.760

0.608 -3.502

0.958 -3.883

2.713 -3.942

0.56

0.0

0.001

0.003

0.01

0.03

0.1

0.55

0.001

0.529 -1.802

0.555 -3.738

0.557 -3.840

0.558 -3.904

0.557 -3.934

0.585 -3.953

0.01

0.522 -1.639

0.530 -3.698

0.535 -3.824

0.531 -3.897

0.537 -3.930

0.556 -3.952

0.1

0.507 -1.273

0.516 -3.618

0.521 -3.787

0.528 -3.877

0.538 -3.919

0.554 -3.952

1.0

0.514 -0.542

0.518 -3.323

0.533 -3.630

0.550 -3.810

0.545 -3.882

0.596 -3.947

10.0

0.521 -0.058

0.548 -2.222

0.552 -3.202

0.573 -3.675

0.577 -3.795

0.568 -3.912

0.0

0.001

0.003

0.01

0.03

0.1

0.54

dyn

0.53

unif

0.52

Test Loss

0.001

0.607 -3.155

0.51 0.50

Figure 16: Sweep of λdyn and λunif . Cell color indicates the RNN test loss for each setting. Top number in each cell is the RNN test loss. Bottom number in each cell shows the Lunif . Lunif varies from 0 (collapsed latent space) to −4 (fully uniform latent space). memory tokens. The encoder is 8 layers deep, while the decoder is 4 layers deep. The RNN is also 8 layers deep, and its readout function is 4 layers deep. We use a batch size of 128 sequences.

Additional Experiments

Figure 16 shows the results of ablating the λdyn and λunif . Results show the optimal RNN performance requires a moderate dynamics loss, paired with a very low uniformity loss. However, a little uniformity is critical for avoiding memory space collapse. Figure 17 shows more samples of generations from Figure 4. Samples generated by the BPTT RNN (Transformer backbone) seem to only pick up on short range context and act accordingly: either output large streaks of white or black based on the current row. BPTT RNN (GRU backbone) improves this significantly, but still fails to capture the nuanced structure of digits. SMT→DMT RNN (Transformer backbone) is able to capture this structure quite well. Figure 18 shows more samples of generations from Figure 5. These generations are often not fully interpretable, but do capture the stroke structure of human-drawn sketches. Capturing this stroke structure is itself a difficult problem, given the long-horizon nature of pixel sequence modeling. Figure 19 provides an analysis of our Sketchy RNN as it “reads” a sequence corresponding to the classic Attneave’s cat image [5]. The memory sequence does not seem to be fully interpretable, but does show significant structure. Figure 20 provides samples of generations when conditioned on partial context of Attneave’s cat image. Figure 21, 22, 23, show additional analysis of the RNNs on MNIST and Sketchy data.

Compute Resources Used

All individual training runs were conducted on one H200 GPU within 48 hours. The synthetic experiments comprised more than 375 small-scale training runs, while the real-data experiments required 144 large-scale runs.

MNIST Dataset Samples

Samples Generated by BPTT RNN (Transformer Backbone)

Samples Generated by BPTT RNN (GRU Backbone)

Samples Generated by SMT DMT RNN (Transformer Backbone)

Figure 17: Additional MNIST Samples. Here we give more examples of samples of MNIST images generated by the various methods. SMT→DMT RNN outperforms BPTT, even when BPTT is applied on a GRU architecture, in processing long-horizon information, which is required for pixel modeling.

Sequence to Set Reframing

As described in Section 2.2, consider a hypothetical oracle memory-encoding model Q that takes as input the sequence of tokens and outputs an effective compressed memory. Here we show that Q does not have to be a recurrent function over the sequence of tokens, but can instead be represented as a permutation-invariant function over a set of timestamped tokens. Claim Let xseq = [x0 , x1 , . . . , xt ] be the original sequence of tokens. We define the set xset = {(x0 , 0), (x1 , 1), . . . , (xt , t)}. Assume that Q is a recurrent function over xseq . In other words, m = Q(xseq ) = f (. . . f (f (m∅ , x0 ), x1 ), . . . , xt ) with m∅ = 0 for some function f . For any such Q, ∃ a function g such that g(xset ) = Q(xseq ). Proof We construct g explicitly. Define g(xset ) as follows: given the input set xset = {(x0 , 0), (x1 , 1), . . . , (xt , t)}, sort the elements in ascending order of their timestamp to recover the sequence xseq = [x0 , x1 , . . . , xt ], then apply Q to this sequence. This is well-defined because the timestamps {0, 1, . . . , t} are distinct integers, so the sort order is unique. The resulting sequence is identical to the original xseq , and therefore g(xset ) = Q(xseq ) = m. Moreover, g is permutation-invariant: any permutation of the elements of xset yields the same sorted sequence and thus the same output. Since Q was arbitrary, this construction applies to every recurrent Q, completing the proof. ■ Implication This result implies that any sufficiently expressive permutation-invariant set model can in principle exactly model a recurrent memory function. Because sets are unordered, time-parallel processing naturally follows. In particular, Transformer-based architectures can be interpreted as operating over sets of timestamped tokens rather than strictly ordered sequences. Notably, the proof is constructive: g recovers the sequential computation by sorting the timestamps and implicitly applying the recurrent update rule f up to t times. Consequently, when implemented with bounded-depth architectures such as Transformers, the required depth may need to scale with sequence length, consistent with prior work on sequential depth and time-parallel training discussed in Section 4. Scaling depth with sequence length seems to present a major theoretical limitation. However, our empirical results suggest that even relatively shallow Transformer encoders can learn highly effective memory representations for both synthetic and natural tasks. Thus, despite lacking full theoretical expressivity, this sequence-to-set reframing may still provide a practical strategy for memory pretraining. For full expressivity, some light-weight post-training may be required.

Encoder Markovian Training

SMT consists of two primary objectives: future predicting with Ldec and dynamics modeling with Ldyn . The dynamics objective serves two purposes: (1) training the RNN to predict the next memory state from the current one, and (2) encouraging the encoder to produce memory states that are predictable from one another, i.e. approximately Markovian. In this section, we show that the predictive state objective Ldec alone is sufficient for learning Markovian memories, implying that Ldyn is theoretically unnecessary, though still practically useful. Claim Let x = [x0 , x1 , . . . , xT ] and y = [y0 , y1 , . . . , yT ] be input and output sequences. For each timestep t, define xctx t = [x0 , . . . , xt ],

xfut t = [xt+1 , . . . , xT ],

ytfut = [yt , . . . , yT ],

fut fut with memory state mt = Eϕ (xctx t ) and reconstructed future ŷt = Dψ (mt , xt ). If mt is an optimal ctx fut fut minimal sufficient statistic of xt for predicting yt given xt at every t, then the memory sequence (mt ) is Markovian:

mt+1 ⊥⊥ xctx mt , xt+1 . t Proof

fut fut By optimality of mt , it is a minimal sufficient statistic of xctx t for predicting yt given xt : fut H ytfut | mt , xfut = H ytfut | xctx t t , xt .

fut fut Note that yt+1 ⊆ ytfut and xfut t+1 ⊆ xt , so optimality of mt at time t implies it is also sufficient for fut fut yt+1 given xt+1 : fut fut ctx fut fut ctx fut H yt+1 | mt , xt+1 , xfut t+1 = H yt+1 | xt , xt+1 , xt+1 = H yt+1 | xt+1 , xt+1 . fut fut By optimality of mt+1 , it is a minimal sufficient statistic of xctx t+1 for predicting yt+1 given xt+1 . ctx Minimality means mt+1 retains no information from xt+1 beyond what is predictively necessary. Since (mt , xt+1 ) already constitutes a sufficient statistic for this same prediction task—as shown above—minimality of mt+1 forces it to be a function of (mt , xt+1 ):

mt+1 = f (mt , xt+1 ) for some measurable f . Therefore mt+1 is determined entirely by (mt , xt+1 ), and conditioning on these renders it independent of all earlier context: H(mt+1 | mt , xt+1 ) = 0, which is equivalent to mt+1 ⊥ ⊥ xctx t | mt , xt+1 . Hence (mt ) is Markovian. ■ Implication This result establishes that under ideal conditions—sufficient encoder and decoder capacity, infinite future horizon, and exact optimization—the memory states learned by Eϕ form a Markov chain driven only by the previous state and the incoming token. In other words, the encoder implicitly learns a Markovian memory representation: mt+1 can be predicted from only (mt , xt+1 ). In practice, finite capacity and approximate optimization relax this property, leaving mt+1 with residual dependence on xctx t beyond (mt , xt+1 ). This gap motivates jointly training the dynamics loss Ldyn alongside Ldec to explicitly encourage Markovian structure in the learned memory sequence. Relatedly, Teoh et al. [122] provide a proof that one-step RNN dynamics with a Tf = 1 encoder also induce a predictive-state memory representation.

Sketchy Dataset Samples

Samples Generated by SMT DMT RNN (Transformer Backbone)

Figure 18: Additional Sketchy Samples. Here we give more examples of samples of Sketchy images from the dataset and generated by SMT→DMT. Even in this hard sparse domain, SMT→DMT can capture the overall stroke structure, which requires integrating information over hundreds of pixels.

Data as Image

Memory as Image (t-SNE 3D)

Memory Space (t-SNE 2D)

1024

Timestep 0

Data as Sequence (partial) Memory (t-SNE 3D) as Sequence (partial) Figure 19: Analysis on Attneave’s Cat. We apply the SMT→DMT-trained RNN on Sketchy and evaluate it on the classic image of Attneave’s cat. The RNN reads the image pixel-by-pixel in raster scan order. Top Left: Input image presented in its original 2D form. Top Middle: 3D t-SNE projection of the RNN memory state, visualized as RGB values over time, showing the evolution of memory throughout sequence processing. Top Right: 2D t-SNE projection of the memory state trajectory over time. Middle: The same image presented as a flat token sequence. From the RNN’s perspective, the task resembles modeling a barcode-like sequence, requiring long-range associations between distant tokens and highlighting the difficulty of pixel sequence modeling. Bottom: 3D t-SNE projection of the memory state visualized along the flattened sequence.

Image Context

Samples generated by SMT DMT RNN

Figure 20: Generations of Attneave’s Cat. We apply the SMT→DMT-trained RNN on Sketchy and apply it to generate part of the image of Attneave’s cat. Given more of the image context, the RNN seems to understand the image better and make somewhat more plausible predictions.

MNIST Memory Visualization (PCA) Data as Image

Memory as Image (PCA 3D)

Memory Space (PCA 2D)

Data as Image

Memory as Image (PCA 3D)

Memory Space (PCA 2D)

Figure 21: RNN Memory Evolution on MNIST (PCA). We analyze the memory evolution of our SMT→DMT MNIST RNN. Left: Input image presented in its original 2D form. Middle: 3D PCA projection of the RNN memory state, visualized as RGB values over time, showing the evolution of memory during processing. Right: 2D PCA projection of the memory state trajectory over time. MNIST Memory Visualization (t-SNE) Data as Image

Memory as Image (t-SNE 3D)

Memory Space (t-SNE 2D)

Data as Image

Memory as Image (t-SNE 3D)

Memory Space (t-SNE 2D)

Figure 22: RNN Memory Evolution on MNIST (t-SNE). We analyze the memory evolution of our SMT→DMT MNIST RNN. Left: Input image presented in its original 2D form. Middle: 3D t-SNE projection of the RNN memory state, visualized as RGB values over time, showing the evolution of memory during processing. Right: 2D t-SNE projection of the memory state trajectory over time.

Data as Image

Memory as Image Memory Space Memory as Image Memory Space (t-SNE 3D) (t-SNE 2D) Data as Image (t-SNE 3D) (t-SNE 2D)

Figure 23: RNN Memory Evolution on Sketchy (t-SNE). We analyze the memory evolution of our SMT→DMT Sketchy RNN. Left: Input image presented in its original 2D form. Middle: 3D t-SNE projection of the RNN memory state, visualized as RGB values over time, showing the evolution of memory during processing. Right: 2D t-SNE projection of the memory state trajectory over time.

Record · ID 259457 · SHA-256 74d72cbad3e45522
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.