Pretraining Recurrent Networks without Recurrence

Pretraining Recurrent Networks without Recurrence

Training recurrent neural networks (RNNs) is traditionally a slow and difficult process. These models must process information step-by-step, much like reading a sentence one word at a time. This sequential nature creates a massive hurdle for modern hardware. GPUs thrive on doing many things at once. A new paper proposes Supervised Memory Training (SMT). This method uses a "teacher" model to show the RNN exactly what its memory should look like at each step. This allows the RNN to learn in parallel, similar to a Transformer. Yet, it retains the efficient, fixed-size memory of traditional recurrent architectures.

Decoupling What to Remember from How to Update

The core challenge in training recurrent models is credit assignment (the process of attributing a loss signal to specific past actions). When a model makes a mistake at step 100, calculating which action at step 5 caused that error is incredibly complex. In a recurrent system, the "credit" signal must pass backward through every intermediate step.

The authors argue we can bypass this struggle by splitting the learning process. First, the model decides what to remember by creating a summary of the past. Second, it decides how to update that summary when new information arrives. By treating the first task as supervised learning, the researchers turn a sequential puzzle into a pattern-matching exercise.

The Bottleneck of Backpropagation Through Time

To understand why SMT is necessary, one must understand Backpropagation Through Time (BPTT). In BPTT, an RNN is "unrolled" in time. This turns the recurrence into a very long, deep chain of operations. As shown in, the gradient signal must travel backward through this entire unrolled chain.

This approach suffers from two fatal flaws. First, it is inherently sequential. You cannot easily calculate step 10 until step 9 is finished. This prevents the model from using the massive parallelism of modern GPUs. Second, as signals travel through many nonlinear layers, they tend to vanish (become too small) or explode (become too large). This makes it difficult to learn long-range dependencies (associations between distant pieces of information).

Learning from a Teacher's Memory

The SMT framework replaces BPTT with a teacher-student architecture. The researchers employ a Transformer to act as an "oracle" or teacher. This Transformer is trained to compress a sequence into a "predictive state." This is a memory representation containing only the information needed to predict the future.

A key mathematical insight makes this possible. The authors prove that any recurrent memory function can be represented as a permutation-invariant function over a set of timestamped tokens. This "sequence-to-set" reframing allows the teacher to process the past in parallel.

Once the teacher has constructed these optimal memory states ($m^*t$), the RNN's job simplifies. Instead of figuring out what is important, the RNN mimics the teacher's transitions. It is trained on one-step labels: given the current memory and the next input, can the RNN predict the teacher's next state? This is expressed as $(m_t, x$.}) \rightarrow m_{t+1

Because the teacher sees the whole sequence at once, it provides these labels for every timestep simultaneously. This enables time-parallel training. As shown in, the gradient path in SMT is $O(1)$. This means the distance the signal travels to assign credit does not grow with the sequence length.

Practitioners should note the importance of the objective function [Equation 5]. The authors found that optimal performance requires specific balances. Specifically, they report that $\lambda_{dyn}$ (the dynamics loss weight) should be 0.1. They also recommend a very low $\lambda_{unif}$ (uniformity loss weight) of 0.001 to prevent memory collapse.

Breaking the Recency Bias

The implications are visible in complex, long-horizon tasks. In synthetic experiments, SMT-trained RNNs significantly outperformed BPTT in state tracking and in-context learning .

While BPTT struggled as sequences grew longer, SMT remained robust.

This advantage is striking in "Attneave’s task," which involves modeling pixel sequences from images. A model must remember a white pixel from a previous row to understand a shape in the current row. Traditional BPTT-trained RNNs often fail here. They suffer from "recency bias" (a tendency to only rely on the most recent inputs) .

In contrast, SMT-trained models successfully reconstruct MNIST digits and human-drawn sketches .

The researchers also identified a new way to scale these models. They found that increasing training compute allows for higher memory state compression. This means the model learns to represent vast experience in increasingly tiny, efficient states .

Limits of the Imitation

SMT is not a complete replacement for recurrent training. The authors note a phenomenon called "drift." Small errors in the RNN's own predictions can accumulate during actual use (inference). Eventually, the memory trajectory wanders away from the teacher's ideal path.

To fix this, the researchers introduced DAgger Memory Training (DMT). Unlike the parallel SMT phase, DMT is an "on-policy" fine-tuning step. In this phase, the RNN uses its own predicted memories to move forward. It is then corrected when it deviates from the teacher .

While DMT solves the drift problem, it is not time-parallel. Finally, because the RNN is a student of a Transformer, it may inherit the teacher's expressivity limits. Extra post-training may be required to exceed the teacher's capabilities.

Figures from the paper