将大模型注意力机制改为低阶复杂度,大幅降低长文本推理开销。
Lizard: An Efficient Linearization Framework for Large Language Models

- 用可学习模块实现动态内存控制的线性化注意力机制。
- 在5次提示的MMLU上性能提升9.4至24.5分,接近原始模型表现。
- 适配硬件优化,解决门控注意力数值不稳问题,加速训练。
我们提出Lizard,一种将预训练Transformer架构的大型语言模型转换为亚二次复杂度的新框架。传统Transformer在处理长序列时面临计算与内存瓶颈,源于softmax注意力的二次复杂度以及随上下文长度增长的键值缓存导致的内存受限问题。Lizard通过引入近似softmax注意力的亚二次注意力机制,同时保持模型质量。不同于以往固定、非自适应结构的线性化方法,Lizard采用紧凑可学习模块,实现自适应内存控制和强泛化能力。此外,我们设计了一种面向硬件的算法,解决门控注意力中的数值不稳定性,加速训练过程。大量实验表明,Lizard在教师模型性能上近乎无损恢复,在5次提示的MMLU基准测试中显著优于此前方法,最高提升达9.4至24.5分,并展现出优异的联想记忆能力。
原文摘要 · Abstract (English)
We propose Lizard, a linearization framework that transforms pretrained Transformer-based Large Language Models (LLMs) into subquadratic architectures. Transformers faces severe computational and memory bottlenecks with long sequences due to the quadratic complexity of softmax attention and the growing Key-Value (KV) cache that makes inference memory-bound by context length. Lizard addresses these limitations by introducing a subquadratic attention mechanism that closely approximates softmax attention while preserving model quality. Unlike prior linearization methods constrained by fixed, non-adaptive structures, Lizard augments the architecture with compact, learnable modules that enable adaptive memory control and robust length generalization. Moreover, we introduce a hardwareaware algorithm that solves numerical instability in gated attention to accelerate training. Extensive experiments show that Lizard achieves near-lossless recovery of its teacher model's performance, significantly outperforming previous methods by up to 9.4 - 24.5 points on the 5-shot MMLU benchmark and demonstrating superior associative recall.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。