Titan-Llama bolts external neural memory onto a frozen Llama: long-document QA accuracy rises from 23.87% to 40.51%
Synopsis
Titan-Llama converts the first four layers of Llama-3.2-1B-Instruct to blockwise attention and attaches a fixed-size neural memory module to each KV head, using test-time training to compress historical key-value bindings while the backbone stays frozen. Against segmented attention windows alone, accuracy rises from 23.87% to 40.51% on QuALITY and from 14.50% to 40.75% on LongHealth. In the evaluated configurations, prefill is more than 20% faster, decode 37% faster at medium context lengths, and steady-state decode memory 24% lower.
Interpretation
A pre-trained attention model can gain fixed-size long-context memory without retraining from scratch: Titan-Llama post-trains external weights amounting to about 1% of the total parameter footprint, making history cut off by segmented attention usable again. Earlier routes either kept the full KV cache, making long-context inference increasingly expensive, or switched to linear attention or state-space models, which generally requires training a new model rather than reusing an existing attention backbone. On Llama-3.2-1B-Instruct, compared with a baseline using segmented attention windows only, QuALITY accuracy rises from 23.87% to 40.51% and LongHealth from 14.50% to 40.75%.
Historical attention is split in two: the recent window keeps exact attention, while older keys and values are read from the neural memory module and added to the local output through a per-channel learned gate. The neural memory module is a three-layer network with hidden dimension 64 and GELU activations, one per KV head, totaling 17.23M trainable parameters; it is updated by test-time training, meaning gradient descent on the module weights at inference. After each fixed-size token chunk, the module is updated once by a squared-error key-value reconstruction loss using learned momentum, decay and step size; evaluation runs on a single 80GB NVIDIA A100 with training sequences of at most 16K tokens.
Perspective
The demonstrated setting is Llama-3.2-1B-Instruct: the first four layers use blockwise attention with one neural memory module each, the remaining layers keep full attention, training sequences are at most 16K tokens, and inference is measured on a single 80GB A100 at batch size one.
Key-value reconstruction is a surrogate for useful memory, not a guarantee of exactly reproducing historical attention; performance declines at 32K and beyond the 16K training length, multi-value retrieval stays below full attention at every length, and chunk size affects RULER results differently by task and length.