想象你拥有一座庞大且极其聪明的图书馆(即大型语言模型),它知晓一切。它非常卓越,但体积过于巨大,以至于向一位顾客运送一本书都需要一辆巨大而昂贵的卡车。你希望构建一个更小、更便宜的图书馆版本,使其能装进背包,同时仍要足够聪明以正确回答问题。
这就是蒸馏(Distillation)的问题:尝试在不损失其智慧的前提下,将一颗巨大的大脑压缩进一颗更小的头脑中。
旧方法:“附加组件”问题
以往的方法试图通过在那些巨大而厚重的书籍(即密集骨干网络)旁边添加一个小型、高效的“笔记本”(称为LoRA)来缩小图书馆。
- 问题所在:那些巨大的书籍(密集骨干网络)依然存在,占据了所有空间和重量。你虽然在“训练”笔记本上节省了成本,但在实际“使用”图书馆(推理)时,你仍然不得不搬运这些沉重的书籍。运送卡车依然太大。
新理念:“预算化 LoRA"
作者提出了一种思考这一问题的新方式。他们不再仅仅添加一个笔记本,而是将整个图书馆视为一个拥有严格预算的施工现场。
想象你是一位建筑师,手中只有有限量的混凝土(计算能力)可用于建造最终的图书馆。你有两种类型的材料:
- 实心混凝土块(密集路径):这些是模型中沉重、可靠但昂贵的部分。
- 轻质钢架(低秩路径):这些是新型、高效且灵活的部分。
预算化 LoRA就像一位聪明的工头,他说:“我们只有足够建造原建筑 40% 的混凝土。让我们弄清楚哪些墙壁必须保持实心,哪些可以替换为轻质钢架。”
工作原理(三个步骤)
预算旋钮:你设定一个“预算”(0 到 1 之间的数字)。
- 如果设定较高,你会保留更多沉重的混凝土。
- 如果设定较低,你就被迫用更多钢架替换混凝土。
- 这个旋钮控制着你图书馆的最终规模和速度。
智能替换:在图书馆建造(训练)过程中,系统会自动决定:
- “这面特定的墙(数学的一部分)作为混凝土保留太昂贵了。让我们将其替换为钢架。”
- “这面墙对于理解复杂故事至关重要。让我们将其保留为混凝土。”
- 它并非随机移除内容,而是学习在哪里替换材料以保持图书馆的聪明才智。
最终清理:一旦训练完成,系统会进行最后的清扫。
- 任何几乎未被使用的混凝土都会被完全移除。
- 某些仍然需要但体积较小的混凝土会被粉碎成微小、高效的钢架版本。
- 最终结果是一个物理上更小、运送更快的图书馆,但它依然知道如何完成工作。
他们的发现
研究人员通过将一个 12 层的“巨型”图书馆压缩至 6 层的“背包”版本来测试这一方法。
- 速度与智慧:他们发现了一个“最佳平衡点”。
- 在中等预算下,他们获得了一个运送速度快1.74 倍的图书馆,且智慧几乎没有损失。
- 在激进预算下(允许的混凝土极少),他们获得了一个速度快4 倍的图书馆,智慧仅有小幅下降。
- “功能”测试:他们在特定任务上测试了这些图书馆,例如“阅读此列表并选出第三项”或“将这些单词改为大写”。
- 旧方法(仅添加笔记本)随着教师模型变得更聪明,在这些任务上的表现反而变差。
- 预算化 LoRA在这些“功能”技能上保持了更好的表现。似乎通过仔细决定如何替换材料(混凝土与钢架),图书馆即使在大幅缩小后,仍保留了遵循指令的能力。
核心结论
该论文认为,提高 AI 效率不仅仅在于计算你拥有多少参数(砖块)。关键在于你如何分配你的建筑预算。
通过将模型视为重材料与轻材料的混合体,并在训练过程中强制实施预算约束,你可以构建出一个不仅训练成本更低,而且在现实世界中运行结构速度更快的学生模型。这之间的区别在于:试图通过仅仅拆掉轮子来把沙发塞进车里,与重新设计整车,使其成为一辆紧凑、高效的滑板车,尽管形态不同,却仍能带你到达相同的目的地。
技术摘要:预算化 LoRA
问题陈述
大型语言模型(LLM)通过扩展规模实现了强劲的性能,但这带来了巨大的部署成本。虽然知识蒸馏(KD)通过将知识从大型教师模型迁移到小型学生模型提供了一种解决方案,但现有的流程往往无法生成在推理时具有结构效率的模型。
主要局限性在于参数高效微调方法(如 LoRA,低秩自适应)的广泛使用。虽然 LoRA 通过冻结稠密骨干网络并添加低秩更新来减少可训练参数的数量,但它保留了稠密骨干网络不变。因此,基于 LoRA 的蒸馏产生的学生模型训练成本更低,但并未降低占主导地位的计算成本,因为稠密矩阵乘法依然完整存在。当前的方法通常针对参数量或训练效率进行优化,而未明确控制学生模型架构在推理期间如何分配计算资源。
方法:预算化 LoRA
作者提出了预算化 LoRA,这是一个将模型压缩重新定义为结构化计算分配问题的蒸馏框架。该方法不固定学生模型架构,而是引入一个全局计算预算,规定最终保留的稠密计算比例。
该方法分为三个阶段运行:
预算感知门控 LoRA 模块:
学生模型将选定的线性投影(例如注意力机制中的 q/k/v/o 和 MLP 中的 $up/gate/down$)替换为门控模块。这些模块计算如下:
y=d⋅Wx+rmaxαB((Ax)⊙g)
其中:
- W 是冻结的稠密权重。
- d∈[0,1] 是由预算调度器控制的稠密路径保留系数。
- A,B 是 LoRA 因子,g 代表学习到的每秩门控。
- 该公式允许模型在不同模块中动态调整稠密路径和低秩路径的活跃容量。
预算调度与控制器:
全局预算参数 F∈[0,1] 指定训练结束时目标保留的稠密计算比例。控制器通过以下方式强制执行该预算:
- 使用余弦调度定义目标稠密成本轨迹 b(t)。
- 基于活跃稠密投影的 MAC(乘加运算)计数计算总保留稠密成本 C(t)。
- 贪婪地降低模块级保留系数(d)以满足目标预算,优先移除成本较低的模块。
- 这在训练过程中实现了从稠密路径到低秩路径的受控计算转移。
训练后压缩:
训练结束后,模型经历一个整合步骤以实现实际的推理节省。根据最终保留系数 d,每个模块被映射到三种部署形式之一:
- 情况 1(丢弃稠密): 如果 d<ϵzero,则完全移除稠密路径,仅保留剪枝后的低秩组件。
- 情况 2(SVD 近似): 如果 ϵzero≤d<ϵlr,则通过低秩 SVD 近似剩余的稠密贡献,将其与 LoRA 因子融合。
- 情况 3(合并稠密): 如果 d≥ϵlr,则保留稠密路径并将其与 LoRA 更新合并为标准稠密线性层。
该流程生成了一族由单一预算旋钮 F 参数化的学生模型,实现了效率与性能之间的直接权衡。
主要贡献
- 新公式: 本文提出将蒸馏视为全局预算约束下的结构化计算分配问题,而非固定架构优化。
- 预算化 LoRA 框架: 引入了一种结合控制器强制的稠密保留与学习到的低秩分配的方法,随后进行训练后压缩,以创建部署高效的学生模型。
- 可控权衡: 证明单一预算参数可产生可控的质量 - 效率权衡,在训练和推理成本之间形成帕累托前沿。
- ICL 保留: 表明在函数式上下文学习(ICL)探针上,预算化 LoRA 比基于标准 LoRA 的蒸馏保留了显著更高的性能。
结果
作者使用从 0.13B 模型初始化的 6 层学生模型评估了预算化 LoRA,该模型从不同规模(0.13B、0.5B、1B)的教师模型进行蒸馏。
质量 - 效率权衡:
- 在中等预算(F=0.4)下,预算化 LoRA 在压缩模块推理中实现了1.74 倍的加速,并将训练计算量减少了0.59 倍,同时保持了与标准 KD-LoRA 相当的困惑度。
- 在激进预算(F=0.0)下,该方法在压缩模块推理中实现了4.05 倍的加速,训练计算量为0.38 倍,仅伴随适度的困惑度下降。
- 该方法描绘了一条清晰的帕累托前沿,与标准 LoRA 相比,提供了具有受控困惑度下降的更低成本工作点。
上下文学习(ICL)保留:
- 在函数式 ICL 探针(改编自 Function Vectors 套件)上,预算化 LoRA 显著优于标准 KD-LoRA,特别是在从更大规模教师模型蒸馏时。
- 例如,使用 1B 教师模型时,KD-LoRA 的宏观平均 ICL 得分为 2.7,而预算化 LoRA(F=0.0)达到 5.0,(F=0.4)达到 5.6。
- 这表明保留高层行为不仅取决于匹配困惑度,还取决于通过预算化稠密保留来维持底层计算结构。
意义与主张
本文主张,在计算受限的蒸馏中,保留行为的关键不在于匹配困惑度或简单地移除参数,而在于控制稠密计算如何转移到低秩路径。
作者认为,标准 LoRA 未能实现推理节省,因为它保留了完整的稠密骨干网络。相比之下,预算化 LoRA 将稠密路径和低秩路径视为竞争资源。结果表明:
- 结构化分配至关重要: 与突然移除相比,逐步、预算化的稠密容量减少允许能力更平滑地转移到低秩路径。
- 超越困惑度: 当训练过程显式管理计算结构时,高层行为(如 ICL)能得到更好的保留,即使最终困惑度与标准 LoRA 相似。
- 部署效率: 该方法通过训练后压缩成功地将训练时的预算信号转化为实际的推理时节省,提供了一系列从稠密到完全低秩的模型。
这项工作将自己定位为迈向将效率视为不仅仅是参数量问题,而是训练和蒸馏过程中计算分配问题的关键一步。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。