Effective Distillation to Hybrid xLSTM Architectures
作者: Lukas Hauzenberger, Niklas Schmidinger, Thomas Schmied, Anamaria-Roberta Hartl, David Stap, Pieter-Jan Hoedt, Maximilian Beck, Sebastian Böck, Günter Klambauer, Sepp Hochreiter
分类: cs.LG
发布日期: 2026-03-16
💡 一句话要点
提出一种有效的蒸馏流程,将大型语言模型提炼到混合xLSTM架构,实现性能匹配甚至超越。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 知识蒸馏 xLSTM 大型语言模型 模型压缩 线性化架构
📋 核心要点
- 现有蒸馏方法难以将二次复杂度LLM无损地迁移到亚二次复杂度模型,导致下游任务性能下降。
- 提出一种针对xLSTM架构的蒸馏流程,包含专家模型合并阶段,旨在实现学生模型性能与教师模型对齐。
- 实验表明,该流程能有效蒸馏Llama、Qwen和Olmo模型,在部分任务上甚至超越教师模型性能。
📝 摘要(中文)
本文旨在将基于二次复杂度注意力机制的大型语言模型(LLM)蒸馏到亚二次复杂度线性化架构中。尽管已有大量研究,但蒸馏模型在各种下游任务上的性能通常无法与教师LLM相媲美。本文设定了无损蒸馏的目标,即学生模型和教师模型在一系列任务上的容错胜负率达到一致。为此,我们提出了一种针对基于xLSTM的学生模型的有效蒸馏流程。我们增加了一个额外的合并阶段,将单独线性化的专家模型组合成一个单一模型。通过从Llama、Qwen和Olmo系列中蒸馏基础模型和指令调优模型,证明了该流程的有效性。在许多情况下,基于xLSTM的学生模型恢复了教师模型的大部分性能,甚至在某些下游任务上超过了它。我们的贡献是朝着更节能、更具成本效益的Transformer LLM替代品迈出的重要一步。
🔬 方法详解
问题定义:现有的大型语言模型(LLM)通常基于二次复杂度注意力机制,计算成本高昂。将这些LLM蒸馏到亚二次复杂度模型,如线性化架构,可以降低计算成本,但现有的蒸馏方法往往无法保证学生模型的性能与教师模型相匹配,导致在下游任务上的表现不佳。因此,需要一种有效的蒸馏方法,能够在降低计算复杂度的同时,保持甚至提升模型性能。
核心思路:本文的核心思路是通过一种专门设计的蒸馏流程,将教师LLM的知识有效地迁移到基于xLSTM的学生模型中。该流程的关键在于引入了一个额外的合并阶段,将多个单独线性化的专家模型组合成一个单一模型。这种方法允许学生模型学习到教师模型不同方面的知识,从而提高整体性能。
技术框架:该蒸馏流程包含以下主要阶段:1) 教师模型选择:选择Llama、Qwen或Olmo等大型语言模型作为教师模型。2) 学生模型构建:构建基于xLSTM的线性化架构作为学生模型。3) 专家模型训练:训练多个独立的线性化专家模型,每个模型专注于学习教师模型的部分知识。4) 模型合并:将训练好的专家模型合并成一个单一的学生模型。5) 性能评估:在各种下游任务上评估学生模型的性能,并与教师模型进行比较。
关键创新:本文最重要的技术创新点在于引入了专家模型合并阶段。传统的蒸馏方法通常直接将教师模型的输出作为学生模型的训练目标,而本文通过训练多个专家模型,并将其合并,使得学生模型能够学习到教师模型更全面的知识。这种方法能够有效提高学生模型的性能,使其在某些任务上甚至超越教师模型。
关键设计:具体的模型合并方法未知,论文中可能没有详细描述。损失函数的设计可能包括模仿教师模型输出的损失项,以及鼓励专家模型多样性的损失项。xLSTM的具体结构和参数设置也需要根据具体的任务和数据集进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明,使用该蒸馏流程训练的基于xLSTM的学生模型,在许多情况下能够恢复教师模型的大部分性能,甚至在某些下游任务上超过了教师模型。具体性能数据和对比基线在论文中给出,但摘要中未明确提及具体的性能提升幅度。
🎯 应用场景
该研究成果可应用于各种需要低计算成本和高效率的自然语言处理任务,例如移动设备上的语言模型部署、边缘计算环境下的实时翻译、以及资源受限场景下的文本生成等。通过将大型语言模型蒸馏到更小的xLSTM架构,可以降低能源消耗和硬件成本,从而推动LLM的广泛应用。
📄 摘要(原文)
There have been numerous attempts to distill quadratic attention-based large language models (LLMs) into sub-quadratic linearized architectures. However, despite extensive research, such distilled models often fail to match the performance of their teacher LLMs on various downstream tasks. We set out the goal of lossless distillation, which we define in terms of tolerance-corrected Win-and-Tie rates between student and teacher on sets of tasks. To this end, we introduce an effective distillation pipeline for xLSTM-based students. We propose an additional merging stage, where individually linearized experts are combined into a single model. We show the effectiveness of this pipeline by distilling base and instruction-tuned models from the Llama, Qwen, and Olmo families. In many settings, our xLSTM-based students recover most of the teacher's performance, and even exceed it on some downstream tasks. Our contributions are an important step towards more energy-efficient and cost-effective replacements for transformer-based LLMs.