Training lets looped language models reuse inference caches
BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing · arXiv:2610.31633
In plain words
Some language models repeat the same layers to do more computation, but each repeat usually stores its own attention cache. This paper introduces a training method that helps earlier repeats use the cache from the last repeat, reducing the need to keep several caches.
The method first trains with varied repeat counts, then teaches the first repeat to work with the last one’s cache. In tests, this often preserved quality, though varied-depth training alone did not fix every model.
The research
Looped language models can add computation by repeating the same layers without adding parameters, but their inference cache grows with each repeat. Can they keep that quality-scaling approach while using less cache memory, without harming the model’s language quality?
The authors pair stochastic-depth pretraining with BeNice!, which trains a model’s first recurrence to use the final recurrence’s cache while preserving language quality.
The authors first trained models with stochastic recurrence depth, then added BeNice! to teach the first recurrence to use the final recurrence’s cache while retaining the standard language-modeling loss. They compared shared-cache and unshared-cache evaluations, including against standard fine-tuning.
On a public 770M model, shared-cache validation perplexity was 11.86 after BeNice! and 21.61 after standard fine-tuning, while unshared-cache scores were nearly equal. This supports the claim that BeNice! can preserve quality when caches are shared in this tested model.
The inference speed and request-time results come from simulated replays of production-derived requests using proxy models, not live serving.
The convergence proof also depends on contractive recurrent dynamics and Lipschitz key and value projections, conditions the paper does not guarantee for every looped architecture.
What the evidence shows
- The memory problem
Looped language models repeat the same layers, allowing recurrence depth to grow without adding model parameters. But conventional inference keeps a separate key-value cache for every recurrence, so cache memory grows with the number of repeats.
The paper asks how these models can reuse the final recurrence’s cache across earlier ones while preserving quality.
- A first training step
The authors’ two-stage method combines stochastic-depth pretraining with BeNice! post-training.
In their size-matched 770M Parcae comparison, stochastic-depth training produced a shared-minus-unshared validation-perplexity difference of +0.01 points, while static-depth training produced +3.10 points. The models shared the same mean recurrence depth, data, and backpropagation depth; the result associates stochastic depth with shareability in this setup.
- Why another step was needed
Stochastic-depth pretraining did not solve the problem for every tested model. With eight recurrences on ClimbMix validation data, the Parcae 1.3B model’s final hidden representations differed by a mean relative 177.5% between shared-cache and unshared-cache evaluation.
That gap indicates the two evaluations approached different points in this example, helping explain why a further correction was needed.
- Correcting cache mismatch
BeNice! trains the first recurrence to use the final recurrence’s cache while retaining the usual language-modeling loss.
Across four newly trained Parcae size classes, shared-cache scores improved by up to 8.35 points on Core, a 22-task language suite, and up to 4.56 points on Core-Extended, a 46-task suite. These are maximum gains across model classes.
In a separate 770M public-checkpoint test, BeNice! also brought shared-cache validation perplexity close to the unshared result.
- A bounded workload result
In a simulated replay of production-derived requests at two requests per second, sharing reduced the 90th-percentile wait for the first output token from 139.2 to 23.2 seconds and increased throughput by 23.6% against separate caches.
This supports a performance benefit under that simulated high-load condition. It does not establish a live-serving result.
Deep
- Original paper authors
- Hayden Prairie, Jacob M. Ryan, Taylor Berg-Kirkpatrick, Daniel Y. Fu
- Submitted to arXiv
- License
- CC BY 4.0