PaperlayerView exact revision on arXiv

Source-linked AI summary

You Only Cache Once: Decoder-Decoder Architectures for Language Models

Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, Furu Wei

arXiv:2405.05254v2cs.CL

TL;DR

Long-context language models face severe KV-cache memory and prefill-latency bottlenecks. YOCO uses a self-decoder and cross-decoder to reuse one global KV cache while preserving autoregressive behavior, and it reports competitive performance with orders-of-magnitude inference-efficiency gains, including near-perfect 1M-context needle retrieval. The paper’s “once” claim applies specifically to global KV caches, while self-decoder caches remain constant-sized.

  • Problem

    Long-context language-model inference is constrained by growing KV-cache memory and high prefill latency.

  • Method

    YOCO stacks a cross-decoder on a self-decoder, whose global KV caches are reused through cross-attention so the model caches global KV pairs once.

  • Results

    YOCO achieves competitive Transformer-level performance while improving inference efficiency by orders of magnitude, including 71.8× faster prefill at 1M context and near-perfect needle retrieval.

  • Takeaways & Limitations

    YOCO supports long-sequence language modeling with substantially lower inference memory and latency across model sizes and context lengths.

  • Takeaways & Limitations

    “Once” refers to the global KV cache; self-decoder caches remain, though their size is bounded to a constant and negligible for long sequences.

Abstract

from arXiv · show

We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once. It consists of two components, i.e., a cross-decoder stacked upon a self-decoder. The self-decoder efficiently encodes global key-value (KV) caches that are reused by the cross-decoder via cross-attention. The overall model behaves like a decoder-only Transformer, although YOCO only caches once. The design substantially reduces GPU memory demands, yet retains global attention capability. Additionally, the computation flow enables prefilling to early exit without changing the final output, thereby significantly speeding up the prefill stage. Experimental results demonstrate that YOCO achieves favorable performance compared to Transformer in various settings of scaling up model size and number of training tokens. We also extend YOCO to 1M context length with near-perfect needle retrieval accuracy. The profiling results show that YOCO improves inference memory, prefill latency, and throughput by orders of magnitude across context lengths and model sizes. Code is available at https://aka.ms/YOCO.

1 Introduction

YOCO addresses the memory and latency bottlenecks of long-context language-model inference with a decoder-decoder design that caches global KV pairs once. It retains decoder-only autoregressive behavior while improving inference efficiency and maintaining competitive language-modeling performance.

  • Motivation: The single-cache design targets long-context deployment because KV caches can otherwise consume 86GB for 512K tokens in a 65B model.A 7B model requires about 110 seconds to prefill 450K tokens and 380 seconds for 1M tokens on four H100 GPUs.
  • Architecture: YOCO stacks a cross-decoder on a self-decoder: the self-decoder creates global KV caches that cross-attention reuses.Both components use causal masking, and the architecture behaves externally like an autoregressive decoder-only Transformer.
  • Experimental results: YOCO achieves results on par with prominent Transformer language models when scaling a 3B model to trillions of training tokens.Scaling curves from 160M to 13B show competitiveness with Transformer models.
  • Inference efficiency: 80× lower KV-cache memory is reported for 65B models, while 3B inference memory falls by more than 9× at 1M tokens.At 32K tokens, overall inference memory for the 3B model is reduced two times.
  • Inference efficiency: Prefill speed improves by 71.8× at 1M context and 2.87× for 32K input, reducing 512K-token latency from 180 seconds to under six seconds.The computation flow permits prefill to exit early before entering the self-decoder.
  • Architecture: The “once” claim applies to the global KV cache; self-decoder caches remain bounded to a constant and become negligible for long sequences.This qualification is stated explicitly in the paper’s footnote.

2 You Only Cache Once (YOCO)

YOCO stacks a self-decoder beneath a cross-decoder so global KV caches are produced once, reused across layers, and support decoder-only autoregressive generation with lower inference cost.

  • 2 You Only Cache Once (YOCO): YOCO uses self-decoder layers to produce shared global KV caches, which subsequent cross-decoder layers reuse through cross-attention.The architecture behaves externally like a decoder-only model while reusing the same caches across cross-decoder layers.
  • 2 You Only Cache Once (YOCO): Efficient self-attention keeps the self-decoder's inference cache bounded to a constant, such as a sliding-window size rather than input length.This supports the architecture's single global-cache design for long sequences.
  • 2 You Only Cache Once (YOCO): YOCO requires O(N + CL) caches, compared with N × L keys and values for Transformer decoders, yielding roughly L-times lower KV-cache memory.N is sequence length, C is a constant cache size, and L is the number of layers.
  • 2 You Only Cache Once (YOCO): Shared caches reduce the inference memory bottleneck, allowing more tokens to be served and larger batch sizes that improve throughput.The architecture is designed to reduce serving costs and improve inference performance.
  • 2 You Only Cache Once (YOCO): Early exit before the cross-decoder during prefilling reduces forward computation to the self-decoder while preserving the final output.Only half the layers are needed for forward computation, giving at least half prefilling latency reduction.
  • 2 You Only Cache Once (YOCO): YOCO reduces prefilling latency from 180 seconds to less than 6 seconds at 512K context and provides about three times speedup at 32K length.The comparison uses a Transformer with optimized inference, including Flash-Decoding and kernel fusion.

3 Design Choices of Self-Decoder

YOCO’s self-decoder can use gated retention or sliding-window attention to encode reusable global information with efficient inference memory. Gated retention offers parallel, recurrent, and chunkwise computations, while sliding-window attention bounds cache growth by a fixed window.

  • Gated Retention: Gated retention unifies parallel, recurrent, and chunkwise recurrent computation while producing equivalent results.Parallel or chunkwise forms support training; recurrent computation maintains an intermediate state during autoregressive inference.
  • Gated Retention: Gated retention adds data-dependent, head-wise decay gating to retention, supporting training parallelism and low-cost inference.The temperature term encourages the decay toward 1 for better memorization, while head-wise decay supports tensor-core utilization.
  • Gated Retention: Chunkwise retention divides computation into inner-chunk and cross-chunk parts for linear memory complexity on long sequences.The chunkwise representation combines parallelism within chunks with recurrence across chunks.
  • Gated Retention: Multi-head gated retention applies gated retention independently to heads, combines their outputs with GroupNorm, and adds a swish gate.The projection and normalization components are learnable and normalize each head before combination.
  • Sliding-Window Attention: Sliding-window attention restricts each query to a fixed window, reducing inference KV-cache complexity from O(N) to O(C).Unlike vanilla Transformer decoders, it does not attend to all previous tokens.

4 Experiments

Experiments evaluate YOCO across training-token and model-size scaling, 1M-context retrieval and language modeling, and deployment efficiency. YOCO matches Transformer-based baselines on reported scaling evaluations while reducing memory and latency and improving throughput.

  • Scaling training tokens: YOCO achieves comparable performance to Transformer language models across downstream tasks when trained with 1T and 1.6T tokens.The results also indicate that YOCO scales with the number of training tokens.
  • Scaling model size: YOCO obtains comparable validation-loss performance to the Llama-optimized Transformer from 160M to 13B parameters, while YOCOgRet outperforms Transformer and YOCOSWA.The comparison uses models trained with the same data and settings and evaluates validation loss across parameter counts.
  • Long-context modeling: YOCO-3B-1M passes the 1M-token needle-in-a-haystack test with near-perfect accuracy and achieves comparable or better multi-needle results than several long-context models.YOCO-3B-1M matches LWM-1M-text with half the model size and outperforms MiniCPM-128K and ChatGLM3-128K in the reported comparison.
  • GPU memory: At 1M context, YOCO uses 12.4GB of inference memory while Transformers use 9.4× more GPU memory.At 32K context, YOCO requires about 2× less memory than Transformer; its KV-cache design also supports 128K tokens with 1GB GPU memory in the reported example.
  • Prefilling latency: YOCO reduces Transformer prefilling from 180 seconds to less than 6 seconds at 512K context and is 2.87× faster at 32K.YOCO prefill grows linearly with sequence length, and early exit before the cross-decoder provides speedup even for shorter contexts.
  • Throughput: At 512K queries, YOCO reaches 43.1 token/s versus Transformer’s 4.5 token/s, achieving 9.6× higher throughput.The throughput comparison includes both prefilling and generation time.

5 Conclusion

YOCO is presented as a decoder-decoder architecture with competitive language-modeling performance and substantially improved inference efficiency. The conclusion also outlines extensions for deployment, multimodal modeling, and optimized KV-cache mechanisms.

  • Conclusion: YOCO achieves competitive performance with Transformers across scaling in training tokens, model size, and context length, while improving inference efficiency by orders of magnitude.The reported context-length scaling reaches 1M tokens.
  • Future perspectives: YOCO combined with BitNet is proposed to reduce deployment costs by shrinking both KV-cache memory and model-weight memory.The proposal targets systems such as Groq, where memory capacity limits model size and input-token count.
  • Future perspectives: The YOCO layout is proposed for multimodal language models because multiple self-decoders support multimodal fusion and causal streaming-video processing.The stated applications include asynchronous multimodal systems for real-time tasks such as robotics.
  • Future perspectives: Future KV-cache mechanisms could combine compression, indexed retrieval, and pre-caching for native retrieval-augmented generation and search.Cache reuse allows maintaining one index rather than creating an index for each layer.

A Chunk Parallelism for Long-Sequence Training of YOCO

YOCO uses chunk parallelism to distribute long sequences while limiting communication through local self-decoder dependencies and a single cross-decoder KV-cache gather. Its chunkwise gated-retention formulation preserves equivalent computation while improving training and prefill efficiency.

  • A Chunk Parallelism for Long-Sequence Training of YOCO: Chunk parallelism partitions long sequences across devices, reducing communication demands while accelerating long-sequence YOCO training.Cross-decoder disentanglement removes repeated layer-wise communication while preserving modeling capability.
  • A Chunk Parallelism for Long-Sequence Training of YOCO: The cross-decoder gathers keys and values only once, whereas self-decoder communication is limited to adjacent devices through local dependencies.The self-decoder examples include gated retention and sliding-window attention.
  • B Chunk-wise Representation of Gated Retention: Chunkwise recurrent gated retention splits computation into inner-chunk and cross-chunk parts for sequences whose positions satisfy n = kB + r.The formulation uses B as the chunk size and represents computation chunk by chunk.
  • B Chunk-wise Representation of Gated Retention: The chunkwise recurrent representation is equivalent to the parallel and recurrent representations, producing the same computation results.The formulation combines parallelism with recurrence rather than changing the gated-retention computation.

C Hyperparameters for YOCO-3B

The YOCO-3B configuration uses a 3072-dimensional hidden state, 26 layers, grouped-query attention, and a 4M-token training batch. Training uses AdamW with a 5T-token learning-rate schedule.

  • C Hyperparameters for YOCO-3B: The YOCO-3B model has a 3072 hidden dimension, 26 layers, 24 query heads, 8 key/value heads, and 2.83B parameters excluding embeddings.The attention configuration uses grouped-query attention.
  • C Hyperparameters for YOCO-3B: Training uses a 4M-token batch, 4096-token training length, AdamW with β = (0.9, 0.95), and a 3.2 × 10^-4 learning rate.The schedule includes 1000 warmup steps.
  • C Hyperparameters for YOCO-3B: The YOCO-3B learning-rate schedule is set for 5T training tokens.The supplied hyperparameter passage introduces this schedule after specifying the optimizer, warmup, and learning rate.

D Hyperparameters for Scaling Curves

The scaling-curve experiments vary model size from 160M to 13B under a common 2048-token training setup. They use AdamW, size-dependent learning rates, linear decay, and 40k training steps totaling 10B tokens.

  • D Hyperparameters for Scaling Curves: Scaling-curve training uses 2048-token sequences, 0.25M-token batches, AdamW, and 40k steps totaling 10B training tokens.AdamW uses β1 = 0.9 and β2 = 0.98, with weight decay 0.05.
  • D Hyperparameters for Scaling Curves: Learning rates are 1.5 × 10^-4 for 160M–1.4B models and 7.5 × 10^-5 for 2.7B–13B models, followed by linear rate decay.The schedule uses 375 warmup steps.

E Hyperparameters for Length Extension

The paper extends YOCO to 1M-token contexts through staged length extension and compares multiple efficient-attention architectures under matched parameter settings.

  • Length extension: Context length is progressively extended through 64K, 256K, and 1M-token stages.Documents longer than the training length are up-sampled, with stage-specific RoPE θ and learning-rate settings.
  • Length extension: The length-extension setup uses distinct RoPE θ and learning rates at each training stage.
  • Efficient attention: Gated retention supports parallel, recurrent, and chunkwise computation paradigms that produce equivalent results.Chunkwise retention combines parallel computation within chunks with recurrent computation across chunks and has linear memory complexity for long sequences.
  • Model comparison: The evaluation compares YOCO variants with Transformer, H3, RetNet, Mamba, and gRetNet using 160M-parameter models.The compared models use 12 layers and hidden dimension 768, with shared word-embedding and softmax-projection weights.

G.1 Fine-Grained LM Perplexity Results

The fine-grained language-modeling results report perplexity on both an overall validation set and diagnosis subsets that separate recall from regular language modeling.

  • Perplexity is reported on both the overall validation set and fine-grained diagnosis sets.
  • The AR-Hit diagnosis set evaluates associative recall capability.
  • The First-Occur diagnosis set indicates regular language-modeling performance.

G.2 Long-Context Evaluation

YOCO is evaluated for long-context modeling across four ZeroSCROLLS tasks, with input lengths extending to 16,384 tokens. Across tasks and lengths, YOCO and Transformer consistently perform better than the other evaluated architectures.

  • Evaluation setup: The long-context evaluation covers four tasks from the ZeroSCROLLS benchmark.The 160M models are further trained with 2B tokens at sequence length 16,384.
  • Evaluation setup: Perplexity is evaluated at input lengths of 4,096, 8,192, 12,288, and 16,384 tokens.
  • Long-context results: YOCO and Transformer consistently perform better than the other architectures across tasks and input lengths.Figure 12 reports answer perplexity for the long-context evaluation.
Loading 2405.05254v2…