On-the-Fly Adaptive Distillation of Transformer to Dual-State  Linear Attention for Long-Context LLM Serving
Authors: Yeonju Ro
Zhenyu Zhang
Souvik Kundu
Zhangyang Wang
Aditya Akella
Venue: ICML 2025 poster (2025)
Keywords: LLM
Efficient serving
Linear attention
adaptive inference
dynamic serving
Primary area: deep_learning->large_language_models

Abstract: Large language models (LLMs) excel at capturing global token dependencies via self-attention but face prohibitive compute and memory costs on lengthy inputs. While sub-quadratic methods (e.g., linear attention) can reduce these costs, they often degrade accuracy due to overemphasizing recent tokens. In this work, we first propose *dual-state linear attention* (**DSLA**), a novel design that maintains two specialized hidden states—one for preserving historical context and one for tracking recency—thereby mitigating the short-range bias typical of linear-attention architectures. To further balance efficiency and accuracy under dynamic workload conditions, we introduce 
DSLA-*Serve*, an online *adaptive distillation* framework that progressively replaces Transformer layers with DSLA layers at inference time, guided by a sensitivity-based layer ordering. 
DSLA-*Serve* uses a chained fine-tuning strategy to ensure that each newly converted DSLA layer remains consistent with previously replaced layers, preserving the overall quality. Extensive evaluations on commonsense reasoning, long-context QA, and text summarization demonstrate that 
DSLA-*Serve* yields **2.3×** faster inference than Llama2-7B and **3.0×** faster than the hybrid Zamba-7B, while retaining comparable performance across downstream tasks. Our ablation studies show that DSLA’s dual states capture both global and local dependencies, addressing the historical-token underrepresentation seen in prior linear attentions.

On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention for Long-Conter LLM Serving
Yeonju Ro Zhenyu Zhang' Souvik Kundu Zhangyang Wang !
Aditya Akella
Abstract
1. Introduction
Large language models (LLMs) have become indispensable for tasks such advanced reasoning. code generation, and multi-turn dialogue. However, long-sequence inference re- mains a major bottleneck due to the ((?)) complexity of self-attention (Vaswani, 2017), which evaluates pairwise relationships among all tokens in a sequence. Alongside heavy compute cost, storing key-value (KV) caches for ex- tended contexts forces memory usage to grow linearly with the number of tokens, significantly challenging scalability.
Large language models (LLMs) excel at captur- ing global token dependencies via self-attention but face prohibitive compute and memory costs lengthy inputs. While sub-quadratic meth- ods (e.g., linear attention) can reduce these costs, they often degrade accuracy due to overempha- sizing recent tokens. In this work, we first pro- pose dual-state linear attention (DSLA), a novel design that maintains two specialized hidden one for preserving historical context and one for tracking recency- thereby mitigating the short-range bias typical of inear-attention archi- tectures. To further balance efficiency and ac- curacy under dynamic workload conditions, we introduce DSLA-Serve, an online adaptive dis- rillation framework that progressively replaces Transformer layers with DSLA layers at inference time, guided by a sensitivity- layer order- ing. DSLA-Serve uses a chained fine-tuning strat- egy to ensure that each newly converted DSLA layer remains consistent with previously replaced layers, preserving the overall quality. Extensive evaluations on commonsense reasoning, long- context QA, and text summarization demonstrate that DSLA-Serve yields 2.3x faster inference than Llama2-7B and 3.0x faster than the hybrid Zamba-7B, while retaining comparable perfor- mance across downstream tasks. Our ablation studies show that DSLA's dual states capture both global and local dependencies, addressing the historical-token underrepresentation seen prior linear attentions. Codes are available at: https: / / thub. www/ctt////c/// utsll
To reduce inference complexity, prior work explores both system-level optimizations (c.g.. cache management and GPU kernels (Ye et al., 2025; Dao et al., 2022; Kwon et al., 2023b)) and algorithmic approaches (Leviathan et al., 2023; Cai et al., 2024a; Kudugunta et al., 2023; Xiao al., 2023b; Zhang al., 2023). Among algorithmic solutions, linear attention has gained traction (Peng et al., 2023; Wang et al., 2020; Gu & Dao, 2023), maintaining fixed-size hidden state to update token embeddings in (TT) time. This makes such models appealing for large-context processing.
Despite their efficiency, purely linear attention models often underperform on tasks requiring careful long-range interac- tions (e.g., large-scale text retrieval) (Waleffe al., 2024; Gu & Dao, 2023; Lieber et al., 2024). Their fixed hidden states impose capacity bottleneck, especially when the task demands rich cross-token relationships. Hybrid archi- tectures have thus emerged, mixing a few self-attention and linear-attention layers (Lieber et al., 2024; Ren et al., 2024; hug). However, these hybrids typically rely on a static de- sign, which cannot flexibly trade off compute cost versus accuracy on the fly. Meanwhile, recent attempts to distill transformers into linear attention (Bick et al., 2024; Zhang et al., 2024a; Wang et al., 2024a) either incur significant ac- curacy drops when fully replacing self-attention, or remain constrained by bottleneck layers if only partially converted. Consequently, current approaches do not adapt well to the variable demands of real-world inference pipelines.
'The University of Texas at Austin Intel Labs. Correspondence to: Yeonju Ro <yro@es cs. utcxas.ccco edu>.
A naturally emerging, intuitively promising solution is then to dynamically adapt architectures at inference time, convert- ing self-attention layers into linear-attention analogs only under memory or latency pressure. Yet achieving such "on-
Proceedings of the nnnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
the-fly" adaptive distillation is non-trivial. First, we must de- cide how many layers can be replaced without unacceptably harming accuracy. Second, mismatches between training and inference architectures easily degrade performance if not carefully addressed. Third, since linear-attention meth- ods often exhibit a strong bias toward recent tokens (see $3), naive distillation may fail to capture the global attention patterns that long contexts demand.
state can overemphasize the most recent context, limiting deeper cross-token learning.
Gated Linear Attention (GLA), To mitigate underper- formance in standard SSMs, GLA (Yang et al., 2023) intro- duces data-dependent gating, adding a learnable forget gate G. E Raxd that modulates the previous hidden state S,-1:
S = G. o S + +k, V..
Our Contributions. In this work, address these chal- lenges through:
where G2 formed from outer products of vectors at 3,, and k;, V, are the key/value embeddings time 1. Although this gating can help retain some historical information, GLA still fundamentally relies on one hidden state, often losing mid- or far-context representations.
Dual-State Linear Attention (DSLA). We first ana- lyze how single-state linear attention overemphasizes recent tokens and propose new module with two spe- cialized states: one for preserving historical context and the other tracking recency. Driven by a constrastive regularization, this design alleviates the short-range bias and better emulates full self aelf-attention.
Knowledge Distillation into Linear Attention. To nar- row performance gaps between quadratic self-attention and linear methods, some works apply distillation (Wang et al., 2024a; Zhang et al., 2024a). However, these static conver- sions do not address dynamic inference scenarios, where the model must flexibly trade accuracy for reduced memory cost under changing request loads or context lengths.
Adaptive Distillation via DSLA-Serve. We then in- troduce DSLA-Serve, a framework that adaptively con- verts Transformer layers into DSLA at inference time, guided by a sensitivity- based layer ordering. This "on- the-fly" approach allows different workloads or prompt lengths to flexibly trade off accuracy and efficiency.
Hybrid Architectures Another approach is to mix self- attention with linear layers different layers or heads (Wal- effe et al., 2024; Lieber et al., 2024; hug). While hybrid models can mitigate the drawbacks of purely linear or purely quadratic methods, their fixed attention layouts remain un- suitable when inference loads fluctuate.
Chained Fine-Tuning for Consistency. To minimize performance drops during layer conversion, DSLA- Serve employs a progressive distillation procedure. Layers are distilled in precisely the order they will be converted at inference, ensuring each newly replaced layer is compatible with previously converted ones.
By contrast, our work aims to (i) improve the typical single- state design of linear attention through a dual-state mech- anism, thereby alleviating recency bias, and (ii) enable on- the-fly conversion of self-attention layers so that the model can dynamically trade off accuracy and efficiency based on system constraints. We demonstrate how single-state de- signs exhibit short-range bias through an attention analysis, then show how dual-state linear attention (DSLA) success- fully mitigates this issue. We finally (iii) incorporate an adaptive inference framework that progressively replaces self-attention with DSLA layers under high resource load, ensuring scalable and efficient LLM serving without sacri- ficing crucial long-context performance.
Extensive Evaluations. Across multiple datasets, DSLA-Serve achieves a 2.29x and 3.0x reduction in end-to-end latency compared to Llama2-7B and Zamba-7B baselines, respectively, while maintaining comparable performance on tasks spanning reasoning, code, and long-context understanding.
2. Background
Notations. We denote matrices with bold uppercase letters (e.g.. X), vectors with bold lowercase letters (e.g., x), and scalars with lowercase letters (e.g.. 2).
3. Motivational Analyses
Linear Attention. A primary driver behind sub-quadratic complexity is linear attention, which replaces the (TT) ) all- pairs interactions of self-attention with recurrences or kernel- based factorizations (Child al., 2019; Beltagy al., 2020; Zaheer et al., 2020; Choromanski et al., 2020; Katharopou- los et al., 2020). Recent approaches often adapt stace-space models (SSMs) (Gu& & Dao, 2023; Dao & Gu, 2024) or RNN gating (Peng et al., 2023) to maintain a fixed-size key-v (KV) state that updates incrementally across tokens. While this design yields (T) computation, a single recurrent
In this section, we present two key motivating observations:
how single-state linear attention tends to overempha- size recent tokens at the expense of historical context;
why LLM serving systems would desire to dynami- cally adapt to changing resource constraints rather than relying on static architectures.
Attention Score Recency Analysis. Gated linear attention (GLA) (Yang et al., 2023) updates the hidden state and
2

---

Onlthe-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Head
Layer33 Head
0.000
Mermry Roquirement (GB)
120
Node Node:
Node
0.004
80-
0.002
0.000
1000
4000
3000
40-
Postloan
Posttio
Amterinn Score
Allaniion Scrres
1800
14.00
16
Time
Figure 2. Illustration of resource usage spikes (memory) over time and across nodes. Long-lived or bursty sessions can heavily strain certain GPUs.
1000
2000
1000
3000
Posion Index
Popibon S1004
Figure 1. Attention scores for Llama2-7B (top) vs. GLA (bottom). The latter is distilled from Llama2-7B with using IB tokens and KL divergence w.r.t. attention oulputs. We examine several randomly selected attention heads/layers on 16 samples from C4 (sequence length 4096). The horizontal axis is the token index, and the vertical axis shows the attention scores to past tokens from the most recent query. GLA focuses heavily on near tokens, while the Transformer maintains more balanced coverage of earlier positions.
multiple nodes. In production scale deployments, load fluc- tuations arise due to:
Temporal variability: Sporadic bursts of lengthy re- quests or multiple concurrent sessions can dramatically increase memory usage when each session accumulates a large KV cache.
output at time step t via:
Spatial imbalance: Session affinity in multi-turn chat systems can repeatedly assign requests from the same user to single GPU, causing that node to saturate while other nodes remain underutilized.
S: = G. @ S + k, Vt,
(1) (2)
0, = q St:
where G. is a learnable gating term in Radd, and k,, Vi, qi are the key, value, and query representations. From these, the final query qr produces
A common fallback is to over-provision GPU capacity, which is both costly and under-utilized most of the time. Conversely, fully replacing self-attention with linear lay- ers reduces memory but can degrade performance on tasks demanding nuanced cross-token relationships. Hence, flexible mechanism is needed to partially convert layers on the fly, balancing efficiency and accuracy in response to real-time resource pressure. Existing static architectures are ill-suited for such dynamic scenarios, motivating our framework that progressively substitutes a Transformer's 's self-attention layers with linear-attention analogs under high load, while retaining more accurate (i.e., quadratic) layers when resources are sufficient.
T OT 2O.V a: (qr
O k;)
II +=
G.). (3)
illustrating how the gating products affect attention to each past token Ut.
Figure compares attention scores for qr across multiple heads in Llama2-7B (top) vs. GLA (bottom) on sequences of length 4096. In many Transformer heads, the attention distribution extends across earlier tokens (including middle or far context). By contrast, GLA consistently places dispro- portionate weight on recent tokens. We hypothesize that this stems from GLA's reliance on a single compressive state S, with a fixed dimension, which must simultancously encode all prior context. Over long sequences, the gating G, of- ten pushes the model "forget" older segments (Ben- Kish et al., 2024; Azizi et al., 2025). In next-token prediction settings, focusing on local tokens often suffices for short contexts, but can be detrimental when crucial information resides far back in the input. We notice a concurrent work (Wang et al., 2024b) has drawn similar "recency bias" ob- servations on state-space models via theoretical analysis.
4. Methodology
Motivated by the recency bias single-state linear attention ($3). this section introduces our proposed Dual-State Linear Attention (DSLA), which maintains both "history" and "re- cent" contexts. We then describe an adaptive LeM-serving framework, DSLA --SSevee,, that progressively converts self- attention layers into DSLA layers in response to dynamic workload demands. Figure 8 gives a high-level overview.
4.1. Dual-State Linear Attention (DSLA)
Need for Adaptive Models. We also highlight the ne- cessity of adaptable inference architectures through a real serving trace (Appendix C), shown in Figure 2. The y y-axis indicates resource usage, measured over time (x-axis) across
Gated Linear Attention (GLA) (Yang et al., 2023) uses a single hidden state S, updated by a forget gate, which causes earlier tokens to be truncated over long sequences. To mitigate this recency hias, we propose Dual-State Linear
3

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Attention (DSLA), where each layer maintains Iwo states:
gains: one state tends to preserve global or historical con- text, while the other focuses on recency or local patterns (cf. $5.5), Empirically, using more than two states showed diminishing returns, increasing training overhead and model complexity without sizable accuracy improvements. Once again, the key is to properly specialize the two states.
S = G.O.S + k vt, i Gl + k, Vi. (4)
Here, G. G? E Rdxd serve as data-driven forget mecha- nisms; k, Vi are the key/value vectors at time t. We then blend these two states in the output:
O, q: (r- s + (1)
(5)
Contrastive Regularization. To drive specialization, we combine a distillation term with a contrastive penalty. Let o, denote DSLA output at position t, and (gt) be the self-attention output from the teacher Transformer. Defnn:
where q: is the query, and is a learnable coefficient that determines the relative weight of each state. By default, we initialize G (the "history" gate) to be closer to the identity matrix, helping preserve older context, while G? (the "re- cency" gate) is randomly initialized (e.g., from (0, o²)) to have broader forgetting effect. This dual-stara design offers more capacity to store mid- or far-range tokens, while still benefiting from the temer-time update rule.
Laist = (o:((ol)) 1f )
Lcont = sim (G), G?),
where sim(.,)) measures cosine similarity between the two gating matrices G G? at time t. Minimizing Ldist aligns DSLA's outputs with the teacher's, while minimizing Coont reduces the similarity of gating functions across the two states, encouraging them to attend to different aspects of the sequence. The total loss is:
Notably, attention patterns vary significantly across layers in a deep model (Ren et al., 2024; Waleffe et al., 2024). Early layers often focus on local context, middle layers tend to retrieve from mid-range segments, and later layers aggregate more global representations. To accommodate these differing roles, we introduce a learnable coefficient Y per layer, which adjusts the relative emphasis on our two states (history vs. recent). Formally, at the final token T, the DSLA output is:
C Edist + X Lcont,
(10)
with a hyperparameter À balancing teacher alignment and state differentiation. During training, the contrastive term ensures that the two gates do not collapse to similar solu- tions, while the distillation term aligns the overall DSLA output with the teacher's 's self-attention. In practice, we ob- serve that DSLA converges in roughly the same number of steps as single-state GLA, indicating minimal overhead.
T Or
[(7o)
Σ 1=0
+ (1- vc.
(6)
where
of = (qr 1 kt)
G o G ... o 1)
Interpretation. Qualitative attention-score analyses (Fig- ure 5) show that the two gates yield distinct patterns: S retains capacity for earlier tokens, whereas S adaptively highlights recent context. Together, they more closely repli- cate full Transformer attention while retaining the linear- time update rule. We next explain how we use DSLA at inference time in a partial or fully converted manner ($4.3), enabling efficient serving across diverse workloads.
a. = (qr 2 kt)
G o G
G2 )
By learning alongside the gating parameters, each layer can balance long-range and local information differently. We interpret o + (1 -) ofll2 as the attention score to token v, for the final query qr, reflecting a superposition of history and recency signals tailored that layer's role.
4.3. Adaptive Layer-wise Conversion with DSLA-Serve
4.2. Specializing the Two Hidden States
While converting all self-attention layers DSLA maxi- mizes speedups and memory savings, it also carries higher risk of accuracy degradation. In practice, the optimal trade- off between efficiency and accuracy can vary with both workload and application constraints. To handle this in fine granularity, we propose an automatic conversion framework, DSLA-Serve that progressively replaces Transformer layers with DSLA at inference time (Figure 8).
While two hidden states S!. S3 expand the capacity beyond single-state GLA, we also seek to ensure that each state specializes to capture different parts of the context. If S and SP converge to the same gating focus on overlapping token spans, the dual-state advantage diminishes.
Why Two States Suffice? In principle, multiple additional states could be introduced, each with its own gating matrix. However, in our experiments, we observe that splitting the capacity into two distinct states already provides significant
Layer Ranking via Sensitivity. Before deployment, we distill each layer from the teacher model (c.g., a Trans- former) into a DSLA counterpart, then measure how much
4

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention
replacing that layer alone impacts validation accuracy (or perplexity). For more fine-grained indicator, compute attention entropy:
Senstiivity
2.8
T Σ ==0
Entropy(A)
AT, log AT,
(11)
Full Transformer
Full Lincar Attntion
Figure 3. Chained Fine-Tuning, We replace layers in ascending
where A is the attention score matrix for the final query Qr. Lower entropy often indicates a layer focusing on fewer
sensitivity order, freezing all other layers. Upon complction of each stage, we commit the newly trained DSLA layer so that subsequent stages see the updated architecture.
tokens. making it "less sensitive" to lincarization. We sort layers in ascending order of sensitivity.
layers have already been converted. To address this, adopt a chained fine-tuning approach that ensures consis- tency between training and deployment architectures. As illustrated in Figure 3, we process with the following recipe:
Adaptive Conversion Logic. During runtime, DSLA- Serve monitors memory load and query lengths. When the system detects high GPU pressure, begins converting layers from the least sensitive end of the queue to DSLA form, thereby reducing memory usage by removing the large key-value (KV) cache footprint of those layers. This process continues until either:
Layer Ranking, Identify and order the layers by their sensitivity (ef. $4.3).
Replace and Finetune. Convert the first (lowest- sensitivity) layer to DSLA, freeze all other layers, and fine-tune until convergence.
1.
Memory Relief Achieved: Enough layers are con- verted that GPU usage falls below a threshold; or
3. Update Baseline. Permanently replace that layer with its DSLA counterpart, so subsequent layers see the updated partial-linear architecture.
2. Accuracy at Risk: A certain service-level objective
(SLO) on quality (e-g., perplexity or BLEU score) is at risk, so DSLA-Serve halts further conversion.
Iterate. Move to the next layer in the queue, repeating the same procedure.
Prioritizing Long Contexts. Many real-world requests are short and do not stress memory resources. By contrast, large prompts and multi-turn sessions can rapidly inflate the KV cache. Thus, DSLA-Serve focuses first on converting layers for those long queries, yielding outsized gains in memory and speed. Our experiments (Tables 1-2) show that partial DSLA substitution for long prompts retains strong performance while substantially reducing latency.
By the time we finish training the most sensitive layer, the model can handle any partial combination of DSLA and Transformer layers without unexpected degradations. This yields family of progressively converted models whose partial configurations are consistent at test time.
This approach prevents train-test mismatch that would oth- erwise arise if multiple DSLA layers used in deployment but were never jointly seen during training. Empirically, our experiments confirm that chained fine-tuning mitigates accu- racy drops and ensures a smooth tradc-off between memory usage and performance.
Connection to Real-World Load. As discussed in $3, production workloads exhibit both temporal and spatial load imbalances. By flexibly converting growing fraction of layers to DSLA, DSLA-Serve can quickly address transient spikes in GPU usage, then stop once the system stabilizes. This "pay as you go" strategy avoids static overprovisioning or fully switching to linear layers (which may excessively degrade accuracy). We next describe how DSLA-Serve en- sures consistency across partial and fully converted configu- rations ($4.4), followed by batching considerations ($4.5).
4.5. Batched Inference Considerations
Batch decoding (Yu et al., 2022; Kwon et al., 2023a) is a common practice for high-throughput LLM services, where multiple generation requests (each corresponding to a dif- ferent user prompt) are processed in parallel. However, partial layer conversion can create branching paths withs-some requests may need DSLA layer while others still rely on self-attention if they have not (or do not) trigger conversion. These requests cannot be fused into a single GEMM pass be- cause DSLA parameters differ from those of self-attention.
4.4. Chained Fine-Tuning for Seamless Conversion
Although adaptively converting layers during inference of- fers flexible memory-accuracy trade-off, naively training DSLA replacements for each layer in isolation can lead to inconsistent behavior when multiple layers are simultane- ously converted. In other words, replacing layer i while assuming all other layers are Transformers at training time will not match actual scenarios where one or more other
Re- DeeBBatching Overhead. When the model encounters a layer where some requests use DSLA while others still use self-attention. DSLA-Serve separates them into two sub- batches. After that layer, the sub-batches may be merged
5

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
again for subsequent shared layers or re-split depending on conversion status. In principle, this re-batching can add scheduling overhead. However, we find that:
and the hybrid model (Zamba-7B), achieving up 72.23% % performance improvement (e.g., compared to Mamba-7B TriviaQA). This substantial gain is likely attributed to the distillation process, which effectively transfers knowledge from the LJama-2-7B teacher model to the converted model. Notably, in multi-document QA tasks, our converted models even surpuss pure self-altention models, achieving accuracy gains of 3.96 to 5.44 for the 25% converted model. Such conversion also offers practical benefits, reducing memory consumption by 1GB for a 4K context and 2GB for an 8K context per request.
KV Cache Savings Dominate. Converting even frac- tion of layers significantly reduces the cumulative KV cache, lowering overall memory pressure and boosting throughput in the rest of the pipeline.
Empirical Minimal Impact. Our experiments ($5.4) show that the time spent splitting and merging batches is relatively small compared to the overall inference cost, especially for large prompt lengths.
Thus, although batched requests may follow diverging paths, the net system throughput remains high. Combined with our chained fine-tuning strategy, DSLA-Serve enables a flexible and efficient inference pipeline that handles heterogeneous requests while maintaining strong accuracy.
Furthermore, we evaluate our DSLA on general language modeling tasks, reporting perplexity in Table 1, as well as on summarization tasks Table 2. For summarization, we utilize the widely used CNN/DailyyMii dataset (Hermann et al., 2015; See et al., 2017) and the XSum dataset (Narayan et al., 2018), reporting performance using ROUGE met- rics (Lin, 2004). Our results indicate that DSLA effectively retains performance comparable to Llama2-7B and Zamba- 7B while outperforming GLA-7B and Mamba-7B This sug- gests that the DSLAA retains the flexibility and adaptability needed for long-context understanding, all while providing substantial memory and computation savings.
5. Exxerimments
Experimental Setup. In this section, we present the experi- mental analysis of the proposed method. We used 1.6 billion tokens sampled from the SlimPajama dataset (Shen et al., 2024) to fine-tune the model. We distilled our model from Llama2-7B (Touvron et al., 2023). For additional setup, please see Appendix B.
5.2. Short-Context Benchmarks
Baselines. We compare the adaptive distilled models gen- erated via DSLA with representatives tfte-offthhe-rr from linear-complexity models (RetNet-6.7B (Sun et al., 2023b), GLA-7B-20B (Yang et al., 2023), Mamba-7B (Gu & Dao, 2023)), hybrid model (Zamba-7B (Glorioso et al., 2024)),
Although the primary goal of DSLA is to enhance long- context understanding in efficient models, we also ensure that it maintains strong short-context performance with- out any degradation. Table 3 presents the performance of DSLA on seven common-sense reasoning tasks (lm.eval), where it demonstrates competitive performance against other bascline models, including pure self-attention models, hybrid and linear- complexxty models. Additional results compared with an alternative linearization method (Zhang et al., 2024a) are provided in Table 7 in Section 5.5.
and quadratic model (Llama2-7B (Touvron et al., 2023)), respectively. We report the performance of DSLA versions with 25%, and 50% layers converted to linear attention al- ternatives.
5.1. Long-Contee Understanding
5.3. Inference Efficiency
Setup. We evaluate our DSLA model on tasks requiring careful long-range interactions (Bai et al., 2024), includ- ing: (i) Multi-Document QA, where the model must answer questions using information spread across least two doc- uments; (ii) Code Understanding, which require reasoning across multiple code files to answer complex questions; (iii) Few-shot learning requires the model to comprehend a lim- ited number of examples and effectively perform various tasks, including classification, summarization, and reading comprehension. Additionally, we report the perplexity on WikiText-2 and Lambada and compare the DSLA model, with 25% and 50% converted layers, against other models at 7B-scale. Training cost is reported in terms of the number of tokens processed.
In this section, we evaluate the inference latency and per- request memory usage during the prefill and decoding stages. Figure 4 shows the prefill latency and decoding latency normalized by generated length.
In the prefill stage, the full transformer shows lower latency for prompts under 2K tokens, while the converted model excels for prompts over 2K tokens. This improvement is attributed to the DSLA layer's design, which utilizes single linear layers for gating to minimize latency overhead while retaining the gating mechanism's 's benefits. Performance gains are particularly notable for context lengths up to 8K tokens.
During decoding, transformer latency increases with the length of generated tokens due to the self attentio mecha-
Results. As shown in Table 1, our DSLA significantly outperforms other lontar-cost models (GLA-7B, Mamba-7B)
6

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Table Results of long context understanding.
We report 25%, 50% converted layers of DSLA. DSLA is distilled from Llama2-7B.
Multi-doe QA
Code Understanding LCC1
Fow-shot Learning
Methods
Cost
WikiTexx-2 4 Lambada
HotporQA 1 2W1kiMQA
Repobench
TREC1 f
Samsum
TriviaQA f
Llama2-7B
2T
8.79
4.13
5.63
10.24
69.83
56.88
59.00
39.1
86.19
GLA-7B
20B
NaN
4.98
3.61
6.89
41.26
44.24
28.50
16.94
57.68
Mamba-7B
N/A
10.55
4.05
1.23
0.80
17.56
10.54
11.0
4.55
15.23
Zamba-7B
10.25
3.74
7.90
7.97
40.70
43.20
64.0
37.74
82.19
DSLA [25%]]
1.6B
9.26
4.14
11.07
14.20
66.91
51,53
55.0
38.66
87.46
DSLA [50%]
1.6B
9.89
6.19
10.61
13.64
61.68
49.84
46.0
37.28
81.99
Table 2. Comparison results on text summarization tasks.
11.
The KV cache size per layer is given by:
Mail
XSUM
Avg.
Rouge-1 1
Rouge-L.1 L.
Rouge- 1
Rouge-L.1
Llama2-7B
30.13%
27.63%
30.07%
24.51%
28.00%
KVCache [B] 2 Texnxbxse
dA Bogte
GLA-7B
18.55%
17.62%
24.65%
20.92%
20.44%
Mamba-7B
23%13%
21.04%
27.92%
23.49%
23.65%
Zamba-7B
31.34%
28.67%
28.41%
23.09%
27.88%
Here, is the number of layers, b is the batch size, s is the context length, h is the number of heads, d, is the dimension per head, and Boyte is for the data type size (e.g., for BF16, for FP64). Thus reducing the number of transformer layers significantly reduces memory usage.
DSLA [25%)
30.011
27.10%
30.10%%
24.70%
27.98%
DSLA [50%)
29.29%
26.42%
29.88%
23.73%
27.33%
Table 3. Comparison results 1m_eval 1 tasks.
war
PIQA
ARC-E.1
ARC-CC
MMLU1
LagiAA
Aug.
1.amm2271
69.3%
57.19
70.05
16.4%
45.31
30.6%
57.2%
77.8%
73.3%
26.1%
5.4. End-to-end System Performance
GLA-7B
69.5%
57.0%
79.2%
34.6%
44.8%
28.45
28.8%
54.5%
Ma
71.9%
58.6%
79.0%
77.5%
42.8%
33.9%
23.7%
55.59
Zumbu 7B
69.9%
57.0%
78.8%
74.2%
42.2%
48.9 16
27.0%
56.9%
In this section, we evaluate the end-to-end system per- formance by simulating a real-world workload using aug- mented Azure inference traces (Appendix. C). To simulate this workload, we replayed user requests from these traces DSLA- and measured the normalized end-to-end latency of the requests.
DSLA [25%]]
60.9%
$6.9%
78.7%
74.9%
43.2%
46.2 5
24.9%
56.3%
DSLA [509]
68.1%
35.4%
75.0%%
71.6%
40.4%
43.2 9
25.4%
54.29
Thansla mor
Sarcyy (meil
4830
Corvncte 25%
By detecting the load, based on the prompt lengths of re- quests seen, the DSLA-Serve automatically adjusts the con- version rate and progressively converts self-attention layers to DSLA layers, ranging from 25% to 50%. Table presents the prompt length distribution from the trace and the corre- sponding conversion rates we applied.
Corvettd 30%
Panslamer
325
250
515
1557
2022
4338
1132
Promttooyth
Gonerated Taken Length
Figure 4. Prefill (left) and per-token decoding (right) latency.
nism's 's reliance on past token interactions, resulting in linear cost growth. Conversely, the converted model maintains a constant latency increase by updating only two states for his- torical and recent information, significantly reducing overall latency.
The result shows a significant improvement in the system's 's efficiency. The average end-to-end latency, normalized by token length. improved by 2.29 x, demonstrating the effec- tiveness of DSLAA-Sevve in adapting to dynamic system load while maintaining performance.
During decoding, we observed an interesting phenomenon: the latency fluctuation is significantly higher for the trans- former model and decreased as we converted more layers into DSLA layers. Notably, the latency spikes consistently occurred at the same position (Appendix. H). Through pro- filing, we observed that memory allocation scudaMalloc) occasionally took 300-1000 ms, while a single inference typically takes only 30-60 ms. Since this issue is related to memory allocation, the fluctuations reduce more self-attention layers to DSLA layers. This is because DSLA reduces memory usage by eliminating the necessity for a key-value (KV) cache.
Table 4. Prompt length distribution and maximum conversion rate during the experiment. DSLA-Serve improves end-to-end per-token generation latency by 2.29
Prompt Length
Distribution
Max Conv. Rate
convert
seq len
64.68%
12.5%
2k<
16. 16%
25%
4k<
16.03%
37.5%
Ske
3.1%
50%
Latency (before) Latency (after)
63.64ms 40.83ms
DSLA avoids maintaining a KV cache, whose size grows linearly with the context length and the number of layers

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
5.5. Further Investigation and Ablation Studies
ing the number of hidden states. The resulting models re-trained on the CoL.A task recover performance. For configurations with more than two states, we compute the Lcont between consecutive states and average the values.
Importance of specialized hidden states. As discussed in $3, the primary limitation of linear attention lies in its tendency to focus on the most recent tokens while neglecting earlier ones. In Figure 5, we compare the attention patterns of standard self-attention (left) and DSLA (right). Unlike the single-state GLA 1), the recency state and history state in DSLA are specialized to attend to different regions of the input. This specialization enables the history state to effectively capture information from earlier tokens.
Table 5. Ablation study of the number of hidden states. Results are reported in accuracy on the validation set.
States
8
Baseline
Lluma-3.2-1 1.2-IB-tnttttt 78.97
81.19
81.44
80.96
82.83
The results, presented in Table 5, show that our DSLA sig- nificantly improves the ability to replicate the original self- attention mechanism, achieving 2.22 accuracy increase over the single-state GLA. Moreover, increasing the number of hidden states beyond two yields diminishing returns, as the dual-state setup is sufficient to effectively capture both historical and recent information.
Self Attention
Layer
Aetention Scores
85
2,5
0.3
2.0
0.2
42
-.0
011
22-
6.5
0.0
1200
1500
3000
13:00
2000
Position Index
Figure 5. Attention score analysis.
Comparison with hybrid architecture. We measured the latency of Zamba-7B to compare its efficiency. For fair evaluation, we used the author-provided checkpoint and code, adhering to the instructions for environment setup. On average, Zamba-7B is 1.8x slower in prefill across all tested context lengths, 3.0x slower in end-to-end (including generation) latency than the transformer-basee architecture Llama2-7B on an A100 machine. This performance gap is primarily attributed to Zamba-7B's extensive number of activated parameters and self-attention layers. Specif- ically. Zamba-7B employs 13 shared self-attention layers that are repeatedly loaded and utilized, exacerbating the memory bottleneck unless executed on state-of-the-as hard- ware such as the H100 (used in the original paper (Glorioso et al., 2024)), which offers 1.5x higher memory bandwidth compared to A A100 that we used (H100; A A100).
History vs recency per layer. We visualize the y value used as weighting factor that determines how the two hidden states contribute to the output (o: (7 + +1 (1-7) )))) in Figure 6.
weighting factor
0,7
history
recency
0.6
0.5
0.0
0.3
10
15
20
25
30
Layer Index
Figure 6. Hidden states weighting factor visualization. In the figure, the state with the higher value dominates each layer. Specifically, the recency state is predominant in the initial layers (0-3) and the final layers (30-31), whereas the history state takes precedence in the middle layers (4-29).
Effectiveness of sensitivity metric. Here, we discuss the effectiveness of attention entropy ($4.3) as a sensitivity met- ric. For the ablation study on other potential metrics (outlier percentage, and downstream task accuracy), please refer to Appendix F. Figure 7 visualizes the correlation between the metric and its performance impact, where the X-axi: represents the layer index at which self-attention layer is replaced with DSLA layer, and the primary Y-axis shows the resulting performance impact, measured as perplexity on the WikiText-2 dataset. The metric is plotted on the secondary Y-axis. A strong correlation between a metric and performance impact suggests its usefulness in guiding model modifications.
This observation aligns with previous research emphasizing the distinct roles played by different layers in model (Jawa- har et al., 2019; Liu et al., 2024a). Fine-grained control across layers can significantly enhance language model- ing performance. By learning the parameter 7 on layer- by-layer basis, the model dynamically adjusts the relative importance of the two hidden states at each layer. This adaptability enables the model to focus on specific attention patterns or temporal dependencies as required by each layer, ultimately improving overall performance.
Is dual-state enough? We conduct an ablation study to assess the impact of the number of hidden states in GLA. These experiments are performed using Llama-3.2 2- In-Ituct on the CoLA task (Wang, 2018). Specifically, we first perform supervised fine-tuning on CoLA, then re- place the self- layers with GLA layers while vary-
Comparison with different linearization method In Ta- ble 7, we compare our method with LoLCATS (Zhang et al., 2024a) on lm.eval tasks. For a fair comparison, we ap- plied our approach to the same teacher model (Llama3-8B) used by LoLCATs. While LoLCATs leverages low-rank adaptation to recover performance lost during linearization,
8

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Table Results of long context understanding. We report 25%, 50% converted layers of DSLAA. DSLA distilled from -lama2-7B---- BB-chat, an instruction tuned model.
Multi-doc QA
Code Understanding
Few-shot Learning
Methods
Cost
HotpotQA
1
2WikiMQA 1
LCC
Repobench t
TREC
Samsum
TriviaQA
Llama2-7B-chan
2T
29.79
27.15
59.71
48.96
58.33
39.03
86.11
GLA-7B
20B
3.61
6.89
41.26
44.24
28.50
16.94
57.68
Mamba-78
N/A
1.23
0.80
17.56
10.54
11.0
4.55
15.23
7amba-7B-Innsin
IT
24.15
29.96
47.06
48.96
70
40.35
88.89
DSLA [25%]
800M
35.37
32.33
58.11
45.17
55.0
40.07
85.56
DSLA [50%]
800M
33.28
30.71
53.55
41.75
44.0
37.64
73.04
Aftention Entropy
10.0
Attuztiun Lniooyy
Specifically, Table reports results using the Llama-2-7B: chat model. The main observations discussed $5.1 remain consistent, confirming that our method performs robustly across varying backbone models.
Paplexity
8.5
8.0
5
20
30
Limitations To enable fast inference without adding model loading latency to the critical path, DSLA-Serv loads both architectures (i.e., Transformer layers and DSLA layers) into memory. This approach may require additional memory to store the model weights. To mitigate this, we can consider offloading layers and prefetching layers, as the layer replace- ment order is known in advance. Offloading and prefetching of weights at inference time have been explored in prior work (Aminabadi et al., 2022; Cai et al., 2024a; Eliseev & Mazur, 2023), and are regarded as effective techniques for reducing memory Saage-especially if the server overlap prefetching with computation or communication, thereby hiding loading latency.
Layer Index
Figure? 7. Attention entropy as sensitivity metric.
DSLA employs a dual-state mechanism to separately re- tain historical and recent information. Additionally, DSLA- Serve enables dynamic adaptation to accuracy-latency trade- offs. On average, DSLA achieves superior performance. We report the standard deviation of our measurements for reference.
Table 7. Comparison of different lincarization methods. DSLA is distilled from Llama-3 8B.
WG
HS
ARC-E
ARC-C1
MMLU'T
8D
24.1%
80.1%
In terms of accuracy, we empirically observed that perfor- mance begins to degrade when more than 75% of the model is converted, consistent with observations from (Wang et al., 2024a). However, DSLA-Serve dynamically adjusts the conversion rate to meet the service level agreement (SLA). In general, the performance of distilled model can be improved by employing stronger teacher model, as demon- strated in Table 6 and Table 7.
LoLCATs
80.9%
79.7%
80.9%
81.7%
$4.9%
52.8%
70.7%
DSLA
23.2%
77.5%
79.8 12
80.7%
$8.9%
68.9%
727%
Application to different models We evaluated the pro- posed method on a smaller-scale model (1.5B) to demon- strate its scalability. Specifically, we distilled the Microsoft Phi-1.5 (Li et al., 2023) model and compared it with the baseline Transformer architecture, the 1.5B SSM model Phi-Mamba (Bick al., 2024), and our hybrid model. Con- sistent with the results on larger models (Table T-Table 3), our hybrid model outperforms all baselines. These results confirm that our method generalizes well across different model scales.
6. Conclusion
In this work, we proposed DSLA-Serve, an adaptive infer- ence system that distills quadratic complexity self-attention layers to linear-cost DSI.Alayers to reduce inference cost on the fly. For the linear-cost model, we use DSLA a dual- state linear attention highlighting two hidden states that separately hold history and recent context. With the spe- cialized states, DSLA overcomes the limitations of previous linear attention models, which primarily focus on localized context. With DSLA-Serve we achieve 2.3x x end-to-end improvement over the transformer model and 3.0x reduc- tion over the hybrid model while maintaining competitive performance downstream tasks.
Table Comparison of 1.5B-scale models 1m Amveval tasks.
DSLA is distilled from Phi-1.5.
Architocture
Cost
Wat
ARC-Et
ARC-CI
PIQAI
HS*
Transformer
1508
73.4
75.0
48.0
62.6
Phi- Phi-Mamiba
SSM
3.08
71.7
74.0
44.1
75.5
60.2
DSLA
DSLA
800M
72.9
74.5
44.5
75.9
60.5
We provide additional results in Table 6 to demonstrate the generalizability of our method across different architectures.

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Acknowledgements
Ben-Kish, A., Zimerman, I., Abu-Hussein, S., Cohen, N.,
Globerson, A., Wolf, I.., and Giryes, R. Decimamba: Exploring the length extrapolation potential of mamba. arXiv preprint arXiv:2406. 14528, 2024.
We thank the anonymous reviewers for their constructive feedback. Ro and Akella are supported by NSF grants CNS- 2105890 and CNS-2232135, as well as by Cisco Research and Meta. Ro is also supported by IBM PhD Fellowship. This work was supported in parts from Intel (single PI SRS grant).
Bick, A., Li, K. Y., Xing, E. P., Kolter, J. Z., and Gu, A. Transformers to ssms: Distilling quadratic knowledge to subquadratic models, 2024. URL https: htRp::/:: //arxiv.. org/ abs /2408 10189.
Impact Statement
Cai, R., Ro, Y., Kim, G.-W., Wang, P., Ehteshami Bejnordi, B., Akella, A., Wang, Z., et al. Read-me: Refactorizing Ilms as router-decoupled mixture of experts with system co-design. Advances in Neural Information Processing Systems, 37:1 17:116126-116111 116148, 2024a.
This work aims improve the efficiency of inference of LLMs. It has partieular negative social or ethical impact beyond typical efficient LLM serving systems.
References
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10774, 2024b.
Bamba: Inference-Efficien Hybrid Mamba2 Model hug- gingface.co. https: http///gggggggagg/ co/blog/ bamba. [Accessed 05-01-2025).
hug-
A100, N.
nvidia.com.
https://wwt
Child, R., Gray, S., Radford, A., and Sutskever, I. Gen- erating long sequences with sparse transformers. arXiv preprint urXivv 10509, 2019.
nvidia. htt.:///ooteee//////emaaaaa
ooolations//atta--Cntteee///////
wwwss-viiii----1 pdf. [Accessed 26-01-2025).
wridia-al00-daten
4 Choromanski, K., Likhosherstov, V., Dohan, D., Song. X.,
Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Aminabadi, R. Y., Rajbhandari, S., Awan, A. A., Li, C.,, Li, D., Zheng, E., Ruwase, O., Smith, S., Zhang, M., Rasley, J., et al. Deepspeed-inferenee inference: enabling efficient inference of transformer models unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15. IEEE, 2022.
Clement, C. B., Bierbaum, M., O' O'Keeffe, K. P., and Alemi, A. A. On the use of arxiv as dataset, 2019.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:4455.21060 2024.
Azizi, S., Kundu, S., Sadeghi, M. E., and Pedram, M. Mambaextend: A training-free approach to improve long context extension of mamba. In The Thirteenth International Conference on Learning Representations, net/ forum?
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io- awareness, 2022. URL https: http///axxvvvvo/// oba/ 2205.14133
2025. URL httpp::///ppee/e/e////e Ld=IgzRooRp1S.
Eliseev, A. and Mazur, D. Fast inference of mixture-of- experts language models with offloading. arXiv preprint arXiv:2312.17238, 2023.
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and J. LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 3119-3137, Bangkok, Thailand, August 2024. Association for Com- putational Linguistics. doi: 10. 10.18653/ 172. URL https:/ http://ccaaatthoogooooo org// acl-long. 172.
Excahnge, S. Stack exchange data dump, 2024. URL https://archae http:://aachiie..orgdeetttaa org/details/ stackexchange.
Foundation, W. Wikimedia downloads. URL https//
dumps wi mediaa org.
Frantar, E. and Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. In Inter- national Conference on Machine Learning, pp. 10323- 10337. PMLR, 2023.
Beltagy, I. Peters, M. E., and Cohan, A. Long- former: The long-document transformer. arXiv preprint arXiv:2004.00 /5000 2020.
10

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre- trained transformers. arXiv preprint arXiv:2210.17333 2022.
with pagedattention. In Proceedings of the 29th Sym- posium on Operating Systems Principles, pp. 61 611-626, 2023a.
K Kwon, W. Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H. Gonzalez, J. E., Zhang, H., and Stoica, Efficient memory management for large language model serving with pagedattention, 2023b. URL https: /etps://arxiv.. org/ wtw./233/5 06180.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling, arXiv preprint arXiv:2101.00027 2020.
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Inter- national Conference on Machine Learning, pp. 19274- 19286. PMLR, 2023.
Glorioso, P., Anthony, Q. Tokpanov, Y., Whittington, Pi- lault, J., Ibrahim, A., and Millidge, B. Zamba: A compact 7b ssm hybrid model. arXiv preprint arXiv:2405. 2024.
Li, Y., Bubeck, S., Eldan, R., Del Giorno, A., Gunasekar, S., and Lee, Y. T. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.0544 202.000990534,, 2023.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A.,
H100, N. NVIDIA H100 Tensor Core GPU Datasheet resources..vidiiacoom
Ye, H., Cai, T., Lewis. P., and Chen, D. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404. 14469, 2024.
nvidia-tensee coreessoo---o--roe wwteee-puuddatssaae cessed 26-01-2025].
[Ac-
Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedi- gos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev- Shwartz, S., et al. Jamba: A hybrid transformer- language model. arXiv preprint arXiv:2403. 1.Xi::2403.19887., 2024.
Hermann, K. M., Kocisky, T., Grefenstette, E., Espeholt, L., Kay, W., Suleyman, M., and Blunsom, P. Teaching machines to read and comprehend. Advances in neural information processing systems, 28, 2015.
Lin, C.-Y. Rouge: A package for automatic evaluation summaries. Text summarization branches out, pp.
Hinton, G. Distilling the knowledge in a neural network. arXiv preprint arXiv: 1503.0251,, 2015.
74- 2004.
Lin, J., Tang, J., Tang, H., Yang, S., Chen, W.-.., Wang. W.-C., Xiao, G., Dang, X., Gan, C., and Han, S. Awq: Activation-aware weight quantization for on-device Ilm compression and accelcration. Proceedings of Machine Learning and Systems, 6:87-100, 2024.
Jawahar, G., Sagot, B., and Seddah, D. What does bert learn about the structure of language? In ACL 2019-57th Annual Meeting of the Association for Computational Linguistics, 2019.
Kasai, J., Peng, H., Zhang, Y., Yogatama, D., Ilharco, G., Pappas, N., Mao, Y., Chen, W., and Smith, N. A. Fine- tuning pretrained transformers into mns. arXiv preprint arXiv:2103.13076 2021.
Liu, Z., Kong, C., Liu, Y., and Sun, M. Fantastic semantics and where to find them: Investigating which layers of generative Ilms reflect lexical semanties. arXiv preprint 24Xiv:2403001500 2024a.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are mns:: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156 5165. PMLR, 2020.
Liu, Z., Zhao, C., landola, F., Lai, C., Tian, Y., Fedorov, I., Xiong. Y., Chang, E., Shi, Y., Krishnamoorthi, R., et al. Mobilellm: Optimizing sub-billion parameter lan- guage models for on-device use cases. arXiv preprint 14945 2024b.
Kudugunta, S., Kusupati, A., Dettmers, T., Chen, K., Dhillon, I., Tsvetkov, Y., Hajishirzi, H., Kakade, S., Farhadi, A., Jain, P., al. Matformer: Nested transformer for elastic inference. arXiv preprint arXiv22310.07707, 2023.
Loshchilov, I. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05000 2017.
Ma, X., Fang, G., and Wang, X. LIm-pruner: On the struc- tural pruning of large language models. Advances in neural information processing systems, 36:21702-2117 2023.
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Mercat, J., Vasiljevie, I., Keh, S., Arora, K., Dave, A., Gaidon, A., and Kollar, T. Linearizing large language models. arXiv preprint arXiv:2405.000 06640, 2024.
Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikei, B., and Han, S. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406. 10774, 2024.
Narayan, S., Cohen, S. B., and Lapata, M. Don't give me the details, just the summary! topic-aware convolutional neu- ral networks for extreme summarization. arXiv preprint arXiv: 0808, 2018.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302. 2023.
Patel, P., Choukse, E., Zhang, C., Shah, A., Goiri, f., Maleki, S., and Bianchini, R. Splitwise: Efficient generative Ilm inference using phase splitting. In 2024 ACM/IEEE S1st Annual International Symposium on Computer Architec- ture (ISCA), pp. 118-132. TEEE, 2024.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Waleffe, R., Byeon, W., Riach, D., Norick, B., Kor- thikanti, V., Dao, T. Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M., and Catanzaro, B. An empiri- cal study of mamba-based language models, 2024. URL 2400.07887.
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung. M., Grella, M., et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint 2r0.v /3448, 2023.
https: http///axvvvv.. org/
Raffel, C.,, Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
Wang, A. Glue: A multi-task benchmark and analysis plat- form for natural language understanding. arXiv preprint arXiv: 1804.07461, 2018.
Wang. J., Paliotta, D., May, A., Rush, A. M., and Dao, T. The mamba the llama: Distilling and accelerating hybrid models, 2024a. URL httts:://axxiii abs 22///2408.20 15237.
Ren, L., Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.0r 07522, 2024.
Wang, P., Cai, R., Wang, Y., Zhu, J., Srivastava, P., Wang, Z., and Li, P. Understanding and mitigating bottlenecks of state space models through the lens of recency and over- smoothing. arXiv preprint arXiv:2501 100058, 2024b.
Ro, Y., Zhang, Z.., Chidambaram, V., and Akella, A. Opti- mizing transformer inference with selective distillation: Layerwise conversion to linear attention. 2nd Workshop on Hot Topics in System Infrastructure (HotInfra 2024), 2024.
Wang. S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprini arXiv:2006.047668 2020.
See, A., Liu, J., and Manning, C. D. Get to the point: Summarization with pointer-generator networks, arXiv preprint arXiv: 100:: 704.04368, 2017.
Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. International Conference on Machine Learning, pp. 38087-38099. PMLR, 2023a.
Shen, Z., Tao, T., Ma, L., Neiswanger, W., Liu, Z., Wang, H., Tan, B., Hestness, J., Vassilieva, N., Soboleva, D., and Xing, E. Slimpajama-de: Understanding data combi- nations for Ilm training. 2024. URL http::
org/ brg/abs/e abs/ 2309033022 10818.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17453, 2023b.
Sun, H., Chen, Z., Yang, X., Tian, Y., and Chen, B. Tri- force: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint arXiv:2404. 19/2, 2024.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficint train- ing, arXiv preprint arXiv:2312.06663. 2023.
Sun, M., Liu, Z., Bair, A., and Kolter, J. Z. A simple and effective pruning approach for large language models. arXiv preprint (tf::2306.11999) 2023a.
Ye, Z.. Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikci, B., Grover, V., Krishnamurthy, A., and Ceze. L. Flashinfer: Efficient and customiz- able attention engine for Ilm inference serving. arXiv preprint arXiv:2501.01005. 2025. URL https:// a arxiv. www.:oogg/bbss//0/
Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.0862/ 2023b.
12

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
Yu, G.-I,, Jeong, J. S., Kim, G.-W., Kim, S., and Chun, B.- G. Orca: A distributed serving system for (Trranformer- Based] generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pp. 521-538, 2022.
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Al- berti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q. Yang, L., et al. Big bird: Transformers for longer se- quences. Advances in neural information processing 3 13:17283-17297, 2020.
Zhang, M., Arora, S., Chalamala, R., Wu, A., Spector, B., Singhal, A., Ramesh, K., and Ré, C. Lolcats: On low- rank linearizing of large language models. arXiv preprinl 2xX1:24100112544 2024a.
Zhang, M. Bhatia, K., Kumbong, H., and Ré, C. The hedgehog the porcupine: Expressive linear attentions with softmax mimicry, 2024b. URL https: //appsi//xxxv... arrvvv.
org/abs/24022 2402. 04347.
Zhang, Z., Sheng. Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 2023.
Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang. S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
13

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
A. System Overview
V
layer id: 31, sensitvity: 8.78
Step I. Layer- Wise Distillation
{layer senstiiiity 8.88)
Offline Distillation
{ layer jd: 25, sensiivity: 8.89)
Step II. Sensitivity Evaluation
Attentonn
Dual-State Gl.A
[layerjdt sensilivity 10:44)
Step IV. Retrieve distiiled layer with min sensitivity
Req 7
[aayer_k (X: sensilivity 8.78
Online Serving
Req T-1
impor
load threshold && accuracy target ace?
True
Req_1
Tpad-Tprompt_P len
Step III. Start serving with base transformer model
Step Convert transformer layer into Multi-state GLA layer
Figure Overview DSLA-Serve.
We provide visual overview of DSLA-Serve in Figure 8.
We first perform a layer-wise fine-tuning (distillation) offline (Top), generating DSLA substitutes for each Transformer layer and ranking layers by their sensitivity to conversion.
During inference (Bottom), DSLA-Serve automatically replaces layers in ascending order of sensitivity as resource pressure increases, allowing the system adaptively reduce memory and compute overhead with minimal accuracy loss.
B. Detailed Experimental Setup
In this appendix, we provide additional details on the fine-tuning process of the DSLA model. We initialize the query, key, value, and output projection layers by reusing their pre- trained parameters. The history gating layer initialized to close the identity matrix, while the recency gate is randomly initialized using (o, d?). The biases of the gating layers are set zero.
For fine-tuning, et al., 2020), Common Crawl, C4 (Raffel et al., 2019), Github, Wikipedia (Foundation), and StackExchange (Excahnge, 2024). We randomly sample from the first chunk of the training dataset, truncating inputs to a maximum of 4K tokens. Tokenization is performed using LLaMA 2-7B's tokenizer (Touvron et al., 2023).
the Slimpajama dataset (Shen et al., 2024), which comprises Arxiv (Clement et al., 2019), Books (Gao
Fine-tuning conducted using the AdamW W optimizer (Loshchilov, 2017), with a weight decay of 0.01. The learning rate linearly warmed up to le-4 over the first 5% of training steps, followed by cosine decay schedule down 5e-5, We leverage open-source Triton kernel for fine-tuning (Yang et al., 2023).
For fine-tuning the 7B model and measuring latency (Fig.4), use A100 GPUs with 80GB of memory. End-to-end performance evaluation (Table4), downstream task evaluations (Table 1-3), and ablation studies (Table 5) are conducted on A6000 GPUs with 49GB of memory.
terms of fine-tuning cost, fine-tuning single layer of a 7B model on tokens takes approximately hours on 4xA 100 GPUs (80GB). Compared to full pretraining (858 days (Touvron al., 2023)), our fine-tuning costs just 0.07%- one-time expense.
C. Serving Trace Information
To replay real-world serving scenarios, we augmented Azure LLM inference trace 2023 (Patel et al., 2024). The original trace contains a timestamp, the number of context tokens (prompt tokens), and the number of generated tokens. Since the original dataset does have a text prompt and session, we generated text prompt and added session enhance the
14

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
original dataset and use it to replay the real-world chat serving scenario. Figure 9 shows the distribution of chats by number of turns per session and the number of requests per second.
20-
Pecertoge (60
244
Number Tums Smaton
Figure 9. (a) Number of turns per session follows Poisson distribution (A = 3). (b) Varying load defined as the number of user requests real world serving scenario. Across time, the load the fluctuates, indicating the need adaptive de-stressing solution for server.
When we are replaying the conversation, we concatenate history of the session to simulate the multi-turn chat. By doing this, resulting trace's s prompt length shows following distribution.
6000
Frequency
5000
3000
1000
5000
10000
15000
20000
Prompt Length
Figure 10. Prompt length distribution
D. Serving of Mixed Batch
Assume there are four requests waiting in the queue from R1 to R4 (Figure 11). These are from different users. Due to differences in prompt length or other factors, the model used the previous generation step for each request may differ. For instance, R2 with short prompt used transformer, while R4 with very long prompt used 3-layeer-converted model.
If we want to batch two generation requests, RI and R2, together for the next token generation, path divergence could occur at Layer 2. RI needs Dual-state GLA (Layer 2-1), and R2 needs Self Attention (Layer 2-2). The diverged path from different needs resembles a off-uue--ff-experts architecture. As explained in the main manuscript, this can cause a delay because cannot be batched for the operations in Layer 2. However, requests RI and R2 can still be batched together at Layer 1, Layer and Layer 4, thereby maximizing throughput. Additionally, the memory savings achieved at Layer2-1 by conversion are substantial enough to compensate for the delay.
User request used in the Figure is from Chatbot Arena Dataset (Zheng et al., 2023).
E. Latency Fluctuation
To analyze the fluctuation in decoding latency, we ran 2048 token generation tasks multiple times with the same prompt. Interestingly, we observed that the latency fluctuation occurs in the same location across multiple runs. With profiling with
15

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
8
R1: the difference OpenCL. CUDA?
Model Used
each
request
8
R2: hollo
8
hypotenuse triangle divides the hypotenuse into two segments with lengthe the ratio 1:22.
R1
R2
R3
R4
Batch__siee size=2
8
help train concentration and memory/nint. Mecitation: Meditation powerfu reduce stress and improve focus and concentration. help you learn the and redirect your attention breathing:
techniques
Layer1
Layer2-1
Layer2-2
wanders, Deep
R
Generation Request
Layer3
[eriited]
Salf Attention Layer
Make as you progress.
regularly
with
Layer4
Dual-state GLA
Figure 11. Serving of Mixed Batch
Run1
Run2
- Run3
250
187.5
Latency
125
62.5
0
200
400
600
800
1000
1200
1400
1800
1800
2000
Generated Token Length
Figure Lateney fluctuation
pytorch profiler, we observed that this is due to the cudaMal loc operation, which takes 100ms-1000ml intermittently. This unexpected delay can be mitigated by converting more transformer layers into linear layers, which reduces the memory usage by removing the KV cache.
For this experiment, we used an A6000 GPU and llama2-7b model, so the order of lateney could differ from the numbers reported in the main manuscript.
F. Further Discussion on Sensitivity Metric: An Ablation Study
In this section, we compare three sensitivity metrics ($4.3): (a) attention entropy, (b) outlier percentage, and (c) downstream task accuracy. The goal is to determine whether these metrics correlate with performance impact, making them effective sensitivity measures.
For the interpretation of figures, each sensitivity metric is plotted on the secondary Y-axis. A strong correlation between metric and performance impact suggests its usefulness in guiding model modifications.
Other than attention entropy that we used, downstream task accuracy in Figure 13(c) worth further discussion. While the downstream task accuracy provides more direct measure of sensitivity, it can be only known after the distillation is done so could be less practical.
Additionally, its utility varies across tasks, making it less generalizable when a model is used for multiple applications.
16

---

On-the-Fly Adaptive Distillation of Transformer to Dual-State Lincar Attention
10.00
Entropy
10.00
10.00
Accuazyy
0.300
9.75
Poplouy
9.50
9.25
A Sitroyy
Peplesty
9.50
9.25
Raiio
9.50
9:25
0.225
0.250
8.75
875
8:75
0.115
8.50
8:25
X.00
8.00
8.00
0.100
Layer
Layer
Layer
Figure 13. Different possible sensitivity metrics: (a) attention entropy, (b) outlier percentage, and (c) downstream task accuracy. We also observed that attention entropy tends to perform reasonably well across variety of downstream tasks.
Figure 13(c) illustrates this limitation: the bar plot represents Wikitext perplexity, while the line plot shows downstream task performance (Rouge-1 score for summarization). These metrics do not always align, suggesting that task-specific accuracy is not always a suitable proxy for sensitivity. However, if model is fine-tuned for single task, using task accuracy as sensitivity metric can be advantageous and improve performance (Ro et al., 2024).
G. Additional Related Work
G.1. Efficient Inference Techniques
Large language models (LLMs) are known for their high deployment costs, primarily due to the substantial parameter counts and the increasing overhead of the KV cache during generative decoding. To mitigate the parameter-related overhead, numerous studies have been conducted. Broadly, some research has focused developing smaller models with carefully optimized training recipes (Liu et al., 2024b), while others have addressed parameter reduction through sparsity (Frantar & Alistarh, 2023; Ma et al., 2023; Sun et al., 2023a) and quantization (Frantar et al., 2022; Xiao et al., 2023a; Lin et al., 2024). To reduce the KV cache cost, which often memory-boun -buudd during the iterative generation process, a series of works have focused on KV cache compression techniques (Zhang al., 2023; et al., 2024; Tang et al., 2024) to alleviate memory I/O limitations. Additionally, other research has explored optimizing the decoding process itself, such as speculative decoding (Sun al., 2024) and multi-head decoding (Cai al., 2024b). In contrast to these efforts, this work investigates elastic inference through layer-wiis wise conversion of sub- -ubd-uudrttic attention to support long-context processing.
G.2. Knowledge Distillation
Knowledge distillation (Hinton, 2015) an effective technique for transferring knowledge from large models to smaller ones, improving inference efficiency. To address the quadratic complexity of transformers, several initial efforts have focused on distilling transformers into linear attention models (Mercat et al., 2024; Kasai et al., 2021). Recently, (Bick et al., 2024) proposed stage-wise distillation pipeline that effectively distills transformer- based Phi-3 models into hybrid of transformer and Mamba models, with only a slight performance degradation. In addition, (Wang et al., 2024a) further improved efficiency by integrating speculative decoding. (Zhang et al., 2024b) explored the expressiveness of linear attention and introduced a simple polynomial approximations to enhance its capabilities. Previous approaches typically used knowledge distillation as a plug-amt and-play technique, successfully distilling quadratic transformers into sub-quadratic linear attention models. However, in this work, systematically analyze the layer-wiss and head-wise properties of the misalignments in attention patterns between self-attention and linear attention. And then we propose simple and effective scaling strategy to mitigate the suboptimal distillation process caused by these misalignments.
17


Reviews:
[{"review_id": "gO0f3Gsnz5", "reviewer": "Reviewer_9L4x", "summary": "This paper introduces Dual-State Linear Attention (DSLA), a novel attention mechanism designed to mitigate the short-range bias of traditional linear attention methods while maintaining the efficiency benefits necessary for long-context LLM serving. The key idea is to maintain two specialized hidden states, one for capturing historical context and another for tracking recency, allowing DSLA to balance global and local dependencies better than prior linear attention methods. To efficiently integrate DSLA into Transformer architectures, the paper proposes DSLA-Serve, an adaptive on-the-fly distillation framework that selectively replaces Transformer layers with DSLA layers during inference. This conversion is guided by sensitivity-based layer ordering, ensuring that DSLA layers are introduced progressively without degrading performance. The chained fine-tuning strategy ensures that each newly converted layer remains compatible with the previously replaced ones, preserving model consistency across partial conversions. Experiments demonstrate the effectiveness of DSLA.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive comments! Please see below. \n\n**A1. Comparison with RetNet and Mamba**\nWe added more baselines including RetNet [1] and Mamba [2], and measured zero-shot performance on challenging tasks including PiQA, ARC-challenge, Hellaswag (HS), MMLU, Winogrande.  On average, DSLA (ours) outperforms both RetNet and Mamba. \n\nTable A. \n| Model | PiQA | ARC | HS | MMLU | WG | Avg |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| RetNet 6.7B | 77.8 | 39.9 | 72.9 | 26.1 | 66.1 | 56.56 |\n| Mamba 7B | 78.3 | 42.8 | **77.8** | 33.0 | **71.9** | 60.76 |\n| Ours | **78.7** | **43.2** | 75.4 | **46.1** | 69.9 | **62.66** |\n\n**A2. FLOPs Report**\nThanks for the great suggestions. Please note that we primarily reported end-to-end latency in our paper, as FLOPs reflect only the total computation per inference and do not account for dynamic load, memory bandwidth, or parallelism. Following table shows FLOPs measured with  torchprofile [3] by feeding a fixed input of length 12 and performing a single token generation. Since FLOPs scale linearly with the number of activated parameters, Zamba-7B exhibits higher FLOPs due to its repeated activation of shared attention layers. \n\nTable B. \n| Model | Architecture | GFLOPS |\n|---|---|---|\n| Llama-2-7B | Transformer | 158.6 |\n| Mamba 7B | SSM | 125.6 |\n| Zamba 7B | Hybrid | 277.7 |\n| DSLA 7B (100% converted) | Dual-state Linear Attention | 163.5 |\n\nWe will reorganize the figure’s location on our revised version. Thanks for your thoughtful comments!\n\n\n**Reference** \n\n[1] Sun, Yutao, et al. \"Retentive network: A successor to transformer for large language models.\" arXiv preprint arXiv:2307.08621 (2023).\n\n[2] Gu, Albert, and Tri Dao. \"Mamba: Linear-time sequence modeling with selective state spaces.\" arXiv preprint arXiv:2312.00752 (2023).\n\n[3] https://github.com/zhijian-liu/torchprofile", "claims_and_evidence": "Claim: Canonical Single-state linear attention overemphasizes recent tokens.\nEvidence: observation in Fig. 1\n\nClaim: DSLA enables model to flexibly trade off accuracy and efficiency.\nEvidence: Tab. 1,2,3,4", "methods_and_evaluation_criteria": "The normal benchmarks for LLMs are adopted.", "theoretical_claims": "Not applicable.", "experimental_designs_or_analyses": "I haven't spotted any problems with experiment designs. It uses a normal llm setting and compare with well-known baselines like GLA and Zamba.", "relation_to_broader_scientific_literature": "In contrast to previous one-stage linaer attention, this paper proposes a two-stage solution that is flexible for use and efficient.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n\n1. Novelty: the method is novel compared with existing linear attention works.\n2. Practical: since the method does not require re-training, it could be adapted to existing transformers.\n3. Clearity: charts and plots demonstrate ideas clearly.\n\nWeaknesses:\n\nThe authors are advised to add more baselines for comparison, including RetNet and Mamba.", "other_comments_or_suggestions": "The system overview in the appendix should be moved to the main paper for better understanding.", "questions_for_authors": "Apart from latency, is it okay to report and compare FLOPs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LljQJPkDLZ", "reviewer": "Reviewer_RVY3", "summary": "This paper introduces an online distillation framework that dynamically converts Transformer layers to dual state linear attention (DSLA) during inference to improve efficiency for long-context LLM serving. DSLA uses two specialized hidden states to better preserve both historical context and recent information, addressing the short-range bias typical in linear attention architectures. Through a sensitivity-based layer conversion ordering and chained fine-tuning strategy, DSLA-Serve achieves 2.3-3.0× faster inference compared to baselines while maintaining comparable accuracy across various downstream tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive comments! Please see below. \n\n**A1. Intuition behind dual state and model capacity**\nThis work stems from our observation (Fig 1) that a single-state linear attention struggles to capture the full range of contextual information handled by self-attention. To address this, we introduce a dual-state mechanism that separately captures historical and recent context (Fig 5). Similar observations on foundation models’ memorization was also made in contemporary pre-training works (e.g. [6]). Our ablation (Tab 5) shows that more than two states yield limited gains. Despite the shift from self-attention, model capacity remains largely intact due to the dual states' expressiveness, as evidenced by strong results across LongBench (Tab 1), summarization (Tab 2), and Harness (Tab 3).\n\n**A2. Generalization: Application to a different linear architecture**\nDue to the space limit, please check A2 and Table C in our response to Reviewer a2iq.\n\n**A3. Scalability: Application to different model size**\nAs requested, we evaluated our method on both 1.5B and 7B models. Results for the 7B model are in Tables 1–3 of the paper. Table A shows 1.5B results, where our distilled hybrid model outperforms the baseline Transformer (microsoft/phi-1_5) and 1.5B SSM model. Similar trends are observed at the 7B scale.\n\nTable A. \n| Model                | Token | Wino | ARC-e     | ARC-c     | PiQA      | HS |\n|---|---|---|---|---|---|---|\n| Transformer                 | 150B  | 73.4       | 75.0      | 48.0      | 76.6      | 62.6      |\n| SSM [1]                        | 3.0B  | 71.7       | 74.0      | 44.1      | 75.5      | 60.2      |\n| DSLA              | 800M  | **72.93** ± 0.01 | **74.49** ± 0.08 | **44.45** ±0.01 | **75.90**  ± 0.01 | **60.48** ± 0.0049 |\n\n**A4. On performance gain on Code understanding and Multi-doc QA**\nTasks where DSLA performs better: Our method outperforms Zamba on code understanding, likely due to the strong coding understanding of our teacher model (LLaMA-2-7B). For Multi-doc QA, we even surpass the full-attention teacher—possibly because dual-state linear attention alleviates the \"lost in the middle\" issue. To support this, we ran a needle-in-a-haystack task and observed better retrieval when answers appeared between positions 1k–2k (e.g., DSLA retrieved 80% of answers at depth=19%, compared to 40% for the Transformer), Since Multi-doc QA often involves mid-context retrieval, DSLA is particularly effective. \n\nTasks where DSLA performs worse: The performance drop may stem from approximation limitations in certain layers. We view deeper task-specific analysis as promising future work and will highlight these findings in the paper.\n\n**A5. Different deployment scenarios**\nWhile our method targets large-scale serving with multiple users, it also benefits single-user settings like PCs or edge devices. While the need for adaptiveness in large-scale systems arises from temporal variability and spatial imbalance, as explained in Sec. 3.2, the need for adaptiveness in edge devices stems from dynamic changes in power requirements [4], to meet a specific SLO. Static attention models may underperform under such variability, but DSLA adapts to changing hardware limits. To simulate an edge scenario, we ran a 340M model on a P100 GPU (12GB) and saw a 1.2× latency improvement over full attention when generating 1024 tokens from a 128 token prompt.\n\n**A6. Discussion on conversion rate**\nImportantly, our method highlights *progressive and selective conversion* to minimize accuracy degradation. Additionally, our framework allows setting a maximum conversion threshold to ensure accuracy is not significantly impacted. \n\n**A7. On performance deterioration**\nEmpirically, we observed that performance begins to deteriorate when more than 75% of the model is converted, similar to [2]. However, this degradation can be mitigated through additional instruction tuning or Lora as in [3]. \n\n**A8. Training overhead details**\nTraining a single layer of a 7B model on 1B tokens takes ~5 hours on 4×A100 GPUs (80GB). Compared to full pretraining (858 days [5]), our fine-tuning costs just ~0.07%—a one-time expense.\n\n**A9. Statistical error information**\nWe report the statistical error of End-to-end latency in Table B and lm harness result on Table A. We will update other experiments in final version.\n\nTable B. \n| Model     | Latency      |\n|---|---|\n| Llama2-7B | 93.64 ± 5.8  |\n| Zamba-7B  | 122.52 ± 8.6 |\n| DSLA-7B   | 40.83 ± 2.2  |\n\n[1] Transformers to ssms: Distilling quadratic knowledge to subquadratic models\n\n[2] The mamba in the llama: Distilling and accelerating hybrid models\n\n[3] LoLCATs: On Low-Rank Linearizing of Large Language Models\n\n[4] Dynamic-ofa: Runtime dnn architecture switching for performance scaling on heterogeneous embedded platforms\n\n[5] Llama: Open and efficient foundation language models\n\n[6] Leave no context behind: Efficient infinite context transformers with infini-attention", "claims_and_evidence": "Well supported claims:\n1. *Dual-state architecture*: The claim that it maintains both historical and recent context is supported through ablation analysis compared to single-state architecture and attention/weighting factor visualizations. \n2. *Framework adaptability*: The framework is shown to be adaptive on a real-world workload, demonstrating that the conversation rate is adapted to the prompt length which leads to about 2.27x reduction in latency. \n3. *Inference speed*: The authors provide evidence for the greater than 2x speed claims using two different base large language models, Llama2 and Zamba 7B models. In addition, they show latency improvements in a real-world workload scenario. \n\nPartially supported claims:\n\n4. *Comparable quality to baselines*: The results do support that results are comparable on average but the individual results are sometimes worse and are dependent on the conversion rate. This could be problematic as it may non-trivial to find a conversion rate that works well across all datasets.", "methods_and_evaluation_criteria": "Overall, the methods and evaluation criteria are well-chosen and appropriate for validating the paper's claims about improving LLM serving efficiency while maintaining performance. However, I have the following concerns: \n\n1. The authors compare only against a single type of single-state baseline, which is quite simplistic. The gated-based attention models, for example, even though they are single-state could be capturing the historical context better. Ideally, the paper would be stronger if it was providing evidence that the dual state approach is useful for different types of recurrent attention mechanisms. \n\n2. Studying the scaling effect to the performance of the proposed method is missing. There is no evidence that the approach works as the model size increases; the evaluation should demonstrate the relative effect even with limited budget (e.g. 1B, 4B, 7B, 12B). \n\n3. Performing a qualitative analysis would be useful to better understand why there is degradation or improvement in certain datasets. It is somewhat surprising that the proposed method works better than full attention especially in code understanding and multi-doc QA.", "theoretical_claims": "The paper is empirical in nature and doesn't make any explicit theoretical claims. It only uses prior theoretical findings in the complexity analysis part. There are no formal proofs or theoretical arguments regarding the effect of conversion to model capacity and optimality of dual states.", "experimental_designs_or_analyses": "The experimental design has been well-designed for the most part: baseline comparisons, evaluation task diversity, ablation studies, real-world workload results. There a few parts that can be improved in the exposition and analysis:\n\n1. Adding some discussion on the training overhead that is introduced when training the dual-state attention. \n2. It would be useful to experiment with different deployment scenarios to identify in which ones the proposed approach is most beneficial. \n3. A few technical details that are missing from the experiment section such as what inference optimization techniques were used and what are the details of the inference workflow used by each method. \n4. The experimental analysis would benefit by showing some statistical error information to better understand the performance differences and how consistent they are.", "relation_to_broader_scientific_literature": "C1: Dual-state linear attention\n\nThe literature in this area has focused mainly on single state mechanisms (e.g. Mamba, GLA). Using a dual state seems to be novel in this area and should be applicable to different types of methods in principle. Other than that this work makes use of well-established techniques such as GLA and theoretical (complexity) or empirical (recency bias) findings from prior work. \n\nC2: On-the-fly layer conversion and efficient serving\n\nThere are offline methods focusing on reducing the complexity of attention through fine-tuning or layer-wise training (e.g. https://arxiv.org/abs/2103.13076). The on-the-fly conversion idea is interesting and is adaptive to different workloads but it is not compared to and positioned well with respect to other distillation methods (https://arxiv.org/abs/2305.02301), model compression methods (https://openreview.net/pdf?id=MxF0IKJtKW), or other inference optimization methods (flash attention, paged attention, KV caching).", "essential_references_not_discussed": "See my reply above.", "other_strengths_and_weaknesses": "See above.", "other_comments_or_suggestions": "See above", "questions_for_authors": "1. What is the training overhead cost of the proposed approach? Please mention how long does it take to train each layer and the total time it would take to convert the whole model.  \n2. Could the authors elaborate what is the intuition behind the dual attention state. Why it is sufficient and how layer conversion affects model capacity?\n3. After which level of conversion does the quality of the model start to deteriorate?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "wPpIMMptFp", "reviewer": "Reviewer_a2iq", "summary": "This work presents a robust approach to deploying gated linear attention (GLA) in practical production environments, addressing key limitations through two primary innovations.\n\nFirst, the authors address GLA’s strong recency bias, which impairs long-context performance. To mitigate this, they propose a dual-state mechanism. One hidden state models local context with a randomly initialized forget gate, while the other models long-term dependencies with a forget gate initialized closer to one (preserving earlier information). By introducing data-dependent dynamic interpolation and a contrastive penalty loss, the model effectively balances reliance between these two states.\n\nSecond, the paper proposes DSLA-Serve, an adaptive distillation framework that selectively converts Transformer layers into DSLA layers during inference. To guide this conversion, the authors introduce an entropy-based sensitivity metric that identifies less critical layers. The framework dynamically substitutes these layers based on system conditions, ensuring performance stability. The chained fine-tuning strategy is introduced to maintain consistency after layer conversion.\n\nThe experimental results demonstrate that DSLA effectively mitigates recency bias and achieves competitive performance on benchmarks, offering improved efficiency compared to other hybrid models. Additional discussions cover batched inference strategies for further deployment optimization.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful and constructive comments! Please see below. \n\n**A1. Comparison to Other Baselines**\nThanks for the great suggestions. We compare our 1.5B-scale model with the Phi-Mamba-1.5B [1] and 7B-scale model to distilled Mamba [2]. As demonstrated in table A and B, our method consistently outperforms these baselines as well.\n\nFor the 1.5B-scale model, we used the same teacher model– Microsoft's phi-1_5 model and used 800M tokens. We used the same environment setting and dataset reported in the main draft. \n\nTable A. \n\n| Method                   | Token | Winogrande | ARC-e     | ARC-c     | PIQA      | Hellaswag |\n|---|---|---|---|---|---|---|\n| Transformers to SSMs [1] | 3.0B  | 71.7       | 74.0      | 44.1      | 75.5      | 60.2      |\n| DSLA (Ours)              | 800M  | **72.93**  | **74.49** | **44.45** | **75.90** | **60.48** |\n\n\nFor the 7B-scale model, note that we used Llama2-7B while [2] used Llama3 to start with. \n\nTable B. \n\n| Method | PiQA | ARC-c | Hellaswag | MMLU | Winogrande |\n|---|---|---|---|---|---|\n| distilled Mamba [2] | **78.7** | **52.4** | **77.7** | 42.4 | 64.8 |\n| DSLA (Ours) | **78.7** | 42.8 | 74.1 | **48.5** | **67.8** |\n\nAlso, thanks for bringing up Nvidia’s 8B Mamba Hybrid model. It's an interesting baseline for hybrid architectures. We chose Zamba-7B as one of the representative hybrid models because it is an open-source model of similar scale (7B) that can run on our machines. While both Zamba-7B and Nvidia’s Hybrid Mamba share similarities in terms of model architecture, both have limitations due to its fixed architecture, whereas DSLA can dynamically adapt to variability of load during inference.\n \n\n**A2. Generalization: Application to a different linear architecture**\nThanks for the great suggestion! Unfortunately we were not able to get the parameters for GatedDeltaNet as it was not public. Instead, we applied the proposed method to an alternative sub-quadratic model, Mamba, to demonstrate its generalizability. Mamba shares similarities with GLA that we picked in the paper due to the presence of \"selective gate\". We extended Mamba’s selective gate to dual gate and fine-tuned the gate parameters with 400M tokens using the same setting used in our paper. Due to the time constraints, we were only able to convert 25% of layers. Table C shows a zero-shot performance of single state Mamba and dual state Mamba. \n\nTable C. \n\n| Model                  | Winogrande | ARC-e    | ARC-c    | PIQA     | Hellaswag |\n|---|---|---|---|---|---|\n| Mamba-1.3B             | 54.1       | 59.0     | 28.2     | 72.2     | 40.1      |\n| Dual-state Mamba [25%] | **54.4**   | **60.1** | **28.7** | **72.2** | **47.5**  |\n\nThis highlights the generalizability of the DSLA method across various architectures.\n\nReference.\n\n[1] Bick, Aviv, et al. \"Transformers to ssms: Distilling quadratic knowledge to subquadratic models.\" Advances in Neural Information Processing Systems 37 (2024): 31788-31812.\n\n[2] Wang, Junxiong, et al. \"The mamba in the llama: Distilling and accelerating hybrid models.\" Advances in Neural Information Processing Systems 37 (2024): 62432-62457.", "claims_and_evidence": "Figures 3 and 5 provide evidence that the dual-state mechanism effectively enhances long-context capabilities.\n\nTable 4 demonstrates that converting multiple layers to linear attention can maintain performance.\n\nRegarding the efficiency claim, the evaluation is limited to comparisons with Llama2-7B and Zamba-7B, omitting a comparison with other hybrid architectures such as NVIDIA's 7B hybrid Mamba2.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria are appropriate for the problem at hand. The paper evaluates performance on both short- and long-context benchmarks, complemented by efficiency comparisons, aligning well with the intended application.", "theoretical_claims": "This is not a theoretical paper.", "experimental_designs_or_analyses": "Yes, the experimental design and analyses appear generally sound.\n\nAblation studies are relatively thorough, but the impact of chaining order during fine-tuning is not fully explored.", "relation_to_broader_scientific_literature": "This paper leans more toward an industrial focus; however, the proposed dual-state mechanism presents a meaningful methodological contribution to long-context modeling in linear attention. The dual-state design effectively addresses recency bias, a known limitation in linear attention models, by combining specialized hidden states for local and long-term dependencies. This innovation aligns with broader research on improving linear attention mechanisms and contributes to the ongoing effort to enhance efficiency in hybrid softmax/linear attention architectures. By enabling more effective deployment of such hybrids in long-sequence settings, this approach opens new opportunities for scalable long-context modeling.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Overall, I believe this is a solid industrial-style paper that aims to improve efficient long-context modeling. If ICML is receptive to such contributions, I would recommend acceptance.\n\nHowever, the paper lacks comparisons with key baselines, particularly other distilled hybrid models such as LOLCATs, Transformers to SSMs, and Mamba-in-LLaMA. Including these comparisons would provide a more comprehensive evaluation of the proposed method’s effectiveness.", "other_comments_or_suggestions": "No", "questions_for_authors": "Have you explored more recent linear attention variants, such as Gated DeltaNet [ICLR '25]? It would be interesting to investigate whether the proposed dual-state technique is applicable to other types of linear attention. Demonstrating its generalizability could significantly strengthen the paper’s contribution.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zKzJ2Tt1pl", "reviewer": "Reviewer_Gcze", "summary": "This work intends to reduce LLM inference complexity without unacceptable serving quality degradation. The contributions are two folds: (1) dual-state linear attention (DSLA): a variant of gated linear attention with two hidden states for history and recency (2) DSLA-Serve: an online adaptive distillation framework that uses an offline chained fine-tuning recipe and an inference time Transformer layer to DSLA layer replacement strategy guided by a sensitivity-based layer ordering. Evaluations on commonsense reasoning, long-context QA, and text summarization demonstrate that DSLA-Serve yields 2.3× faster inference than Llama2-7B and 3.0× faster than the hybrid Zamba-7B, while retaining comparable performance across downstream tasks", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive comments! Please see below. \n\n**A1. Serving system details**\nWe implemented our inference system on top of the DeepSpeed inference serving framework (MII). In the autoregressive process, for relatively short prefill lengths, our method may be slower than Transformers due to the additional linear projection layers. In addition to the standard K, Q, V, and O projections, DSLA introduces two extra layers to compute G1 and G2 (see Eq. 4 in our paper). This overhead impacts latency when memory is not the bottleneck. However, as we generate more tokens, GPU memory becomes the main bottleneck, and our method begins to outperform Transformers in terms of latency.\n\n**A2. KV cache**\nThank you for the question. To clarify, the KV cache is not entirely “removed” but rather “implicitly integrated” into the hidden states of the DSLA layers. In the original attention layers, the full KV cache is retained, where each KV pair can be considered a hidden state, and the output depends on all such states. In contrast, the DSLA layer maintains only two hidden states, into which the information from all KV pairs is incrementally aggregated (please check the equation 4). This design significantly reduces the memory footprint, thereby improving efficiency.\n\n**A3. Effectiveness on post-trained models**\nThanks for the interesting question! To evaluate our method's effectiveness on reasoning tasks, we tested it on the Qwen/Qwen2.5-Math-1.5B-Instruct model, which is fine-tuned with SFT followed by RL using GRPO. During distillation, we used open-source reasoning datasets [1,2], distilled 10% of the layers, and evaluated them on GSM8K and AIME24. Our model matched the teacher’s accuracy on GSM8K, showing that our method works on post-trained models. \n\nTable A. \n| Model                           | GSM8K |\n|---------------------------------|-------|\n| Qwen/Qwen2.5-Math-1.5B-Instruct | 85.0  |\n| DSLA-Qwen2.5-Math-1.5B-Instruct | 84.9  |\n\nReference.\n\n[1] https://github.com/huggingface/open-r1 \n\n[2] https://huggingface.co/datasets/HuggingFaceH4/numina-deepseek-r1-qwen-7b\n\n[3] Muennighoff, Niklas, et al. \"s1: Simple test-time scaling.\" arXiv preprint arXiv:2501.19393 (2025).", "claims_and_evidence": "1. Claim: DSLA can preserve historical context, which overcomes the short-range bias observed in single-state GLA.\n\nEvidence: Figure 5 shows that the recency state and history state in DSLA are both effective.\n\n2. Claim: Attention entropy is an effective sensitivity metric.\n\nEvidence: Figure 7 shows a strong correlation between the attention entropy of each layer and its performance impact. Appendix F also compares other metrics.\n\n3. Claim: DSLA-Serve achieves significant inference speed-ups while maintaining competitive task performance compared to both full Transformer and hybrid models.\n\nEvidence: Section 5 (supported by Tables 1–3 and Figure 4) reports that the adaptive model runs 2.3× faster than Llama2-7B and 3.0× faster than", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-tailored to the problem. The authors use a broad range of benchmark datasets (e.g., Multi-Document QA, Code Understanding, WikiText-2, Lambada, and summarization tasks like CNN/DailyMail and XSum) that stress both long-context understanding and short-context performance. Hidden states weighting factor clearly shows the tradeoff between history and recency. LLM serving metrics such as memory usage and latency show the system advantage.", "theoretical_claims": "There are no theoretical claims to check.", "experimental_designs_or_analyses": "I have checked the soundness of the experimental designs and analyses.\n1. Section 5.1 evaluates the DSLA models of two conversion rates on perplexity and the performance on long-context understanding.\n2. Section 5.2 evaluates the DSLA models on short-context benchmarks\n3. Section 5.3 evaluates the inference latency and memory usage during prefill and decoding stages. However, it is unclear what the prompt length is for the decoding latency experiment.\n4. Section 5.4 evaluates the end-to-end system performance on open-source LLM serving traces.\n5. Section 5.5 does ablation study on the ratio between history and recency, number of states, efficiency compared with Zamba-7B, and the effectiveness of sensitivity metric.", "relation_to_broader_scientific_literature": "DSLA provides new insights to modeling memory and forget mechanisms. DSLA-Serve provides a new recipe for adaptive inference, which is key to understanding the extremely efficient human reasoning process.", "essential_references_not_discussed": "I didn’t identify essential references not discussed in the paper.", "other_strengths_and_weaknesses": "The experiments are comprehensive. However, the effectiveness of sensitivity is unclear after all the post-training stages including RL.", "other_comments_or_suggestions": "I don’t have other comments or suggestions.", "questions_for_authors": "1. What LLM serving system are you using? Could you explain why shorter generation length can take longer time in the autoregressive process? GPU memory allocation can explain the burst latency but cannot explain my concern.\n2. If I understand correctly, the KV cache of the replaced layer is removed without being incorporated into the hidden states of the new DSLA layer. Have you tried to use the KV cache to further improve the model quality?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents a novel and practically effective approach to improving long-context LLM serving efficiency.
While baseline comparisons and scalability analysis are limited, the empirical results and methodological contributions are strong.
I would like to recommend conditional acceptance pending the following revisions:
1. add comparisons with RetNet, Mamba, and other recent hybrid architectures.
2. Clarify training overhead (e.g., time/resource costs for fine-tuning DSLA layers).
3. Discuss scalability limitations and potential solutions (e.g., GPU memory constraints).
Decision: Accept (poster)