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

Authors: yutao  sun
li  dong
yi  zhu
shaohan  huang
wenhui  wang
shuming  ma
quanlu  zhang
jianyong  wang
‡  furu
wei  †⋄
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 17:06:41

Abstract: 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.

Main paper: 
 introduction: 
 Decoder-only Transformer [40] has become the de facto architecture for language models. By caching the previously computed key/value vectors, the model can reuse them for future generation steps. The key-value cache avoids encoding the history again for each token, greatly improving the inference speed. The compelling feature establishes the decoder-only architecture as the standard option.However, as the number of serving tokens increases, the key-value (KV) caches occupy a lot of GPU memory, rendering the inference of large language models memory-bounded [29] . For the example of a 65B-size language model (augmented with grouped-query attention [1] and 8-bit KV quantization), 512K tokens occupy about 86GB GPU memory, which is even larger than the capacity of one H100-80GB GPU. In addition, the prefilling latency of long-sequence input is extremely high. For instance, using four H100 GPUs, the 7B language model (augmented with Flash-Decoding [6] and kernel fusion) requires about 110 seconds to prefill 450K tokens, and 380 seconds for 1M length. The above bottlenecks make it difficult to deploy long-context language models in practice.In this work, we propose a decoder-decoder architecture, YOCO, for large language models, which only caches KV pairs once. Specifically, we stack cross-decoder upon self-decoder. Given an input sequence, the self-decoder utilizes efficient self-attention to obtain KV caches. Then the cross-decoder layers employ cross-attention to reuse the shared KV caches. The decoder-decoder architecture is conceptually similar to encoder-decoder, but the whole model behaves more like a decoder-only model from the external view. It naturally fits into autoregressive generation tasks, such as language 3 You Only Cache Once (YOCO)The proposed architecture, named YOCO, is designed for autoregressive modeling, such as large language models (LLMs). As shown in Figure 1 , the decoder-decoder architecture has two parts, i.e., self-decoder and cross-decoder. Specifically, YOCO is stacked with L blocks, where the first L 2 layers are self-decoder while the rest modules are cross-decoder. Given an input sequence x = x 1 • • • x |x| , the input embeddings are packed intoX 0 = [x 1 , • • • , x |x| ] ∈ R |x|×dmodel ,where d model is hidden dimension. We first obtain contextualized vector representations X l = Self-Decoder(X l-1 ), l ∈ [1, L 2 ], where X L /2 is used to produce KV caches K, V for cross-decoder. Then we compute X l = Cross-Decoder(X l-1 , K, V ), l ∈ [ L 2 + 1, L] to get the output vectors X L . After obtaining X L , a softmax classifier performs next-token prediction over the vocabulary.Both self-and cross-decoder follow a similar block layout (i.e., interleaved attention and feedforward network) as in Transformer [40] . We also include pre-RMSNorm [48] , SwiGLU [32] , and grouped-query attention [1] as improvements. The difference between the two parts lies in attention modules. Self-decoder uses efficient self-attention (e.g., sliding-window attention). In comparison, cross-decoder uses global cross-attention to attend to the shared KV caches produced by the output of the self-decoder. Notice that the whole model behaves like a decoder-only model. The tokens generated by cross-decoder are also fed back to self-decoder. Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification: The abstract and introduction is carefully written.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Limitations are discussed in the paper.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
self-decoder: 
 Self-decoder takes embeddings X 0 as input and compute intermediate vector representation X L /2 : Y l = ESA(LN(X l )) + X lEQUATIONwhere ESA(•) represents efficient self-attention, SwiGLU(X) = (swish(XW G ) ⊙ XW 1 )W 2 , and RMSNorm [48] is used for LN(•). Causal masking is used for efficient self-attention.The key property of the efficient self-attention module is O(1) inference memory, i.e., constant number of KV caches. For example, the cache size of sliding-window attention [4] depends on the window size instead of the input length. More design choices (e.g., gated retention) of the efficient self-attention module are detailed in Section 4.First, the output of the self-decoder X L /2 generates global KV caches K, V for cross-decoder:EQUATIONwhere W K , W V ∈ R d×d are learnable. Then, cross-decoder layers are stacked after self-decoder to obtain the final output X L . The KV caches K, V are reused by all the L 2 cross-decoder modules:Q l = LN(X l )W l Q Y l = Attention(Q l , K, V ) + X l X l+1 = SwiGLU(LN(Y l )) + Y l (3)where Attention(•) is standard multi-head attention [40] , and W l Q ∈ R d×d is a learnable matrix. Causal masking is also used for cross-attention. Because cross-attention is compatible with group query attention [1] , we can further save the memory consumption of KV caches.(Skipped)Self-Decoder KV CachePre-filling context and then generate then generate new In addition to competitive language modeling results, YOCO significantly reduces serving costs and improves inference performance. We report detailed inference comparisons in Section 5.4.Saving GPU Memory and Serving More Tokens. Reducing Prefilling Time and Improving Throughput. As shown in Figure 2 , because the cross-decoder reuses the outputs of self-decoder, we can exit early before entering the cross-decoder during the prefill stage. The intriguing property of computation dependency greatly accelerates the prefilling speed. First, only half the layers are needed for forward computation, i.e., at least half prefilling latency reduction. Second, the efficient attention modules of the self-decoder are usually fast. For the example of 512K context length, we can decrease the prefilling latency from 180 seconds (Transformer with optimized inference, such as Flash-Decoding) to less than 6 seconds (Figure 9 ). Even for 32K length, YOCO has about three times speedup in terms of prefilling time. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
design choices of self-decoder: 
 We can choose various efficient self-attention methods for self-decoder. As long as the module only requires constant inference memory, the self-decoder's cache memory complexity depends on the number of layers. Moreover, a good module choice improves both training and deployment costs. In this work, we use sliding-window attention (Section 4.1) or gated retention (Section 4.2).Sliding-window attention [4] restricts the attention range into a fixed window size C. In contrast, vanilla Transformer decoders attend to all previous tokens. During inference, the KV cache memory complexity can be reduced from O(N ) to O(C), i.e., the memory usage is constant rather than increasing with sequence length. Similar to multi-head self-attention [40] , we compute the output of sliding-window attention via:Q = XW Q , K = XW K , V = XW V head i = softmax(Q [i] K ⊺ [i] + B)V, B ij = 0, i -C < j ≤ i -∞, otherwise SWA(X) = Concat(head 1 , • • • , head h )W O (4)where W Q , W K , W V , W O ∈ R d×d are learnable matrices, and the window causal mask B controls each query only attends to the previous keys whose distances are less than C. The pre-normalization and residual connection are also applied to the module.Gated retention (gRet, aka gRetNet) augments retention [35] with a data-dependent gating mechanism. We use gRet as the default efficient self-attention module. The method unifies the parallel, recurrent, and chunkwise recurrent computation paradigms, which are equivalent and can obtain the same computation results. The training process usually uses the parallel or chunkwise recurrent paradigms, while the inference stage can employ the recurrent paradigm for constant KV memory.The Parallel Representation The gated retention is defined as:EQUATIONwhere W Q , W K , W V ∈ R d×d and W γ ∈ R d×1 are learnable weights, and the temperature term τ encourages γ to 1 for better memorization [46] . The data-controlled decay is head-wise [19] rather than element-wise so that the computation can fully utilize NVIDIA tensor cores. Refer to [35] for more details about the other designs.The Recurrent Representation Being equivalent to Equation (5), the output of gated retention can be computed recurrently. For the n-th timestep, the output is obtained via:EQUATIONwhere Q, K, V, γ are the same as in Equation (5) . During auto-regressive inference, the self-decoder maintains S n as the intermediate state for an efficient generation.The Chunkwise Recurrent Representation The chunk-wise representation is a unified formulation of recurrent and parallel representations. Given chunk size B, the outputs are computed chunk by chunk. The computation is divided into inner-chunk and cross-chunk parts. Denote [i] as the i-th chunk, i.e., x [i] = x (i-1)B+1 , • • • , x iB , we compute the i-th chunk as:β (i-1)B+j = (i-1)B+j k=(i-1)B+1 γ k , D [i] (j, k) = β (i-1)B+k β (i-1)B+j if j ≤ k else 0 R i = K ⊺ [i] (V [i] ⊙ β iB β [i] ) + β iB R i-1 , β [i] (j, k) = β (i-1)B+j gRet(X) = (Q [i] K ⊺ [i] ⊙ D [i] )V [i] Inner-Chunk + (Q [i] R i-1 ) ⊙ β [i]Cross-Chunk (7) where R i is the intermediate state of the i-th chunk, and β summarizes the data-controlled decay γ. Appendix B proves the equivalence between the computation paradigms. The chunkwise paradigm combines the best of parallelism and recurrence, i.e., saving FLOPs compared to fully parallel computation and reducing iterations compared to recurrent computation. During the training and prefill stages, the chunk-wise representation increases throughput and reduces GPU memory consumption. [39] . YOCO-3B-1M is extended to the context length of 1M tokens.Multi-Head Gated Retention Similar to multi-head attention [40] and multi-scale retention [35] , we apply gated retention to each head and combine the outputs together:head i = gRet(X) Y = GroupNorm h (Concat(head 1 , • • • , head n )) MHGR(X) = (swish(XW G ) ⊙ Y )W O (8)where W G , W O ∈ R d×d are learnable matrices, and GroupNorm [42] normalizes each head [41] . We also apply swish gate to increase non-linearity [35] . 
experiments: 
 We evaluate YOCO for large language models from the following perspectives. First, we follow the setting of StableLM-3B-4E1T [39] to scale up training tokens (Section 5.1). Second, we present the scaling curves of the proposed architectures (Section 5.2). Third, we scale up the YOCO model to 1M context length and evaluate its long-sequence modeling capability (Section 5.3). Fourth, we analyze the deployment advantages, including GPU memory footprint, serving capacity, prefilling time, and throughput (Section 5.4). Experimental results show that YOCO achieves competitive performance in various evaluation metrics and significantly reduces the inference cost.We train a 3B-size YOCO language model by scaling up the number of training tokens. Then we compare the checkpoints with strong Transformer-based language models. We use a similar training recipe to that in StableLM-3B-4E1T [39] . Detailed hyperparameters are described in Appendix D.Results Table 3 compares YOCO with OpenLLaMA-v2-3B [12] , StableLM-base-alpha-3B-v2 [38] , and StableLM-3B-4E1T [39] . We use LM Eval Harness [10] to evaluate zero-shot performance on various downstream tasks. OpenLLaMA-v2-3B and StableLM-base-alpha-3B-v2 are trained with 1T tokens. The intermediate numbers of StableLM-3B-4E1T are taken from its technical report [39] . Experimental results indicate that YOCO achieves comparable results with previous well-tuned Transformer language models. Both the checkpoints trained with 1T tokens and 1.6T tokens obtain a consistent trend. Moreover, the results show that YOCO is scalable in terms of training tokens.We compare the scaling curves between Llama Transformer [40, 37] , YOCO with gated retention (YOCO gRet ; Section 4.2), and YOCO with sliding-window attention (YOCO SWA ; Section 4.1). We train language models of various sizes (i.e., 160M, 400M, 830M, 1.4B, 2.7B, 6.8B, and 13B) using the same training data and settings. We augment the Transformer architecture with Llama [37] improvements, such as RMSNorm [48] , SwiGLU [32] , and removing bias. The sliding window size of YOCO SWA is 1,024. The training batch size is 0.25M tokens with a 2k sequence length. We train the models with 40k steps, i.e., 10B tokens. In practice, we find that the setting is effective for loss convergence, and scaling laws can be well-fitted. More hyperparameters are detailed in Appendix E. Results Figure 3 reports the validation loss with various parameter counts. We also fit the scaling curves as in [17] . YOCO obtains comparable performance from 160M to 13B compared to the Llama-optimized transformer architecture. The findings demonstrate that YOCO scales effectively with respect to model size. Moreover, YOCO gRet outperforms Transformer and YOCO SWA . The gains come from hybrid architectures of attention and retention, whose inductive biases tend to be complementary to each other. Recent hybrid architectures [21] also confirm similar findings.We extend the context length of YOCO-3B (Section 5.1) to 1M tokens. We continue the model training with longer lengths progressively. The length schedule is 64K, 256K, and 1M tokens. Training data is up-sampled according to sequence length [9] . For a fair comparison, we do not use long-instruction tuning data. More training details are described in Appendix F.The pressure test evaluates whether models can retrieve "needles" from a long document [16] . We follow the evaluation setting of Gemini 1.5 [30] and LWM [24] . The needles are constructed as a city with a magic number. We run 10 times at the same depth and length. The average accuracy is reported. Figure 4 shows that YOCO-3B-1M passes the Needle-In-A-Haystack test with near perfect accuracy. The results indicate that YOCO has strong long-context modeling capability.Multi-Needle Retrieval Besides single-needle retrieval, we conduct a multi-needle evaluation. We compare YOCO-3B-1M with previous long-context language models, including MiniCPM-128K [15] , ChatGLM3-128K [47] , YaRN-Mistral-128K [28] , and LWM-1M-text [24] . The evaluation is conducted in 128K sequence length, because most previous models are tuned with this length. Table 4 presents accuracy results with N needles. LWM-1M-text and YOCO-3B-1M are trained with a 1M context length, while the others are of 128K length. Although LWM-1M-text continues training of Llama-2-7B, YOCO-3B-1M can still achieve comparable performance with half the model size. Moreover, the 7B-size YaRN-Mistral-128K [28] obtained by position interpolation lags behind the other models. Compared to MiniCPM-128K and ChatGLM3-128K, YOCO-3B-1M also outperforms these well-trained language models.SizeN = 1 N = 2 N = 4 N = 8YaRN-Mistral-128K [ Perplexity over Long Sequences Figure 5 shows the cumulative average negative log-likelihood (NLL) as a function of context length. We evaluate both book and repository-level code data. We follow the setting of [30] and filter validation data that are longer than 1M tokens. NLL decreases consistently with longer sequence length. The results indicate that YOCO can effectively utilize long-distance dependency for language modeling. We also observe that the NLL-length curves tend to fit the power law, where the gaps are affected by the noise within the validation examples.We analyze inference efficiency from various perspectives, such as GPU memory footprint, prefilling latency, throughput, and serving capacity. We show that YOCO reduces the deployment cost by orders of magnitude, especially for long-sequence inference. More importantly, the user experience (such as latency) is improved while maintaining good performance and reducing expenses.We compare YOCO gRet with Transformer. The default model configuration follows Section 5.1.Notice that Transformer uses grouped-query attention [1] , Flash-Decoding [6] , and kernel fusion for a fair comparison. As described in Section 4.2, gated retention uses the chunk-recurrent representation in the prefill stage, and the recurrent representation in the generation stage. The chunk size is set to 256. We implement a Triton [36] kernel for gated retention. The evaluation sequence length ranges from 32K to 1M. The last 1,024 tokens are supposed to be generated, while the previous tokens are given input context. The experiments are conducted with H100-80GB GPU cards. The inference memory consumption is made up of three parts, namely model weights, intermediate activation, and KV cache. Figure 6 presents the breakdown memory profiling results. Along with an increase in context length, the main memory bottleneck becomes KV caches, while model weights consume constant memory. The results show that YOCO gRet alleviates the activation cost and KV cache memory footprint.As shown in Figure 7 , the memory cost is significantly reduced using YOCO. Moreover, the memory consumption of YOCO increases slowly along the sequence length. For example of 1M length, the overall inference memory usage is only 12.4GB, while Transformers occupy 9.4× GPU memory. YOCO makes it feasible to deploy long-sequence modeling on customer-level GPUs. Even with a 32K sequence length, YOCO requires about 2× less memory than Transformer. Although we compare 3B-size models here, the reduction ratio becomes larger as the number of layers increases. For example, YOCO can serve 128K tokens with 1GB GPU memory, while Transformer with GQA [1] can only support 1.6K tokens at 65B model size.Prefilling Latency In the prefill stage, the model encodes input tokens in parallel. As shown in Figure 9 , the prefilling latency is a pain point of user experience for long-context models. For 512Kand 1M-length input sequences, Transformer needs about 180 seconds and 300 seconds, respectively. The computational complexity of Transformer is O(N 2 ), which requires a large number of FLOPs for long context. In contrast, YOCO's prefilling time is O(N ), growing linearly (Section 3.3) along the sequence length. Figure 9 shows that YOCO reduces the Transformer prefilling time from 180 seconds to less than 6 seconds for 512K context. As described in Section 3.3, the prefill stage can early exit before entering cross-decoder. So, there is at least two times speedup of prefilling latency even for short context. For example, YOCO is 2.87× faster than Transformer for 32K length.Throughput The throughput indicates how many tokens the model can process per second, involving both pre-filling and generation time. Figure 10 shows that YOCO achieves higher throughput across context lengths compared to Transformer. For the example of 512K queries, Transformer's throughput is 4.5 token/s while YOCO reaches 43.1 token/s, i.e., achieving 9.6× speedup. The throughput is improved for the following reasons. First, YOCO decreases the time required for prefilling as previously demonstrated. Second, as the memory consumption is reduced, we can use larger batch size for inference, which also contributes to the throughput improvement.We compare YOCO gRet and YOCO SWA with Transformer and other variants, including H3 [5] , RetNet [35] , Mamba [13] , and gRetNet (Section 4.2). All models have 160M parameters with 12 layers and a hidden dimension of 768. The weights of word embedding and softmax projection are shared. For Mamba, we follow the details in [13] , where double-SSM layers are implemented instead of "SSM + SwiGLU". For H3, the experiment uses a hybrid version following the original paper [5] , where two attention layers are inserted after the first and L 2 -th layers. Fine-Grained LM Perplexity Table 5 reports the fine-grained validation perplexity for language modeling. Following Zoology [2] , we divide the perplexity into "Ar-Hit" and "First-Occur". Specifically, "Ar-Hit" considers the predicted tokens that are bigrams previously seen in the previous context, which evaluates the associative recall capability. "First-Occur" considers the tokens that cannot be recalled from the context.Long-Context Evaluation Figure 11 reports the answer perplexity with varying context length (ranging from 4K to 16K) on the ZeroSCROLLS [31] benchmark. We continue training the above models in 16,384 length with 2B tokens. The rotation base scaling [44] is used for length extension.For sparse Transformer, we use the context window of 2,048 and keep RoPE θ unmodified. As shown in Figure 11 , YOCO and Transformer consistently outperform other methods across tasks and lengths, which is consistent with the findings in Section 5.3. Moreover, the results highlight the importance of global attention for long-context modeling. Notice that the 12K and 16K results in Qasper are similar because the lengths of most documents are shorter than 16K.Valid. Set↓ AR-Hit↓ First-Occur↓ as the default layout. Future work can refine a scaling law to guide the choice of layer ratio. Second, the setting "Unstacked YOCO [1:1] " uses word embeddings X 0 as input to the cross-decoder, rather than stacking cross-decoder upon self-decoder (i.e., using X L /2 in Equation (3)). Third, the model "Interleaved & Hybrid" is a hybrid architecture that interleaves gRetNet and Transformer layers.YOCO [1:1] 3. 
conclusion: 
 In this work, we propose a decoder-decoder architecture (YOCO) for large language modeling. YOCO achieves significantly better inference efficiency and competitive performance compared to Transformers. Experimental results demonstrate that YOCO achieves favorable results for large language models in various settings, i.e., scaling up the number of training tokens, scaling up model size, and scaling up context length to 1M tokens. Profiling results also show that YOCO improves inference efficiency by orders of magnitude, especially for long-sequence modeling. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We work on fundamental research that has no direct societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not pose safety risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We carefully follow the licenses of open-source code, data, and models.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: We introduce chunk parallelism for YOCO to reduce the communication frequency, accelerating long-sequence training in Section 5.3. Dividing long sequences into different devices is essential when the training length is extremely long [20, 8] . However, the overall throughput tends to be bounded by GPU communication [23] . Cross-decoder disentangles self-attention dependency while preserving modeling capability, bringing intriguing advantages to distributed long-sequence training.𝑋 𝑋 1 𝑋 2 Split Self-Decoder 𝐾 𝑉 𝑂 1 𝑂 2 Cross-Decoder 𝑀 1 𝑀 2 Project 𝑄 1 𝐾 1 𝑉 1 𝐾 2 𝑉 2 𝑄 2All-Gather Once In self-decoder, the dependency only exists in the adjacent devices. For example, gated retention only requires the hidden state S n in Equation ( 6), and sliding-window attention attends to tokens within the context window. Therefore, the communication amount of self-decoder is relatively small. In the cross-decoder, the all-gather operation is only triggered once for the KV cache, rather than communicating in each layer. The hardware-friendly architecture gives more flexibility to distributed long-sequence training.We illustrate the equivalence between recurrent representation and chunkwise recurrent representation of gated retention. For the output O n , n can be split as n = kB + r where B is the chunk size:EQUATIONwhereΓ i = n k=i+1 γ i , ζ [c] (j, k) = cB i=(c-1)B+j+1 γ i , α i = iB j=(i-1)B+1 γ j , [i] indicates the i-th chunk, i.e., x [i] = [x (i-1)B+1 , • • • , x iB ].R n is written as a recurrent function:EQUATIONDenote [i] as the i-th chunk, i.e.,x [i] = [x (i-1)B+1 , • • • , x iB ], β (i-1)B+j = (i-1)B+j k=(i-1)B+1 , β [i] (j, k) = β (i-1)B+j , We concatenate the output in a block together:O [n] = [n] m=kB+1 β [n] Q [n] K ⊺ m V m + kB m=1 β [n] Q [n] n i=m+1 γ i K ⊺ m V m [n] m=kB+1 β [n] Q [n] K ⊺ m V m = (Q [n] K ⊺ [n] ⊙ D [n] )V [n] , D [n] (j, k) = β (n-1)B+k β (n-1)B+j if j ≤ k else 0 kB m=1 β [n] Q [n] n i=m+1 γ i K ⊺ m V m = β [n] Q [n] R i-1 , R i = K ⊺ [i] (V [i] ⊙ β iB β [i] ) + β iB R i-1 , O [n] = (Q [n] K ⊺ [n] ⊙ D [n] )V [n] Inner-Chunk + (Q [n] R n-1 ) ⊙ β [n]Cross-Chunk (11) Finally, we show that the chunkwise recurrent representation of gated retention is equivalent to the other two representations.We present pseudocode for the three computation paradigms of gated retention (Section 4.2). Parallel implementation enables training parallelism to fully utilize GPUs. The recurrent paradigm enables low-cost inference. Chunkwise retention combines the above advantages (i.e., parallel within each chunk and recurrent across chunks), which has linear memory complexity for long sequences.def ParallelRetention( q, # bsz * num_head * len * dim k, # bsz * num_head * len * dim v, # bsz * num_head * len * dim gt): # bsz * num_head * len retention = q @ k.transpose(-1, -2) causal_mask = torch.full([q.shape[-2], q.shape[-2]], float("-inf"), device=q.device). triu (1) .type_as(q) gt = F.logsigmoid(gt). We adjust the head dimension to 128 instead of 80 as in StableLM for better kernel support. To keep the model size unchanged, we set the hidden size to 3072 and the number of layers to 26. Grouped-query attention [1] is used, where the number of query heads is 24, and the number of key-value heads is 8. We train YOCO with gated retention (Section 4.2). The non-embedding parameter count is 2.8B. In comparison, StableLM-3B-4E1T is 2.7B and OpenLLaMA-v2-3B [12] is 3.2B. The training sequence length is 4096. The batch size is 4M tokens. We use the AdamW [26] optimizer with β = 0.9, 0.95. The maximal learning rate is 3.2e-4 with 1000 warmup steps and linear decay to 1.28e-5. The total schedule is set to 5T tokens. Given the resource budget, we train the model with 400k steps (1.6T tokens). The curated training corpus is similar to [39] . We use tiktoken-cl100k_base as the tokenizer. The hidden dimension is set to 3072. The number of layers is 26. The number of query heads is 24, and the number of key/value heads is 8 with grouped-query attention [1] . The total number of parameters without embedding is 2.83B. The training batch size is 4M tokens. We use 4096 training length. The optimizer is AdamW [26] with β = (0.9, 0.95). The learning rate is 3.2 × 10 -4 with 1000 warmup steps. We set a 5T-token learning rate schedule with linear decay to 1.28 × 10 -5 . Table 8 reports the hidden dimension, number of layers, and number of heads used for different model sizes. The head dimension of gated retention is set to 256. To align the number of parameters, the FFN size for Transformer is 8 3 d while the FFN size for YOCO is 3d. The training length is set to 2048. The batch size is set to 0.25M tokens. We use the AdamW [26] optimizer with β 1 = 0.9, β 2 = 0.98. The learning rate is 1.5 × 10 -4 for 160M to 1.4B sizes and 7.5 × 10 -5 for 2.7B to 13B sizes. The warmup step is 375 with linear rate decay. The weight decay is set to 0.05. We train the models with 40k steps, i.e., 10B tokens.Hidden We progressively extend the context length to 1M tokens in Section 5.3. The length schedule is 64K, 256K, and 1M. We up-sample the documents that are longer than the training length [9] . The word "once" refers to global KV cache. Strictly, self-decoder also needs to store a certain number of caches. As the self-decoder utilizes an efficient attention module, the cache size is bounded to a constant, which can be ignored compared to global caches when the sequence length is large. 


Decision: This paper proposes a novel decoder-only architecture that uses only one layer global KV cache to improve inference efficiency. Reviewers all agree this is a highly novel paper with clear writing and thorough experiments (similar performance with vanilla transformer, high memory efficiency, low latency, and scaling to 1m context). Overall this is a great contribution to the community.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: Please refer to the weakness section.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: The paper does not explicitly discuss any limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please refer to **Weaknesses**
limitations: limitations: Please refer to **Weaknesses**
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: N/A

Comments: title: Thanks for the rebuttal
comment: I thank the authors for the clarifications. I will keep my rating.
Thanks for addressing most of my concerns, I will keep my score.
I appreciate the authors for replying to my questions. Considering the impressive performance and hybrid architecture of this work, I would like to increase the final rating to 6.
Thanks for replying to my question, I have no other questions and keep my score.