Supervised Memory Training (SMT) trains nonlinear recurrent neural networks by teaching them what to remember and how to update that memory separately. MIT researchers Akarsh Kumar and Phillip Isola use a Transformer encoder to compress past inputs into a state that helps predict future inputs. The RNN learns individual transitions between those states, allowing pretraining across time steps in parallel. Paper
Figure 1: BPTT and SMT. Credit: Akarsh Kumar and Phillip Isola, Pretraining Recurrent Networks without Recurrence, CC BY 4.0.
Standard backpropagation through time sends learning signals through a chain of recurrent computations, where gradients can vanish or explode. SMT uses supervised memory targets to shorten that path. At inference, the RNN processes inputs sequentially with fixed-size memory.
A second stage, DAgger Memory Training, corrects errors that accumulate when the RNN uses its own predicted memories. This fine-tuning stage requires sequential execution.
Tests cover synthetic memory tasks, TinyStories character prediction, and pixel sequences from MNIST and Sketchy. SMT followed by fine-tuning improves long-range learning and sequential-compute efficiency over conventional training for the tested Transformer- and MLP-based RNNs. The tested GRU variant fails because its memory representations collapse. On TinyStories, larger models, context windows and memory states improve performance. Methods and results
