akarshkumar0101/smt

TL;DR: SMT (Supervised Memory Training), from Akarsh Kumar and Phillip Isola at MIT, trains nonlinear RNNs without ever backpropagating through time. A Transformer teacher is trained to compress the past into memory tokens that suffice to predict the future; those memory states become supervised labels, and the RNN is trained on one-step transitions (mₜ, xₜ₊₁) → mₜ₊₁ — fully time-parallel, with an O(1) gradient path between any two tokens. The repo is a clean, minimal PyTorch implementation with MNIST-pixels and TinyStories experiments; the paper claims SMT beats BPTT on these tasks, though the README itself carries no numbers.

The idea

BPTT has two structural problems: it is sequential in time (no parallelism across the unroll) and its gradient path length grows with sequence length (vanishing/exploding gradients, poor long-range credit assignment). Transformers dodged both by abandoning recurrence — at the cost of O(T) inference state. SMT tries to keep recurrent inference but train like a Transformer.

The trick is to decouple what to remember from how to update memory:

1. Teacher (what to remember). A Transformer encoder maps a context window to a fixed set of memory tokens; a decoder must predict future tokens from those memory tokens alone. This is a predictive-state objective — the memory is forced to retain exactly the information about the past that's useful for the future. The encoder can see the whole context at once, so it's trained with ordinary parallel attention, no recurrence.

2. Student RNN (how to update memory). Given teacher memories mₜ at every position, RNN training collapses to plain supervised regression on one-step transitions: predict mₜ₊₁ from (mₜ, xₜ₊₁). Every timestep is an independent training example — trivially parallel, and no gradient ever flows through more than one step of recurrence.

3. DMT (fixing distribution shift). Pure one-step imitation suffers the classic compounding-error problem: at inference, the RNN feeds on its own imperfect memories, which drift off the teacher's distribution. The DAgger Memory Training variant rolls the RNN out for a stretch, then has the teacher relabel targets from the RNN's own visited states — the standard DAgger fix, here applied to latent memory rather than actions. In the provided configs, DMT is a short fine-tune (1,500 iters) on top of a long SMT run (150,000 iters).

2. Student RNN: supervised one-step transitions (all t in parallel) mₜ xₜ₊₁ RNN cell (Transformer) m̂ₜ₊₁ teacher label mₜ₊₁ regression loss No unrolling: gradient path between any two tokens is O(1). DMT relabels the RNN's own rollout states (DAgger).

What's in the repo

• src/model.py — teacher (encoder+decoder) and RNN, SMT/DMT forward passes, generation. The RNN cell is itself a small Transformer (d_model 256, depth 8, 8 heads, rotary embeddings) whose state is 16 memory tokens (mem_size 4096).
• src/train_smt.py, src/train_dmt.py — training scripts with full CLI configs (tyro).
• main.ipynb — walkthrough of configs, launching runs, analyzing the RNN. tutorial.ipynb is "coming soon."
• Dataset code for MNIST pixel sequences and TinyStories (byte-level, vocab 256), stored as one long token array — easy to swap in your own data.

No pretrained weights, no benchmark tables in the README, no reported wall-clock or loss numbers — the "SMT outperforms BPTT" claim lives in the paper (the arXiv/project-page links in the README are present but the header is partly commented out). Apache-2.0, actively updated.

Assessment

The framing is genuinely different from the recent linear-RNN wave (Mamba, RWKV, linear attention), which buys parallel training by restricting the recurrence to be linear/associative. SMT keeps the recurrence fully nonlinear and instead changes the training signal — closer in spirit to predictive state representations and target-propagation than to architecture design. It's also conceptually related to distilling Transformers into RNNs, but here the teacher is trained jointly and specifically to emit imitable memory states.

Caveats to hold onto: everything shown is small (≈256-dim, byte vocab, MNIST/TinyStories), so scaling is unproven. Final RNN quality is capped by the teacher, and the teacher's encoder only sees a finite context window (T=320 in the configs) — so it's unclear how the RNN learns dependencies longer than what the teacher itself can see. The DMT fine-tune suggests one-step imitation alone isn't enough, and whether DAgger-style relabeling holds up at very long rollouts is exactly the open question. Compute-wise this is single-GPU friendly.

Try it

git clone https://github.com/akarshkumar0101/smt && cd smt
pip install numpy pandas xarray matplotlib tqdm einops einop tyro \
  jupyterlab ipywidgets x_transformers huggingface_hub transformers datasets
cd src
python my_datasets/mnist.py          # or tinystories.py
python train_smt.py --seed=0 --dataset="mnist" --n_iters=150000 \
  --save_dir="./runs/smt_mnist_0/"   # defaults match the README config

Then open main.ipynb for the guided walkthrough, and run train_dmt.py pointed at the saved rnn.pkl/teacher.pkl for the DAgger fine-tune.