Pretraining Recurrent Networks without Recurrence

Akarsh Kumar
Phillip Isola
MIT

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 $\mathcal{O}(1)$order one (compared to $\mathcal{O}(T)$order 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 $(m_t,x_{t+1})\to m_{t+1}$m sub t and x sub t plus one mapping to m sub t plus one. SMT acquires these memory labels by training a Transformer-based 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 $\mathcal{O}(1)$order one 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.

GRAY

Project Page: akarshkumar.com/smt
Source Code: github.com/akarshkumar0101/smt

1 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 [90].

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 [103, 132]. Although conceptually well-motivated, BPTT is sequential in time and suffers from unstable high variance gradients that may vanish or explode [98]. 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 $\mathcal{O}(T)$order 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^{\ast }$m star sub t. Then, RNN training reduces to learning the one-step update $(m_t^{\ast },x_{t+1})\to m_{t+1}^{\ast }$m star t and x t plus one to m star t plus one 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 [76]. 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 non-sequential 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 $\mathcal{O}(1)$order one gradient path for long-range associations.

Indeed, Transformers solved time-parallelism and credit assignment in the same way [125], 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, 65]. 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 [122]. 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 [68, 40, 39, 24]. But, because their transition function is linear, the class of functions they can represent is fundamentally constrained [85], which can lead to failure on important sequential tasks such as state tracking [84, 78].

SMT aims to combine the best of all worlds: time-parallel training, stable $\mathcal{O}(1)$order one 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 [82]. 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$t to state at time $t+1$t plus one).

2 Methods

2.1 Background

Causal Conditional Sequence Modeling Let $x=[x_0,\dots ,x_T]$x and $y=[y_0,\dots ,y_T]$y denote input and output sequences. The objective is to learn a model of the conditional distribution $p(y\mid x)$p of y given x. We assume each output $y_t$y sub t depends only on $x_0,\dots ,x_t$. Formally we model this distribution with ${\prod }_{t=0}^T{p}_{\theta }(y_t\mid {\mathbf{x}}_{\le t})$the product from t equals zero to T of the probability of y sub t given x up to time t. Autoregressive sequence modeling is a special case when $x_t=y_{t-1}$x sub t equals y sub t minus one.

Recurrent Neural Networks (RNNs) An RNN models this problem using a fixed-size latent state, $m_t$m sub t, that summarizes past inputs. At each timestep, this state is updated according to:

m sub t plus one equals f of m sub t and x sub t plus one

where $f_{\theta }$f is the transition function. The predicted output token distribution is then $p_{\theta }(y_t\mid {\mathbf{x}}_{\le t})=\text{softmax}(g_{\theta }(m_t))$the probability of y sub t given x up to time t is the softmax of g of m sub t, where $g_{\theta }$g is the readout function. Ideally, $m_t$m sub t “remembers” important information from the past and intentionally “forgets” unimportant information, i.e., $m_t$m sub t is a memory.

Backpropagation Through Time (BPTT) Traditionally, RNNs are trained with BPTT [103, 132]. In the forward pass, $f_{\theta }$f is recurrently unrolled over the sequence. The input sequence $x_t$x sub t is provided via teacher forcing, while the memory sequence $m_t$m sub t is generated by the RNN’s transition and is used to compute the output predictions. Conceptually, the computation graph takes the form:

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 $\mathcal{O}(T)$order T steps. Depending on the singular values of the Jacobian of $f_{\theta }$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 [97].
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 [101]. When gradients explode, the induced dynamical system is chaotic, causing training instability [98].

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$Q, that takes as input the sequence of tokens up to timestep $t$t, ${\mathbf{x}}_t^{\text{ctx}}=[x_0,x_1,\dots ,x_t]$x context t, the sequence from x zero to x t, and outputs an effective compressed memory for that timestep $m_t^{\ast }=Q({\mathbf{x}}_t^{\text{ctx}})$m star t equals Q of x context t. This memory retains all information from the past input that is relevant for predicting the future output, ${\mathbf{y}}_t^{\text{fut}}=[y_t,\dots ,y_T]$y future t, the sequence from y t to y T, 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$Q at different points along the sequence produces a corresponding sequence of memory labels $[m_0^{\ast },m_1^{\ast },\dots ,m_T^{\ast }]$m star zero through m star T. With $Q$Q, the RNN’s problem of learning a temporal update collapses to standard supervised learning on oracle memory transitions labels $(m_t^{\ast },x_{t+1})\to m_{t+1}^{\ast }$m star t and x t plus one leading to m star t plus one. Our key insight is that $Q$Q does not have to be a recurrent function over $[x_0,x_1,\dots ,x_t]$, but can instead be represented as a permutation-invariant function over the set $\{(x_0,0),(x_1,1),\dots ,(x_t,t)\}$ (details in Appendix E).

In practice, SMT approximates $Q$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 [76].

Formulation Formally, we have the RNN $f_{\theta }$f, bidirectional encoder ${\mathcal{E}}_{\varphi }$E, and causal decoder ${\mathcal{D}}_{\psi }$D. ${\mathcal{E}}_{\varphi }$E and ${\mathcal{D}}_{\psi }$D are time-parallel Transformer architectures. Given $\mathbf{x}=[x_0,x_1,\dots ,x_T]$x and $\mathbf{y}=[y_0,y_1,\dots ,y_T]$y, we consider, for each timestep $t$t, a decomposition into the past and future:

The encoder maps each context to a memory state with $m_t={\mathcal{E}}_{\varphi }({\mathbf{x}}_t^{\text{ctx}})$m t equals E of x context t. Then, the decoder predicts the future output distribution using the memory of the past and teacher forced future inputs:

the probability of y future t given x context and x future is the product of the probabilities of each future y given the memory and future x, which equals D of m t and x future t

The future decoding loss for timestep $t$t is (CE denotes the sequence level cross-entropy loss):

We have the RNN predict the next memory given the current memory and the next input with ${\hat{m}}_{t+1}=f_{\theta }(m_t,x_{t+1})$m hat sub t plus one equals f theta of m t and x sub t plus one. This prediction is supervised with the next timestep’s memory:

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. $m_{t+1}$m sub t plus one is predictable solely from $(m_t,x_{t+1})$m t and x sub t plus one).

We add a uniformity loss [131] to prevent the memory space from collapsing:

The full objective is a weighted sum of all three losses:

where the $\lambda$lambda terms control the trade-off between memory representation, dynamics, and collapse.

Practice Theoretically, it should be enough to train ${\mathcal{E}}_{\varphi }$E phi and ${\mathcal{D}}_{\psi }$D psi with only ${\mathcal{L}}^{\text{dec}}$L dec, and separately train $f_{\theta }$f theta with only ${\mathcal{L}}^{\text{dyn}}$L dyn (proof in Appendix F). However, in practice we find it beneficial to jointly train all models in one stage with ${\mathcal{L}}^{\text{smt}}$L smt, since that explicitly optimizes $m_t$m t to be Markovian, and provides additional temporal credit propagation benefits described in Section 3.6.

For experiments, we truncate ${\mathbf{x}}_t^{\text{ctx}}$x ctx to a context length $T_c$T sub c and ${\mathbf{y}}_t^{\text{fut}}$y fut to a future length $T_f$T sub f. For computational efficiency, we estimate the expectation in ${\mathcal{L}}^{\text{smt}}$L smt by randomly sampling a single timestep $t$t, rather than computing all timesteps in the sequence. This yields SMT a smaller training memory footprint than BPTT: $\mathcal{O}(M+T)$order M plus T instead of $\mathcal{O}(MT)$order M times T, where $M$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, $\mathcal{O}(1)$order one credit path for long-range associations, and gradient stability.

2.3 DAgger Memory Training (DMT)

Figure 2: SMT vs DMT. SMT trains the RNN with behavior cloning on the encoder-generated memory states (off-policy imitation learning). DMT unrolls the RNN with its own memory states and then imitates the encoder trajectory (on-policy imitation learning). Figure design inspired by Jacobs et al. [60]Jacobs and colleagues.

After SMT, the RNN achieves low one-step error in predicting $(m_t,x_{t+1})\to m_{t+1}$m t and x sub t plus one leading to m sub t plus one when $m_t$m t comes from the encoder. However, at evaluation time, the model is unrolled autoregressively, using its own predicted memories rather than the encoder memories as input. This train–test mismatch causes small prediction errors to accumulate over time, leading to a growing drift between the RNN-generated memory trajectory $[{\hat{m}}_0,\dots ,{\hat{m}}_T]$m hat zero through m hat T and the encoder trajectory $[m_0,\dots ,m_T]$m zero through m T, even with teacher forced input tokens. This drift is quantified as ${\delta }_t=\text{MSE}({\hat{m}}_t,m_t)$delta t equals the mean squared error between m hat t and m t.

We introduce DAgger Memory Training (DMT), a finetuning phase that corrects this drift via on-policy imitation learning [102]. By exposing the RNN to its own induced memory state distribution, DMT trains the RNN to autocorrect its errors to stay aligned with the encoder trajectory (Figure 2).

Concretely, given $\mathbf{x}$x, we first compute the encoder trajectory $[m_0,\dots ,m_T]$m zero through m T using only ${\mathcal{E}}_{\varphi }$E phi and then the RNN trajectory $[{\hat{m}}_0,\dots ,{\hat{m}}_T]$m hat zero through m hat T using $f_{\theta }$f theta. Instead of training on SMT labels $(m_t,x_{t+1})\to m_{t+1}$m t and x sub t plus one leading to m sub t plus one, we train on DMT labels $({\hat{m}}_t,x_{t+1})\to m_{t+1}$m hat t and x sub t plus one leading to m sub t plus one. Equivalently, the training loss is:

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 $x_t$x t 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, $m_t$m t. 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.

3 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 [85, 38].

“BPTT RNN” denotes the BPTT baseline. “SMT Encoder*” generates memories $m_t$m t with the SMT-trained 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 $\to$ 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 [119]. As a more challenging problem, we test our method on
raster-scan order pixel sequence modeling of sparse images from MNIST [72] and Sketchy [106].
This is a hard problem for RNNs [71, 124]. 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 $T_c=T_f=T$T sub c equals T sub f equals T and train all timesteps in the ${\mathcal{L}}^{\text{smt}}$L S M T
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 $\to$ 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 $\mathcal{O}(T)$order T credit path length, compared to SMT’s $\mathcal{O}(1)$order one. Further analysis in
Section 3.7 confirms the difference in gradient stability in both methods.

Figure 3: Synthetic Task Experiments. We evaluate BPTT, SMT, and SMT$\to$DMT using five synthetic tasks with various settings to probe different properties of the algorithms. $\ast$ 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$\to$DMT outperforms BPTT, signaling that SMT has better gradient properties, memory utilization, state tracking, associative recall, and in-context learning than BPTT.

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$\to$DMT. Figure 5 shows images generated by an SMT$\to$DMT RNN trained on Sketchy.

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$\to$DMT captures these dependencies with a non-gated RNN architecture. More samples are in Appendix Figure 17.

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.

3.3 Sequential Compute and Data

We now evaluate BPTT and SMT across real domains and various RNN architectures. Each method is allowed $N$N optimization steps on token batches of shape $B\times T$B by T (number of sequences $\times$ sequence length). We sweep $N,B,$N, B, and $T$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 ($\sim$roughly 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.

Figure 6 shows the results. In sequential compute, SMT Encoder and SMT$\to$DMT RNN are significantly more efficient than BPTT with the Transformer and MLP backbones. In data, SMT Encoder and SMT$\to$DMT RNN has approximately

Figure 6: Sequential Compute and Data Efficiency. We sweep training hyperparameters for BPTT, SMT, and SMT$\to$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$\to$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$\to$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 [121]. SMT$\to$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$\to$DMT along three axes: context length, memory state size, and model parameter count. For the first two, we logarithmically sweep $T_c$T sub c 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$\to$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$\to$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.

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

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.

3.5 Compression as a Scaling Axis

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.

Neural scaling laws predict the relationship between a resource (e.g. compute, data) and a desired property (e.g. validation loss, benchmark accuracy) [66, 54]. Can the desired property instead be compression [59]? 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.

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 [114, 70], 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 Ablations

Predictive State and Detached RNN The impact of the predictive state objective (Equation 2) is evaluated by sweeping the future length $T_f$T sub f, while keeping $T_c$T sub c 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 ${\mathcal{L}}_t^{\text{dyn}}$L dyn 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 $m_t$m sub t 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 $T_f$T sub f or joint training. When $T_f$T sub f is large enough, there is a $\mathcal{O}(1)$order one credit path length between the needle and the answer at all timesteps. Interestingly, when $T_f$T sub f is small, there exists no credit path to learn early timestep memories, yet joint training still learns effectively, even when $T_f=1$T f equals one. Credit must be propagating through the RNN dynamics from $m_T$m sub T to $m_{T-1}$m sub T minus one, and so on, to $m_0$m zero. 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 $m_t$m sub t to $m_{t-1}$m sub t minus one through $f_{\theta }$f theta; $T$T such gradients steps sends information $T$T steps back in the sequence. This implies that solving $T$T sequence length credit assignment task when $T_f=1$T f equals one, requires at least $T$T gradient optimization steps. This credit amortization phenomenon is reminiscent of value bootstrapping in RL [120].

$\lambda$Lambda Coefficients The values of ${\lambda }_{\text{dyn}}$lambda dyn and ${\lambda }_{\text{unif}}$lambda unif are swept here to check their effects. Figure 16 shows the results. The best RNNs require ${\lambda }_{\text{dyn}}=0.1$lambda dyn equals zero point one, and ${\lambda }_{\text{unif}}=0.001$lambda unif equals zero point zero zero one. When ${\lambda }_{\text{unif}}=0$lambda unif equals zero, although the RNN performance is preserved, the memory space is collapsed, as indicated by ${\mathcal{L}}^{\text{unif}}$L unif.

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

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 $m_t$, depending on the weight initialization. SMT is non-recurrent and has a $\mathcal{O}(1)$ credit path length, making its gradients agnostic to initialization and time-horizon.

Figure 12: Impact of DMT across many runs with different SMT ${\lambda }_{\text{dec}}$ and ${\lambda }_{\text{dyn}}$ hyperparameters. Left: Applying DMT reduces the drift of the RNN rollout (measured with $1-R^2$ of RNN memory prediction ${\hat{m}}_t$ of encoder ground truth $m_t$). 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 of $m_t$, $\Vert \frac{\partial L}{\partial m_t}\Vert$the norm of the partial derivative of L with respect to m t, at different $t$t for both methods with different model weight initializations on the needle 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$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.

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)$constant rather than $O(T)$linear memory and compute per generated token (Table 1). RNNs also constitute a more expressive class of models [85, 78].

Figure 13: Sequence Length Generalization. An SMT$\to$DMT trained RNN generalizes better than its Transformer teacher when evaluated on sequence lengths longer than training. The task is synthetic state tracking.

Here, we compare an SMT$\to$DMT RNN against a Transformer on the synthetic stack state tracking task. For a fair comparison, we use the SMT encoder as the Transformer baseline, since it imposes the same memory-information bottleneck as the RNN.

Figure 13 shows the Transformer outperforms the RNN on training sequence lengths, but significantly underperforms the RNN on sequence lengths longer than training. Prior work on length generalization reports similar findings [100]. 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.

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.

4 Related Works

Relation with NextLat Highly related and concurrent to our work, Teoh et al. [123]Teoh and colleagues introduced Next-Latent Prediction (NextLat), which co-trains an RNN with memory state supervision from a Transformer. With a particular setting of hyperparameters, NextLat closely resembles DMT (and SMT when the rollout length is 1), though still utilizing BPTT on multi-step rollouts. However, the central focus of NextLat is to regularize the Transformer to learn compact world models, rather than providing a method to train RNNs. Consequently, their analysis focuses primarily on the Transformer. In contrast, our work focuses exclusively on training the RNN in a one-step manner.

Recurrent Neural Networks (RNNs) RNNs were studied extensively early in AI because their recurrence mechanism resembles biological brains [81] and can be applied to any sequential task [28]. Many different algorithms were proposed for learning, including random guessing [109], evolutionary algorithms [88, 3, 116, 105, 107], hebbian learning [45, 56, 87, 93], real-time recurrent learning [133], and other algorithms [94, 10, 71, 8, 64]. BPTT is the only widely adopted algorithm [132].

However, it has repeatedly been shown that BPTT produces unstable gradients that vanish, explode, or exhibit high variance [11, 53, 98, 7]. Several directions address this issue. One direction focuses on architectural modifications, including residual connections [115, 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 [108, 4, 135, 86, 127, 47]. Others explored external memory [99, 37, 41], hierarchical modeling [48, 19, 128, 62], other unique directions [61, 79, 89].

Recently, there has been renewed interest in RNNs in the form of linear state space models [40, 113, 39], linear attention models [68, 118, 24, 138, 137], and even nonlinear RNN models [9, 15, 91, 95]. 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 [125] (unlike prior attention methods [6]), which is crucial for leveraging modern hardware [49, 55] to scale performance [66, 54]. Linear RNNs gained popularity [68, 24, 138, 31] after it was realized they can be parallelized with the associative scan algorithm [13, 80, 130].

A recent line of work attempts to parallelize nonlinear RNNs as well [75, 22]. Rather than computing $[m_0,\dots ,m_T]$m zero through m T with $mt+1=f_{\theta }(m_t)$m t plus one equals f theta of m t, they formulate the forward pass as an iterative optimization procedure. Starting with an initial guess $[m_0^0,\dots ,m_T^0]$m zero superscript zero through m T superscript zero, they construct a system of $T$T equations, $\{m_{t+1}-f_{\theta }(m_t)=0{\}}_{t=0}^T$the set of equations m t plus one minus f theta of m t equals zero, for t from zero to T, and solve this system with Newton’s method [96]. Many works have further built on this approach [23, 35]. Although appealing, this approach approximates BPTT and hence will suffer from its $\mathcal{O}(T)$order 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 $m_t$m t to be a predictive state while satisfying $m_{t+1}\approx f_{\theta }(m_t)$m t plus one is approximately f theta of m t and providing an $\mathcal{O}(1)$order one 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 [104, 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 [83]. 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 [82, 83, 139]. While such models succeed on tasks amenable to parallelization (e.g. parity tracking via associative scan [77, 73])for example, parity tracking via associative scan, they systematically fail on tasks requiring deep sequential computation (e.g. tracking a chess board [84])such as tracking a chess board. Interestingly, the aspect that makes models parallelizable, limits their performance on harder problems [82]. Nonlinear RNNs are one of the few classes of models where its sequential depth grows with the input sequence length [85]. Although these constraints were seen as theoretical, there is growing evidence they affect models in practice as well [78].

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 [76, 111], a representation that is sufficient for optimal decision making [112]. Early works interpreted PSRs as a literal vector of probabilities of future events, but have since been generalized [25]. Belief states are a similar concept, which also defines a sufficient statistic of the past [63, 58].

PSRs have been previously incorporated into RNNs [26, 46]. Venkatraman et al. (2015)Venkatraman and colleagues 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 $\mathcal{O}(T)$order T credit path.

Other Related Work Our work is related to the literature on cross-architecture teacher-student distillation [67, 129, 43, 16, 92], but these works do not address the challenges of training nonlinear RNNs. The Recurrent Transformer is an RNN architecture that attends to all past hidden states, creating an $\mathcal{O}(1)$order one gradient path that stabilizes credit assignment [95]. 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, 136]. A new line of work uses principles from diffusion models to train blocks of a feed-forward network in parallel, avoiding global backpropagation [74, 110].

5 Discussion

In SMT, the teacher model is time-parallel, and is thus constrained in expressivity [82], 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 $m_t$m sub t within a sequence. We found that training on all memories $[m_0,\dots ,m_T]$m zero through m T 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 [75].

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 $\mathcal{O}(1)$order one 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

References