Teaching RNNs to Remember Without Unrolling Time

For three decades, the central obstacle to training recurrent neural networks has not been the architecture itself but the way we assign credit through it. Backpropagation through time (BPTT) remains the default method, yet it carries two well-documented liabilities: it is inherently sequential, and it propagates gradients across the full length of a sequence, where they tend to vanish or explode. The paper "Pretraining Recurrent Networks without Recurrence" by Akarsh Kumar and Phillip Isola proposes a method, Supervised Memory Training (SMT), that questions whether recurrent credit propagation is necessary at all. The answer they arrive at is a qualified no, and the reasoning behind that answer is worth examining carefully.

The Problem SMT Tries to Solve

An RNN maintains a compressed memory state $m_t$ that it updates as it observes new tokens. The training difficulty arises because a piece of information written into memory at time $t$ may only become useful many steps later. BPTT handles this delayed utility by unrolling the network across all timesteps and backpropagating through the entire computation graph. The longest gradient path between two tokens is therefore $O(T)$, which is precisely where instability creeps in. The product of many Jacobians either shrinks toward zero or grows without bound, and the sequential structure prevents the kind of parallelism that made Transformers scale so effectively.

The framing in the paper is elegant: suppose we somehow had access to the optimal memory state $m_t^$ at every timestep. Then training the RNN would reduce to learning the one-step transition $(m_t^, x_{t+1}) \rightarrow m_{t+1}^*$ as a plain supervised regression problem. No unrolling, no temporal gradient chains, just a sequence of independent supervised targets that can be computed in parallel. The entire challenge collapses to a single question: where do these memory labels come from?

The Predictive State Insight

SMT's answer rests on a definition of what memory should be. The authors adopt the notion of a predictive state from the literature on predictive state representations: an effective memory is a sufficient statistic of the past for predicting the future. It should retain everything relevant for future prediction and discard everything else.

The subtle move is in how the past is represented. We habitually treat the past as a sequence, which implies sequential computation. But if each observation is tagged with its timestamp, the past can be represented losslessly as a set of timestamped events rather than an ordered sequence. Under this reparameterization, the optimal memory becomes a permutation-invariant function of the set, and permutation-invariant functions can be computed by models that operate fully in parallel over time. This is the conceptual hinge of the whole method.

In practice, the authors train a Transformer encoder to compress past context into a memory representation, paired with a decoder that uses that memory to predict future tokens. This teacher encoder-decoder pair is trained on the predictive state objective. Once trained, the encoder produces a memory trajectory $m_0, m_1, \ldots, m_T$ across a sequence, and these become the supervision targets for the RNN. The RNN, in turn, only has to learn the one-step update rule that walks along this trajectory.

The result is a clean decoupling:

• What to remember (the memory representation) is a non-sequential problem solved by the parallel Transformer teacher.
• How to update memory (the memory dynamics) is sequential in nature but can be supervised one step at a time, with no backpropagation across steps.

This is why the longest gradient path between any two tokens becomes $O(1)$ rather than $O(T)$. Each one-step transition is trained as an isolated supervised example.

Where SMT Sits in the Landscape

It is useful to position SMT against the three dominant approaches to sequence modeling, because the paper is essentially arguing it captures the strengths of each without their characteristic weaknesses.

Transformers solved time-parallelism and credit assignment together through attention, but they store the entire history of token representations and attend over all of them. Memory grows with sequence length, which is prohibitive for very long or unbounded sequences. Sliding-window variants bound the memory but lose access to anything before the window.

Linear attention RNNs (including the recent state space models and linear recurrence families) achieve parallel training, stable credit assignment, and fixed memory. Their limitation is expressivity: a linear transition function constrains the class of representable functions, and there is a body of theoretical work showing such models fail at certain state-tracking tasks.

Nonlinear RNNs are maximally expressive and have fixed-size memory, but they have historically been trapped with BPTT and its $O(T)$ gradient pathology.

SMT's pitch is to deliver time-parallel training, $O(1)$ credit assignment, fixed-memory inference, and full nonlinear expressivity at once. The empirical results reported on language modeling and pixel sequence modeling show SMT outperforming BPTT for pretraining several RNN architectures, particularly on long-range dependencies. The authors are explicit that SMT is a pretraining method: the RNN learns to imitate the teacher's memory trajectory, then requires lightweight post-training to correct for drift and to surpass the ceiling imposed by the teacher.

My Take

The most interesting aspect of SMT is what it reveals about the nature of the RNN training problem. It has long been treated as fundamentally temporal, but SMT demonstrates that the temporal difficulty is largely an artifact of trying to learn the representation and the dynamics simultaneously through one backward pass. By outsourcing the representation problem to a parallel teacher, the genuinely sequential part of the problem becomes trivially supervisable.

This is conceptually close to distillation and to decoupled neural interfaces like synthetic gradients, but the framing as predictive states gives it a cleaner theoretical grounding. It is also, in a real sense, a method for learning world models: the one-step transition $m_t \rightarrow m_{t+1}$ under input $x_{t+1}$ is exactly the form of a latent dynamics model. I suspect the technique will turn out to be useful well beyond RNN pretraining, in any setting where one wants compact latent state with learned transitions.

There are honest limitations to flag. First, the student RNN can only be as good as its teacher unless post-training pushes it further; the teacher's predictive state objective becomes the ceiling. Second, the teacher is a Transformer, which means SMT does not eliminate the quadratic-context cost at training time. It moves that cost into a teacher whose outputs are then compressed, but you still pay it once. Third, the assumption that the past is losslessly a permutation-invariant set of timestamped events is clean in theory, but the practical fidelity of the learned encoder under finite capacity is an empirical question that deserves scrutiny across diverse task distributions. Finally, drift is a genuine concern: an RNN trained to imitate one-step transitions can accumulate error over long rollouts, and how much post-training is required to control this drift will determine whether SMT is practical at scale.

Open Questions

Several directions seem worth pursuing. Can the teacher itself be replaced by a cheaper architecture, removing the quadratic dependency entirely from the pipeline? How does the quality of the predictive state degrade as memory dimensionality shrinks, and is there a principled way to allocate memory capacity to the most predictively useful features? And perhaps most intriguing: if biological memory really does construct temporally compressed abstractions rather than storing entire histories, does the decoupling of representation from dynamics offer a useful model for how such abstractions might be acquired without backpropagating across a lifetime of experience?

SMT does not retire BPTT, but it reframes the credit assignment problem in a way that makes the hardest part disappear. That reframing alone is a contribution worth taking seriously.