Research map / Architectures and efficiency / Attention and transformers
You Only Cache Once: Decoder-Decoder Architectures for Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes YOCO, a decoder-decoder LLM architecture (self-decoder producing one global KV cache reused by a cross-decoder) that cuts KV-cache memory and prefill latency while matching Transformer quality.
Problem
Decoder-only Transformers store KV caches for every layer, making long-context inference memory-bound and prefilling very slow.
Approach
Stack a cross-decoder on a self-decoder so that only the self-decoder's final output produces one global KV cache, which all cross-decoder layers reuse via cross-attention; the self-decoder uses constant-memory efficient attention (gated retention or sliding-window attention), and prefill can exit after the self-decoder.
Key results
- YOCO-3B trained on 1.6T tokens averages 0.636 across eight zero-shot tasks, comparable to StableLM-3B-4E1T's reported results.
- For 512K context, prefilling latency drops from about 180 seconds for an optimized Transformer to under 6 seconds with YOCO; 32K prefill speedup is 2.87x.
- At 1M context, YOCO-3B-1M uses 12.4GB total inference memory, while Transformer uses 9.4x more; 512K throughput is 43.1 vs 4.5 tokens/s (9.6x).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use YOCO instead of Transformer | 3B language modeling · Eval Harness, 1T tokens | Avg | Helped |
| Increase YOCO training from 1T to 1.6T tokens | YOCO-3B · Eval Harness | Avg: 0.634 -> 0.636 | Helped |
| Extend context length to 1M tokens | YOCO-3B · Eval Harness | Avg: 0.636 -> 0.645 | Helped |
| Use gated retention instead of sliding-window attention | YOCO · Language modeling, validation set | perplexity: 3.553 -> 3.530 | Helped |
| Use gated retention instead of sliding-window attention | YOCO · Language modeling, AR-Hit | perplexity: 1.202 -> 1.199 | Helped |
| Use gated retention instead of sliding-window attention | YOCO · Language modeling, First-Occur | perplexity: 4.094 -> 4.067 | Helped |
| Use YOCO-3B-1M rather than LWM-1M-text | 128K multi-needle retrieval · N=2 | accuracy: 0.90 -> 0.98 | Helped |
| Use YOCO-3B-1M rather than LWM-1M-text | 128K multi-needle retrieval · N=4 | accuracy: 0.76 -> 0.84 | Helped |
| Use YOCO-3B-1M rather than LWM-1M-text | 128K multi-needle retrieval · N=8 | accuracy: 0.62 -> 0.56 | Hurt |
| Use YOCO-gRet rather than Transformer | language modeling · Validation set | perplexity: 3.564 -> 3.530 | Helped |
| Use YOCO-gRet rather than Transformer | language modeling · AR-Hit | perplexity: 1.219 -> 1.199 | Helped |
| Use YOCO rather than Transformer | inference at 512K context · 512K input context | prefill latency: 180 seconds to less than 6 seconds | Helped |
Surprising findings
- YOCO with gated retention outperforms both plain Transformer and YOCO with sliding-window attention in validation loss across scales, attributed to complementary attention/retention inductive biases.
Limitations
- Main training budget was 1.6T of a planned 5T tokens due to resource limits.
- Long-context and efficiency results use a 3B model; reduction ratios for larger models are argued rather than fully measured.
- Multi-needle and long-context evaluations are limited to a few benchmarks and no long-instruction tuning.
Experimental setup
Datasets: curated pretraining corpus similar to StableLM-3B-4E1T book data repository-level code data
Benchmarks: ARC-C ARC-E BoolQ HellaSwag OBQA PIQA Winogrande SciQ Needle-in-a-Haystack multi-needle retrieval ZeroSCROLLS AR-Hit/First-Occur perplexity
Metrics: accuracy validation loss perplexity GPU memory prefilling latency throughput needle retrieval accuracy
Baselines: Llama Transformer OpenLLaMA-v2-3B StableLM-base-alpha-3B-v2 StableLM-3B-4E1T YaRN-Mistral-128K LWM-1M-text MiniCPM-128K ChatGLM3-128K Mamba RetNet Hybrid H3
Builds on: Transformer RetNet StableLM-3B-4E1T Flash-Decoding grouped-query attention
Closest papers by meaning
- Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation (NeurIPS 2025 · similarity 0.88)
- KV Cache is 1 Bit Per Channel: Efficient Large Language Model Inference with Coupled Quantization (NeurIPS 2024 · similarity 0.87)
- Reducing Transformer Key-Value Cache Size with Cross-Layer Attention (NeurIPS 2024 · similarity 0.87)
- Block Transformer: Global-to-Local Language Modeling for Fast Inference (NeurIPS 2024 · similarity 0.87)
- VL-Cache: Sparsity and Modality-Aware KV Cache Compression for Vision-Language Model Inference Acceleration (ICLR 2025 · similarity 0.87)
- SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.