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

Summary

This paper proposes YOCO, a hybrid model that combines gated linear attention with standard attention (SA). The model stacks efficient self attention (ESA) in the first $L/2$ layers, succeeded by another $L/2$ cross-attention layers. Notably, the output of the last ESA is shared across subsequent CA layers, thereby achieving significant parameter reduction and enabling exceptional key-value (KV) cache compression, critical for optimizing inference. Two ESA variants are evaluated: sliding window attention and a novel gated retention method, which incorporates data-driven head-wise decay over retention. Upon scaling YOCO to a 3-billion-parameter model trained on a corpus of 1 trillion tokens, the authors report superior performance relative to Llama-like architectures in language modeling tasks. They also conduct some analysis on long-seq evals and observe near-perfect performance on needle-in-haystack tests and other benchmarks like Qasper.

Strengths

1. YOCO's hybrid structure delivers remarkable results in needle-in-haystack scenarios and demonstrates robust performance on retrieval-centric tasks, marking a pioneering achievement. 2. The proposed data-dependent gated-retention brings great improvement against retention. 3. By facilitating substantial KV cache compression relative to standard attention, YOCO exhibits superior retrieval capabilities compared to existing linear attention models. I'm very glad to see the results of YOCO scaling to larger sizes.

Weaknesses

I see no obvious disadvantages of this paper; however, the manuscript would benefit from: 1) The authors should add more comparions with exisiting linear-time / hybrid models trained on trillions of tokens, e.g., RWKV6 and TransNormer, whose checkpoints are publicly available. 2) Despite concurrent works, I suggest the authors to add discussions with Samba [1] and Mamba2 [2] in their next version. [1] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling [2] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Questions

1. Some notations are confusing: 1) Regarding Eq.7, the usage of $\beta_{iB}$ suggests an accumulation effect from preceding chunks, which may mislead readers. Additionally, the notations of $\beta_{[i]}(j,k)$ appears unused. If I understand correctly, $x_{[i]}$ is a 2-d tensor while $\beta_{[i]}$ is a scalar, it could be better to use another notation to distinguish the two. 2) Eq. 8 should be $\mathrm{head}_1,\dots,\mathrm{head}_n=\dots$ 2. I'm curious if the authors have tried other linear attention variants instead of gRet, e.g., Mamba, and GLA.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

The authors propose a new architecture for language models, where the top half of the transformer layers uses the KV from the bottom layer, while the bottom half applies efficient self-attention. The proposed architecture effectively reduces the KV cache size while maintaining the performance of the model, especially for long-context scenarios. Experiments also show that the method could scale up to 13B parameters.

Strengths

1. The proposed architecture is simple and effective, which could be easily integrated into existing transformer model implementations. 2. The experiments are comprehensive and convincing. The authors prove the effectiveness of the method on a 3B model and up to 1M context.

Weaknesses

1. As opposed to the first strength, the paper does not introduce new techniques or insights, thus limited in novelty. The authors also did not give possible explanations for the effectiveness of the proposed architecture. 2. The paper is lack of sufficient argumentation surrounding the design decisions. Though section 4.5 and 4.6 provide some preliminary analysis, further discussions are required to make the paper more convincing. For example, how the efficient self-attention and decoder-decoder structure affect the model's performance respectively. 3. The paper reports that the model outperforms the baseline transformers, but it remains unclear what contributes to the performance improvement. The main experiment is a partial comparison of the 1T token checkpoint instead of the fully trained model, so it is possible that the model is just easy to optimize (under large learning rates) but not necessarily converge to a better point. Also, the YOCO model has a different hyperparameter setting from the baseline model, with larger intermediate size (the scaling curve), which may also contribute to the performance improvement.

Questions

Please refer to **Weaknesses**

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Please refer to **Weaknesses**

Summary

The paper introduces YOCO, a decoder-decoder architecture designed for large language models. This architecture comprises a cross-decoder stacked upon a self-decoder, efficiently encoding global key-value caches reused by the cross-decoder. YOCO aims to reduce GPU memory demands and improve prefill latency and throughput while maintaining global attention capabilities. Experimental results demonstrate that YOCO achieves competitive performance compared to Transformer models, significantly reducing inference memory and prefill latency, and effectively extending context lengths up to 1M tokens with high retrieval accuracy.

Strengths

- YOCO's design, with its cross-decoder and self-decoder, offers a novel approach to caching key-value pairs, reducing GPU memory consumption. - The architecture significantly reduces prefill latency and improves throughput, addressing critical bottlenecks in long-sequence language model inference. - YOCO demonstrates effective scalability in model size and training tokens, maintaining competitive performance with other leading Transformer models. - Extensive experiments validate YOCO's performance and efficiency gains, showing substantial improvements in memory usage and latency across various model sizes and context lengths.

Weaknesses

- Transformers with flash attention could also scale to 1m tokens (e.g. FlashDecoding, https://crfm.stanford.edu/2023/10/12/flashdecoding.html) any comparison/discussion? Additional complexity with the cross-decoder and self-decoder mechanisms may pose implementation challenges. - While the architecture shows significant improvements in inference efficiency involving very long context lengths, it remains unclear how the fixed-size sliding window size affects the performance versus efficiency tradeoffs.

Questions

- The evaluation primarily focuses on memory and latency improvements. Does YOCO also bring training efficiency gains? - Are YOCO models slower than models in Table 3? Since the context size is usually much smaller, but YOCO used fixed window size of 1024 while most task examples probably contain <1024 tokens.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The paper does not explicitly discuss any limitations.

Summary

The paper introduces YOCO (You Only Cache Once), a novel decoder-decoder architecture for large language models. YOCO uses a self-decoder to generate global key-value (KV) caches, reused by a cross-decoder, reducing GPU memory usage and improving inference efficiency. The architecture achieves comparable performance to full transformers but with significantly lower memory demands. Extensive experiments show YOCO's effectiveness in scaling with more training tokens, larger model sizes, and longer context lengths, up to 1 million tokens. YOCO demonstrates substantial improvements in memory footprint, prefill latency, and throughput, making it a promising model for long-context understanding and multimodal applications.

Strengths

Overall, this is a high-quality paper. Originality: The paper presents a novel architecture that achieves performance comparable to full transformers with only one layer storing global KV tokens. Quality: The paper includes extensive experiments that robustly demonstrate the proposed model structure's ability to maintain excellent scaling performance while achieving good inference efficiency. The experiments are comprehensive and well support the claims made in the paper. Clarity: The paper is well-motivated, clearly stating the problem it aims to solve. The overall model structure is also clearly explained. The experimental section is well-organized, effectively showcasing how the model scales up with more training tokens, larger model sizes, and longer context lengths. It was very enjoyable to read. Significance: I believe this paper highlights the importance of achieving good scaling performance with only a single layer of global KV cache, including strong needle retrieval capabilities. This is a significant contribution, demonstrating the potential for efficiently handling long sequences with such models.

Weaknesses

- The paper should evaluate the in-context learning ability of the new architecture. - I believe more ablation studies on the window size of the sliding-window attention are necessary. The paper could more thoroughly investigate several important model parameters. - I think a significant future application for long context models is long video understanding. While this paper focuses on language modeling, it could benefit from including some discussion on extending the model to multimodal scenarios. - There are a few typos in the paper. For example, in line 36, "early exit before entering the self-decoder" should be "cross-decoder" instead of "cross-encoder."

Questions

- In the ablation study, does Unstacked YOCO refer to the model without the self-decoder? - Therefore, in the new model, will the number of layers and the number of attention heads per layer differ from the standard transformer design?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Please refer to the weakness section.

Thanks for the rebuttal

Thanks for addressing most of my concerns, I will keep my score.