用卷积增强注意力,让小模型高效学会大模型的长文本处理能力。
LAWCAT: Efficient Distillation from Quadratic to Linear Attention with Convolution across Tokens for Long Context Modeling
- 用因果卷积和门控线性注意力,将二次复杂度模型转为线性结构
- 仅用1K长度数据蒸馏,即可在22K上下文保持90%以上准确率
- 适合边缘部署,训练成本低于0.1%,推理速度超FlashAttention-2
尽管Transformer在多个领域达到顶尖性能,但其序列长度相关的二次计算复杂度仍是瓶颈,尤其对延迟敏感的长上下文应用。现有线性复杂度替代方案虽日益强大,但从零训练仍需大量资源。为此,我们提出LAWCAT(Linear Attention with Convolution Across Time),一种新颖的线性化框架,可高效将预训练Transformer的能力迁移至高性能线性注意力架构。LAWCAT引入因果Conv1D层增强局部依赖建模,并采用归一化门控线性注意力提升不同上下文长度下的泛化能力。全面评估表明,仅使用1K长度序列蒸馏Mistral-7B,即可在22K上下文实现超过90%的passkey检索准确率,显著扩展有效上下文窗口。类似地,Llama3.2-1B LAWCAT变体在S-NIAH 1&2&3任务(1K-8K上下文)和BABILong基准(QA2&QA3,0K-16K上下文)上表现优异,所需预训练样本不足预训练模型的0.1%。此外,对于超过8K序列,LAWCAT的prefill速度优于FlashAttention-2。LAWCAT为高性能、长上下文线性模型提供了高效路径,适用于边缘部署,减少对海量长序列训练数据与算力的依赖。
原文摘要 · Abstract (English)
Although transformer architectures have achieved state-of-the-art performance across diverse domains, their quadratic computational complexity with respect to sequence length remains a significant bottleneck, particularly for latency-sensitive long-context applications. While recent linear-complexity alternatives are increasingly powerful, effectively training them from scratch is still resource-intensive. To overcome these limitations, we propose LAWCAT (Linear Attention with Convolution Across Time), a novel linearization framework designed to efficiently transfer the capabilities of pre-trained transformers into a performant linear attention architecture. LAWCAT integrates causal Conv1D layers to enhance local dependency modeling and employs normalized gated linear attention to improve generalization across varying context lengths. Our comprehensive evaluations demonstrate that, distilling Mistral-7B with only 1K-length sequences yields over 90\% passkey retrieval accuracy up to 22K tokens, significantly extending its effective context window. Similarly, Llama3.2-1B LAWCAT variant achieves competitive performance on S-NIAH 1\&2\&3 tasks (1K-8K context length) and BABILong benchmark (QA2\&QA3, 0K-16K context length), requiring less than 0.1\% pre-training tokens compared with pre-training models. Furthermore, LAWCAT exhibits faster prefill speeds than FlashAttention-2 for sequences exceeding 8K tokens. LAWCAT thus provides an efficient pathway to high-performance, long-context linear models suitable for edge deployment, reducing reliance on extensive long-sequence training data and computational resources. Code is released at: https://github.com/zeyuliu1037/LAWCAT
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。