Test-time gradient descent beats neural cache: WikiText-2 perplexity 44.3

2026-08-25

Edinburgh's dynamic evaluation adapts an LM with gradient descent on recent history. On WikiText-2, AWD-LSTM perplexity falls from 66.1 to 44.3, beating neural cache at 52.0.

What problem this solves

RNN language models miss a blunt signal: a word that just appeared in a document is much more likely to appear again, and the hidden state often fails to use that fact. A unigram cache on top of an RNN already lowers perplexity, which is a public admission that the network is bad at recent frequency.

Pointer networks, copy nets, and the neural cache can handle direct repeats, the same symbol coming back. They do not handle indirect repeats: synonyms, related words, topic, style. The neural cache stores past hidden states paired with the next token and, at a new state, boosts words whose stored keys are nearby. It cannot change recurrent dynamics, and it cannot raise the probability of a symbol it has not yet seen in the test sequence. Character-level models make that limitation sharper, because a single character carries almost no meaning on its own.

Training fits a global distribution. At test time each stretch of text comes from a local distribution that drifts as the topic changes, or as benchmarks concatenate many documents into one long sequence. Fixed parameters are right for the average and wrong for the current local.

Method

Dynamic evaluation adapts the trained weights with gradient descent on text the model has just scored, then uses the new weights on the next chunk.

The procedure still yields a valid autoregressive log-probability, because the update uses only tokens already evaluated. Cost is one forward pass, one backward pass, and a cheap update per segment.

Mikolov (2010) and Graves (2013) already ran this idea with SGD at every step. Results were mixed. This paper turns the update into a recipe that holds up. Longer segments give cleaner gradients and fewer updates. A decay term pulls weights back toward the trained parameters so older adaptations fade and recent history counts more, which matches a drifting local distribution. Per-parameter step sizes follow an RMSprop-style scale, but the mean squared gradient is collected on the training set rather than on the test stream, because early in a test sequence there is nothing reliable to average. Decay is scaled the same way, so high-RMS weights decay faster.

The learning rate has to be swept on the validation set. Mixed older results mostly came from skipping that sweep and reusing training-time optimizer settings. Sprechmann et al. later did exactly that and saw almost no test-time gain. One pass over validation per setting is cheap; the learning rate is the knob that matters.

A sparse variant leaves the original weights alone and introduces a zero-initialized matrix M so the hidden state becomes h + Mh, optionally on a subset of units. On Hutter Prize they adapt a 500 by 500 matrix, 250k parameters, about 0.5% of a full-model update.

Results

The base models were the public SOTA of the moment: AWD-LSTM for words, a 2800-unit mLSTM for characters. Dynamic-evaluation hyperparameters are tuned on validation; both static and dynamic numbers are test-set.

SetupPTB PPLWikiText-2 PPL
AWD-LSTM static (rerun)57.766.1
AWD-LSTM + neural cache52.852.0
Traditional dynamic eval (SGD, step 1)53.549.0
Final dynamic eval (RMS + RMS prior)51.144.3

WikiText-2 moves the most. Documents stay in order, the vocabulary is 33k, and cross-article dependence is exactly what an adaptive method should use. On PTB, even the old SGD-every-step rule already reaches 53.5; each change in the recipe shaves a bit more, down to 51.1. On word-level text8 the static model is 87.5, neural cache 75.1, dynamic evaluation 70.3, so the gap survives a larger corpus.

Character-level, same pattern. Hutter Prize: mLSTM static 1.24 bits/char, sparse dynamic eval 1.13, full 1.08. Character-level text8: 1.27 down to 1.19. The advantage shows up after a few hundred characters. On Hutter it peaks around 2k to 3k characters, so the fit is local rather than a new global model of the whole test set. On Spanish European Parliament text, which Hutter barely contains, the gap keeps growing. After 10k Spanish characters, a static model’s continuation snaps back to Wikipedia-like English; the dynamic model keeps emitting Spanish-shaped words for the whole sample.

Indirect repeats show up in a probe on WikiText-2. One update on the word production raises its own log probability by 2.42, construction by 1.29, recording by 1.27, development by 1.14. The vocabulary median is −0.66: related words go up, unrelated words go down.

Why it matters

This is test-time adaptation, not a new architecture. The base model stays put; recent history is used to take gradient steps. In 2018 that was enough to reset the public word-level and character-level language modelling numbers, and to beat a neural cache that can only replay symbols it has already seen.

The transferable piece is the operation, not the LSTM scores. For any autoregressive task, adapting weights on recent context generalizes past a cache, to related words and to style. The sparse form cuts adaptation parameters to 250k, which is the opening for batched inference where each sequence would otherwise need its own copy of the model. A Transformer KV cache stores activations. Dynamic evaluation writes into the weights. Both ideas are still around, usually under names like test-time training or inner-loop adaptation.

Limitations

The authors draw the boundary themselves: autoregressive sequence modelling only, and the true token has to be observed after the prediction. Mini-batching at test time means one parameter copy per sequence, which blows up memory; the sparse variant exists for that. There is an extra backward pass. Segment length, learning rate, and decay all need a validation sweep, so test time is not hyperparameter-free. Get the learning rate wrong and the older “almost no gain” reports come back.

Every experiment is an LSTM or mLSTM. No Transformer, no instruction-tuned model. WikiText-2 is unshuffled; the paper never isolates how much of the gain is “document order is exploitable.” Hutter Prize includes XML and non-Latin bytes, so 1.08 bits/char is a poor comparison to English-only compression figures. The production probe is a single case, not a systematic study; the Spanish samples sit in the appendix with no automatic metric. Usefulness for speech and translation is claimed, not measured.

Terms

Source

What people are saying

All paper explainers