Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Untangling Component Imbalance in Hybrid Linear Attention Conversion Methods

Martin Benfeghoul, Teresa Delgado, Adnan Oomerjee, Haitham Bou Ammar, Jun Wang, Zafeirios Fountas

TL;DR

This work addresses the risk that hybrid linear-attention conversions can neglect the linear path, yielding misleading attribution that the model uses linear attention when it predominantly relies on sliding-window softmax. It identifies and quantifies this component imbalance, introduces a diagnostic framework, and proposes three remedies—Inference-Time Hybridisation, HedgeCATs (attention-weight transfer + LoRA fine-tuning), and Scheduled Sliding-window Dropout (SSD)—to enforce genuine LA usage while preserving efficiency. Across multiple base models and common benchmarks, the proposed strategies recover most of the base-model performance and restore credible LA contribution, with exponential feature maps and larger φ dimensions further boosting results. The findings offer practical routes to reliable, scalable Transformer conversions and ensure attributional validity for hybrid architectures in long-context tasks.

Abstract

Transformers' quadratic computational complexity limits their scalability despite remarkable performance. While linear attention reduces this to linear complexity, pre-training such models from scratch remains, in most cases, prohibitively expensive. Recent post-training linearisation methods convert pre-trained Transformers to linear models efficiently, often using hybrid approaches that combine linear attention with sliding-window softmax. We identify a critical flaw: existing hybrid methods inadvertently bypass the linear component, relying almost entirely on SWA. Component-level diagnostics reveal this previously undetected behaviour stems from overlooked evaluation practices on common-sense benchmarks. We propose three solutions to ensure balanced component usage: (i) inference-time hybridisation of linear-only conversions with sliding-window softmax; (ii) HedgeCATs, combining attention-weight transfer with targeted LoRA fine-tuning; and (iii) Scheduled Sliding-window Dropout (SSD), which stochastically suppresses the softmax branch during training to prevent component collapse. Our methods maintain computational efficiency while recovering most base model performance and ensuring genuine linear attention adoption, restoring the validity of performance attributions in hybrid conversions.

Untangling Component Imbalance in Hybrid Linear Attention Conversion Methods

TL;DR

This work addresses the risk that hybrid linear-attention conversions can neglect the linear path, yielding misleading attribution that the model uses linear attention when it predominantly relies on sliding-window softmax. It identifies and quantifies this component imbalance, introduces a diagnostic framework, and proposes three remedies—Inference-Time Hybridisation, HedgeCATs (attention-weight transfer + LoRA fine-tuning), and Scheduled Sliding-window Dropout (SSD)—to enforce genuine LA usage while preserving efficiency. Across multiple base models and common benchmarks, the proposed strategies recover most of the base-model performance and restore credible LA contribution, with exponential feature maps and larger φ dimensions further boosting results. The findings offer practical routes to reliable, scalable Transformer conversions and ensure attributional validity for hybrid architectures in long-context tasks.

Abstract

Transformers' quadratic computational complexity limits their scalability despite remarkable performance. While linear attention reduces this to linear complexity, pre-training such models from scratch remains, in most cases, prohibitively expensive. Recent post-training linearisation methods convert pre-trained Transformers to linear models efficiently, often using hybrid approaches that combine linear attention with sliding-window softmax. We identify a critical flaw: existing hybrid methods inadvertently bypass the linear component, relying almost entirely on SWA. Component-level diagnostics reveal this previously undetected behaviour stems from overlooked evaluation practices on common-sense benchmarks. We propose three solutions to ensure balanced component usage: (i) inference-time hybridisation of linear-only conversions with sliding-window softmax; (ii) HedgeCATs, combining attention-weight transfer with targeted LoRA fine-tuning; and (iii) Scheduled Sliding-window Dropout (SSD), which stochastically suppresses the softmax branch during training to prevent component collapse. Our methods maintain computational efficiency while recovering most base model performance and ensuring genuine linear attention adoption, restoring the validity of performance attributions in hybrid conversions.

Paper Structure

This paper contains 40 sections, 4 equations, 4 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Contributions
  3. Background & Related Works
  4. Linear Attention
  5. Hybrid Attention
  6. Linearising Pre-Trained Transformers
  7. Training Interventions to Tackle Component Collapse
  8. Identifying the Issues within Hybrid Conversion Methods
  9. Experimental Set-up
  10. Feature Map $\Phi$
  11. Combining Sliding-Window and Linear Attention
  12. Models and Datasets
  13. Training & LoRA Parameters
  14. Ablation Study of LoLCATs Hybrid Attention Modules at Inference time
  15. Reverting Back to Linear Attention-only
  16. ...and 25 more sections

Figures (4)

  • Figure 1: Measuring the performance of HedgeHog conversion at various points during the conversion process. LoRA refers to a LA-only fine-tuning.
  • Figure 2: Adding SWA at inference time at various points of the conversion process.
  • Figure 3: The average performance of HedgeCATs-trained models for difference amounts of LoRA fine-tuning.
  • Figure 4: Performance comparison of dropout and sliding-window size schedules for a different number of fine-tuning epochs.