MVA: Linear Attention with High-order Query-Keys Integration and Multi-level Vocabulary Decomposition

We sincerely appreciate the insightful feedback and guidance provided. Based on your suggestions, we have made substantial improvements to our work.

1. Additional Experiments with MVA-SW on Qwen2.5-14B-1M and Qwen2.5-32B
   We have conducted further experiments with MVA-SW on Qwen2.5 models and report the results below.

   Model	MMLU	PIQA	Hellaswag
   Qwen2.5-14B-1M	80.7	85.2	87.3
   → MVA-SW (14B)	77.3	83.8	86.8
   Qwen2.5-32B	83.9	-	85.2
   → MVA-SW (32B)	79.8	82.5	85.0

   We have also included results for Llama3-8B with MVA:

   Model	+Tokens	ARC-C	Winogrande	MMLU	TriviaQA
   Llama3.1-8B	-	79.7	60.5	66.7	77.6
   → MVA (8B)	+4B	60.4	58.2	32.1	63.2

2. LoRA Configuration Details
   As suggested, we explicitly clarify the LoRA settings. Our method only requires replacing Full Attention in Mistral or Llama models with MVA-SW or MVA, followed by fine-tuning. The configurations are detailed in Sections 4.1, 4.2, and 4.3.

   • MVA-SW (General Setting): Fine-tuning QKV and down_proj weights, with rank=128 and α=256 (default α=2×rank).
   • MVA-SW (LoLCATs Setting): Fine-tuning only QKV weights, with rank=8 and α=16 (aligned with LoLCATs).
   • MVA: Fine-tuning QKV and down_proj weights, with rank=256 and α=512.

3. Inference Efficiency and Memory Usage
   While inference efficiency and memory footprint are closely related to training efficiency in linear models, we acknowledge that our previous evidence was insufficient. To address this, we now report inference efficiency and memory usage across different sequence lengths.

   Table: Inference Performance Comparison

   Seq Len	Full Inf Time (s)	Full Inf Mem (GB)	Prefill Time (s)	Gen Latency (ms/token)	Total Mem (GB)
   MVA
   4K	0.14	0.77	0.125	98.8	15.32
   8K	0.26	1.53	0.249	60.3	16.58
   16K	0.51	3.06	0.508	78.8	19.08
   32K	1.08	6.11	1.090	78.8	24.09
   64K	2.25	12.22	2.265	97.3	34.11
   128K	5.08	24.44	7.156	58.1	54.14
   GSA
   4K	0.10	0.76	0.077	93.2	15.30
   8K	0.17	1.53	0.153	40.9	16.55
   16K	0.32	3.06	0.315	48.5	19.06
   32K	0.64	6.11	0.630	63.0	24.07
   64K	1.29	12.22	1.293	90.2	34.08
   128K	2.66	24.44	5.102	38.7	54.11
   Flash-Attention
   4K	0.05	1.26	0.056	21.7	15.79
   8K	0.11	2.53	0.116	27.5	18.54
   16K	0.27	5.05	0.287	46.3	23.55
   32K	0.73	10.11	0.750	92.4	33.55
   64K	2.19	20.22	2.208	220.4	53.57
   128K	OOM	OOM	OOM	OOM	OOM

As shown in the table, the Full Inf Time and Full Inf Mem correspond to inference using the entire sequence without KV cache. In contrast, the Prefill Time refers to inference where the sequence is first prefilled and subsequently decoded using the KV cache. The Gen Latency then represents the per-token decoding time in this cached scenario. There may be minor fluctuations in the Prefill Time results due to the absence of prewarming during testing. However, the Gen Latency of the linear model should be close to a fixed value in theory but not in practice, probably due to some shortcomings of the triton operator in slicing and compiling at different lengths. Overall, our method starts to dominate over flash-attention in all aspects at 64K length.

4. Long Context Experiments
   Following your suggestion, we evaluated our method on LongBench and compared it against GSA and other baselines.

   Model	Qasper	NarrativeQA	QMSum
   Models trained from scratch
   RWKV6	9.2	14.4	1.1
   Mamba	5.6	27.9	0.8
   Mistral	25.8	25.1	5.0
   Finetuned from Mistral 7B on 20B tokens
   RetNet	11.1	0.0	0.0
   GLA	18.4	17.2	9.0
   GSA	18.8	19.2	10.0
   MVA	20.7	20.4	9.58

5. Theoretical Proof Organization
   We greatly appreciate your suggestions regarding the structure of our theoretical proofs. We have now revised the theoretical analysis to follow a structure similar to MetaLA’s appendix:

   • Subsection → Proposition (or Lemma) → Proof → Discussion

Conclusion

We are deeply grateful for your constructive feedback, which has significantly improved our research and manuscript.

We sincerely appreciate your insightful feedback and valuable support, which have greatly benefited us. Below, we provide further clarifications based on your suggestions. Thank you again for your guidance.

1. Scaling to Larger Models
   Due to resource and time constraints, we are unable to conduct fine-tuning experiments on a 70B model. However, to approximate your suggestion, we have conducted experiments with MVA-SW on Qwen2.5-14B-1M and Qwen2.5-32B, as shown in the table below:

   Model	MMLU	PIQA	Hellaswag
   Qwen2.5-14B-1M	80.7	85.2	87.3
   → MVA-SW (14B)	77.3	83.8	86.8
   Qwen2.5-32B	83.9	-	85.2
   → MVA-SW (32B)	79.8	82.5	85.0
   Additionally, we have incorporated MVA into Llama3-8B, and the results are presented below:

   Model	+Tokens	ARC-C	Winogrande	MMLU	TriviaQA
   Llama3.1-8B	-	79.7	60.5	66.7	77.6
   → MVA (8B)	+4B	60.4	58.2	32.1	63.2

Due to time constraints we have only carried out the above experiments, if you require additional experiments or more tests please let us know, we would be happy to talk to you again and receive your guidance, thank you very much!

2. LongBench Evaluation
   For long-sequence tasks, we followed the benchmark selection in the GSA paper and conducted experiments on Qasper, NarrativeQA, and QMSum, all of which are part of LongBench except QuALITY. Consequently, we utilized the LongBench framework to ensure consistency across experiments when comparing against GSA and other methods. However, it appears that GSA may not have used LongBench directly for evaluation. As a result, there are slight discrepancies between our results on Mistral and those reported in the GSA paper. To ensure fair comparisons, we will adopt a normalization strategy: we will adjust the sampling procedure of our Mistral and MVA evaluations such that the resulting Mistral performance aligns more closely with the GSA-reported Mistral results. We will then conduct our MVA experiments under the same conditions.

   Model	Qasper	NarrativeQA	QMSum
   Models trained from scratch
   RWKV6	9.2	14.4	1.1
   Mamba	5.6	27.9	0.8
   Mistral	25.8	25.1	5.0
   Finetuned from Mistral 7B on 20B tokens
   RetNet	11.1	0.0	0.0
   GLA	18.4	17.2	9.0
   GSA	18.8	19.2	10.0
   MVA	20.7	20.4	9.58

3. Loss Curve Analysis
   We have added loss curves from fine-tuning different models in the https://anonymous.4open.science/r/icml25_mva-B2AB provided. The results indicate that MVA exhibits the fastest loss convergence compared to other methods such as GLA, GSA, and MetaLA. Additionally, as training progresses, the loss of MVA gradually approaches that of Mistral, further demonstrating its effectiveness. Please let us know if you need any other convergence curves, thank you very much!

We sincerely appreciate your constructive feedback, which has been instrumental in refining our work. Thank you again for your valuable time and insightful suggestions!

Thank you for your valuable feedback. We sincerely appreciate the time and effort you have taken to review our work. Your insightful suggestions have significantly helped us refine the clarity and presentation of our paper. Below, we address each of your concerns in detail.

1. Clarification on Naming Conventions:

We sincerely apologize for any confusion caused by our last-minute modification of the naming conventions. Our advisor suggested that the abbreviation should follow a three-letter format, similar to GLA and GSA, as MQVA was considered too long. From the paper, it can be derived that MQVA = MVA and MQVSWA = MVA-SW. However, we acknowledge that our original presentation might have been unclear, and we have now carefully revised the notations and paper formatting to ensure clarity.

2. Title:

Thank you for pointing out. We have now corrected the content in \icmltitlerunning to match our paper’s title.

3. Justification for Maintaining or Slightly Improving Performance with Additional Training Tokens:

We appreciate your suggestion to explicitly clarify why fine-tuning with more tokens maintains or slightly improves performance, as this requires a deeper understanding of the method. The reason is as follows:

When combining Sliding Window Attention (SWA) with Linear Attention during fine-tuning, the two branches initially exhibit different learning distributions, with SWA being more dominant. As training progresses, performance first increases and then decreases due to the impact of Linear Attention on the final outcome. However, our method mitigates this issue by either fine-tuning only the introduced additional parameters or selectively truncating gradient information, effectively transforming the declining phase into a slow improvement and highlighting the advantages of our approach.

4. Removal of Unused Equation Numbering:

We acknowledge that numbering equations that are not referenced in the main text may introduce unnecessary complexity. However, we intentionally retained some equation numbers because they facilitate discussions in group meetings and other research communications. That said, we recognize your expertise in academic writing and have followed your suggestion to remove some of the unreferenced equation numbers for a more formal and concise presentation.

5. Notation Refinements:

We have revised our notation for clarity:

• The previous notation has been updated to $\exp(\cdot) R_d$, where $R_d \in \mathbb{R}^{s_w \times 1}$, $s_w$ represents the window size and the values in the $R_d$ matrix are all 1.
• We have replaced sigmoid with $\sigma(\cdot)$ and provided its mathematical definition for consistency.
• The subscripts in our MetaLA equations have been adjusted following your recommendations, resulting in a much more concise and readable formula set. We truly appreciate your guidance in improving the mathematical presentation.

6. Clarification of MVA-2 Level:

As noted in Section 3.2, MVA-2 Level employs multi-level vocabulary decomposition, where the vocabulary is recursively decomposed into two hierarchical levels.

7. Explanation of With and Without SWA Comparisons:

We appreciate your suggestion to enhance the explanations in our table captions. We will carefully incorporate this feedback in future revisions. Below, we provide a rationale for the inclusion of both "with" and "without SWA" comparisons:

Currently, there are two primary research directions for converting full-attention models into linear models:

1. Hybrid architectures, which combine sliding window attention and linear attention either at the intra-layer level or by selectively replacing different layers.
2. Pure linear attention models, such as GSA, which do not mix different mechanisms.

Both approaches share a common challenge: the inherent limitations of linear attention. Since our method offers a theoretically superior approximation to full attention, it demonstrates advantages in both hybrid and purely linear settings.

Another important reason for including these comparisons is that SWA preserves the original model’s local distribution within the window, significantly reducing fine-tuning resource requirements while achieving higher performance. As a result, many recent approaches incorporate sliding window attention or global window attention mechanisms. However, researchers focused on linear models, such as those studying GSA, are particularly interested in evaluating the effectiveness of purely linear attention. Thus, our experiments on MVA without SWA serve to highlight the strengths of our linear component and provide valuable insights for researchers working in this domain.

We truly appreciate your constructive feedback, which has greatly contributed to the refinement of our paper. Your recommendation of Handbook of Writing for the Mathematical Sciences has been particularly beneficial.

Thank you for your insightful feedback. Your suggestions have helped us refine our methodology and validation, especially regarding dataset selection. Below, we provide responses to each of your concerns.

1. Essential References Not Discussed:

1. Conclusion: We emphasise the importance of weighting LA with SWA based on attention score correlations in hybrid models. This is explicitly stated in our paper, along with supporting empirical validation.

   Explanation: Our baseline model indeed integrates linear attention with sliding window attention, as seen in prior work like Based and Infini-Attention. However, these methods typically rely on standard gating mechanisms. We argue that a more effective approach is to weight based on attention score correlations—our proposed soft fusion method. Upon further analysis, we recognize that LoLCATs also employ attention score weighting for linear and sliding window attention. However, our study was also conducted very early not later than LoLCATs, and our method extending softmax or sigmoid-based linear models is more natural and logical compared with the ordinary LA of LoLCATs. Notably, if full attention is computed recursively, both current and historical attention terms are weighted by attention scores. While LoLCATs apply this strategy, they do not explicitly highlight its significance. Furthermore, our approach preserves a softmax-like function, making our method more intuitive and theoretically sound.

2. Perceptron Mechanism: We do not claim the perceptron mechanism as a novel contribution, as it has been widely studied. Instead, we employ it due to constraints in our GSA branch. Since the KV sequence is dynamically compressed to a fixed size, applying a simple softmax may no longer be optimal. Prior studies suggest that softmax acts as a relevance filter, enhancing focus on crucial tokens. However, given the severe compression in GSA or our method, more sophisticated information extraction mechanisms are necessary. This motivates our use of a perceptron, which we empirically validate.

3. ** Prior Work:** Our method shares only minimal overlap with existing work, apart from the well-known Taylor expansion of $e^x$, which is a natural and intuitive formulation. Nevertheless, these prior works provide valuable insights, and we would like to include them as references to strengthen the theoretical grounding of our approach.

2. Addressing Weaknesses:

1. Comparison with LoLCATs on Alpaca-Clean Data:
   LoLCATs did not use SlimPajama but rather Alpaca-Clean. We conducted experiments on Alpaca-Clean, and our results are as follows:

   Model	Training Data	PIQA	ARC-e	ARC-c	HellaSwag	Wino-grande	MMLU	Avg.	Avg. (w/o MMLU)
   Mistral-7B (v0.1)	-	82.1	80.9	53.8	81.0	74.0	62.4	72.4	74.4

| → LoLCATs (LoRA rank=8) | AlpacaClean(+40M) | 81.5 | 81.7 | 54.9 | 80.7 | 74.0 | 51.4 | 70.7 | 74.5 | | → LoLCATs (rank=8)| RedPajama(+40M) | 80.1 | 77.6 | 49.0 | 80.3 | 71.7 | 53.2 | 68.6 | 71.7 | | → MVA-SW (rank=32)| AlpacaClean(+20M)| 82.1 | 81.5 | 54.7 | 81.2 | 74.1 | 52.2 | 70.7 | 74.4 | | → MVA-SW (rank=8)| AlpacaClean(+40M)| 82.3 | 81.9 | 57.6 | 80.2 | 74.0 | 51.6 | 71.2 | 75.2| | → MVA-SW (rank=8)| RedPajama(+40M)| 82.5| 81.5 | 55.7 | 79.7 | 72.9 | 52.4 | 70.8 | 74.5|

These results indicate that, unlike LoLCATs, which require a two-stage process, our method maintains strong performance across datasets with a single-stage fine-tuning approach. In particular, our approach to AlpacaClean fine-tuning performance is much better than LoLCATs.

2. Comparison of GLA/GSA:
   We ensured a fair comparison in all evaluations. The advantage of our MVA (without sliding window) over GSA stems from the incorporation of information across different frequency scales and the extending of memory capacity. Additionally, when applying the method described in Section 3.4, we extended it to both GLA and GSA for a fair comparison.

3. MetaLA vs. Softmax:
   The key limitation of MetaLA compared to softmax is the absence of a K matrix, which is replaced by a gated matrix. This substitution hinders the effective use of the original attention’s K matrix and introduces additional instability in training. Although our paper proposes mitigations, this deviation from standard attention mechanisms poses inherent challenges in fine-tuning convergence.

Last but not least, thank you very much for your guidance and comments, so that we can improve our work!