Lizard: An Efficient Linearization Framework for Large Language Models

📄 arXiv: 2507.09025v3 📥 PDF

作者: Chien Van Nguyen, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Viet Dac Lai, Haoliang Wang, Jayakumar Subramanian, Ryan A. Rossi, Trung Bui, Nikos Vlassis, Franck Dernoncourt, Thien Huu Nguyen

分类: cs.CL, cs.LG

发布日期: 2025-07-11 (更新: 2025-10-09)

备注: 13 pages


💡 一句话要点

Lizard:一种高效线性化框架,用于加速和优化大型语言模型

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 线性化 亚二次注意力 自适应内存控制 长度泛化 硬件感知 MMLU基准测试

📋 核心要点

  1. Transformer模型在处理长序列时,softmax注意力和KV缓存导致计算和内存瓶颈。
  2. Lizard通过可学习模块增强架构,引入亚二次注意力机制,实现自适应内存控制和长度泛化。
  3. 实验表明,Lizard在MMLU基准测试中显著优于现有方法,并具有卓越的关联召回能力。

📝 摘要(中文)

我们提出了Lizard,一个线性化框架,可以将基于Transformer的预训练大型语言模型(LLM)转换为亚二次架构。由于softmax注意力的二次复杂度以及不断增长的Key-Value(KV)缓存导致推理受限于上下文长度,Transformer在处理长序列时面临严重的计算和内存瓶颈。Lizard通过引入一种亚二次注意力机制来解决这些限制,该机制在保持模型质量的同时,能够紧密地近似softmax注意力。与受限于固定、非自适应结构的先前线性化方法不同,Lizard通过紧凑、可学习的模块来增强架构,从而实现自适应内存控制和强大的长度泛化能力。此外,我们还引入了一种硬件感知算法,可以解决门控注意力中的数值不稳定性,从而加速训练。大量实验表明,Lizard几乎无损地恢复了其教师模型的性能,在5-shot MMLU基准测试中,显著优于以前的方法,提升高达9.4 - 24.5个百分点,并展示了卓越的关联召回能力。

🔬 方法详解

问题定义:大型语言模型(LLMs)基于Transformer架构,在处理长序列时,softmax注意力的计算复杂度和Key-Value (KV)缓存的内存占用呈二次方增长,导致计算和内存瓶颈。现有的线性化方法通常采用固定的、非自适应的结构,限制了模型的性能和泛化能力。

核心思路:Lizard的核心思路是通过引入一种亚二次注意力机制来近似softmax注意力,同时保持模型质量。该机制通过可学习的模块来增强架构,实现自适应的内存控制和长度泛化。此外,Lizard还提出了一种硬件感知算法,用于解决门控注意力中的数值不稳定性,从而加速训练。

技术框架:Lizard框架主要包含以下几个关键模块:1) 亚二次注意力机制:用于近似softmax注意力,降低计算复杂度;2) 可学习的模块:用于自适应地控制内存,提高模型的泛化能力;3) 硬件感知算法:用于解决门控注意力中的数值不稳定性,加速训练。整体流程是将预训练的Transformer模型转换为Lizard架构,然后进行微调。

关键创新:Lizard的关键创新在于其自适应的内存控制和长度泛化能力。与现有的线性化方法相比,Lizard通过可学习的模块来动态地调整内存的使用,从而更好地适应不同的输入序列长度。此外,Lizard的硬件感知算法能够有效地解决门控注意力中的数值不稳定性,提高了训练的稳定性和效率。

关键设计:Lizard的关键设计包括:1) 亚二次注意力机制的具体实现方式,例如采用低秩分解或核方法等;2) 可学习模块的结构和参数设置,例如采用MLP或卷积神经网络等;3) 硬件感知算法的具体实现细节,例如采用梯度裁剪或混合精度训练等;4) 损失函数的设计,例如采用交叉熵损失或对比学习损失等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Lizard在5-shot MMLU基准测试中,显著优于以前的方法,提升高达9.4 - 24.5个百分点。实验结果表明,Lizard能够几乎无损地恢复其教师模型的性能,并具有卓越的关联召回能力。这些结果验证了Lizard框架的有效性和优越性。

🎯 应用场景

Lizard框架可应用于各种需要处理长序列的自然语言处理任务,例如文档摘要、机器翻译、问答系统和对话生成等。通过降低计算复杂度和内存占用,Lizard可以使大型语言模型在资源受限的设备上运行,并提高模型的推理速度和效率。该研究的实际价值在于能够推动大型语言模型在更广泛的领域得到应用,并为未来的模型优化提供新的思路。

📄 摘要(原文)

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.