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 2024 oral (2024)
Keywords: Decoder-Decoder
Model Architecture
Primary area: deep_learning_architectures
TL;DR: We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once.

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.

You Only Cache Once: Decoder-Decoder Architectures for Language Models
Yutao Sun Li Dong" Yi Zhut Shaohan Huang
Wenhui Wang Shuming Ma² Quanlu Zhang!
Jianyong Wang' Furu Weit
Tsinghua University t Microsoft Research
htts://aka.ms/GGeeralAll
Abstract
We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once. consists of two components, i.c., a cross- decoder stacked upon self-decoder. The self-decoder efficiently encodes global key-value (KV) caches that are reused by the cross-decoder via cross-altention. The overall model behaves like 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- perfept 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.
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 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 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 H1 0000-0GG 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 1 10 seconds to prefill 450K tokens, and 380 seconds for 1M length. The above bottlenecks make it difficult 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, stack cross-decoder upon self-decoder. Given input sequence, the self delf-decoder utilizes efficient self-attention obtain KV caches. Then the cross-decoder layers employ cross cotss-attnniinn to reuse the shared KV caches. The decoder- architecture is conceptually similar to encoder-decoder, but the whole model behaves more like decoder-only model from the external view. It naturally fits into autoregressive generation tasks, such language
Equal contribution. Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

modeling, First, because YOCO only caches once". the GPU memory consumption of KV caches significantly reduced. Second, the computation flow of the decoder-decoder architecture enables prefilling to early exit before entering the self-decoder. The nice property speeds up the prefill stage dramatically, improving user experience for long-context language models. Third, YOCO allows for more efficient system design for distributed long-sequence training. In addition, we propose gated retention for self-decoder, which augments retention [35] with a data-controlled gating mechanism.
We conduct extensive experiments to show that YOCO achieves favorable language modeling perfor- mance and has many advantages in terms of inference efficiency. Experimental results demonstrate that YOCO can be scaled up with more training tokens, larger model size, and longer context length. Specifically, we scale up the 3B YOCO model to trillions of training tokens, attaining results on par with prominent Transformer language models, such as StableLM [39]. Moreover, the scaling curves ranging from 160M to 13B show that YOCO are competitive compared to Transformer. We also extend the context length of YOCO to 1M tokens, achieving near-perfect needle retrieval accuracy. In the multi-needle test, YOCO obtains competitive results even compared to larger Transformers.
In addition to good performance on various tasks, the profiling results show that YOCO improves the GPU memory footprint, prefill latency, throughput, and serving capacity. In particular, the memory of KV caches can be reduced by about 80x for 65B models. Even for a 3B model, the overall inference memory consumption can be reduced by two times for 32K tokens and by more than nine times for IM tokens. The prefill stage is speeded up by 71.8x for the IM context and 2.87x x for the 32K input. For example, for a 512K context, YOCO reduces the Transformer prefilling latency from 180 seconds to less than six seconds. The results position YOCO as a strong candidate model architecture for future large language models with native long-sequence support.
2 Related Work
Numerous efforts have been made to reduce KV caches for inference. Efficient attention mechanisms are proposed, such as sparse Transformer [4], linear attention [18], and recurrent modeling [27, 13, 46, 3, 19]. Another strand of research drops KV caches to achieve sparsity [49, 43, 11]. In comparison, we keep one global KV cache and still conduct full cross-attention for better long-context modeling. Moreover, some previous methods are complementary to our proposed architecture. For example, multi-fgrouped--queee query attention [33, I] and multi-latent attention [7] can be used in YOCO. Low-bit KV quantization [14, 25, 34] can also be used together to reduce memory consumption. In addition, the intriguing property of YOCO greatly speeds up the prefill stage.
3 You Only Cache Once (YOCO)
The proposed architecture, named YOCO, is designed for autoregressive modeling, such 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 & layers are self-decoder while the rest modules are cross-decoder. Given an input sequence r = 21...-p0 T|al+ the input embeddings are packed into X [x1, Iz R Xdmand, where dmodel is hidden dimension. We first obtain contextualized vector representations X Self Delf--ecoder(x"-)) e [1, 4).. where 2/2 is used to produce KV caches R, V for cross-decoder. Then we compute x = Cross-Decoder((x"-1 V). E 16 + 1, L. to get the output vectors XL. After obtaining X softmax classifier performs next-token prediction over the vocabulary.
Both self- and cross-decoder follow a similar block layout (i.c., interleaved attention and feed- forward 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
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 decoder-only model. The tokens generated by cross-decoder are also fed back to self-decoder.
'The word "once" refers to global KV cache. Strietly, self-decoder also needs to store a certain number of caches. As the deco-de utilizes an efficient attention module, the cache size is bounded to constant, which be ignored compared global caches when the sequence length is large.

---

You
Only
Cache Once
t
Output
Feed Forward
Layer
KV Cache
Cross-Attention
Q
K
Feed Forward
Layer
Efficient Self-Attention
Q
Embedding
You
Only Cache
Figure 1: Overview the decoder-decoder architecture. Self-decoder generates the global KV cache. Then cross-decoder employs cross-attention to reuse the shared KV caches. Both self-decoder and cross-decoder use causal masking. The overall architecture behaves like decoder-only Transformer, autoregressively generating tokens.
3.1 Self-Decoder
Self-decoder takes embeddings X as input and compute intermediate vector representation 1/:;;
Y' ESA(LN(X')) + x
(I)
X SviGLU +
where ESA(.) represents efficient self-attention, SwiGLU(X) (swish(X c XWi)W,,, 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- attention [4] depends on the window size instead of the input length. More design choices (c.g., gated retention) of the efficient self-attention module are detailed Section
3.2 Cross-Decoder
First, the output of the delf-decoder 4/2 generates global KV caches K, V for cross-decoder: LN(X
K LN(X*>>WKx
where WK. Wv E dddd are learnable. Then, cross-decoder layers are stacked after self-decoder to obtain the final output X The KV caches K, V are reused by all the & cross-decoder modules:
Q' LN(X')W3
Y Attention(Q'' K, V) + x
/+1 SWiGLU(LN(Y")) + Y
where Allention(-) is standard multi-head attention [40], and W R is learnable matrix. Causal masking is also used for cross-attention. Because cross-attention compatible with group query attention [1]. we can further save the memory consumption of KV caches.

---

Prefilling
Generation
then generate new
KV Cache Memory
Cross-Decode (Skipped)
Cross-Decoder
Transformer
O(LND) O( +
YOCO
KV Cache
Self-Decoder
Table Inference memory complexity KV caches. N, L, D are the sequence length, num- ber layers, and hidden dimension.
t Pre- filling context and then generate
Prefilling Time
Figure 2: YOCO Inference. Prefill: encode input to- kens in parallel. Generation: decode output tokens one by one. The computation flow enables prefill- ing to early exit without changing the final output, Table 2: Prefilling time complexity of attention thereby significantly speeding up the prefill stage. modules. N, L, D are the same as above.
Transformer YOCO
O(LN2D) O(LND)
3.3 Inference Advantages
In addition 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. Table compares the memory complexity between Transformers and YOCO. Specifically, because global KV caches are reused and efficient self-attention needs constant caches, the number of caches is O(N + CL). where N the input length, C' is a constant (c.g., sliding window size), and L is the number of layers. For long sequences, CL is much smaller than N, about (NN) caches are required, i.c., you only cache once. In comparison, Transformer decoders have to store N x I. keys and values during inference. So YOCO roughly saves L times GPU memory for caches compared to Transformer. Because the bottleneck of inference capacity becomes KV caches, our method enables to serve many more tokens without being out of GPU memory. The increased batch size is also beneficial to inference throughput.
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.c., 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- Even for 32K length, YOCO has about three times speedup terms of prefilling time. Table compares prefilling time complexity of attention modules between Transformer and YOCO.
to less than 6 seconds (Figure 9).
4 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 the number of layers. Moreover, good module choice improves both training and deployment costs. In this work, we use sliding-window attention (Section 4.1) gated retention (Section 4.2).
4.1 Biiding-Window Attention
Sliding- window attention [4] restricts the attention range into a fixed window size C. contrast, vanilla Transformer decoders attend to all previous tokens. During inference, the KV cache memory complexity can be reduced from O(N) to OCC), i.e., the memory usage is constant rather than increasing with sequence length. Similar multi-head self-attention [40], compute the output of

---

Mlding-windoww attention via:
= = K = =XW,, V ==
[0
i-C
head; softmax(Q1a + B)V, Bij -00
(4)
otherwise
SWA(X) Concat(head,,
head,) Wo
where We, WK, Wv, Wo Rdxd 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.
4.2 Gated Retention
Gated retention (gRet, aka gRetNet) augments retention [35] with 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:
Q (XWQ) O, = G
V XWv,
O inl
%
n 2 m
= sigmoid ( D Dam
i=m+1
(5)
0.
n <
gRet (X) (QKT @ D)V
where Wo. WK. Wy Rddd and W,, E gdx1 are learnable weights, and the temperature term
encourages to for better memorization [46]. The data-controlled decay is head-wise [19] rather than element-wise that the computation can fully utilize NVIDIA tensor cores. Refer [35] for more details about the other designs.
The Recurrent Representation Being equivalent Equation (5), the output of gated retention be computed recurrently. For the n-th timestep, the output is obtained via:
S, = In +TVV gRet(X,)) Qn 1=
(6)
|æ|
where Q, K, V are the same in Equation (5). During regressegee inference, the self-decoder maintains S, as the intermediate state for an efficient generation.
The Chunkwise Recurrent Representation The chunk-wise representation is 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] the i-th chunk, i.c., TPp = (1- ((-1))++1+" .. T4B, we compute the i-th chunk as:
(i- 1)+++ II
Dalisk) = B if j k else 0 The
B =
B1--11B+5
k=
BIB
R; K7)(11 (Vig
+ BIBRI-1 Bpq(i,)) B
311
gRel(X) (Q KT @ Dia)
(:11 2 Big
Inoer-Chunk
Cross-Chunk
where R; is the intermediate state of the i-th chunk, and | summarizes the data-com cata-connrolled decay Appendix B proves the equivalence between the computation paradigms. The chunkwise paradigm combines the best of parallelism and recurrence, i.e., saving FL.OPs compared to fully parallel com- putation and reducing iterations compared to recurrent computation. During the training and prefill stages, the chunk-witg representation increases throughput and reduces GPU memory consumption.
5

---

Model
ARC-C
ARC-E
BoolQ
Hellaswag
OBQA
PIQA
Winogrunde
SeiQ
Avg
Training with IT tokens
OpenLLaMA-33B- [12]
0.339
0.676
0.657
0.700
0.260
0.767
0.629
0.924
Stahlel M --aapha--8--2 [38]
0.324
0.673
0.646
0.686
0.264
0.760
0.621
0.921
0.612
StableLM-33B-45 [39]
0.666
0.768
0.632
0.914
YOCO-3B
0.379
0.731
0.645
0.689
0.298
0.763
0.639
0.924
0.634
Training with 1.6T tokens
StableLMM33B-ie [39] YOCO-3B
0.688
0. 0.762
0.627
0.913
0.396
0.733
0.644
0.698
0.300
0.764
0.631
0.921
0.636
Extending context length to IM tokens
YOCO-3B-1M
0.413
0.747
0.638
0.705
0.300
0.773
0.651
0.932
0.645
Table Eval Harness [10] accuracy compared with well-trained Transformer language models. We scale the 3B model to 1.6 trillion training tokens. The IT and 1.6T results of Stablel.M-3B-4EITT taken from its technical report [39]. YOCO-3B-IM is extended to the context length of 1M tokens.
Multi-Head Gated Retention Similar multi-head attention [40] and multi-scale retention [35], we apply gated retention to each head and combine the outputs together:
head, gRet(X)
Y GroupNorm, (Conca((caaaa head,))
(8)
MHGR(X) (swish(XWa) @ Y)Wo
where We, Wo E Rddd 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, follow the setting of MaableLM-3B-4E111 [39] to scale up training tokens (Section 5.1). Second, we present the scaling curves of the proposed architectures (Section 5.2). Third, 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 metries and significantly reduces the inference cost.
5.1 Language Modeling Evaluation
We train a 3B-size YOCO language model by scaling up the number of training tokens. Then we compare the checkpoints with strong Transformer-baeee language models. We use similar training recipe to that in StableLM-3B MaAtdeMM-B--4444 3ETTT [39]. Detailed hyperparameters are described in Appendix D.
Results Table 3 compares YOCO with OpenLLaMA-vv-338 [12], Bhablel-M-bass-aalpa-3BBB- and StableLMM33--44F [39]. We use LM Eval Harness [10] evaluate zero-shot performance on various downstream tasks. BaenLLaMA-v2-38 and StableLM-baan
[38],
tokens. The intermediate numbers of StableLM-3B-4FE MabbleLM-baae-allhaa be-baaallphaaa----e-- are trained with IT
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 IT tokens and 1.6T tokens obtain a consistent trend. Moreover, the results show that YOCO is scalable in terms of training tokens.
5.2 Scalability Compared with Transformers
We compare the scaling curves between Llama Transformer [40, 371, YOCO with gated retention (YOCOgRet Section 4.2), and YOCO with sliding-window attention (YOCOswA: 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 YOCOswA is 1,024. The training batch size is 0.25M tokens with 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.

---

3.6
Transformer
18
3.4
SwA
27
g
Loss
YOCO geel
45
0.6 0
3.2
54
63
0,40
72
3.0
81
0.2
90
100
0.0
128K
256K
384K
512k
640K
768K
896k
1M
100
10a
# Parameters (B)
Context Length
Figure 3: LM loss decreases along with scaling up the model size (ranging from 160M to 13B).
Figure 4: Needle-in-a-haysttcc results in IM length.
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, YOCOgRes outperforms Transformer and YOCOswA- 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.
5.3 Long-Conterr Evaluation
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 IM 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 described in Appendix F.
Needle In Haystack with 1M Context 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 city with a magic number. run 10 times at the same depth and length. The average accuracy is reported. Figure 4 shows that YOCO-3B-IM 1M passes the Ncedle-In- Inddle-In-A-HEe 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 multi-neede evaluation. We compare YOCO-3B-1M with previous long-contens language models, including MiniCPM- 128K [15], ChatGL.M3 [47], YaRN-Mistral-1288 [28], and L.M--1Mtext [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- and YOCO-3B-IM trained with a IM context length, while the others are of 128K length. Although LWM-IM-ter continues training of Llama-2-7B, YOCO-3B-IM 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 MiniCPM- 128K and ChatGLM3 -2-.MM--1KKK YOCO-3B- also outperforms these well-trained language models.
Model
Size
N = 1
N = 2
N =
N
8
YaRN-Mistral-1225 128K [28]
7B
0.02
0.12
0.08
0.20
L.M-IIM-tett [24]
7B
1.00
0.90
0.76
0.62
MiniCPM-128K [15]
2.4B
1.00
1.00
0.54
0.56
ChatGLM3--12K [47]
6B
0.94
0.72
0.52
0.44
YOCO-3B-IM
3B
0.98
0.98
0.84
0.56
Table 4: Multi-needle retrieval accuracy. N indicates the number of needles. N = 1 is single-needle retrieval used as reference, and N > indicates the multi-needle test. The evaluation conducted in 128K length, because most previous long-context models are tuned with this length.

---

Negative Log Log- LggLLikllhood (Lower Better)
Negative Log Li-kelhhlood (Lower Better)
10
100
1K
10K
100K
1M
10
100
1K
10K
100K
1M
Sequence Position
Sequence Position
(a) Book data.
(b) Repository-Ieve code data.
Figure 5: Cumulative average negative log-likelihood on book and repository-level code. We filter the validation examples that are longer than 1M tokens. YOCO achieves improved performance with longer context, i.e., utilizing long-distance information for language modeling.
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 gaps are affected by the noise within the validation examples.
5.4 Inference Advantages
We analyze inference efficiency from various perspectives, such 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 geet with Transformer. The default model configuration follows Section 5.1. Notice that Transformer uses grouped-query attention [1], Daas-DDeodiing [6], and kernel fusion for fair comparison. As described in Section 4.2, gated retention uses the chunk-recurrent representation in the prefill stage, and the recurrent representation the generation stage. The chunk size set to 256. We implement& 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.
GPU Memory The inference memory consumption is made up of three parts, namely model weights, intermediate activation, and KV cache. Figure 6 presents the breakdown memory profil- ing results. Along with increase in context length, the main memory bottleneck becomes KV caches, while model weights consume constant memory. The results show that YOCO alleviates the activation cost and KV cache memory footprint.
KV Cache Weight
GPU Memory (GB)
100
Other
75
9.38x
50
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 IM length, the overall inference memory usage is only 12.4GB, while Transformers occupy 9.4x GPU memory. YOCO makes feasible to deploy long- long-sequence modeling on customer-level GPUs. Even with 32K sequence length, YOCO requires about 2x less memory than Transformer. Although we compare 3B-size models here, the reduction ratio becomes larger as the number of layers increases.
25
Transformer
YOCO
Figure 6: Breakdown memory sumption in 1M context length.
Figure reports the GPU memory consumption of KV cache for each token. As YOCO only caches one layer of global kaluv pairs, it needs roughly L times less memory compared to Transformer.

---

120 30
Transformer
600
Transformer
GPU Memory (GB)
100
15
1.95 2.32x
3.01x
YOCO
KV Cache Memory
YOCO
IKB / Token)
80
400
32K 64K
128K
9.38x
60
40
6.39x
200
64X
20
4.16x
32x
40x
256K
512K
1M
1.2B
6.4B
13B
30B
65B
Length
Model Size
Figure 7: Inference memory of Transformer and YOCO across various lengths.
Figure 8: GPU memory of KV cache for each token with different model size.
300
Transformer
600
2.72x
Transformer YOCO
Prefilling Time (s)
8.36x
YOCO
Throughput (token/S)
5.05x
500
$2.87
200
32K 64K
128K
71.82x
2.57x
300
100
30.3x
15.55x
200
2.77x
32K
256K
512K
1M
100
4.37x
9.56x
Length
32K
128K
Context Length
Figure 9: Prefilling latency for different lengths. Transformer's time grows quadratically while YOCO's grows linearly.
Figure 10: Inference throughput of Transformer and YOCO varying the context length.
For example, YOCO can serve 128K tokens with 1GB GPU memory, while Transformer with GQA [1] can only support 1.6K tokens 65B model size.
Prefilling Latency the prefill stage, the model encodes input tokens in parallel. As shown in Figure the prefilling latency is pain point of user experience for long-context models. For 512K- and IM-length input sequences, Transformer needs about 180 seconds and 300 seconds, respectively, The computational complexity of Transformer ((N2), which requires large number of FLOPs for long context. In contrast, YOCO's prefilling time is (NN), 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 even for short context. For example, YOCO is 2.87 faster than Transformer for 32K length.
least two times speedup prefilling latency
Throughput The throughput indicates how many tokens the model process per second, involving both pre-filling and generation time. Figure 10 shows that YOCO achieves higher throughput across context lengths compared Transformer. For the example of 512K queries, Transformer's 's throughput is 4.5 token/s while YOCO reaches 43.1 token/s, i.c., achieving 9.6x 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 the throughput improvement.
5.5 Comparisons with Transformer Variants
We compare YOCO gRe gRet and YOCOswA 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 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.

---

GovReport
QMSum
Qasper
NarrativeQA
Answer PPL
3.0
4.0
4
6
3.8
5
3
2.5
8K 12K 16K
4K
8K 12K 16K
4K
8K 12K 16K
4K
12K 16K
Length
Length
Length
Length
Mamba
Sparse TRM
Hybrid H3
Transformer
YOCO geot
Figure 11: Long sequence task perplexity decreases along with the increasing input length.
Valid. Set.
AR-Hit
First-Oceur/
Fine-Grained LM Perplexity Table 5 re- ports the fine-gras fine-grained validation perplexity for language modeling. Following Zoology [2], we divide the perplexity into "Ar-Hit" and "First-Occur" Specifically, "Ar-Hit" consid- ers the predicted tokens that are bigrams previ- ously in the previous context, which eval- uates the associative recall capability. "First- Occur" considers the tokens that cannot be recalled from the context.
Mamba [13]
3.645
1.555
4.126
RetNet [35]
3.633
1.466
4.131
Hybrid H3 [S]
3.591
1.251
4,130
gRetNet
3,600
1.354
4.116
Transformer
3.564
1.219
4.104
YOCOsWA
3.553
1.202
4.094
YOCO gRet
3.530
1.199
4.067
Table 5: Fine-grained LM perplexity results.
Cong-Context Evaluation Figure 11 reports the answer perplexity with varying context length (ranging from 4K 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 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-com cong-contxtt modeling. Notice that the 12K and 16K results in Qasper are similar because the lengths most documents are shorter than 16K.
5.6 Ablation Studies
As shown in Table 6, we explore dif- ferent layout configurations for YOCO. First, compare the ratio of self-decoder to cross-decoder layers. For example, YOC is the default setting, where each module contains I/2 layers. The results show that YOCO|1:1] compara- ble to (0:0 and outperforms both YOCO [:3] and /OCO We use [1:1] as the default layout. Future work can re- fine a scaling law to guide the choice of layer ratio. Second, the setting "Unstacked YOCO " uses word embeddings X as input to the cross-decoder, rather than stacking cross-decoder upon self-decoder (i.e., using X 4/2 in Equa- tion (3)). Third, the model "Interleaved & Hybrid" is hybrid architecture that interleaves gRetNet and Transformer layers.
Valid. Set.
AR-Hit, First-Occur]
YOCO [:a]
3.530
1.199
4.067
YOCO [a
3.526
1.207
4.060
YOCO [:3]
3.565
1.230
4.102
YOCO 0:1
3,898
1.827
4.374
Unstacked YOCO [1:1]
3.531
1.188
4.071
Interlcaved & Hybrid
3.542
1.204
4.081
Table 6: Fine Fine-grained LM perplexity results. "[s:c]" is the ratio self-decoder to cross-decoder layers.
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 1M tokens. Profiling results also show that YOCO improves inference efficiency by orders of magnitude, especially for long-sequence modeling.
10

---

Acknowledgement
We would like acknowledge Ben Huntley for maintaining the GPU cluster. The long-sequence training utilizes CUBE, which is an internal version of [22]. We implement the Triton kernel of gated retention based on FLA [45]. This work was supported in part by National Key Research and Development Program of China under Grant No. 2020YFA0804503, National Natural Science Foundation of China under Grant No. 62272264, and Beijing Academy of Artificial Intelligence.
References
[1] J. Ainslie, J. Lee-Thory M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. Training generalized multi-quen transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13225, 2023.
[2] S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré. Zoology: Measuring and improving recall in efficient language models, arXiv preprint arXiv:2312.04997, 2023.
[3] M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer. J. Brand- stetter, and S. Hochreiter. xLSTM: Extended long short-term memory. arXiv preprint arXiv:2405.04517, 2024.
[4] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse Transformers. URL http:///wwnuii.oooooooo/////e//p
2019.
[5] T. Dao, D. Y. Fu, K. K. Saab, A. W. Thomas, A. Rudra, and Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint 1rXiv:2212.14003 2022.
[6] T. Dao, D. Haziza, F. Massa, and G. Sizov. Flash-Decoding for long-context inference. https: //crfm. stanf..ssaoorrrr edu/2023/ /t////////aaaccodiiingggg html, 2023.
[7] DeepSeek-AL. Deepseek-v2 v2: A strong, economical, and efficient mixture-or -ore--of---etee-e-eeeeee language model. 2024.
[8] J. Ding, S. Ma, L. Dong, X. Zhang, S. Huang, W. Wang, N. Zheng, and F. Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307 02446, 2023.
[9] Y. Fu, R. Panda, X. Niu, X. Yue, H. Hajishirzi, Y. Kim, and H. Peng. Data engineering for scaling language models to 128k context. ArXiv, abs/2402.10177 2024.
[10] L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac"h, h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. A framework for few-shot language model evaluation, 12 2023.
[11] S. Ge, Y. Zhang, L. Liu, M. Zhang, J. Han, and J. Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. arXiv preprint arXiv:2310.01801, 2023.
[12] X. Geng and H. Liu. LaMALLaMA: An open reproduction of LLaMA. openIm-researe hwtttaereeeerrchhhhoeeeee /oenn_ 11 11ama, 2023.
com/
[13] A. Gu and T. Dao. Mamba: Linear- time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.000722 2023.
[14] C. Hooper, S. Kim, H. Mohammadzadeh, M. W. Mahoney, Y. S. Shao, K. Keutzer, and A. Gholami. Kvquant: Towards 10 million context length LLM inference with KV cache quantization. arXiv preprint arXiv:2401.18077 2024.
[15] S. Hu, Y. Tu, X. Han, C. He, G. Cui, X. Long, Z. Zheng, Y. Fang, Y. Huang, W. Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06630 44.4096939, 2024.
[16] G. Kamradt. Needle in a Haystack pressure testing llms. https://github. com/ homm/kkamrdt// LLITeett_Needleellliiysssssdddd htwp://aaattttcaeeee////imaaaa 2023.

---

[17] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. CoRR, abs/2001 1.0861,, 2020.
[18] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages $156-5165. PMLR, 2020.
[19] T. Katsch. Gateloop:
Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:231 2023.
[20] S. Li, F. Xue, Y. Li, and Y. You. Sequence parallelism: Making 4D parallelism possible. arXiv preprint arXiv:2 /3120, 2021.
[21] O. Lieber, B. Lenz, Bata, Cohen, Osin, Dalmedigos, E. Safahi, S. Meirom, Y. Belinkov, S. Shalev-Shwartz. O. Abend, R. Alon, T. Asida, A. Bergman, R. Glozman, M. Gokhman, A. Manevich, N. Ratner, N. Rozen, E. Shwartz, M. Zusman, and Y. Shoham. Jamba: A hybrid Transformer-Manba language model. CoRR, abs/2403. 19887, 2024.
[22] Z. Lin, Y. Miao, G. Liu, X. Shi, Q. Zhang, F. Yang, S. Maleki, Y. Zhu, X. Cao, C. Li, M. Yang. L. Zhang, and L. Zhou. SuperScaler: Supporting flexible DNN parallelization via unified abstraction, 2023.
[23] H. Liu, M. Zaharia, and P. Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2110008889, 2023.
[24] H. Liu, W. Yan, M. Zaharia, and Abbeel. World model on million-length video and language with ringattention. arXiv preprint arXiv:2402. 08268, 2024.
[25] Z. Liu, J. Yuan, H. Jin, S. Zhong, Z. Xu, V. Braverman, B. Chen, and X. Hu. KIVI: A tuning- free asymmetric 2bit quantization for KV cache. arXiv preprint arXiv:2402.02750, 2024.
[26] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. International Conference on Learning Representations, 2019.
[27] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, H. Cao, X. Cheng, M. Chung, M. Grella, K. K. GV, X. He, H. Hou, Kazienko, Kocon, J. Kong, B. Koptyra, H. Lau, S. I. Mantri, F. Mom, A. Saito, X. Tang, B. Wang, J. S. Wind, S. Wozniak, R. Zhang, Z. Zhang, Q. Zhao, Zhou, Zhu, and R.-J Zhu. Rwkv: Reinventing rnns for the transformer era, 2023.
[28] B. Peng, J. Quesnelle, H. Fan, and Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint
[29] R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, and Dean. Efficiently scaling Transformer inference. ArXiv, abs/2211.05102, 2022.
[30] M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.055300 2024.
[31] U. Shaham, M. Ivgi, A. Efrat, J. Berant, and O. Levy. Zeroscrolls: A zero-shot benchmark for long text understanding. arXiv preprint arXiv:2305. 14....35..119,, 2023.
[32] N. Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 2020.
[33] N. M. Shazeer. Fast transformer decoding: One write-head all you need. ArXiv,
abs/1911.021555 2019.
[34] Y. Sheng, I.. Zheng, B. Yuan, Z. Li, M. Ryabinin, B. Chen, P. Liang, C. Ré, I. Stoica, and C. Zhang. Flexgen: High-throughput generative inference of large language models with a single GPU. In International Conference Machine Learning, pages 31094-31116. PMLR, 2023.
[35] Y. Sun, I. Dong, S. Huang, S. Ma, Xia, J. Xue, J. Wang, and Wei. Retentive network: successor to transformer for large language models. arXiv preprint arXiv:2307.086211 2023.
12

---

[36] P. Tillet and D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
[37] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro. F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397,, 2023.
[38] J. Tow. StableLM Alpha v2 models. Sabblelm-base-alpha--2
http:://huggingffcecccc//saa http:///ggg.....otttaaatttallll.
[39] J. Tow, M. Bellagente, D. Mahan, and C. Riquelme. StableLM 3B 4EIT. batts://aka..nn
StableLM-3B-4EI1
[40] A. Vaswani, Shazeer, N. Parmar, Uszkoreit, L. Jones, N. Gomez, Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pages 6000-6010, 2017.
[41] H. Wang. S. Ma, S. Huang, L. Dong. W. Wang. Z. Peng, Y. Wu, P. Bajaj. S. Singhal, A. Benhaim, B. Patra, Z. Liu, V. Chaudhary, X. Song, and F. Wei. Magneto: A foundation Transformer. In Proceedings of the 40th International Conference Machine Learning, volume 202, pages 36077-36092, 2023.
[42] Y. Wu and K. He. Group normalization. In Proceedings the European conference on computer vision (ECCV), pages 3-19. 2018.
[43] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. arXiv preprini arXiv:2309. 010222399.77755 17453, 2023.
[44] W. Xiong, J. Liu, I. Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint 21X.v:2309.16039 2023.
[45] S. Yang and Y. Zhang. FLA: A Triton-based library for hardware-efficient im- plementations of linear attention mechanism. https://githi: com/ sust sustcsongln/ attention, Lat---iinaaaateetio... 2024.
[46] S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efffiient training. arXiv preprint arXiv:2312.06633 2023.
[47] A. Zeng, X. Liu, 7.. Du, 7. Wang, H. Lai, M. Ding, Z. Yang, Y. Xu, W. Zheng, X. Xia, et al. GLMM-130B An open bilingual pre-trained model. arXiv preprint UrXiv:2210.0244 2022.
[48] B. Zhang and R. Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[49] Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, et al. H2O: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
13

---

A Chunk Parallelism for Long-Sequence Training of YOCO
We introduce chunk parallelism for YOCO 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 distributed long-sequence training.
GPU
GPU
02
O2
V
Cross-Decoder
All-Gather Once
Q1
V,
K₂
V/2
Q2
1
1
Project
M2
Self-Decoder
Split
Chunk parallelism of YOCO training on two GPU devices. The training strategy is to partition the sequence into different chunks. .M denotes the intermediate representation X i.e. the output of self-decoder The keys and values in the cross-decoder are only gathered once.
Figure 12:
In self-decoder, the dependency only exists in the adjacent devices. For example, gated retention only requires the hidden state S, in Equation (6), and sliding-windon attention attends to tokens within the context window. Therefore, the communication amount of self-decoder is relatively small. In the rroosde the all-gather operation is only triggered once for the KV cache, rather than communicating in each layer. The hardware-friendll architecture gives more flexibility distributed long-sequence training.
B Chunk-wise Representation of Gated Retention
We illustrate the equivalence between recurrent representation and chunkwise recurrent representation of gated retention. For the output O,,, n can be split as kB + where the chunk size:
n
O
Σ
II
RQKKTV
Σ
II i=m-1
Σ II
7QrKTV Vm
n
Σ
II
RQKKTV (QrKBBt1M
m=
kB T
11 II
(c+1)B
kB
Σ
II RQKKVV (Qn
Σ
KKT
mteB
V
Vnntn
II
%)
II
%
imm+1
i=kB+1
i=mtcB+1
(=c++))+++
n-1
(Qn
II
(K (V
2 (e))
II a
i=kB+1
(=c+1
(Qn
II k=
R4-1
(9)
14

---

where r, = - ni. Sig (5. = -IT= 1) %o = -I]==((-1)------ Y [ indicates the i-th
chunk, i.e., [7(]] = [()11))+11" -.i]- R, is written recurrent function:
R4 = ( o (tj) + +,
(10)
I (1-1)8+)
Denote [i] as the i-th chunk, i.e. 2| 1184111 TBB), P6-1)8+5 Pl (i, k) B(c-1)B+5) We concatenate the output in a block together:
(k=(i-1)B+1
[n]
&B
On1
Σ
Qm) Km mm +
Σ
1l Qml
II
aKT Vm
m=kB+1
m=1
im++1
In]
((n-1)B+k
Bi KT Vm (Q c Din) (in(i,k)
if S k else 0
B(n-11B+5
1++1
kB
BiB
Σ m2
II
Wk - Qn
R3 =
R=K (% o
+ BBB Ri-1=
2mm
[(n] (Q O Pm + (Q Rm-1) 3 1mn)
Inner: Chunk
Cross-Chunk
(11)
Finally, show that the chunkwise recurrent representation of gated retention is equivalent to the other two representations.
C Pseudo Code of Gated Retention
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, baz head len dim k, baz num head len dim
baz num head len dim
gt): bsz t-a
head len
retention causal mausal_makk torch. shape [-2], q. s.shape [-2]], float (M=inf""" (evice=q.device).
triu(1) (yuu()).tpp__as(( gt F=loosignoid(gt). cumsum(m1) gate_lose hal___ogittttommmaaae
mask (g[..., None] g(..., None, causal_mask exp()
retention retention mask output retention output group_norm norm(ouput)) return output
def RecurrentRetentionn
nun head dim num head dim gt baz num head ):
9, k, v, past_kv,
gt logeigmoidd(tt) / gate_ gate_logit, poggit_normaiizer current kv gt. exp() * past kv k. (-1) (-1) * v. unsqueeze(-2)
output torch. sum (q. unsqueeze output group_ nor_ (output)
current kv, dim=-2)
return output, current_kr
15

---

ChunkriseRereation)
9, k, v, baz num num_head chunk_size dim past_ky, baz * num num_head dim * dim
gt); bsz * num head * Saunk_sizo
gt F. F.logeigmoid(gt).. .ulluttt..cmmmmm1)) / gatt__oggt_normatte cross_ reessxeeention (q past_kv) gtl..., .exp)
k, v, gt) retention inner retention cross_ retention output group_ pep__normmeeee (ort(rettt)ion)
inner rnnerreeenttin Parallelketentionn
value_decas (-gt gtl:, :, -1, Nonel) exp() [..., None]
chunk, decay gt ....... -1] -i].exp()
current chunk decay past_ ky trtraaapose -2) (v value value_decay)
return output, current_ir
D Hyperparameters for YOCO-3B (Section 5.1)
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 26.
Grouped-query attention [1] is used, where the number of query heads is 24, and the number of key-value heads is We train YOCO with gated retention (Section 4.2). The non-embedding parameter count is 2.8B. In comparison, MableeLM-3--4EJJ is 2.7B and OpcnLLaMA-v2-33 [12] is 3.2B. The training sequence length is 4096. The batch size is 4M tokens. We use the AdamW [26] optimizer with 8 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 [39]. We use tiktoken-cli00l 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 3 = (0.9, 0.95). The learning rate 3.2 x 10 with 1000 warmup steps. We set 5T-token learning
rate schedule with linear decay 1.28 10
Params
Values
Layers
26
Hidden size
3072
FFN size
8192
Vocab size
100,288
Heads
24
Key-value heads
8
Adam /
(0.9, 0.95)
LR
3.2 x 10
Batch size
4M
Warmup steps
1000
Weight decay
0.1
Table 7: Hyperparamters used for the YOCO-3B model Section 5.1.
E Hyperparameters for Scaling Curves (Section 5.2)
Table 8 reports the hidden dimension, number of layers, and number of heads used for different model
sizes. The head dimension gated retention is set to 256. To align the number of parameters, the
FFN size for Transformer is şd while the FFN size for YOCO is 3d. The training length is set to 2048. The batch size set to 0.25M tokens. We use the AdamW [26] optimizer with 3 0.9, 83 0.98.
The learning rate is 1.5 1 for 160M to 1.4B sizes and 7.5 x 10 for 2. 7B to 13B sizes. The
16

---

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.
Size
Hidden Dim.
#Layers
#Heads
160M
768
12
12
400M
1024
24
830M
1536
24
12
1.4B
2048
24
2.7B
2560
32
20
6.8B
4096
32
32
13B
5120
40
40
Table 8: Model size and parameeassmeters used for scaling curves in Section 5.2.
Hyperparameters for Length Extension
We progressively extend the context length 1M tokens in Section 5.3. The length schedule is 64K, 256K, and I M. We up-sample the documents that are longer than the training length [9]. Table shows that we use different RoPE and learning rate for each stage.
Training Length
65,536
262, 144
1,048,576
Learning Rate
8 x 10-5
4 x 10
2 x 10
RoPE
640K
5M
80M 1.5B
Training Tokens
6B
4B
Table 9: Hyperparamters used for length extension in Section 5.3.
17

---

NeurIPS Paper Checklist
1. Claims
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. 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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
18

---

Justification:
There is no theoretical result in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
The experiment can be easily reproduced based on the model description, hyperparameter, and any well-known pre-training corpus.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [Yes]
Justification: Code will be released in camera-ready version. All of the data we use is aublic-available
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https: // /nipp..ccocc cttps/cccppbbcc//gceeeecec gww.//////eeeeeeeee / ssi ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Hyperparameters are attached in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: For large language models, the variance between different runs is negligible. Moreover, the evaluation pipeline is deterministic.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
20

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The corresponding resources are stated in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttea://////ss/////eeee httpss//aurrpp..ccpppbb hwabbliic/thiisuuiiiio cscuidel ines?
Answer: [Yes]
Justification:
We follow the NeurlPS Code of Ethics in the research.
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS 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., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 fundamental research that has 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 socictal impact 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.
21

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not pose safety risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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
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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "kB559NRact", "reviewer": "Reviewer_CYFA", "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.", "questions": "- In the ablation study, does Unstacked YOCO refer to the model without the self-decoder?\n\n- 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": "Please refer to the weakness section.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Overall, this is a high-quality paper.\n\nOriginality: The paper presents a novel architecture that achieves performance comparable to full transformers with only one layer storing global KV tokens.\n\nQuality: 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.\n\nClarity: 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.\n\nSignificance: 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.\n\n- 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.\n\n- 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.\n\n- 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.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review and insightful feedback.\n\n>Q1: In the ablation study, does Unstacked YOCO refer to the model without the self-decoder?\n\nA1: The input of Unstacked YOCO's cross-decoder is the output of **embedding layer**. In comparison, the input of YOCO's cross-decoder is the output of **self-decoder**. The model without the self-decoder is `YOCO_[0,1]` in Table 6, where the whole model is stacked with cross-decoder and the shared KV cache is word embedding.\n\n\n---\n>Q2: The model hyper-parameters such as the number of layers and the number of attention heads.\n\nA2: We keep most of them the same as the standard Transformer to ensure fair evaluation, where both the model size and depth are comparable.\n\n\n---\n>Q3: a significant future application for long context models is long video understanding\n\nA3: We consider multimodal scenario as one of the most important future directions. Thanks for your suggestions.\n\n\n---\n>Q4: Ablation studies on the window size of SWA.\n\nA4: We find that the window size from 1024 to 4096 achieves similar end performance in our early experiments. Since a larger window size affects inference latency, we keep the default as 1024.\n\n\n---\n>Q5: a few typos\n\nA5: We will fix them in the camera-ready version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FUX4KeSD4e", "reviewer": "Reviewer_fhCA", "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.", "questions": "- The evaluation primarily focuses on memory and latency improvements. Does YOCO also bring training efficiency gains?\n\n- 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": "The paper does not explicitly discuss any limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- YOCO's design, with its cross-decoder and self-decoder, offers a novel approach to caching key-value pairs, reducing GPU memory consumption.\n- The architecture significantly reduces prefill latency and improves throughput, addressing critical bottlenecks in long-sequence language model inference.\n- YOCO demonstrates effective scalability in model size and training tokens, maintaining competitive performance with other leading Transformer models.\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: The comparison and discussion with FlashDecoding.\n\nA1: Flash-Decoding and kernel fusion have been used in comparison (described in L240, L121, L25), i.e., the Transformer results have been based on FlashDecoding. The contributions of YOCO and FlashDecoding are orthogonal. YOCO optimizes the pre-filling complexity and KV cache memory from the perspective of architecture design, while FlashDecoding optimizes the implementation. We can directly utilize FlashDecoding for cross-decoder without rewriting the kernel.\n\n\n---\n>Q2: The performance versus efficiency tradeoffs of Efficient Self-Attention (ESA)\n\nA2: Table 3/4 and Figure 3/4 show that ESA does not harm the end-to-end performance under the YOCO architecture. We find that the window size from 1024 to 4096 of self-decoder (SWA) achieves similar end performance in our early experiments.\n\n\n---\n>Q3: The training efficiency in YOCO\n\nA3: The training efficiency of YOCO and Transformer is comparable when the training length is small. When the training length becomes long, YOCO training is faster compared with Transformers because of the cost saving of self-decoder, with a speedup ratio between 1x and 2x.\n\n\n---\n>Q4: The efficiency comparison when the token length is very short\n\nA4: Even for short sequences, there is still 2 times prefill speedup with YOCO. As described in Figure 2 and Line 115, we can still exit early before entering the cross-decoder during the prefill stage. The YOCO models are not slower than Transformers in Table 3.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D9noSazh8a", "reviewer": "Reviewer_qC7V", "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.", "questions": "Please refer to **Weaknesses**", "limitations": "Please refer to **Weaknesses**", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed architecture is simple and effective, which could be easily integrated into existing transformer model implementations.\n2. 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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: insights and explanations for the effectiveness of the proposed architecture\n\nA1: The key insights are summarized as follows. First, KV cache can be shared across layers without significantly affecting language modeling performance. Most previous work focuses on compressing KV cache along with the sequence dimension, while YOCO improves the cache issues from another perspective, i.e., the layers. Second, the hybrid design is competitive. After the NeurIPS submission, there were several concurrent works indicate this insight, such as Samba, and character.ai's architecture blog. Third, the early exit insight (as described in Figure 2 and Line 115) dramatically improves the prefill speed. All the above insights and advantages make YOCO novel and go beyond conventional KV compression methods, which improves deployment and user experience.\n\n\n---\n>Q2: How the efficient self-attention and decoder-decoder structure affect the model's performance respectively.\n\nA2: \n\n- The comparisons between decoder-decoder and decoder-only architectures are presented in Table 6, i.e., the settings `YOCO_[1:1]` and `Interleaved & Hybrid`, where the interleaved model is a decoder-only architecture with hybrid layers. The results show that the two layouts achieve similar performance. \n\n- For different self-decoder choices, we conducted experiments with sliding-window attention and gated retention. Both representative design choices work well as shown in Figure 3 (i.e., model size scaling up experiments) and Table 5 (i.e., ablation studies).\n\n- Different ratios between self-decoder and cross-decoder are also compared in Table 6.\n\n- In order to comprehensively inspect how the proposed architecture affects performance, we conducted evaluation from diverse perspectives, including scale up training tokens (Section 4.1), scaling curves of the proposed architectures (Section 4.2),  scale up the YOCO model to 1M context length and evaluate its long-sequence modeling capability (Section 4.3), compare with Transformer variants (Section 4.5), and ablation studies on various design choices (Section 4.6).\n\n\n---\n>Q3: what contributes to the performance improvement?\n\nA3: The improvements mainly come from the hybrid-style architecture. Multiple recent works confirmed this point, such as Samba[1], and Jamba[2]. The trends are consistent across different learning rate schedules. Because YOCO saves the key and value projection matrices, for fair comparisons, we accordingly increase the FFN part in order to keep the overall parameter count similar across models. For example, as shown in the page 47 of the Llama 2 paper[3], this is a common practice for fair comparisons across design choices. Besides, instead of performance, we focus more on the improvements in terms of inference memory, prefill latency, and throughput.\n\n[1] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling\n\n[2] Jamba: A Hybrid Transformer-Mamba Language Model\n\n[3] Llama 2: Open Foundation and Fine-Tuned Chat Models\n\n\nWe hope the above explanation clarifies the rationale behind our experiment designs. Thank you again for the valuable feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fJFkVUFabR", "reviewer": "Reviewer_hQmr", "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. \nNotably, 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.\n\nTwo ESA variants are evaluated: sliding window attention and a novel gated retention method, which incorporates data-driven head-wise decay over retention. \nUpon 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. \nThey also conduct some analysis on long-seq evals and observe near-perfect performance on needle-in-haystack tests and other benchmarks like Qasper.", "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$\n2. I'm curious if the authors have tried other linear attention variants instead of gRet, e.g., Mamba, and GLA.", "limitations": "N/A", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. The proposed data-dependent gated-retention brings great improvement against retention.\n3. 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:\n\n1) 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. \n2) Despite concurrent works, I suggest the authors to add discussions with Samba [1] and Mamba2 [2] in their next version.\n\n[1] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling\n\n[2] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive comments.\n\n>Q1: add results of exisiting linear-time / hybrid models trained on trillions of tokens, e.g., RWKV6 and TransNormer\n\nA1: We focus on the evaluation using the same training data for fair comparisons. So we use OpenLLaMA and StableLM in Table 3. We also compare with other linear-time / hybrid architectures (e.g., Mamba, and hybrid H3) in Table 5 with the same training data and settings. We can include RWKV6 and TransNormer numbers for reference results as suggested.\n\n\n---\n>Q2: Despite concurrent works, I suggest the authors to add discussions with Samba and Mamba2 in their next version.\n\nA2: Samba and Mamba2 were released in arXiv after the NeurIPS submission. These two methods are complementary to this work, which is promising to use them as self-decoder in YOCO. Specifically, the ablation setting `Interleaved & Hybrid` in Table 6 is similar to the hybrid design of Samba, and both Mamba2 and gRetNet follow the similar design principles. We can include the discussions in the camera-ready version.\n\n\n---\n>Q3: suggestions about notations\n\nA3: Thanks for the suggestion. We will optimize the notations in Eq. (7) and (8).\n\n\n---\n>Q4: I'm curious if the authors have tried other linear attention variants instead of gRet, e.g., Mamba, and GLA.\n\nA4: We conducted experiments with sliding-window attention and gated retention in the work. The other linear attention variants are supposed to behave similarly and follow the same trend.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (oral)