LightTransfer: Your Long-Context LLM is Secretly a Hybrid Model with Effortless Adaptation
Authors: Xuan Zhang
Fengzhuo Zhang
Cunxiao Du
Chao Du
Tianyu Pang
Wei Gao
Min Lin
Venue: Accepted by TMLR ()

Abstract: Scaling language models to handle longer contexts introduces substantial memory challenges due to the growing cost of key-value (KV) caches.
Motivated by the efficiency gains of hybrid models and the broad availability of pretrained large transformer backbones, we explore transitioning transformer models into hybrid architectures for a more efficient generation.
In this work, we propose \textsc{LightTransfer},
a lightweight method that transforms models such as LLaMA into hybrid variants.
Our approach identifies \textit{lazy} layers---those focusing on recent or initial tokens---and replaces their full attention with streaming attention. 
This transformation can be performed without any training for long-context understanding tasks or with minimal fine-tuning for o1-like long reasoning generation tasks that require stronger reasoning capabilities.
Experiments across diverse benchmarks and models (e.g., LLaMA, Mistral, QwQ-STILL) demonstrate that, 
even when half of the layers are identified as \textit{lazy},
\textsc{LightTransfer} achieves up to 2.17$\times$ throughput improvement with minimal performance loss ($<1.5\%$ on LongBench) and achieves 53.3\% on math benchmark AIME24 of advanced o1-like long reasoning model QwQ-STILL.

Published in Transactions on Machine Learning Research (09/2025)
Light Transfer: Your Long-Context LLM is Secretly a Hybrid Model with Effortless Adaptation
Xuan Zhang Singapore Management University
Fengzhuo Zhang
fazhang So. mus. edu
National University of Singapore
Cunxiao Du
ducr@sca. com
Sea AT Lab, Singapore
Chao Du Sea AI Lab, Singapore
duchao@sea.com
Tianyu Pang Sea AI Lab, Singapore
tienyupang@sea.com com
Wei Gao
Sreigao@smu.edde
Singapore Management University
Min Lin Sea AI Lab, Singapore
linmin seasaaccom
Reviewed OpenReview: /ttps:///ppnreeviee net/ forum? id= knegvWICro
Abstract
Scaling language models to handle longer contexts introduces substantial memory challenges due to the growing cost of key-value (KV) caches. Motivated by the efficiency gains of hybrid models and the broad availability of pretrained large transformer backbones, we explore transitioning transformer models into hybrid architectures for a more efficient generation. In this work, we propose LIGHTTRANSFER, lightweight method that transforms models such LLaMA into hybrid variants. Our approach identifies lazy layers -hhree focusing on recent or initial tokens and replaces their full attention with streaming attention. This transformation can be performed without any training for long-contess understanding tasks or with minimal fine-tuning for ol-like long reasoning generation tasks that require stronger reasoning capabilitics. Experiments across diverse benchmarks and modols (c.g., LLaMA, Mistral, QwQ-STILL. demonstrate that, even when half of the layers are identified as lazy, LIGHTTRANSFER achieves up 2.17x throughput improvement with minimal performance loss (< 1.5% LongBench) and achieves 53.3% on math benchmark AIME24 of advanced ol-like long reasoning model QwQ-STILL. Our project homepage: https
//sites google. wwwgle.comm/eee wwt.:/cww/wwiiieeeeeeeeeeee wwt.e/ewwwwwttaneeee...
//sites.google....
1 Introduction
Recent advancements in large language models (LLMs) have extended their capacity for handling long context inputs and generating long-form reasoning. For example, LLaMA-3.1 supports context lengths up
"Work doue during Xuan Zhang's associate membership at Sea Al Lab.
i "oorreppnndecce Cunxiao Du.

---

Published in Transactions on Machine Learning Research (09/2025)
Outputs
Outputs
FFN
FFN
Non-lagy
Standard Attn layer
Standard Attn layer
KV
cache
cache D
FFN
FFN
Standard Attn layer
Efficient Atti layer
KV
Lagg
cache
state D
Inputs
Inputs
(a) Transformer
(b) Hybrid model
Figure 1: (a) A standard transformer architecture. (b) A hybrid model in which certain layers of a standard transformer are replaced with more memory- designs. LIGHTTRANSFER identifies lazy layers in (a) and transforms them into more efficient variants, yielding (b).
to 128K (Dubey et al. 2024,, while OpenAl's ol can produce sequences of up to 100K tokens (OpenAl 2024). As the cornerstone of the efficient inference of these models on long context, key-value (KV) cache stores precomputed key and value tensors for each token in the language sequence to avoid recomputing them for each attention layer. However, as the number of model layers and input lengths increases, the memory required for storing the KV cache grows significantly, posing challenges for inference efficiency.
Various methods have been proposed to roduce the KV cache storage by modifying the model architoc- ture (Shazoer 2019 Brandon al.) 2024 Goldstein et al.) 2024] Nawrot et al., 2024 Wang et al. 2024b Yu et al.] 2024) One promising approach is the hybrid model. As shown in Figure 1 in these hybrid models, certain layers of a standard transformer are replaced with more memory--fficient mechanisms such
RNNs (Sherstinsky 2020). Mamba (Gu & Dao 2023 and sliding window attention (Beltagy et al.] 2020 These approaches exploit the notion that different layers can be manually assigned distinct functiona lities, such as using memory-efficient mechanisms for local context processing and standard attention for global context handling (Gemmma et al.j 2024), thereby achieving notable memory savings. Concrete examples of such hybrid architectures include Transformer-Rest Recurreett Neural Network (RNN) designs such as YoCo (Sun et al. 2024), Transformer-Mamba approaches such as Jamba. (Lieber et al. 2024 Team et al. 2024 and Transformer-Sliding Window models like Gemma 2 (Gemma et al. 2024 and YoCo (Sun et al. 2024 However, key limitation is that they require training the entire model from scratch.
Given the substantial efficiency gains offered by hybrid models and the availability of large-scale pretrained transformer backbones, a natural direction is to transition these pretrained modols into hybrid architectures with minimal additional training. A straightforward method is to replace traditional full attention layers with sparse attention, thereby adopting a fixed-size KV cache reduce memory overhead. A representative example is streaming attention (Xiao et al. 2023), which augments the sliding-window mechanism by introducing
sink tokens. However, Table 2 shows, completely substituting all standard attention layers with streaming attention leads to severe degradation in the model's I's ability to process long contexts, thereby undermining its capacity to capture global information. Consequently, when transitioning from a pretrained transformer to a hybrid model, two primary challenges arise. First, it is necessary to retain some standard attention layers to preserve the model's s long-context modeling capabilities, raising the critical question: which layers should remain intact? Second, this transition should idcally be lightweight, cnabling efficient adaptation with minimal data even allowing to be applied entirely at test time. Otherwise, large-scale pretraining data were required, one could simply train a hybrid model from scratch, undermining the value of a transition-based approach.
In response to the above challenges, we examine the atlention patterns in different transformer layers to determine whether each layer exhibits distinct functionalities. We conduct preliminary experiments (Section and identified two key findings: First, certain layers in long-contert LLMss exhábit lazy behavior, primarily focusing on semantically unimportant tokens (e.g., the initial few tokens) and the most recent during answer generation. The properties of lazy layers address our first challenge and enable standard transformer-based
2

---

Published in Transactions on Machine Learning Research (09/2025)
LLMs to operate in hybrid-like manner: identified lazy layers use streaming attention, whereas non-lazy layers retain full attention. Second, after analyzing attention weight patterns, we find that layer behavior is consistent across tokens for given long imput. This insight partially addresses the second challenge and paves the way for a test-time transformation in which selective modifications are applied during the prefilling stage, allowing for elficient adaplation without extensive retraining.
Building upon this insight, we propose LIGHTTRAASFERR a lightweight method that transforms models such as LLoMA, Mistral (Jiang et al.) 2023). and Qu (Qwen 2024b) into their corresponding hybrid variants. Specifically, as shown in Figure I we analyze the attention allocation patterns in each layer to determine whether it can be treated as a lazy layer. In lazy layers, we apply streaming attention, while standard attention is retained in non-lazy layers. The output of the transformer with a reduced KV cache differs from the original output due to the reduced cache size, and this difference is theoretically analyzed in Theorem 5.1 For tasks where the input is sufficiently long (i.e., long-context understanding), we leverage on-the-fly lazy layer identification at the prefilling stage, LIGHTTRANSFER TEST.. In addition, for ol-like long reasoning generation lasks, even though the questions can be relatively short (only few dozen tokens) yel. demand higher model capacity, we surprisingly find that minimal training still enables robust performance SEEHETTRAASEER-TRART In practice, this transition requires only around 5K samples (originally utilized for long-reasoning ability distillation (Min et al. 2024 underscoring the lightweight nature of our approach.
We conduct experiments on four representative LLMs (i.e., LLaMA2-7B-chaa (Tourron et al. 2023 ),
Mistral--BBIIntruueet (Jiang et al., 2023), LLaMA3-SB- and its 70B counterpart (Dubey et al. 2024) )
evaluating them on long-contens context benchmarks including LongBench (Bai et al. 2023) and Needle- Beedle-In-A- -eedle-In-A-Haysstcc (NTAH) (Kamradt, 2023). In addition, we adapt an ol-like long reasoning model QwQ-32B-STILLL and assess its performance on MATH-OAI (Lightman et al.,, ALME24 D and GSMSK (Cobbe et al. Experimental results indicate that hybrid models converted via LIGHT TIIHTTRRASFEER achieve performance on par with their standard transformer counterparts. For example, on long-context understanding tasks, it achieves only a 1.45% performance decline on LongBench. For long reasoning tasks, it achieves performance that is comparable to or even belter on the widely used mathematical benchmark AIME24, reaching 53.3% accuracy, Notably, these results were obtained while half of the model's layers employed streaming attention, yiclding up to 2.17x increase in throughput.
2021
2 Related Works
Memooy-efficient architectures, such as linear RNN-based architectures (e.g., Mamba Gu & Dao 2023 and those employing sparse attention methods (e.g., streaming attention (Xiao et al. 2023 have demonstrated
clear advantages in deployment, including reduced memory usage and higher throughput (Peng et al.
2023
Dao & Gu 2024: Vang et al. 2023: Sun et al. 2024 Lieber et al.) 2024 Gemma et al. 2024 However, a key drawback of these memory-efiine efffccee models is their limited ability to handle extended contexts effectively (Behrouz el al.) 2024 Yuan el al.] 2024). Meanwhile, the inference memory cost of standard transformers (i.e., the storage of the KV cache) grows linearly as the context length increases (Shi al.
Li et al.) 2024c a). To address these challenges, recent rescarch has proposed hybrid architecture: maintaining the strong capabilities of pretrained transformers while selectively substituting certain transformer layers with more memory-efficiett modules, thereby balancing high performance with practical deployment
2024
efficiency (Lieber et al. 2024 Gemma et al.) 2024 Sun et al. 2024 Botev et al. 2024 De et al. 2024). For
example, Jamba (Lieber ct al.
2024 Team et al. 2024) integrates Mamba (Gu & Dao 2023) with transformer,
Gemma 2 (Gemma et al. 2024) alternates sliding- window altention ion with standard altention layers, and
Minimax-01 (Li et al.) 2025) employs lightning attention (Qin et al., (2024) in certain layers. However, a key limitation of these methods is their reliance on training the entire model from scratch. Although recent approaches aim to leverage the capabilities of large-scale pretrained models by converting selected layers into memory-efficient structures to form well-trained hybrid model, they still depend on extensive training data (Wang et al.) 2024a] Ge et al.] 2024). For instance, LongGen (Ge et al. 2024) transforms certain layers in pretrained LLM into sparse attention but requires retraining on over 2TB of data. Differently, our framework is designed to be substantially more lightweight. Despite requiring no additional
https://huggingffcecc////a http:///awaac..octtaattstttttttt wtwp://AAA//////---/aaeeeeee valiioo-oaiddatinnnaaa.a
3

---

Published in Transactions on Machine Learning Research (09/2025)
Layer 0
Layer 20
pattern, Astention weights
sink and recent tokens
68
67
06
04
Prefilling stoge
Decoding stage
Figure 2: Visualization of attention weight distributions LLaMA3-8B. Left: The attention patterns across different layers. Right: Each cell in the heatmap represents the aggregated attention weight from each query token (x-axis) to two loken groups: (1) the initial Winnk tokens, and (2) the most recent. Wrecent tokens, during both the prefilling and decoding stages. Layers predominantly attending to these token groups are highlighted with black outlines. The color intonsity indicates the sum of attention weights, averaged across samples a batch.
training for long-context understanding tasks, and only 5K training examples for more demanding long-text reasoming tasks (as originally used for long-reasoning ability distillation (Min et al. 2024)), it still achieves comes from the identification of each layer's function, whereas LongGen always uses fixed structure (retaining the middle layers for full attention). Some works also attempt to fully transfer transformer models into RNN-like architectures (Kasai et al. 2021 Zhang et al.) 2024b Mercat et al.) 2024 Zhang et al. 2024a; Bick et al. 2024). However, these methods primarily focus on short-context tasks (e.g., QA), whereas our approach targets long-context scenarios.
strong performance
both fronts. The superiority of our
3 Preliminary
Before introducing LIGHTTRANSFER, we provide a brief overview of the generative inference in autoregressive LLMs, which the key background for our method.
Inference stages. The typical generative LLM inference process involves two stages: (1) Prefilling: the autoregressive LLM processes the input prompt X by parallel computing, and also saves the KV cache of tokens in X. The output of the last token in this stage is the first token of the response. (2) Decoding: after the prefilling stage is completed, the LLM generates output tokens one by onc, and saves their KV cache. In each decoding step, a new loken is generated based on the current loken and the KV cache stored from earlier steps, contimuing until a stop criterion met.
4 Observations
In this section, we analyze the attention patterns during inference in long-context LLMs, providing insights that motivate our approach to transform the standard transformer into its corresponding hybrid variant. The study conducted on the LLaMA3-8B-Tnstrus model (Dubey et al. 2024 using a sample from the LongBench Bai et al. 2023 benchmark. Our key findings are follows:
Layer behavior in long-context LLMs during inference. Previous research (Xiao et al. 2023 has shown that large portion of attention in LLMs tends to focus semantically unimportant tokens Xinitial (e.g., the first few tokens) and the most recent tokens Xrecent (i.c., tokens in the sliding window). We refer to this paltern as lazy behavior, likening it to skimming a paper by reading only the first lines and the conclusion. While it is also called attention sink [Xiao et al., 2023 Gu et al.) 2024). we emphasize the shorteut nature by referring to as lazy. Through our analysis, we find that even with long contexts, some layers exhibit more pronounced lazy behavior, which we define as lazy layers. The left panel of 2
presents the attention patterns across different layers. We observe that some layers (e.g., layer 0) do not follow clear pattern in
4

---

Published in Transactions on Machine Learning Research (09/2025)
Outputs
Lazy ratio
Test-time identification
Update KV
1. 1.Update priority queue. Priority
Feedforward
Streaming attn weight
Insert in priority order.
2. Save KV cache for current layer
KTT
Lazy ratio
0.9 0.8
0.7
0.4
LX
Attention layer
Layer idx
K
2. When the queue is full,
1. Replace with reduced
KV cache
Pop the first (i.e., lazy layer).
Full attn weight
Lazy ratio 0.9 0.8 0.7
0.4
for lazy layer
Layer idx
10 0
Inputs
Figure 3: The framework of our LIGHITTRANSPER-TESTT A priority queue is maintained during the prefilling stage to store the lazy ratio and corresponding layer index afler processing each layer. Once the queue reaches its capacity, the layer with the highest lazy ratio identified as lazy layer, and its KV cache is reduced, frocing memory for storing the KV cache of the current layer.
attention weight distribution, while others (e.g., layer 20) show a clear lazy behavior pattern. Consequently, more memory-efficient attention mechanism be employed in these lazy layers by retaining only subset. KV cache of constant size.
Layer behavior remains consistent for a given input. To further explore whether a layer consistently functions as a lazy layer during generation for a fixed prompt, we visualize the attention weights for {Xinitial: Xxecent across all layers for all generated tokens in the right panel of Figure 2 using a randomly selected sample (additional examples are provided in Figure g). Notably, for a given input prompt, layers that exhibit lazy behavior maintain this pattern relatively consistently across tokens. This suggests certain degree of stability in attention dynamics throughout the generation process. In addition, the indexes of these consistent lazy layers vary according to different prompts. This necessitates the test-time algorithm in the following section.
5 Methodology: Light Transfer
In this section, we introduce LIGHTTRANSER TRANSFER, method for converting pretrained transformers into hybrid architectures for replacing full attention with streaming attention. The method has two settings: (1) For tasks like long-context understanding, LIGHTTRANSFER TEST allows for on-the-fly transformation at test time without requiring addit ional training.
efficient generation. LIGHITRA TRANSFER leverages our observation of lazy layers by
(2) For tasks demanding higher model capacity, such as ol-like long reasoning generalion, LIGHTTRaNCE TRANSFER- -AAIN- involves fine-tuning to adapt the model to the hybrid architecture.
5.1 LightTransfer-Tee Test
As shown in Figure g the first step in applying LIGHTTRANSFER-TTE is identifying lazy layers, defined those whose final Whast number of tokens in queries (i.e., X allocate the most attention to Xinitial U X To measure how the model allocates attention at layer i, we define lazy ratio Tz
Wlast 2 (Xinttaal A((2,))
Σ
Σ
(1)
where (., is the averaged attention weight over all heads from a query token è to a key token x at layer Intuitively, higher r; indicates that Xinst focuses more heavily on these particular key sets, thus exhibiting more lazy attention. To ensure that only P layers with the lowest lazy ratios maintain full attention during the prefilling stage and thus reduce peak memory usage, we adopt priority queue. We treat the lazy ratio
5

---

Published in Transactions on Machine Learning Research (09/2025)
Table 1: Torch style code for our lazy ratio calculation with flash attention.
def www.rrrttoclllaatiiiee
9. bs nun heads seq_len head__ _im k, be mun hendheads- seq_len head dim v, be nun headseads- seq_len dim w_ w_last, v_sink, v_recent): attt_out,, lse (q k, V, causal=frue, return_ (eturn__e==True)
q_last = qli, [._last:).pernuter 2, 1, 3) k_comb torch, cat O:w_sink],
,
-R_recent:]], dim=1). permuttee 2, 3,
le__y_yyyrtiioooo logeumexp
Log_lazy____ ratio torch. (q_ last, k_com)
return log_ _og_lay_rratio
Ti as the priority in max-based priority queue of size P. Whenever the queue exceeds capacity, the layer with the highest lazy ratio is popped, labeled lazy, and its standard attention is replaced with streaming attention. Here we do not replace the standard attention with streaming attention in a head-wise manner due the inefficieney, discussed in Appendix B.2 Specifically, for cach lazy layer i, we retain only the KV caches corresponding to (Ximital, Xrecent} and discard others. During decoding, memory usage is naturally reduced because the decoding process relies the already updated (and thus reduced) KV caches from the prefilling stage.
Identification burden. FlashAttention Dao 2023 is widely used to accelerate computations during the prefilling phase, but it does not explicitly expose attention weights. A direct application of our lazy layer identification strategy would thus require recomputing the attention matrix, incurring non-negligible overhead. To circumvent this issue, as shown in Table [ we leverage the log-sum-exp values (i.e., the denominator) of all attention weights produced by FlashAttention. Consequently, we only need to recompute the streaming attention score (a constant-size matrix multiplication); thus climinating the need for full rocomputation. Our identification algorithm mitigates additional latency introduced by full recomputation, resulting in only slight throughput reduction of 0.0058 to 0.0014 relative to baseline of 1 across sequence lengths from 4K to 32K. Notably, longer sequences result in smaller relative thronghput reduction. This occurs because the prefill operation grows with sequence length, whereas our identification process remains 0(1). As a result, when n is large, the identification overhead is overshadowed by the overall prefill cost.
5.2
Light Transfer- Train
For ol-like long reasoning tasks, where the input question typically consists of only few dozen words, the lazy ratio r; is not a reliable indicator of lazy. Because the sliding window is relatively large compared to the input, r; remains at across all layers. To address this, we adopt a pre-selection strategy. Specifically, for each sample in the training set, we feed both the question and the answer as input to the LLM, thereby providing sufficient context for each sample to reveal which layers are lazy. We then compute the frequency for each layer and select those with the highest lazy layer counts. However, frequency-based selection may not be fully optimal for each sample, while ol-like long reasoning tasks are inherently difficult, so additional fine-tuning allows the model to adapt to the new hybrid architecture and re-balance capacity across layers. Therefore, once these layers are identified, we perform supervised fine-tuning (SFT) under a hybrid architecture in which lazy layers employ streaming attention, while non-lazy layers retain standard attention. During inference, we simply rely on the preselected lazy layers, without requiring on-the-fly identification.
5.3 Theoretical Analysis
We first provide theorctical analysis of the approximation error of LIGHTTRANSPER-TESS and then discuss how this analysis implies the performance of LiGHITRANSSPRRTTE lazy layer identification procedures in LIGHTTRANSFER TEST are implicitly optimizing an upper bound the error of the whole network output induced by reducing the KV cache. We denote the set of layer indexes whose KV cache is reduced as I. For any layer i I, we denote the attention score of the discarded KV pairs as s; = Then we have the following upper bound of the error of the network output.
We would like to highlight that our
--Eze(Xmaan (XimstattX ) A.
6

---

Published in Transactions on Machine Learning Research (09/2025)
Theorem 5.1 (Informal). If the Frobenáus norms of all the parameters in L-layer with H-attention heads
trunsformer are upper bounded by B and the activalion function is Tip--iipchhll,, then we have that
Err. of LIGHTTRANSSE TRANSFER in logil.
S2 + Lip + +
Lip B²)
Σ
82.
iEZ
If we denote the error of hidden states at layer as et, then it evolves as e, Cemint+ (2--ee)))2H((Btaa)) B B
where C1 and C2 are quantities related to B, H and Lip--
The formal statement and the proof of Theorem 5.1 are provided in Appendix [F] We note that the error recursive expression consists three terms. The first term represents the error from the previous layer. The second term represents the error from the previous layer amplified by the current layer. Thanks to the layer normalization, this term will be truncated by The last term represents the newly introduced error if we shorten the KV cache at the current layer. By relaxing this recursive formula, derive the upper bound of the error between the logits of our method and the original transformer. This shows that the error is upper bounded by the sum of the attention scores of the removed KV pairs up to an additive constant. We highlight that our algorithm optimizes Equ. equation | which is exactly the upper bound of the error induced by LIGHTTRANSFER in logit up to a constant. We note that this theorem also provides the error analysis of the initial point of this fine-tuning process. The fine-tuning will further decrcase the error induced by LIGNSTRANSFER shown in Theorem 5.1 This theoretical bound is further confirmed in Appendix B.3
6 Experiments
In this section, empirically validate that TRANSTRANSEER can accelerate LLM generation while maintaining long-text capabilities including two scenarios 1) long context understanding, and 2) ol-like long reasoning generation, and uncover several insightful findings.
6.1 Experiments on Long-Context Understanding Tasks
In these experiments, we only apply SIGHTTRANSFRR-TTER As previously discussed, the imput length for these understanding tasks is sufficient to enable on-the-fly lazy-layer detection during the prefilling stage, making additional training unnecessary.
6.1.1 Experiments LongBench
TEST using four widely used LLMs, specifically LLaMA2-7B- chat (Touvrou et al. 2023 Mistral--7B-Instruaat (Jiang et al.] 2023), LLLMA3-8B-Instruce and LLaMA3-70B Instruct (Dubey et al. (024)) on LongBench (Bai et al.] 2023), which is a multi-task benchmark designed assess the long-context capabilities of LLMs. Detailed experimental configurations can be found in Appendix A An ablation study on these hyperparameters is provided in the Appendix C.1
Settings. We evaluate
Since no existing approach can convert transformer into a hybrid model at test time only, layer- level KV cache reduction methods serve as our closest baselines (Detailed discussions on how LIGHTTRANSFER- TEST relates to layer-level KV cache reduction methods are available in Appendix B.1]. Specifically, we compare LiGhtTRansPErR- TEST against the following baselines: 1) Standard: a standard transformer-bussed model in which each layer employs the original self-attention mechanism. 2) Strcaming LLM (Xiao al. 2023): A memory-efficient approach that modifies each attention layer in standard transformer to use only the KV cache for the first few tokens and the most recent tokens. 3) MiniCache (Liu et al. 2024a): An inter-layer KV cache reduction method that merges KV cache of every two adjacent layers after the model's midpoint using spherical interpolation while retaining important tokens to reduce cache storage. 4) SqueezeALtentionn (Wang et al. 2024c): An inter-layer KV cache reduction method that precisely distributes the KV-cache budget across layers.
Baselines.
Results. Table 2 summarizes the performance across various tasks in the LongBench Bai et al. 2023
benchmark. We have the following findings:
7

---

Published in Transactions on Machine Learning Research (09/2025)
Table 2: Performance comparison of LIGHTTRRANSFER-TEST and baseline methods on LLaMA-2-7IH
Mistral-7B. Intruct,
and LLaMA-3-70B-Intten using LongBench. Bold denotes the
best method, and underlined denotes the second best.
Single-Doc. QA Muti.-Doc. QA
Summary
Few-shot
Syn.
Code
NrtvQA
Qasper
MF-en
Hotpot
Musique
DuReader
GovReport
QMSum
MultiNews
TREC
TrivisQA
SAMSum
PCount
Average
PRe
LOG
RB-1
LLaMA 2-7tt
Standard
19.1
21.6
36.9
27.7
8.6
6.5
27.1
20.8
26.0
64.0
83.6
41.3
2.9
7.5
60.6
54.9
31.8
Streaming
13.1
15.2
26.9
23.1
5.5
4.4
21.1
19.9
24.2
61.0
82.8
38.9
2.1
4.0
59.0
52.2
28.3
MiniCache
13.1
13.7
30.3
15.6
4.7
9.8
21.5
20.9
24.3
63.0
83.
35.1
2.2
6.1
53.4
46.5
27.7
AqueezeAtt.
15.9
15.7
27.0
25.5
6.5
4.3
21.9
19.6
23.3
62.0
83.2
39.9
1.9
0.5
60.0
53.5
28.7
LITTRANS
15.8
18.3
30.1
27.3
7.0
4.7
22.7
20.2
25.1
62.0
82.8
39.6
2.1
1.2
59.4
53.6
29.5
Mistral- 7B- Instruct
Standard
29.7
40.5
53.4
50.0
29.1
32.9
34.9
25.4
27.7
76.0
89.1
47.3
98.5
60.4
62.1
47.6
Streaming
22.2
32.1
44.8
41.7
23.0
20.3
24.8
21.3
26.0
65.0
86.7
40.4
3.5
46.0
52.8
47.9
37.4
MiniCache
19.7
30.3
35.6
29.5
15.5
20.3
24.8
21.3
26.0
65.0
86.7
40.4
3.8
45.1
52.8
47.9
35.3
SqueezeAtt.
26.8
30.4
38.4 4
44.3
21.0
18.6
24.9
21.0
26.2
75.5
89.2
46.3
6.5
89.0
60.6
60.6
42.5
LITTRANS
20.0
53.6
50.5
27.5
32.3
25.4
27.3
76.0
89.3
47.3
59.9
61.3
Standard
23.4
32.8
39.6
44.7
22.2
20.1
28.8
23.3
27.0
73.5
90.6
41.9
3.6
72.0
58.1
51.3
40.8
Streaming
19.5
17.5
26.1
36.4
16.1
12.1
22.8
21.4
25.4
66.0
86.4
40.1
3.5
70.7
59.7
54.2
36.1
MiniCache
10.9
11.5
20.1
25.5
74.5
37.4
64.1
48.5
45.3
AquoezeAtt.
20.0
19.6
26.2
37.5
18.7
13.3
23.8
22.0
23.8
72.5
90.0
41.5
6.7
66.0
55.2
47.6
36.5
LITTRANS
23.2
18.3
35.7
43.7
20.9
14.5
24.1
22.3
26.0
71.0
91.1
41.4
6.9
67.0
60.2
53.4
38.7
LLaMA-3 mns-rrt
Standard
25.6
46.4
51.4
49.8
28.8
28.7
32.2 2
22.4
27.6
73.5
92.9
45.7
12.0
68.5
41.6
69.7
44.8
Streaming
25.4
36.2
34.4
44.3
22.7
15.0
20.2
26.2
91.1
43.6
11.5
41.9
67.1
40.0
MiniCache
25.1
45.2
38.4
4
46.2
24.9
17.8
29.1
22.3
27.1
71.0
86.7
41.3
10.1
67.0
35.6
54.4
40.1
SqueezeAtt.
26.3
36.8
34.0
48.1
25.0
17.5
28.0
21.5
25.5
71.5
92.8
44.8
11.5
67.0
41.5
68.5
41.3
LITTRANS
25.8
44.3
46.9
49.3
29.4
20.8
28.4
22.1
26.9
74.0
92.3
43.9
11.5
68.0
43.6
69.8
43.6
LLMs ezhibit redundancy across layers. As shown in the table, although MiniCache has some limitations, both SqueczeAttention and LIGHTTRANSPER- TEST enable the model to handle long-text tasks effectively, incurring only slight performance decrease (an average drop of 4.0% and 1.5%, respectively) when removing the KV cache in 50% of the layers. This finding suggests that LLMs exhibit redundancy in their layer-level KV caches.
The tronsferred hybrid architectures can preserve strong long-conteat understanding capability. LIGHTTRANSPER--
applies streaming attention some layers of a transformer-hhan model while retaining standard self-attention in others, striking an effective balance between computational efficiency and representational capacity. In contrast, MiniCache adopts cross layer attention (CLA) [Brandon et al. 2024 (sharing one KV cache across adjacent layers), and SqueezeAttention allocates distinct KV-cache quotas per layer. Under a higher compression ratio than MiniCache and the same ratio as SqucezeAttention. TRGNSPRR TEST surpasses them by 6.1% and 2.6%, respectively, demonstrating the effectiveness of transitioning transformers into hybrid models for memory-effcient inference. This superiority partially originates from the fact that our algorithm explicitly optimizing the error upper bound in Thcorem In contrast, the optimization methods of MiniCache and SqueezeAttention do control the error induced KV reduction in theoretically plausible manner.

---

Published in Transactions on Machine Learning Research (09/2025)
100%
100%
99%
99%
98%
98%
97%
97%
96%
96%
95%
95%
94%
94%
4K
8K
16K
32K
4K
8K
16K
32K
Standard
LLighTTransr (a) Single-Key NIAH
Standard
LightTransfer
(b) Multiple-Key NIAH
Figure 4: Performance comparison of LIGHTTRANSFER and standard model on NIAH tasks using Mistral-7B. Instruct.
Table 3: Performance comparison of LIGHTTRRANFER TRRNNSFRRRRANN and baseline methods on three mathematical benchmarks using QwQ-32B. Bold denotes the best method, and underlined denotes the second best.
Method
MATH-OAI
AIME24
GSM8K
QwQ-STILL
90.2
46.7
95.6
LongGen
78.2
16.7
95.4
LITTRANS
90.7
53.3
95.5
6.1.2 Experiments NIAH
Settings. We also evaluate whether LIGHTTRANSFER-TE while replacing some standard attention layers into memory-efiiciet streaming attention. The evaluation is conducted on single-key and multiple-key NIAH tasks collected in the Ruler (Hsich et al. 2024) benchmark. We report the performance with input context lengths of 4K, 8K, 16K, and 32K. Detailed experimental configurations can be found in Appendix
preserve in-context retrieval capabilities
Results. Figure summarizes the performance on NIAH tasks, with the context length ranging from 4K 32K. While our MIGHTTRANSPER-TTO replacing select transformer layers with streaming attention reduces memory overhead, strategically retaining original attention mechanisms in deeper layers ensures robust long-range dependency modeling. This explains the maintained performance on single-key tasks (32K: 96.7% vs standard 96.6%) and competitive multi-key results at 32K (78.2% vs 78.9%). The retained standard layers serve as an anchor for cross-token reasoning, which is crucial for in-context retrieval.
6.2 Experiments on ol-like Long Reasoning Tasks
In these experiments, we investigate the effectiveness of TAGHTTRRANSERRTTRAT ol-like long reasoning generation tasks. While these tasks feature relatively short inputs, they demand intricate reasoning, Conse- quently, we SFT the model with approximately 5K training examples to facilitate swift adaptation within the transferred hybrid architecture.
Settings. Experiments conducted on three widely used mathematical benchmarks AIME24, MATH-OAI, and GSMSK. We use greedy decoding evaluate the performance of our model with maximum tokens set to 32K. We aim to experiment with models capable of generating CoT, such QwQ. However, the original training data for QwQ is not publicly available, we adopt the public training set provided by QWQ-STILL (Min et al.] 2024), model demonsurated to achieve comparable performance to QwQ. Specifically, following QwQ-STILL, we apply a simple distillation approach using Qwen2.5-32B. 3-B--sruutttt as our base model. We generally follow the original training of QwQ-STILL, and replace 50% of layers with streaming attention. To mitigate the training complexity of attention, we optimize LIGHTTRRANSFRR. TRAIN training using Flex Attention Dong et al. 2024

---

Published in Transactions on Machine Learning Research (09/2025)
Table 4: Performance comparison of LIGHTTRANSFER and baseline methods that operate without tuning
LLM parameters on mathematical benchmarks using QwQ-32B. Bold denotes the best method.
Method
MathOAI
AIME24
QwQ-STILL
90.2
46.7
MiniCache- Test
20.0
0.0
SqueczeAAteentinn Test
82.6
15.0
DuoAltn.-Train
79.2
33.3
LIGHTTRANSFER- TEEST (OURS)
85.0
40.0
LIGHTTRANSFER TRAIN (OURS)
90.7
53.3
0.95
0.9
0.85
0.8
0.75
0.7
Layers that LongGen assumes unnecessary standard full attention.
0.65
0.6
12 15 18
21
24 27 30 33 36 42 45 48 51
54
57
60
63
-Lazy ratio
Figure Lazy ratio scores across layers in QwQ-32B-STH
Baselines. We compare our LIGHTTTRANSFER-THHH against the following baselines: 1) QwQ-STILL (Min
et al. 2024): distilled model on 1eem22.5322BInnsttrr that achieves performance comparable to QwQ-32B
Preview, whose training data is publicly available. 2) LongGen (Ge et al., 2024h: an approach that assumes the layers at both ends of the model do handle global information and predefines the replacement of those layers with sparse allention.
We further benchmark MIGHTTRANSFER- and (ii) DuoAttention (iiao et al. 2024). Unlike our approach, DuoAttention must first run on separate calibration set to pinpoint. those allention heads whose KV cache should be evicted, whereas operates entirely calibration-free
against (i) existing training-free sparsification baselines,
retains its performance on Math-OAI (+0.5%), AIME24 (+6.6%) and GSMSK (-0.1%). In contrast, LongGen, which assumes its middle layers require standard attention, exhibits drop GSMSK but suffers 30.0% and 12.4% decrease AIME24 and Math-OAI, respectively. While the unchanged GSM8K results for LongGen may indicate that GSM8K poses lower complexity for these models, the broader comparisons nevertheless highlight the strength of our data-driven layer selection. Specifically, our MIGHTTRANNSER-TTE (Figure 5) and replaces those exhibiting the highest, which is proven to be more robust than hand-crafted assumptions. Moreover, our findings underscore the existence layer-level KV cache redundancy even ol-like long reasoning models, emphasizing the promise of hybrid transformer architectures. Meanwhile, shown in Table 4
Results. Table shows that LIGHTTRANSFER-THHE
calculates each layer's lazy ratio
surpassos all other training-free free basclines both MathOAI and AIME24, highlighting its suilability for long reasoning generation tasks.
6.3 Ablation Studies Analysis
Standard layer retention ratio vs. model performance. As shown Figure t we systematically vary the fraction of layers that use original attention from 0.25 to 0.5, up 0.75, for both SaaMAA-8B-Inntrrae and ATaMMA370B-Instrriie LongBench benchmark. As expected, higher retention ratios consistently
10

---

Published in Transactions on Machine Learning Research (09/2025)
(a) LLaMA 3-8B
SqueezeAttn
LightTranster
34%
35%
36%
37%
38%
39%
40%
41%
Performance
(b) LLaMA 3-70B
SqueezeAttn
LightTransfer
38%
39%
40%
41 %
42%
43%
44%
45%
Performance
ratio of standard attention layers:
25%
50%
755%
Figure 6: Effect of retaining standard attention in more layers on LongBench.
Table 5: Relative token- generation throughput at different sequence lengths (4K, 8K, 16K, and 32K) compared to the Full baseline. Bold denotes the best method.
Method
4K
8K
16K
32K
SqueezeAtten.
1.03x
1.09x
1.12x
1.04x
MiniCache
1.26x
1.29x
1.52x
1.41x
LIGHTTRAL TRANSFER
1.44x
1.78x
2.17x
1.75x
yield improved model performance. However, this comes at the cost of increased memory consumption, highlighting the tradc-off between efficiency and accuracy. Notably, across all compression settings examined, surpasses the st rongest, baseline on that benchmark (i.e., underscoring the benefit transitioning standard transformers to hybrid models via strategical designs for more efficient generation.
thereby
Throughput of token generation. To evaluate how these memory optimizations impact loken-generation throughput, we conduct experiments with Mistral-7B the Ruler benchmark under maximum batch-size configurations. Input sequence lengths of 4K, 8K, 16K, and 32K were tested while retaining 50% of the standard attention layers. As shown in Table 5 LICHTTRANNFER-TTT TEST consistently achieves the highest throughput compared with other training-free test time inter-layer KV cache reduction methods. In contrast, SqueezeAttention, despite having the same compression ratio, fails to reduce peak memory usage during the prefilling phase, since it must complete prefilling for all layers before applying compression. This constraint limits the feasible batch size, restricting potential throughput. Meanwhile, MiniCache exhibits lower throughput due to its smaller compression ratio (i.c., removing KV caches in at most 25% of layers). These findings underscore the effectiveness of LIGIIT TRANSFER in balancing memory usage and computational efficiency.
Effect of different layer replacement stratogics. As shown in Figure 7 (a-d) we experiment with four different layer replacement strategies for integrating memory-effciint streaming attention into transformers, with consistent replacement counts (except Standard). The results shown Figure 7 indicate noticeable reductions for Pyramid and Random strategies, suggesting that the predefined expectations about each layer's 's function may not fully align with their actual roles. Moreover, the performance of our LIGHTTRANSFER is effective in reducing memory usage while
surpasses other rategie,,, suggesting that maintaining performance. Additional results with two extra basclines are provided in Appendix C.5
11

---

Published in Transactions on Machine Learning Research (09/2025)
Layer
lozy
Layer 16
lazy
Layer
(a) Full
(b) Pyramid
(c) Random
(d) (annffssferr)urs))
24%
50%
24%
MuSiQue
HotpotQA
NarrativeQA
23%
45%
23%
22%
40%
22%
21%
35%
21%
20%
30%
20%
Standard
+ LightTransfer
0 Pyramid
O Random (0,32)
0 Random Random 6,32)
Figure 7: Different layer replacement strategies and their performance on -Lastuct--Instruct:: 1) Standard: Use standard attention in all layers. 2) Our TRGHPTRANSPE:: Dynamically identify lazy layers the fly, and replace their attention mechanism accordingly. 3) Pyramid: Replace each layer with memory-efficient attention; the budget decreases with depth, forming pyramid-like structure. 4) Random: Randomly replace layers with memory-effciiett attention within the ranges [0, 16), [16, (2), or [0, 32). We keep same number of replaced layers, except Standard.
Instruct:
7
Conclusion
We present LIGHTTRANSFER, a lightweight framework for transforming standard transformers into hybrid models for more efficient generation by identifying lazy layers and replacing their full-attention modules with strcaming attention. Extensive experiments show that even when half of the transformer layers are replacod with streaming attention, LIGHTTRANSPER delivers up to 2.17x increase throughput while incurring less than 1.5% performance drop on LongBench. For advanced long reasoning generation tasks like ATME24, our method achieves these gains without any performance degradation on QwQ-STILL.
References
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, ct al. Longbench: bilingual, multitask benchmark for long context understanding. 2023.
arxiv preprint arXiv:2308.1455
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. arxiv preprint arXiv:2501.0066; 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan, Longformer: The long-document transformer, arXiv preprint arXiv:2004.05150 055550 2020.
Aviv Bick, Kevin Y Ti, Fric P Xing, J Zico Kolter, and Albert Gu. Transformers ssms: Distilling quadratic knowledge to subquadratic models. arXiv preprint arxiv:2408. 10189, 2024.
Aleksandar Botev, Soham De, Samuel L Smith, Anushan Fernando, George-Cristian Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sossa, Robert Dadashi, et al. Recurrentgemma: Moving past transformers for efficient open language models. arXiv preprint arXiv:2/04. (xx::22040078333 2024.
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint 0rXiv:2405.1298 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichhro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint or Xin:21 10. 14168, 2021.
12

---

Published in Transactions on Machine Learning Research (09/2025)
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint xrXiv:2307.06691, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:405.21106 2024.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru Albert Gu, Ruba Haroun, Lconard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local atlention for efficient language models. arXiv preprinl Xivv:2202. 19427, 2024.
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:24 12.054.0044 12.05496, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint 2if::2407.2178.. 2024.
Benjamin L Edehman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pp.
PMLR, 2022.
Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. Not all heads matter: head- level kv cache compression method with integrated retrieval and reasoning. arXiv preprint Xrxiv:24 19258, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2810.0180; 2023.
Suyu Ge, Xihui Lin, Yunan Zhang, Jiawei Han, and Hao Peng. little goes a long way: Efficient long context. training and inference with partial contexts. arxiv preprint 2rXiv:2410.01445 2024.
Team Gemma, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at a practical size. arXiv preprint arXiv:2408. 000.0000 2024.
Daniel Goldstein, Fares Obeid, Erie Alcaide, Guangyu Song, and Eugene Cheah. Goldfinch: High perfor- mance twkv trarffrrrooor hybrid with linear pre-fill and extreme kv-cache compression. arXiv preprint arXiv:2407.12077 12077, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXio preprint 2.Ni:23312.00722 2023.
Xiangming Gu, Tianyu Pang, Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, and Min Lin. When attention sink emerges in language models: An empirical view. arXiv preprint arXiv: 10.:2440.11000 10781, 2024.
Cheng- Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arxiv preprint arXiv:244400654,, 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Monsch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arxiv:2910. 06825, 2023.
Gregory Kamradt. Needle in a haystack pressure testing 1lms. http:::/ggithub.commkaaaaea 2023. GitHub repository.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilaarco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into mns. arXiv preprint arx.w.100..100% 2021.
13

---

Published in Transactions on Machine Learning Research (09/2025)
Aonian Li, Bangwei Gong, Bo Yang. Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprinl 2wX0:22501.0813,, 2025.
Haoyang Li, Yiming Li, Anxin Tian, Tianhao Tang, Zhanchao Xu, Xuejia Chen, Nicole Hu, Wei Dong, Qing Li, and Lei Chen. survey on large language model acceleration based on kv cache management. arXiv preprint (rXtt:2412.1944 19442, 2024a.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. arXie preprint WrXiu:2404.14435 2024b.
Zonggian Li, Yinhong Liu, Yixuan Su, and Nigel Collier. Prompt compression for large language models: A survey. arxiv preprint orXiv:2 10. 2388, 2024c.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Mcirom, Yonatan Belinkov, Shai Shalev-Shwaid Shwartz, et al. Jamba: A hybrid canssfrmer-mamleee language model. Xix preprinl or Xiv:24403. 19887, 2024.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Πya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations.
Akide Liu, Jing Liu, Zizheng Pan, Yefei Hc, Gholamreza Haffari, and Bohan Zhuang. Minicache: Ky cache compression depth dimension for large language models. NeurIPS, 2024a.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024b.
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. arXiv preprint arXiv:2405. 06640, 2024.
Yingqian Min, Zhipeng Chen, Jinhao Jiang, Jie Chen, Jia Deng, Yiwen Hu, Yiru Tang, Jiapeng Wang, Xiaoxne Cheng, Huatong Song, et al. Imitate, explore, and self-improve: A reproduction report on slow-thinking reasoning systems. arXiv preprint arXiv:2412.09) 09418, 2024.
Piotr Nawrot, Adrian Laácucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dynamic memory compression: Retrofitting lims for accelerated inference. arXiv preprint ar Xrx::2240330000 09636, 2024.
Team OpenAI. Introducing openai ol. https://openai.comm
2024.
Bo Peng, Eric Alcaide, Quentin Authony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing rmns for the transformer era. arXie preprint 230: 13048, 2023.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv preprint arx:::440.004658 2024.
Team Qwen. Qwen1.5-moe. Matching 7b model performance with 1/3 activated parameters', February 2024a. URL
Team Qwen. Qwq: Reflect deeply on the boundaries the unknown. -uq-3bb-revveew 2024b. Accessed: November 28, 2024.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. primer in bertology: What we know about how bert works. Transactions of the association for computational linguistics, 8:842 866, 2021.
14

---

Published in Transactions on Machine Learning Research (09/2025)
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint 2019.
Alex Sherstinsky. Fundamentals of recurrent neural network (rnn) and long short-term memory (Istm) network. Physica D: Nonlinear Phenomena, 404:132306, 2020.
Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, and Hai Zhao. Keep the cost down: A review methods to optimize Ilm's s kv-cache consumption. arxiv preprint or arXiv:2407. 18003, 2024.
Ken Shoemake. Animating rotation with quaternion curves. In Proceedings of the 12th annual conference on Computer graphics and interactive techniques, pp. 245-254, 1985.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. arXiv preprint WrXiv2245.0525. 2024.
Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, and Gongyi Wang. Razorattention: Efficient kv cache compression through retrieval heads. arXiv preprint arxiv:2407. 15891, 2024.
Jamba Team, Barak Lenz, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, et al. Jamba-1.5: Hybrid transformer-mamly models at scale. ar preprint Xio:24408 12570, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXie preprint or arXin:2307.09955 2023.
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. arXiv preprint 2rXiv:248..12237, 2024a.
Zheng Wang. Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. Model tells you where to merge: Adaptive kv cache merging for 11ms on long-context tasks. arXiv preprint arXiv:2407.08455 08454, 2024b.
Zihao Wang, Bin Cui, and Shaoduo Gan. Squeezeattention: 2d management of kv-cache in Ilm inference via layer-wise optimal budget. arXiv preprint arXiv:2404.0477 0479,, 2024c.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiu preprint arXiv:2009. 17458, 2023.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context Ilm inference with retricval and streaming heads. arXiv preprint arXiv:2410.10814 2024.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid kv cache compression for higu-throughput Ilm inference. arXiv preprint orXiv:2405.12555 12532, 2024.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-cfficient training. arXiv preprint srXio:2312.06633 2023.
Hao Yu, Zelan Yang, Shen Li, Yong Li, and Jianxin Wu. Effectively compress kv heads for 1lm. arXie preprint 0rXiv:2406.07056 2024.
Jiayi Yuan, Hongvi Liu, Shaochen Zhong, Yu-Neng Chuang, Songchen Li, Guanchu Wang, Duy Le, Hongye Jin, Vipin Chaudhary, Zhaozhuo Xu, et al. Kv cache compression, but what must we give in return? a
comprehensive benchmark of long context capable approaches. arXiv preprint 0ii::2007.01227. 2024.
Fengzhuo Zhang, Boyi Liu, Kaixin Wang, Vincent Tan, Zhuoran Yang, and Zhaoran Wang. Relational reasoning via set transformers: Provable efficiency and applications to marl. Advonces in Neural Information Processing Systems, 1335825-35833 2022.
15

---

Published in Transactions on Machine Learning Research (09/2025)
Michael Zhang, Simran Arora, Rahul Chalamala, Alan Wu, Benjamin Spector, Aaryan Singhal, Krithik Ramesh, and Christopher Ré. Lolcats: On low-rank linearizing of large language models. ar preprinl. arXiv:2410.10220 2024a.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the poreupine: Expressive linear attentions with softmax mimicry. arXie preprint orXiv:2402.0/4 04347, 2024b.
Yang Zhang, Yanfei Dong, and Kenji Kawaguchi. Investigating layer importance in large language models. Xixiv preprint arxiv:24 09. 14881, 2024c.
Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjic Hu, Wen Xiao, al. Pyramidkv: Dynamic kv cache compression based pyramidal information funneling. arXiv preprint ar Xiv:2400 02069, 2024d.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint orXin:2305. 19420, 2023.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024e.
A Settings.
We adopt a generative format where answers are produced using greedy decoding for all tasks. All the experiments are conducted using NVIDIA A100 using PyTorch and HuggingFace Transformers replaced with bass_attentionn_e with kvcache to accelerate. All model weights, activations, and KV caches use kicccchc
BF16 precision, with no quantization applied. We set the sink token num Wsink = and the window size Wreeent 1020.
A.1 Settings on LongBench
The input context window sizes of L.AMA2-7B-chal Mistral-7B--Isttruut, I.L.MA--8B-Isstruu and LLaMA3- 70B-Instruct are 4K, 8K, and 32K, with average tokenized sequence lengths approximately 13K, 12K, 10K, and 10K in LongBench. For evaluation, constraints, only include the performance of 16 randomly selected lasks out of the 21 LongBench tasks. For MiniCache, as the code was not open-sourced before our submission, we reimplemented it based on the original paper and the SLERP (Shoemake) 1985) code it references. We followed all the hyper-parameters outlined in the paper, except for the number of retention tokens. SqueezeAttention and our LIGHTTRANSFER TEST TIME are both set the same compression ratio, equivalent to removing KV caches from 50% of the layers (i.e., P is set to 50% of the total number of layers), whereas MiniCache is set 25% (i.e., its maximum possible compression).
the metrics recommended by LongBench. Duc to space
A.2 Settings on NIAH
The evaluation is conducted using the metrics recommended by Ruler. Because synthetic in-context retrieval tasks in the Ruler benchmark require more extensive global context, use slightly lower removal ratio here than the one applied to LongBench. In our LIGHSTRANSPER-TES TIME setup, we remove the KV caches from 25% of the layers.
Discussions
B.1
Relationships with Test-time KV Cache Reduction.
Some techniques (iao al. 2023
Li al.
2024b Wang et al. 2024b Zhang al., 2024e Liu al. 2024d identify redundant. tokens within each attention layer and
2024b Yang et al. 2024 Zhang al.
16

---

Published in Transactions on Machine Learning Research (09/2025)
evict their associated KV cache at test time, thereby effectively lowering memory usage. Within this line of research, the approach most closely aligned with our MIGHTTRRASPPRR-TEE. TIME specifically targets layer-level KV cache redundancies during inference, aiming to further optimize memory consumption by examining how different layers store and reuse keys and values. However, current mothods only consider the relationships of KV caches across layers from a relatively coarse perspective for reducing KV caches across layers. For example, MiniCache Liu et al. (2024a) focuses on the similarity of KV caches between layers, while AtueezeAttention Wang et al. (2024c) optimizes cache usage without a detailed investigation into the internal mechanisms of transformers. In contrast, our LIGHTTRANSPER approach goes further by examining how each layer functions and selectively replacing certain layers with more memory-effiiient architectures.
B.2 Why We Do Not Adopt a Head-Wise Hybrid Model
Prior studies typically do not consider a dise hybrid design Lieber et al. (2024); Gemina et al. (2024); Sun al. (2024): Botev et al. (2024); De et al.] (2024). One practical reason is that LLMs often employ tensor parallelism (TP) to distribute computation across multiple GPUs. In this setup, a single layer generally contains multiple attention heads (e.g., eight heads per layer), and each head is handled by separate GPU. If different heads in the same layer maintain different KV cache sizes, GPUs with smaller caches must wait for those with larger caches to finish. This synchronization bottleneck cancels out any latency benefits gained from compressing only cortain heads, making a head-wisc hybrid approach incfficient in real-world deployments.
In addition, when using TP only, we observe that the head-wise hybrid model struggles to improve throughput. To make fair comparison, we also implement both data-parallel attention + TP feod-forward network (DP+TP) method based on the state-of-the-arl inference framework SGLang, and evaluate this strategy with pure TP on LLAMA3-70B using an A100 40G node.
Despite idealized settings (identical input lengths and perfect batch distribution with BSZ%GPU 0), DP+TP significantly reduces throughput-onn 0.0735 x that of TP-amd and reduces the maximum supported sequence length only 1/128x. The root cause is that TP shards attention layer parameters across GPUs, while DP+TP replicates them, consuming an additional 157.5 GB of GPU memory. Moreover, TP shards the KV cache per reguest, allowing it to support 8x longer context lengths.
Table 6: Throughput comparison (tokens/s) for L.LMA3-70B under different sequence lengths. DP+TP suffers from severe memory issues.
Length
512k
32k
16k
8k
4k
DP+TP (tokens/s)
00M
OOM
OOM
OOM
145.6
TP (tokens/s)
22.6
360.0
772.1
1188.6
1981.9
Another critical issue with DP+TP is its difficulty in achieving load balance, concern also raised in the DeepSeckV3R1 Inference Report. In extreme imbalance cases, throughput under DP+TP decreases lincarly with GPU count. While increasing batch size may alleviate this in short-contexl scenarios, it does not hold for long-contens settings, where KV cache memory constraints restrict batch size (e.g., only 16 sequences at 32k tokens, even with TP).
Furthermore, in production deployments with prefill-dcodde seporation (PD), DP+TP schedules workload balancing based on prefill lengths. We argue this heuristic fails in many real-world settings:
Response lengths in later turns are highly unpredictable, violating assumptions of balanced prefill length.
Multi-turn conversations:
Chain-of- Thought models may generate over 16K tokens even from short prompts, again breaking static prefill balancing.
Long-form generation:
17

---

Published in Transactions on Machine Learning Research (09/2025)
Finally, head-wise hybrid methods are hard to integrate with existing inference systems like vLLM and SGLang due to KV cache granularity constraints. In contrast, industry-leading pretrained hybrid models (e.g., Gemma-2/3) adopt layer-wise hybrid structure, validating our design decision.
Table 7: Performance under different hyperparameters.
(a) Window size Wlecent
(b) Sink token count Wsink
Window size
508
1020
2044
Sink num
4
Performance
39.5
39.8
39.8
40.1
Performance
26.5
39.8
39.8
39.9
(c) Wlast
Whast
8
16
32
64
Performance
39.9
39.8
39.9
39.7
B.3 Additional Evidence for the Usefulness of the Bound
We measure the divergence between the output distributions of the standard and LIGHTTRANSFER methods after prefilling with identical real-world samples. We employ the Kullback-Leiin (KL) divergence as a principled metric to quantify distributional discrepancy.
Table 8: KL divergence on seven LongBench datasets using LLaMA3-BB-Instmie Lower values indicate closer aligument.
Dataset
No layer reduction 50% layer reduction
2WikiMultihoQQA
4 x 10-6
-9 x 10-6
Hotpot
-6 x 10-6
-1 x 10-5
MultiFicldQA-EN
-4 X 10-6
-8 x 10-6
MulttiielddQA-ZH
-1.1 x 10-5
-1.6 x 10-5
MuSiQue
-6 x 10-6
-1 10-5
NarrativeQA
-1.1 x 10-5
-1.5 x 10-5
Qasper
-6 x 10-6
-1x 10-5
Average
-7.1 x 10
-1.2 x 10
The extremcly small KL values, consistently approaching zero across diverse datascts, confirm the practical tightness of our theoretical bound. Furthermore, the KI. divergence values increase with the reduction of layers, consistent with our theoretical predictions. Such negligible divergence indicates that LIGHTTRANSFER faithfully preserves the original token-level predictive distribution during prefilling.
B.4 Limitations
Firstly, acknowledge that the performance of the LightTrrasse method in the test setting does not match its training performance. Improving the effectiveness of this Tightrrrasfor-Tess remains an important direction for fulure research. Secondly, we recognize that our training method constitutes fine-tuning, and thus does not fully exploit the potential the hybrid model architecture. Due to constraints in available GPU resources and data, we were unable to comprehensively explore a fully native hybrid model. We believe that future work with enhanced computational resources could unlock additional capabilities inherent in this promising hybrid model structure.
18

---

Published in Transactions on Machine Learning Research (09/2025)
C Additional Experiment Results.
C.1 Impact of Hyperparameters
We adopt these hyperparaneterrs either directly from StreamingL..M Xiao et al. (2023 (i.e., Wsink and Wrecent ) ensuring consistency with established practices in the field, or through preliminary experiments (i.e., Wlast). We also conducted additional experiments to analyze the impact of hyperparameters (Wsink: Wrecents and Wlatt) model performance. As shown in Table 7 the variation in performance remains within one percenlage point across different configurations, demonstrating the robustness of our approach to hyperparameter choices.
C.2 Combination with Intra-layer KV Cache Reduction Methods
To illustrate the orthogonality between our MIGHTTRANSPER-TEES and intra-layer KV cache compression methods, we conduct additional experi- ments that combine LIGHTTRANSPER-TRSS with SnapKV (a cutting- edge method for intra-layer KV cache reduction). In these experi- ments, SnapKV is applied to compress the KV cache in non-lazy lay-
378
1.8
87.4
ers, while MRAIS'RRANSEERR- TEEST remains active for lazy layers. We use Qwen2.5-3B-chat -32kat for this analysis. As shown in Figure g leveraging SAGHYTRRANSPRR-TES. alongside an intra-layer KV cache compression method can further reduce KV cache size while preserving model perfor- mance, underscoring LIGHTT RRANSFER RAASPER-TEST's orthogonality to cxisting methods focused on intra-layer redundancies.
36.6
08
86.6
0.
86
Staplty Paformanne
Stapity
Compress
Figure Comparison of SnapKV and LLiht Transfer.
C.3 Comparison with Head-wise KV Cache Reduction Methods
We use 50% sparsity across all methods and evaluate on MaAMMa-----nstrnc---nnii-------- to align with DuoAutention's released checkpoint.
Under the same sparsity level, our method on par with DuoAttention (Xiao et al. 2024 which requires training and is head-wise. FastGen (Ge ct al. 2023) will be OOM on most datasets because it is incompatible with flash attn (Xiao et al. 2024). RazorAttn (Tang et al. 2024) and HeadKV (Fu el. al. 2024) are highly sensitive to the calibration set. Although we follow their reported calibration setup, the results are hard to reproduce. Therefore, we use the full LongBench datasct (input only) for their calibration. Unlike (Xiao) et al. 2024 Ge et al. 2023 Tang et al.) 2024 Fu et al.) (204)), our search strategy only needs one additional matmul operation, which fast. Thus, Trahs.Transfer does not need any calibration set and can perform on-the-fly search, thereby avoiding sensitivity issues.
C.4 Performance MoE architectures
We conducted additional experiments using the Qpeen1.5.M6E-44.BBB2.2. (Qwen 2024a) model on tasks within LongBench. As shown in Table 10 when replacing 50% of its layers with streaming attention, LightTransfer exhibits performance drop of less than 1% the MoE architecture, outperforming other layer-wise KV cache pruning baselines. This observation is consistent with our findings on non-MoF, transformer models, thereby further confirming the robustness and effectiveness of our approach.
C.5 Comparasion with other layer layer replacement strategies
To further study the effectiveness of our approach, we benchmark it against the following two baselines
(i) Shapley value-based Zhang et al. 2024c) which replaces layers according to their estimated Shapley contribution;
(ii BERTology Rogers et al. 2021 which iteratively substitutes the least influential layers.
19

---

Published in Transactions on Machine Learning Research (09/2025)
Table 9: Comparison with -ead-wise KV Cache Reduction Methods on
Dataset
Razor
HeadKV
Ours
DuoAttn
Training Free
Yes
Yes
Yes
No
qasper
19.69
29.97
26.85
27.02
multiqa_ en
27.62
30.93
48.26
53.69
hotpotqa
23.98
26.96
37.07
35.52
2wiki
24.83
25.70
30.41
28.08
multinews
25.84
27.31
26.50
27.76
trec
55.50
58.50
66.00
69.00
triviaqa
63.74
78.54
87.68
87.32
samsum
40.10
39.70
41.04
41.13
pcount
2.26
0.39
2.00
2.00
Icc
32.19
34.54
41.20
39.24
repo-p
32.15
36.06
40.02
40.04
Average
31.63
35.33
40.64
40.98
Table 10: Performance on Qxenn
Dataset
Standard
MiniCache
SqueezeAttn
LitTrans
qasper
30.19
18.66
23.48
26.92
multifieldqa_on
38.47
24.14
29.58
37.23
hotpotqa
10.17
5.57
8.56
9.12
2wikimqa
11.92
7.64
11.35
12.78
multi news
24.77
20.09
20.94
24.63
troe
68.00
66.00
64.50
65.50
triviaqa
86.35
70.44
84.80
85.35
samsum
38.09
24.87
38.22
37.74
passage_count
1.67
2.04
3.30
3.10
lec
48.33
33.38
44.90
48.68
repobench-p
40.89
22.09
36.62
38.33
Average
36.26
26.81
33.30
35.43
20

---

Published in Transactions on Machine Learning Research (09/2025)
00
(a) Example 0
(b) Example
Figure 9: Additional examples of layer behavior across tokens.
Table summariscs the performance on three tasks from LongBench. Our Light strategy attains the best performance all datasets, surpassing the next-best baseline by up to +3.7% (HotpotQA) while maintaining the same inference budget.
Table 11: Comparison of different layer-replacement strategies on three datasets on LLaMa3-ShP 8B-Insruut.
Method
HotpotQA
MuSiQue
NarrativeQA
Shapley value based
39.97
18.18
19.63
BERTology
42.48
18.13
20.91
Light Transfer
43.70
20.90
23.20
D More Examples
D.1 Examples about Layer Behavior across Tokens
Additional examples of layer behavior across tokens for a given input can be found in Figure 9 The examples are randomly chosen from LongBench benchmarks. The analysis is conducted using I.aMAB-SB-Instruo..
E
Notation
For positive integer N N, define the set [N] (1. N}. For vector Hd adopt ||p to denote the Ep norm of vectors. For matrix X = [x] Fdd IT E Rdixds, where 2; E Rdd for = 1,... d1, we define the -oorm of X as || [1[[[=a pp.... ... II pplll iee, we first apply &p norm in a row-wise manner and then apply fg norm. The Frobenius norm || ||2. is also denoted as ||. For a matrix X E Roxb i-th row and i-th column are denoted as [X],, and [X]. respectively. The element at i-th row and j-th column X is denoted [X].]-
F Theoretical Analysis
In this section, provide the theoretical analysis of the proposed method. We first define the transformer structure we analyze in this paper. In fact, we analyze the LLaMA-type structure (Dubey et al. 2024), i.e., the transformers that adopt the pre-norm and the res-link. The imput of the transformer is the embedding of the tokens X E Nxd, where N is the number of tokens, and d is the dimension of the token embedding. We consider a L-layer transformer, i.e,, there are transformer blocks in the network. Each transformer block consists of Multi-Heal Attention (MHA) and Feew-rd (FF) module. The MHA module is a combination of multiple causal solf-attention modules. Each causal self-attention module is defined as
attn(X, X, Wq WK. Wv) softmax(X WT X M)XWw,
where X Nxd is the input. Wº, Wx dx and Wy IRdd the weights of the self-attention module, and M NxN, is the causal mask. The causal mask is defined as
ifi
[M]s -00
otherwise.
21

---

Published in Transactions on Machine Learning Research (09/2025)
The MHA with H heads is defined as
mha (x, (W WKA WaAL
Σ
M))
where X Nxd the input, Wo,,, WKA Rdxds and Wyn E Rdxdd are the weights the h-th head of MHA. Here we just merge the parameter Wo into Wy for case of notation. Our analysis can be directly applied to the parameterization that explicitly includes Wo as weight. The FF module applies transformations to X in row-wise manner, which can be defined as
ffn(X, WA.1, WA,2) = (XWA.))
where WA,1. WA,2 E Rdxd weights of FF module, and of-) is element-wise activation function. For example, can be ReLU function. We require that is Lipschitze function.
Assumption F.1. The activation function o-)) is Lipp -ipschite, i.e., lo(x) ovv) s Lip |r yl for any 2, E R.
We note that this assumption is satisficd by all the popular activation functions, including ReLU, sigmoid, ELU, and GELU. The imput of the transformer is denoted as the oulput of the 0-th layer, i.e., x(O) X. Then the i-th block processes in the input (i-1)) as
y()) x(=1) + mha (LN(X(I-"), (W w w ( Q.a K, WYNK=1
(2)
x = y(v) ffn(L)(Y W A.1. w(())
(3)
where the superscript (i) denotes the parameters and hidden states at layer i, and LN is the row-wise normalization of the input. To simplify the mathernttccl calculation, we defined LN as
x if ||]|| $ LN(x) = {{ x/||=|12 otherwise
Our analysis can be directly applied to the LayerNorm function of PyTorch. For case of notation, we will abbreviate mha(;, (W()), W K,k: wr(0) ) and ffn(., ((() w( W. The output logits of the transformer is
mha ( and fin (.) in the following.
A.1 A.2)
x = x W unemb:
where Wunemb E Rdddhhocah is the unembedding matrix. We would like to adopt the last of ((++1) to decode the next token. The parameters of the whole transformer is denoted as ( W( w L.H (W()) ( LL U (Wanembb) Then the whole transformer denoted as A.1
Q.h KAs V.Jihh= U
(L+1) transformer(XX 8).
In our method, will apply mask the MHA in some layers, where only remain the first and last several tokens. This can be described by defined the masked indexes set M, C [i] for i-row for i E [N]. The corresponding mask Mazy can be defined as
0 if M₁
MMszylis
-00
otherwisc.
For example, in our experiments, we Mi as the first and the last 1020 tokens. Then we denote the corresponding MHA as
h
mha x, (W WKA. Wrrl)=e))
softmax(XW K + X
22

---

Published in Transactions on Machine Learning Research (09/2025)
(0) The mha module at i-th layer will be denoted as mha The FF module will remain the same in the method. We denote the set of indexes of the layers that apply this mask as I. Then our method can be expressed as
= x + I} mha 1 (LN(X((-)))) Li I} mha (
(LN((F(11)),
where we denote all the hidden states with our method applicd as X and Y, and I[-} is the indicator function. The output of the whole nelwork denoted
X transformer(X,0,.2)
To derive the theoretical analysis of the error, we need to delineate the norm of the transformer parameters. In fact, all the transformers the real life have bounded parameters due to the calculation and storage requirements of the computer.
Assumption F.2. The Frobenius norms of all the parameters of the transformer is upper bounded by B > O, i.e., [| Q.A || & B, R,All,, S B, Wl < B, | & B, [[ 4,1 || z B, [Wunemblle S B for E [M] and E [L].
To state our main result, we define the maximal sum the original attention scores of the discarded tokens at layer I as 81, which formally defined as
exp (ILN(X"-D)) O (..T LN(X"(11))
b.: QA
max N
h=1
1 M.M, k= exp ([LN(XU-1)),
W O.T
QN QA K.
[LN(X
exp ([LN(X (( W. T [LN(X(-I))T]
K.A.
max c
1=1jEM EB exp ([LN(XV-D)], w W.T
[LN(XV-D)T].. QA K.
Then the main result is follows.
Theorem F.3. We define the difference of the hidden states of our method and the original transformer at layer E [L] as e = | 12,000 Under Assumptions F.1 and F.2. this error involves as
B² + 4H min {2. [1 H (1-1) (HB BBp ex
x- + 2B + Lipp L e I)ss (4)
The error between the logits generated by our method and the original transformer can be upper-bounder as
(transforme((x.0.2 transformer(xx,))
S 2LLB2 (H Lip B + 4H1B) + +2HB2 + Lip L1p
B2)
8i- (5)
2,00
IEI
We note that the error recursive expression consists of three terms. The first term represents the crror from the previous layer. The second term represents the error from the previous layer amplisied by the current. layer. Thanks to the layer normalization, this term will be truncated by 2. The last term represents the newly introduced error if we shorten KV cache at the current layer. By relaxing this recursive formula, we derive the upper bound of the error between logits of our method and the original transformer. This shows that the error is upper bounded by the sum of the attention scores of the removed KV pairs up to an additive constant.
Proof of Theorem F.3 We derive the error analysis of our analysis in three steps.
The error decomposition of the whole network.
Bound each term in the error decomposition.
Conclude the proof.
23

---

Published in Transactions on Machine Learning Research (09/2025)
Step 1: The error decomposition of the whole network.
We derive the error decomposition of the whole network in recursive manner. In fact, for the i-th layer, we have that
|| 12,00 1 [|8(0) 112 + fr ( (LN(F10)) ffn ffn (LN(X()))
2,00
||8000 - - S 118 x 1|,00
(6)
+ |l $ z)- mha ( (LN(X((-1))) + Ii E Z} mha (1)
(LN(X((-))) mha (LN(X
2,00
where the inequalities follow from the triangle inequality. In addition, we have that
transformer(X, (X, 0,I) transformer(X X, 0)
W Wunemblle || ||x(()) - -
(8)
2,20
where the inequality results from Lemma G.2
Step 2: Bound each term the error decomposition
We will bound each term in the right- hand side of Equ. equation G and equation the FF module, we have that
For the term related to
frn (9 (LN(P(()) ( (LN(Y
2,00
I Lip
| [|LN()(0)) - )l
A.2
$ Lip
|I (1)
F
|| min {2. (1/
y [|,,0
Lip B². min -2.1000-Y((111. 11,,00)))
where the first incquality results from Lemma G.2 the second inequality results from the definition of In(.), and the last inequality results from Assumption side of Eqn. equation 7 we have that
For the term related to MHA module in the right-hand
[[| $ I} mha (0) (LN(X(())) +Ifi Z} mhaa " (LN((((-))) mha (LN(X((-))))
2,00
1i $ z}
|maa ( (LN( mha (LN(X
2,00
+ +[(i mha (LN( mha ( (LN(X((-1)))
2,00
< I1 Z)
mha ( (LN((((-1)) mha (3) (LN(X
+ +[(i (l (3) (LN(}((-))) mha ( (LN(
2,00
+ mha ( (LN(X - mha (9) (LN(X (11))))
2,00
< H B(1 + ((1+4B²)||LL (B2))BBB(((((111 - - LN((X(1))
2,00 +lfi I} 12.2 8s
< H. B( + AB²) min (2, - (1--1)
11,000 Ifi E I) 2BH
(10)
where the first inequality results from the triangle inequality, the second inequality results from Lemma G.4 (1)
Define the error ex (+x x00 ||,,00 with d O. Combining Eqn. equation 6 equation
equation
and equation
10
have that
ex (t-1)
HB( + 4Ba) ein ex (-")) + Ili E 1)2BHs:
+ B2 min {2, x (i-1)
+ +HBBI + min ex (-")) + Hie I]2BHs:}.
(11)
24

---

Published in Transactions on Machine Learning Research (09/2025)
Step 3: Conclude the proof.
We derive the recursive expression of the hidden state error by relaxing the right- hand side of Equ. equation I1 as follows.
S (1-1) HB(1 4B²) min m2n () + [(i 2]2BHs
+ LBp min {2, [1+ HB((+ x + 1[i E 2)2L1n H3: Hs;
I x (1-1) + (HB + B + +4HB3) min {2, [1 + HB(1 + AB²) ex (i-1)) 2H BI E Z)8,
This proves the recursive formula. By summing this inequality from = 1 to i L, we have that
(1) S 2L + L + AHB3) + + Lin Bd) Σ x
8,.
(12)
IEI
Combining Equ. equation and equation
12
we have that
transformer(XX 0) transformer(X, 0)
S 2LBB (H + Lip B + 4HBB) + +2 H1 + LipB3
Σ
84.
2,00
IEI
Thus, conclude the proof of Theorem
F.3
G
Supporting Lemmas
Lemma G.1 (Corollary A.7 in Edelman et al. (0222) ). For any v, y E Rdd, we have
|softmax (x) softmax(y)lll 22 yllxo
Lemma G.2 (Lemma 17 in Zhang et al. 2022 ). Given any two conjugate numbers u, u E [1, ie., 1 == 1, and S p S oo, for any I xe and Re, we have
| S ||AT I and ||xx|p S
2023 For any X,X R and any Wo WK, E Lemma G.3 (Lemma 1.8 in Zhang al. Rdxda Wy.k Raxd for h E H of [X 2,00 11 Bx, [Wo I BQ, [[KK,lle, Bx, [Valle By for h [H], then have
(ma((X, (WQ, WK. Wv - mha(X, (WQ.. WK.s. Wyn)ken)
2,00
H Bv 4Bx BoBk) 11
Lemma G.4. For query vector q Rd, and two sets of key-value pairs K, Rixd K₂ Rexd Vi e Nixd and V/2 Naxd, We define attention scores softmax(q" K2] and softmax(q KTT) as
softmax(q [K1, K2]") |sT sT and softmax(g K = 3T.
Then we have that
|oftmaxx V softmax(q [K1, K2] [ V [|, s 2111s211 max V2l1220c)
Proof of Lemma G. In fact, we have that
softmax(g [K,, MVT TTT Io V1 + J.J V2, and softmax(g" (T V.V..
25

---

Published in Transactions on Machine Learning Research (09/2025)
Further, the difference between $1 and 81 can be upper bounded as
||
ll
N1
(KG).)
exp(a [Kile.)
(xp(aa[[i))) E exp(oT(K2kk) ΣE exply [K2l5.) ==1
Ni
Ni
exp(1 _ [Ki]..) Eni exple T [K2l..)
Ni
+ E exp(a' [K2]..)) Ejn1 exply [Kilja)
j=n
S2 1:
where the first equality results from the definition of softmax(+), and the last equality results from the definition of 82 Then we have that
softmax(q K. V[ softmax(q [K1, K2] IM V2 12
= |s| V1 + $ V/2 $1 Vill2
< || 3111 22,0 + || 1V2112.00
< 22|11|1 max V/ 2, [2212,000
Thus, we conclude the proof of Lemma
G.4
26


Reviews:
[{"review_id": "PUPUjgFQce", "reviewer": "Reviewer_9JNz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The paper's method is simple to use and does not add too much overhead\n- The eval benchmarks presented are relevant for scenarios where KV cache compression hurts the most (LongBench and NIAH)\n- The training method is simple, it relies on simple SFT and also uses open datasets for recovering model quality.\n- The paper provides sound theoretical analysis for their method and also the writeup in appendix B.2 \n- The reuse of attention recomputation is good as it minimizes overhead.\n\n**Weaknesses**\n- Some of the parts of the method was unclear. The authors say their method is calibration free for training, but then they pass the SFT data through the model to figure out lazy layers?\n- Some of the parts of the paper are difficult to follow. Many tables such as Tables 3, 4, 7 are missing baseline numbers and Table 5 provide only speedup factors without providing details for actual numbers.\n- There is minimal novelty in the method of replacing attention layers in transformers.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a method to convert transformers to hybrid models (with different attention styles in different layers). The method LightTransfer has two methods, `test` for inference-time transformations and `train` for reasoning models, with a single-shot transformation. For the hybrid layer, they use streaming attention layers and show initial analysis supporting that streaming attention patterns emerge naturally for long context tasks. The `test` method, used for non-reasoning models, an on-the-fly transform is applied per sample. They define the lazy ratio as the metric to figure out which layers to set as lazy (ie, convert to hybrid). Using attention weight recomputation, they reduce the cost to compute the lazy ratios across the model. The `train` method, used for reasoning models, uses the data for prompt + completion to identify the layers in an offline manner, and then run SFT to recover accuracy. The method is shown to have good results on long context tasks (LongBench and NIAH) using the `test` method and math reasoning tasks (AIME / MATH-OAI / GSM8k).", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Some of the wording around the experimentation section is unclear. It takes multiple reads to figure out which model was being used in the `train` method, and towards this, it will be better to be more upfront. [**strengthen the work**]\n- Some more details on benchmarking will be useful. What inference engine setup, HW etc, since those details matter for inference performance. [**critical to secure recommendation**]", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "n4sbWOoZVS", "reviewer": "Reviewer_pnfF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I need to clarify first that I am not a full expert in this domain.\n\nStrengths:\n* The paper is well-organised and easy to read.\n* The comprehensive experiments have been conducted.\n\nWeaknesses:\n* The idea of replacing the attention mechanism of some layers in Transformers with a more efficient one is not novel.\n* The authors should also provide the experiments to verify the bound shown in Theorem 5.1 and should explain why this bound is useful.\n* In *Effect of different layer replacement strategies* part, the authors should also compare other replacement strategies. \n\nReferences:\n* [Investigating Layer Importance in Large Language Models](https://aclanthology.org/2024.blackboxnlp-1.29/) (Zhang et al., BlackboxNLP 2024)\n* [A Primer in BERTology: What we know about how BERT works](https://arxiv.org/abs/2002.12327) (Rogers et al. 2020)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a method to identify the Transformers layers that pay great attention to unimportant tokens and the most recent ones, and replace the full attention mechanism of these layers with streaming attention. On top of four representative LLMs, the authors conduct comprehensive experiments on the long context understanding task and 1-like long reasoning generation tasks.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See Weakness.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "saN7It08KH", "reviewer": "Reviewer_tEVS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- Efficiency without retraining\nLightTransfer delivers substantial memory savings and up to a 2.17× increase in token‐generation throughput with no or minimal additional training, making it straightforward to drop into existing pretrained transformers.\n- Strong empirical performance\nAcross multiple models (LLaMA 2, Mistral, LLaMA 3) and long-context benchmarks, it maintains accuracy within 1.5 % of full self-attention while more than doubling throughput on sequences up to 32 K tokens.\n- Broad applicability\nValidated on diverse architectures (e.g. LLaMA, Mistral, QwQ-STILL) and tasks (zero-shot understanding, fine-tuned reasoning), demonstrating robustness across model sizes and use cases.\n- Insightful analysis\nClear visualizations of layer-wise attention patterns expose which layers are “lazy,” and a principled method leverages these observations to decide where to substitute streaming attention.\n- Solid theoretical underpinning\nAn explicit error bound (Theorem 5.1) guides layer replacement decisions and distinguishes LightTransfer from purely heuristic approaches.\n\n# Weaknesses\n\n## Major\n- Ambiguity around “streaming attention” setup\n\nThe original streaming-attention work uses learnable sink tokens and requires retraining. Here, however, the paper treats the first few prompt tokens as “sink” tokens without retraining. The authors should (a) clearly describe their implementation (are sink tokens learned or fixed?), and (b) compare performance against a fully-trained streaming-attention baseline to quantify any trade-offs.\n- Unclear use of the theoretical error bound\n\nSection 5.3 derives an upper bound on the error introduced by removing KV caches, but it isn’t clear how to translate that bound into practice. How can one use Eqn 1 to predict model quality drop for a given layer replacement ratio? Concrete guidance or examples would elevate the theorem from an observation to an actionable design tool.\n\n## Minor\n-  Overhead measurement in Section 5.1\n\nThe statement “throughput reduction of 0.0058 to 0.0014 relative to a baseline of 1” is opaque. Does this correspond to a 0.14 – 0.58 % overhead? Is it measured only during prefilling or across both prefilling and decoding?\n\n- Figure 2 (left panel) clarity\n\nThe circular highlight around the border can be mistaken for the plot boundary. Consider adding a distinct rectangular border or subtle shading behind the highlight to separate it from the axes.\n\n- Figure 2 (right panel) generation details\n\nIt’s unclear whether the heatmap reflects attention to all early and recent tokens or just one representative token. Please specify how many tokens were aggregated, whether the color scale is an average or maximum, and how the data points were computed (e.g. per-token mean over a batch).\n\n- No limitation discussed", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces LightTransfer, a lightweight framework that transforms pretrained transformer LLMs into hybrid models by:\n- Identifying “lazy” layers—those that predominantly attend only to the first few or most recent tokens—and replacing their full self‐attention with memory‐efficient streaming attention.\n- Zero‐shot adaptation for long‐context understanding tasks (no additional training required) and minimal fine‐tuning (≈5 K examples) for advanced long‐reasoning generation tasks.\n- Extensive empirical evaluation across multiple LLMs (LLaMA 2 7B‐chat, Mistral 7B‐Instruct, LLaMA 3 8B/70B) on benchmarks including LongBench, Needle‐In‐A‐Haystack, and AIME24. Results show up to 2.17× token‐generation throughput improvement with < 1.5 % performance drop on LongBench and 53.3 % accuracy on AIME24—matching or exceeding baselines while replacing half the layers .\n- Ablation studies demonstrating the trade‐off between the fraction of retained full‐attention layers and performance, and superiority over other inter‐layer KV‐cache reduction methods .\n- Theoretical grounding, optimizing an explicit error upper bound (Theorem 5.1) to guide layer replacement decisions, and analysis of attention patterns in long‐context inference", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Clarify experimental setup, particularly around sink tokens\n- Extend the theoretical analysis further\n- Discuss limitations and address the ambiguities in the plots and text", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WTgvsnoycR", "reviewer": "Reviewer_qcXM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Please find the strengths below:\n- The paper defines a clear metric to pinpoint layers that ignore most of the context, revealing a previously underexplored structural property of long-context Transformers.\n- The paper introduces LightTransfer-Test and LightTransfer-Train, which improve inference throughput without harming performance, even on reasoning tasks.\n- The paper provides a logit-level approximation guarantee linking the discarded attention mass $s_i$ directly to the worst-case output deviation.\n\nPlease find the weaknesses below:\n- The lazy-ratio metric was only evaluated on a limited set of datasets and tasks, so its behavior and computation overhead across different domains remain unclear.\n- Experiments use contexts up to tens of thousands of tokens only; multi-document or web-scale inputs are untested and may yield different results.\n- The theoretical error bound lacks any experimental or ablation study to validate its tightness or practical relevance.\n- There is no systematic investigation of how many layers can be converted before accuracy drops sharply, leaving users without clear guidance on setting conversion thresholds.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper analyzes the behavior and weight structure of transformer on relatively long input sequences. The paper discovers that certain layers focus almost only on the initial and most recent tokens, which are called lazy layers in the paper. Inspired by this discovery, the paper proposes LightTransfer-Test, a runtime technique that converts lazy layers from full self-attention to streaming attention, significantly reducing key-value cache usage and improve inference throughput without additional training. Moreover, to handle reasoning tasks sensitive to attention changes, the paper introduces LightTransfer-Train, which fine-tunes only the lazy layers on a few thousand examples to restore performance. The paper then presents a bound on the logit-level approximation error introduced by these methods. Finally, the paper conducts experiment on multiple open-source models (LLaMA, Mistral) and benchmarks (long-context understanding, challenging math tests). The methods demonstrate up to 2.17× speedups, roughly 50% cache reduction, under 1.5% accuracy degradation, and even improved performance on math datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper should validate the lazy-ratio across diverse domains and measure its computation overhead and extend experiments to truly web-scale contexts to ensure generality. To strengthen the work, the paper can include an empirical study of the error bound’s tightness, analyze how accuracy degrades as more layers are converted to guide threshold selection, and quantify domain-specific fine-tuning requirements for LightTransfer-Train.", "broader_impact_concerns": "I do not foresee any ethical issues in this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)