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

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 [2], [3]. Although conceptually well-motivated, BPTT is sequential in time and suffers from unstable high variance gradients that may vanish or explode [4]. 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)\) steps [5]. 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, x_{t+1})\rightarrow 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 [6]. 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)\) gradient path for long-range associations.

Indeed, Transformers solved time-parallelism and credit assignment in the same way [7], and have since revolutionized sequence modeling [8]. However, Transformers do not possess a compressed memory of the past in the way RNNs or human brains do [9], [10]. 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 [11]. 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 [12]. 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 [13].

Linear attention RNN models also exhibit time-parallel training and relatively stable credit assignment, while maintaining a fixed memory size [14]–[17]. But, because their transition function is linear, the class of functions they can represent is fundamentally constrained [18], which can lead to failure on important sequential tasks such as state tracking [19], [20].

SMT aims to combine the best of all worlds: time-parallel training, stable \(\mathcal{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 [21]. 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↩︎

2.1.0.1 Causal Conditional Sequence Modeling

Let \(\mathbf{x} = [x_0, \dots, x_T]\) and \(\mathbf{y} = [y_0, \dots, y_T]\) denote input and output sequences. The objective is to learn a model of the conditional distribution \(p(\mathbf{y} \mid \mathbf{x})\). We assume each output \(y_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}_{\leq t})\). Autoregressive sequence modeling is a special case when \(x_t = y_{t-1}\).

2.1.0.2 Recurrent Neural Networks (RNNs)

An RNN models this problem using a fixed-size latent state, \(m_t\), that summarizes past inputs. At each timestep, this state is updated according to: \[\label{eq:rnn95update} m_{t+1} = f_\theta(m_t, x_{t+1})\tag{1}\] where \(f_\theta\) is the transition function. The predicted output token distribution is then \(p_\theta(y_{t} \mid \mathbf{x}_{\leq t}) = \text{softmax}(g_\theta(m_t))\), where \(g_\theta\) is the readout function. Ideally, \(m_t\) “remembers” important information from the past and intentionally “forgets” unimportant information, i.e., \(m_t\) is a memory.

2.1.0.3 Backpropagation Through Time (BPTT)

Traditionally, RNNs are trained with BPTT [2], [3]. In the forward pass, \(f_\theta\) is recurrently unrolled over the sequence. The input sequence \(x_t\) is provided via teacher forcing, while the memory sequence \(m_t\) is generated by the RNN’s transition and is used to compute the output predictions. Conceptually, the computation graph takes the form: \[m_t = f_\theta(\ldots f_\theta(f_\theta(m_{\emptyset}, x_0), x_1), \ldots , x_t) \quad\text{with}\quad m_{\emptyset}=\mathbf{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 \(\mathcal{O}(T)\) steps. Depending on the singular values of the Jacobian of \(f_\theta\), 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 [22].

2. BPTT often produces unstable high variance gradients [5]. When gradients vanish, the RNN experiences a recency bias, hindering the learning of long-range associations [23]. When gradients explode, the induced dynamical system is chaotic, causing training instability [4].

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.

2.2.0.1 Motivation

Consider a hypothetical oracle memory-encoding model \(\mathcal{Q}\), that takes as input the sequence of tokens up to timestep \(t\), \(\mathbf{x}_t^\text{ctx}=[x_0,x_1,\dots,x_t]\), and outputs an effective compressed memory for that timestep \(m_t^*=\mathcal{Q}(\mathbf{x}_t^\text{ctx})\). 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]\), 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 \(\mathcal{Q}\) at different points along the sequence produces a corresponding sequence of memory labels \([m_0^*, m_1^*, \dots, m_T^*]\). With \(\mathcal{Q}\), the RNN’s problem of learning a temporal update collapses to standard supervised learning on oracle memory transitions labels \((m_t^*, x_{t+1})\rightarrow m_{t+1}^*\). Our key insight is that \(\mathcal{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 10).

In practice, SMT approximates \(\mathcal{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 [6].

2.2.0.2 Formulation

Formally, we have the RNN \(f_\theta\), bidirectional encoder \(\mathcal{E}_\phi\), and causal decoder \(\mathcal{D}_\psi\). \(\mathcal{E}_\phi\) and \(\mathcal{D}_\psi\) are time-parallel Transformer architectures. Given \(\mathbf{x} = [x_0, x_1, \dots, x_T]\) and \(\mathbf{y} = [y_0, y_1, \dots, y_T]\), we consider, for each timestep \(t\), a decomposition into the past and future: \[\mathbf{x}_t^{\text{ctx}} = [x_0, \dots, x_t] \qquad \mathbf{x}_t^{\text{fut}} = [x_{t+1}, \dots, x_T] \qquad \mathbf{y}_t^{\text{fut}} = [y_{t}, \dots, y_T]\] The encoder maps each context to a memory state with \(m_t = \mathcal{E}_\phi(\mathbf{x}_t^{\text{ctx}})\). Then, the decoder predicts the future output distribution using the memory of the past and teacher forced future inputs: \[p_{\phi,\psi}(\mathbf{y}_t^\text{fut} \mid \mathbf{x}_t^\text{ctx}, \mathbf{x}_t^\text{fut}) = \prod_{\tau=t}^{T} p_\psi(y_{\tau} \mid m_t, \mathbf{x}_{t+1:\tau}) = \mathcal{D}_\psi(m_t, \mathbf{x}_t^\text{fut})\] The future decoding loss for timestep \(t\) is (\(\mathrm{CE}\) denotes the sequence level cross-entropy loss): \[\label{eq:loss95dec} \mathcal{L}_t^\text{dec} = \mathrm{CE}\left( \mathbf{y}_t^\text{fut}, p_{\phi,\psi}(\mathbf{y}_t^\text{fut} \mid \mathbf{x}_t^\text{ctx}, \mathbf{x}_t^\text{fut}) \right)\tag{2}\]

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})\). This prediction is supervised with the next timestep’s memory: \[\label{eq:loss95dyn} \mathcal{L}^\text{dyn}_t = \mathrm{MSE} (\hat{m}_{t+1}, m_{t+1})\tag{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. \(m_{t+1}\) is predictable solely from \((m_t, x_{t+1})\)).

We add a uniformity loss [24] to prevent the memory space from collapsing: \[\mathcal{L}^\text{unif} = \log \mathbb{E}_{t_a, t_b \sim [0,\dots,T]} \exp(-2\| m_{t_a} - m_{t_b}\|_2^2)\]

The full objective is a weighted sum of all three losses: \[\label{eq:smt} \mathcal{L}^\text{smt} = \lambda_\text{dec} \mathbb{E}_t\left[ \mathcal{L}^{\text{dec}}_t \right] + \lambda_\text{dyn} \mathbb{E}_t\left[ \mathcal{L}^{\text{dyn}}_t \right] + \lambda_\text{unif} \mathcal{L}^\text{unif}\tag{4}\] where the \(\lambda\) terms control the trade-off between memory representation, dynamics, and collapse.

2.2.0.3 Practice

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

For experiments, we truncate \(\mathbf{x}_t^\text{ctx}\) to a context length \(T_c\) and \(\mathbf{y}_t^\text{fut}\) to a future length \(T_f\). For computational efficiency, we estimate the expectation in \(\mathcal{L}^{\text{smt}}\) 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: \(\mathcal{O}(M+T)\) instead of \(\mathcal{O}(MT)\), where \(M\) is the memory size.

2.2.0.4 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)\) credit path for long-range associations, and gradient stability.

2.3 DAgger Memory Training (DMT)↩︎

After SMT, the RNN achieves low one-step error in predicting \((m_t, x_{t+1})\rightarrow m_{t+1}\) when \(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]\) and the encoder trajectory \([m_0, \dots, m_T]\), even with teacher forced input tokens. This drift is quantified as \(\delta_t = \text{MSE}(\hat{m}_t, m_t)\).

We introduce DAgger Memory Training (DMT), a finetuning phase that corrects this drift via on-policy imitation learning [26]. 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}\), we first compute the encoder trajectory \([m_0, \dots, m_T]\) using only \(\mathcal{E}_\phi\) and then the RNN trajectory \([\hat{m}_0, \dots, \hat{m}_T]\) using \(f_\theta\). Instead of training on SMT labels \((m_t, x_{t+1}) \rightarrow m_{t+1}\), we train on DMT labels \((\hat{m}_t, x_{t+1}) \rightarrow m_{t+1}\). Equivalently, the training loss is: \[\mathcal{L}^\text{dmt}=\mathbb{E}_t \left[ \mathrm{MSE}(\hat{m}_t, m_t) \right]\]

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\) 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\). 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 [9], [18].

“BPTT RNN” denotes the BPTT baseline. “SMT Encoder\(^*\)” generates memories \(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\(\rightarrow\)DMT RNN” denotes the RNN pretrained with SMT and finetuned with DMT, which constitutes our full method.

3.0.0.1 Architectures

We use RNN architectures based on a Transformer, MLP, and GRU [27] backbone.

3.0.0.2 Datasets

We consider character-level language modeling on TinyStories [28] as a naturalistic task requiring long-range memory [29]. As a more challenging problem, we test our method on raster-scan order pixel sequence modeling of sparse images from MNIST [30] and Sketchy [31]. This is a hard problem for RNNs [32], [33]. 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 [34].

More details on architectures, datasets, and experiments are in Appendix 7.

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\) and train all timesteps in the \(\mathcal{L}^\text{smt}\) expected value. Our tasks include the following (details of tasks are in Appendix 7.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\(\rightarrow\)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)\) credit path length, compared to SMT’s \(\mathcal{O}(1)\). Further analysis in Section 3.7.0.1 confirms the difference in gradient stability in both methods.

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\(\rightarrow\)DMT. Figure [fig:sketchy95generations] shows images generated by an SMT\(\rightarrow\)DMT RNN trained on Sketchy.

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\) optimization steps on token batches of shape \(B\times T\) (number of sequences \(\times\) 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 (\(\sim\) 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 [35]. 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 6.

Figure 5 shows the results. In sequential compute, SMT Encoder and SMT\(\rightarrow\)DMT RNN are significantly more efficient than BPTT with the Transformer and MLP backbones. In data, SMT Encoder and SMT\(\rightarrow\)DMT RNN has approximately the same data efficiency as BPTT with the Transformer and MLP backbones on TinyStories. However on MNIST, SMT Encoder and SMT\(\rightarrow\)DMT RNN shows significantly better data efficiency. This result is explained by the short vs long range memory information requirements of natural language [36] vs pixel sequence modeling [37]. SMT\(\rightarrow\)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\(\rightarrow\)DMT along three axes: context length, memory state size, and model parameter count. For the first two, we logarithmically sweep \(T_c\) and the number of memory tokens in the Transformer-based RNN (Figure 12). For model scaling, we vary the width and depth of the RNN, encoder, and decoder. We use the TinyStories domain for these experiments.

Figure 6 shows that SMT\(\rightarrow\)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 [fig:model95scaling] presents the parameter scaling results. The SMT encoder follows a standard power-law-like scaling trend. The SMT\(\rightarrow\)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.

3.5 Compression as a Scaling Axis↩︎

Neural scaling laws predict the relationship between a resource (e.g. compute, data) and a desired property (e.g. validation loss, benchmark accuracy) [38], [39]. Can the desired property instead be compression [40]? 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 7 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 [41], [42], scaling along this compression axis may be a desired direction forward for future sequence models. Notably, Transformers perform no compression of the past [9], which may explain their training efficiency.

3.6 Ablations↩︎

3.6.0.1 Predictive State and Detached RNN

The impact of the predictive state objective (Equation 2 ) is evaluated by sweeping the future length \(T_f\), while keeping \(T_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}\) 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\) to be a Markovian representation.

Figure 8 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\) or joint training. When \(T_f\) is large enough, there is a \(\mathcal{O}(1)\) credit path length between the needle and the answer at all timesteps. Interestingly, when \(T_f\) is small, there exists no credit path to learn early timestep memories, yet joint training still learns effectively, even when \(T_f=1\). Credit must be propagating through the RNN dynamics from \(m_{T}\) to \(m_{T-1}\), and so on, to \(m_0\). 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\) to \(m_{t-1}\) through \(f_\theta\); \(T\) such gradients steps sends information \(T\) steps back in the sequence. This implies that solving \(T\) sequence length credit assignment task when \(T_f=1\), requires at least \(T\) gradient optimization steps. This credit amortization phenomenon is reminiscent of value bootstrapping in RL [43].

3.6.0.2 \(\lambda\) Coefficients

The values of \(\lambda_\text{dyn}\) and \(\lambda_\text{unif}\) are swept here to check their effects. Figure 13 shows the results. The best RNNs require \(\lambda_\text{dyn}= 0.1\), and \(\lambda_\text{unif}=0.001\). When \(\lambda_\text{unif}=0\), although the RNN performance is preserved, the memory space is collapsed, as indicated by \(\mathcal{L}^\text{unif}\).

3.6.0.3 Drift and DMT

As described in Section 2.3, RNN suffers from drift post-SMT. Figure 9 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↩︎

3.7.0.1 Gradient Properties of BPTT and SMT

The fundamental difference between BPTT and SMT in long-range credit assignment is dictated by their gradients. Figure [fig:gradient95vanishing] shows the gradient magnitude of \(m_t\), \(\|\frac{\partial L}{\partial m_t}\|\), at different \(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\), 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.

3.7.0.2 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 [18], [20].

Here, we compare an SMT\(\rightarrow\)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 10 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 [44]. This result reflects the distinct inductive biases of the architectures: Transformers behave like growing lookup tables in context, while RNNs update finite states [9]. The latter is a better inductive bias for generalization.

3.7.0.3 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 11. 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 18 and Figure 19 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.

4 Related Works↩︎

4.0.0.1 Recurrent Neural Networks (RNNs)

RNNs were studied extensively early in AI because their recurrence mechanism resembles biological brains [45] and can be applied to any sequential task [46]. Many different algorithms were proposed for learning, including random guessing [47], evolutionary algorithms [48]–[52], hebbian learning [53]–[56], real-time recurrent learning [57], and other algorithms [32], [58]–[61]. BPTT is the only widely adopted algorithm [3].

However, it has repeatedly been shown that BPTT produces unstable gradients that vanish, explode, or exhibit high variance [4], [5], [62], [63]. Several directions address this issue. One direction focuses on architectural modifications, including residual connections [64], [65] and gating mechanisms [66], culminating in the development of the LSTM [67] and GRU [27]. A parallel direction addressed gradient instability through orthogonal weight parameterizations to prevent exponential growth or decay across time [68]–[73]. Others explored external memory [74]–[76], hierarchical modeling [77]–[80], other unique directions [81]–[83].

Recently, there has been renewed interest in RNNs in the form of linear state space models [15], [16], [84], linear attention models [14], [17], [85]–[87], and even nonlinear RNN models [88]–[91]. Recurrent computation more generally has been reappearing across paradigms including in diffusion [92], looped Transformers [93], and reasoning [94], [95].

4.0.0.2 Time-Parallel Training

Transformers revolutionized sequence modeling [8] largely because they have time-parallel training [7] (unlike prior attention methods [96]), which is crucial for leveraging modern hardware [35], [97] to scale performance [38], [39]. Linear RNNs gained popularity [14], [17], [86], [98] after it was realized they can be parallelized with the associative scan algorithm [99]–[101].

A recent line of work attempts to parallelize nonlinear RNNs as well [102], [103]. Rather than computing \([m_0,\dots,m_T]\) with \(m_{t+1}=f_\theta(m_t)\), they formulate the forward pass as an iterative optimization procedure. Starting with an initial guess \([m_0^0,\dots,m_T^0]\), they construct a system of \(T\) equations, \(\{m_{t+1} - f_\theta(m_t)= 0\}_{t=0}^T\), and solve this system with Newton’s method [104]. Many works have further built on this approach [105], [106]. Although appealing, this approach approximates BPTT and hence will suffer from its \(\mathcal{O}(T)\) credit path length and corresponding gradient instability, along with the added convergence worries of Newton’s method [107]. In contrast, SMT uses an encoder to train \(m_t\) to be a predictive state while satisfying \(m_{t+1}\approx f_\theta(m_t)\) and providing an \(\mathcal{O}(1)\) credit path length.

4.0.0.3 Computation Complexity Class of Models

A model’s architecture determines the problems it can theoretically solve [108]. Some tasks are inherently sequential and cannot be efficiently parallelized [109]; the circuit depth of a task is the minimum number of sequential steps required to solve it on an infinitely parallel computer [110], [111]. 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 [112]. 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 [21], [112], [113]. While such models succeed on tasks amenable to parallelization (e.g. parity tracking via associative scan [114], [115]), they systematically fail on tasks requiring deep sequential computation (e.g. tracking a chess board [19]). Interestingly, the aspect that makes models parallelizable, limits their performance on harder problems [21]. Nonlinear RNNs are one of the few classes of models where its sequential depth grows with the input sequence length [18]. Although these constraints were seen as theoretical, there is growing evidence they affect models in practice as well [20].

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 [116].

4.0.0.4 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 [6], [117], a representation that is sufficient for optimal decision making [118]. Early works interpreted PSRs as a literal vector of probabilities of future events, but have since been generalized [119]. Belief states are a related concept, which also defines a sufficient statistic of the past [120].

PSRs have been previously incorporated into RNNs [121], [122]. [123] 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)\) credit path.

4.0.0.5 Other Related Work

Our work is related to the literature on cross-architecture teacher-student distillation [124]–[128], but these works do not address the challenges of training nonlinear RNNs. In concurrent work, [129] 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, [129] 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 \(\mathcal{O}(1)\) gradient path that stabilizes credit assignment [91]. 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 [89], [130], [131]. A new line of work uses principles from diffusion models to train blocks of a feed-forward network in parallel, avoiding global backpropagation [132], [133].

5 Discussion↩︎

In SMT, the teacher model is time-parallel, and is thus constrained in expressivity [21], 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\) within a sequence. We found that training on all memories \([m_0, \dots, 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 [102].

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)\) 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.

6 Definitions↩︎

6.0.0.1 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 [5]. Under this definition, BPTT has \(\mathcal{O}(T)\) credit assignment path length, whereas Transformers and SMT have \(\mathcal{O}(1)\) path length between any two tokens.

6.0.0.2 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 [35].

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 \(\mathcal{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 \(\mathcal{O}(T)\) per optimization step.

6.0.0.3 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.

7 Experiment Details↩︎

7.1 Architectures↩︎

Our primary architecture used for most experiments is shown in Figure 12.

7.1.0.1 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, \(m_t=[m_t^1,\dots,m_t^M]\). Within a Transformer block, we use rotary position encodings [134] and RMSNorm instead of LayerNorm. We perform RMSNorm on the output memory tokens for stability. Figure 12 shows the encoder architecture.

7.1.0.2 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 [134] and RMSNorm instead of LayerNorm. The output predictions are read out at token positions such that \(\hat{y}_{t+k}\) is a function of only \(m_t\) and \(x_{t+1},\dots,x_{t+k}\). Figure 12 shows the decoder architecture.

7.1.0.3 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.

7.1.0.4 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.

7.1.0.5 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, \(m_t^l\), and outputs a single memory token for the next timestep, \(m_t^{l+1}\). We do not RMSNorm on the output memory tokens, since that would undermine the GRU’s residual structure.

7.1.0.6 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 12 shows the readout architecture.

7.2 Datasets↩︎

7.2.1 Synthetic Tasks↩︎

7.2.1.1 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 \(\mathbf{x}=[3,4,0,2,1,3,1,0]\) the target is \(\mathbf{y}=[\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,2]\), where \(\emptyset\) 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.

7.2.1.2 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 \(\mathbf{x}=[3,4,1,2,0,0,0,0]\), the target is \(\mathbf{y}=[\emptyset,\emptyset,\emptyset,\emptyset,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.

7.2.1.3 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 \(\mathbf{x}=[1,0,2,3,0,1,0,0]\), the target is \(\mathbf{y}=[\emptyset,1,\emptyset,\emptyset,3,\emptyset,1,2]\). By varying the sequence length and state complexity (maximum stack depth), this task evaluates the algorithm’s capacity for state tracking.

7.2.1.4 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 [135], [136]. For example, given \(\mathbf{x}=[b,1,a,3,d,2,4,a]\), the target is \(\mathbf{y}=[\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,\emptyset,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.

7.2.1.5 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 [137], [138]. For each sequence, parameters \(a\) and \(b\) are sampled, then the sequence is presented as \(\mathbf{x}=[x_0,y_0,x_1,y_1,x_2,y_2,x_3,y_3,]\), where \(y_i=(a x_i + b) \mod V\), where \(V\) is the vocabulary size. Then, the target is \(\mathbf{y}=[y_0,\emptyset,y_1,\emptyset,y_2,\emptyset,y_3,\emptyset]\). 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.

7.2.2 Natural Tasks↩︎

7.2.2.1 TinyStories

TinyStories is a curated dataset of short stories generated by OpenAI’s GPT-4 [28]. 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.

7.2.2.2 MNIST

MNIST is a classic image dataset consisting of handwritten digits [30]. Rather than performing classification or 2D image generation, we consider the problem of 1D pixel-sequence modeling. The original 28\(\times\)​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.

7.2.2.3 Sketchy

Sketchy is an image dataset of human-drawn sketches [31]. Rather than performing classification or 2D image generation, we consider the problem of 1D pixel-sequence modeling. The original images are resized to 64\(\times\)​64 using Lanczos resampling, and the pixels are binarized. Non-overlapping 2\(\times\)​2 patches are tokenized, yielding a vocabulary of \(2^{2\times2}=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.

7.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.

7.3.0.1 Synthetic Experiments

We set \(T_c = T_f = T\). To evaluate the expectation in \(\mathcal{L}^\text{smt}\), we compute the loss terms at all timesteps \(t\in[0, \dots, 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\times T\) input contexts to the encoder.

7.3.0.2 Other Experiments

To evaluate the expectation in \(\mathcal{L}^\text{smt}\), we compute the loss term at a single timestep sampled uniformly from \(t \sim \mathcal{U}[0,\dots,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 \(\mathcal{L}^\text{unif}\) over batches of memories from different sequences.

By default for SMT, we set \(T_c = 256\), \(T_f = 64\), \(\lambda_\text{dec}=1.0\), \(\lambda_\text{dyn}=0.1\), \(\lambda_\text{unif}=0.001\) and train for \(150000\) SGD iterations. Unless otherwise specified, models use a hidden dimension of 256 with 16 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.

8 Additional Experiments↩︎

Figure 13 shows the results of ablating the \(\lambda_\text{dyn}\) and \(\lambda_\text{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 14 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\(\rightarrow\)DMT RNN (Transformer backbone) is able to capture this structure quite well.

Figure 15 shows more samples of generations from Figure [fig:sketchy95generations]. 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 16 provides an analysis of our Sketchy RNN as it “reads” a sequence corresponding to the classic Attneave’s cat image [34]. The memory sequence does not seem to be fully interpretable, but does show significant structure. Figure 17 provides samples of generations when conditioned on partial context of Attneave’s cat image. Figure 18, 19, 20, show additional analysis of the RNNs on MNIST and Sketchy data.

9 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.

10 Sequence to Set Reframing↩︎

As described in Section 2.2, consider a hypothetical oracle memory-encoding model \(\mathcal{Q}\) that takes as input the sequence of tokens and outputs an effective compressed memory. Here we show that \(\mathcal{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.

10.0.0.1 Claim

Let \(\mathbf{x}_\text{seq}=[x_0, x_1, \dots, x_t]\) be the original sequence of tokens. We define the set \(\mathbf{x}_\text{set}=\{(x_0, 0), (x_1, 1), \dots, (x_t, t)\}\). Assume that \(\mathcal{Q}\) is a recurrent function over \(\mathbf{x}_\text{seq}\). In other words, \[m = \mathcal{Q}(\mathbf{x}_\text{seq})= f(\ldots f(f(m_{\emptyset}, x_0), x_1), \ldots , x_t)\] with \(m_{\emptyset}=\mathbf{0}\) for some function \(f\). For any such \(\mathcal{Q}\), \(\exists\) a function \(g\) such that \(g(\mathbf{x}_\text{set}) = \mathcal{Q}(\mathbf{x}_\text{seq})\).

10.0.0.2 Proof

We construct \(g\) explicitly. Define \(g(\mathbf{x}_\text{set})\) as follows: given the input set \(\mathbf{x}_\text{set} = \{(x_0, 0), (x_1, 1), \dots, (x_t, t)\}\), sort the elements in ascending order of their timestamp to recover the sequence \(\mathbf{x}_\text{seq} = [x_0, x_1, \dots, x_t]\), then apply \(\mathcal{Q}\) to this sequence.

This is well-defined because the timestamps \(\{0, 1, \dots, t\}\) are distinct integers, so the sort order is unique. The resulting sequence is identical to the original \(\mathbf{x}_\text{seq}\), and therefore \(g(\mathbf{x}_\text{set}) = \mathcal{Q}(\mathbf{x}_\text{seq}) = m\).

Moreover, \(g\) is permutation-invariant: any permutation of the elements of \(\mathbf{x}_\text{set}\) yields the same sorted sequence and thus the same output.

Since \(\mathcal{Q}\) was arbitrary, this construction applies to every recurrent \(\mathcal{Q}\), completing the proof. \(\blacksquare\)

10.0.0.3 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.

11 Encoder Markovian Training↩︎

SMT consists of two primary objectives: future predicting with \(\mathcal{L}^{\text{dec}}\) and dynamics modeling with \(\mathcal{L}^{\text{dyn}}\). 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 \(\mathcal{L}^{\text{dec}}\) alone is sufficient for learning Markovian memories, implying that \(\mathcal{L}^{\text{dyn}}\) is theoretically unnecessary, though still practically useful.

11.0.0.1 Claim

Let \(\mathbf{x} = [x_0, x_1, \dots, x_T]\) and \(\mathbf{y} = [y_0, y_1, \dots, y_T]\) be input and output sequences. For each timestep \(t\), define \[\mathbf{x}_t^{\text{ctx}} = [x_0, \dots, x_t], \qquad \mathbf{x}_t^{\text{fut}} = [x_{t+1}, \dots, x_T], \qquad \mathbf{y}_t^{\text{fut}} = [y_t, \dots, y_T],\] with memory state \(m_t = \mathcal{E}_\phi(\mathbf{x}_t^{\text{ctx}})\) and reconstructed future \(\hat{\mathbf{y}}_t^{\text{fut}} = \mathcal{D}_\psi(m_t, \mathbf{x}_t^{\text{fut}})\). If \(m_t\) is an optimal minimal sufficient statistic of \(\mathbf{x}_t^{\text{ctx}}\) for predicting \(\mathbf{y}_t^{\text{fut}}\) given \(\mathbf{x}_t^{\text{fut}}\) at every \(t\), then the memory sequence \((m_t)\) is Markovian: \[m_{t+1} \perp\!\!\!\perp \mathbf{x}_t^{\text{ctx}} \;\Big|\; m_t,\, x_{t+1}.\]

11.0.0.2 Proof

By optimality of \(m_t\), it is a minimal sufficient statistic of \(\mathbf{x}_t^{\text{ctx}}\) for predicting \(\mathbf{y}_t^{\text{fut}}\) given \(\mathbf{x}_t^{\text{fut}}\): \[H\!\bigl(\mathbf{y}_t^{\text{fut}} \mid m_t, \mathbf{x}_t^{\text{fut}}\bigr) = H\!\bigl(\mathbf{y}_t^{\text{fut}} \mid \mathbf{x}_t^{\text{ctx}}, \mathbf{x}_t^{\text{fut}}\bigr).\] Note that \(\mathbf{y}_{t+1}^{\text{fut}} \subseteq \mathbf{y}_t^{\text{fut}}\) and \(\mathbf{x}_{t+1}^{\text{fut}} \subseteq \mathbf{x}_t^{\text{fut}}\), so optimality of \(m_t\) at time \(t\) implies it is also sufficient for \(\mathbf{y}_{t+1}^{\text{fut}}\) given \(\mathbf{x}_{t+1}^{\text{fut}}\): \[H\!\bigl(\mathbf{y}_{t+1}^{\text{fut}} \mid m_t, x_{t+1}, \mathbf{x}_{t+1}^{\text{fut}}\bigr) = H\!\bigl(\mathbf{y}_{t+1}^{\text{fut}} \mid \mathbf{x}_t^{\text{ctx}}, x_{t+1}, \mathbf{x}_{t+1}^{\text{fut}}\bigr) = H\!\bigl(\mathbf{y}_{t+1}^{\text{fut}} \mid \mathbf{x}_{t+1}^{\text{ctx}}, \mathbf{x}_{t+1}^{\text{fut}}\bigr).\] By optimality of \(m_{t+1}\), it is a minimal sufficient statistic of \(\mathbf{x}_{t+1}^{\text{ctx}}\) for predicting \(\mathbf{y}_{t+1}^{\text{fut}}\) given \(\mathbf{x}_{t+1}^{\text{fut}}\). Minimality means \(m_{t+1}\) retains no information from \(\mathbf{x}_{t+1}^{\text{ctx}}\) beyond what is predictively necessary. Since \((m_t, x_{t+1})\) already constitutes a sufficient statistic for this same prediction task—as shown above—minimality of \(m_{t+1}\) forces it to be a function of \((m_t, x_{t+1})\): \[m_{t+1} = f(m_t,\, x_{t+1})\] for some measurable \(f\). Therefore \(m_{t+1}\) is determined entirely by \((m_t, x_{t+1})\), and conditioning on these renders it independent of all earlier context: \[H(m_{t+1} \mid m_t,\, x_{t+1}) = 0,\] which is equivalent to \(m_{t+1} \perp\!\!\!\perp \mathbf{x}_t^{\text{ctx}} \mid m_t, x_{t+1}\). Hence \((m_t)\) is Markovian. \(\blacksquare\)

11.0.0.3 Implication

This result establishes that under ideal conditions—sufficient encoder and decoder capacity, infinite future horizon, and exact optimization—the memory states learned by \(\mathcal{E}_\phi\) 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: \(m_{t+1}\) can be predicted from only \((m_t, x_{t+1})\).

In practice, finite capacity and approximate optimization relax this property, leaving \(m_{t+1}\) with residual dependence on \(\mathbf{x}_t^{\text{ctx}}\) beyond \((m_t, x_{t+1})\). This gap motivates jointly training the dynamics loss \(\mathcal{L}^{\text{dyn}}\) alongside \(\mathcal{L}^{\text{dec}}\) to explicitly encourage Markovian structure in the learned memory sequence.

Relatedly, [129] provide a proof that one-step RNN dynamics with a \(T_f=1\) encoder also induce a predictive-state memory representation.

References↩︎