On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention

TL;DR

This paper introduces DSLA, a dual-state linear attention mechanism, and DSLA-Serve, an adaptive distillation framework that replaces Transformer layers with DSLA layers during inference, significantly improving efficiency while maintaining accuracy.

Contribution

The work presents a novel dual-state linear attention design and an online adaptive distillation method for transforming Transformers into more efficient models with minimal accuracy loss.

Findings

DSLA achieves better context preservation than traditional linear attention.

DSLA-Serve accelerates inference by over 2x compared to Llama2-7B.

The approach maintains performance across various NLP tasks.

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.3x faster inference than Llama2-7B and 3.0x 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. Codes are available at https://github.com/utnslab/DSLA-Serve.