On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention for Long-Context LLM Serving

Thank you for the detailed and constructive comments! Please see below.

A1. Serving system details We 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.

A2. KV cache Thank 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.

A3. Effectiveness on post-trained models Thanks 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.

Table A.

Reference.

[1] https://github.com/huggingface/open-r1

[2] https://huggingface.co/datasets/HuggingFaceH4/numina-deepseek-r1-qwen-7b

[3] Muennighoff, Niklas, et al. "s1: Simple test-time scaling." arXiv preprint arXiv:2501.19393 (2025).

Thank you for the thoughtful and constructive comments! Please see below.

A1. Comparison to Other Baselines Thanks 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.

For 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.

Table A.

For the 7B-scale model, note that we used Llama2-7B while [2] used Llama3 to start with.

Table B.

Also, 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.

A2. Generalization: Application to a different linear architecture Thanks 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.

Table C.

This highlights the generalizability of the DSLA method across various architectures.

Reference.

[1] Bick, Aviv, et al. "Transformers to ssms: Distilling quadratic knowledge to subquadratic models." Advances in Neural Information Processing Systems 37 (2024): 31788-31812.

[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.

Thank you for the detailed and constructive comments! Please see below.

A1. Intuition behind dual state and model capacity This 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).

A2. Generalization: Application to a different linear architecture Due to the space limit, please check A2 and Table C in our response to Reviewer a2iq.

A3. Scalability: Application to different model size As 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.

Table A.

A4. On performance gain on Code understanding and Multi-doc QA Tasks 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.

Tasks 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.

A5. Different deployment scenarios While 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.

A6. Discussion on conversion rate Importantly, 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.

A7. On performance deterioration Empirically, 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].

A8. Training overhead details Training 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.

A9. Statistical error information We 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.

Table B.

[1] Transformers to ssms: Distilling quadratic knowledge to subquadratic models

[2] The mamba in the llama: Distilling and accelerating hybrid models

[3] LoLCATs: On Low-Rank Linearizing of Large Language Models

[4] Dynamic-ofa: Runtime dnn architecture switching for performance scaling on heterogeneous embedded platforms

[5] Llama: Open and efficient foundation language models

[6] Leave no context behind: Efficient infinite context transformers with infini-attention

Thank you for the detailed and constructive comments! Please see below.

A1. Comparison with RetNet and Mamba We 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.

Table A.

A2. FLOPs Report Thanks 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.

Table B.

We will reorganize the figure’s location on our revised version. Thanks for your thoughtful comments!

Reference

[1] Sun, Yutao, et al. "Retentive network: A successor to transformer for large language models." arXiv preprint arXiv:2307.08621 (2023).

[2] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023).

[3] https://github.com/zhijian-liu/torchprofile