💻 computer science

Where to Begin: Efficient Pretraining via Subnetwork Selection and Distillation

本文提出了一种用于预训练小语言模型（SLMs）的高效框架，该框架结合了用于识别结构化稀疏子网络初始化的进化搜索与来自更大模型的知识蒸馏，在实现与标准基准线相当的性能的同时，显著降低了计算成本。

原作者： Arjun Krishnakumar, Rhea Sanjay Sukthanker, Hannan Javed Mahadik, Gabriela Kadlecová, Vladyslav Moroshan, Timur Carstensen, Frank Hutter, Aaron Klein

发布于 2026-01-15

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Arjun Krishnakumar, Rhea Sanjay Sukthanker, Hannan Javed Mahadik, Gabriela Kadlecová, Vladyslav Moroshan, Timur Carstensen, Frank Hutter, Aaron Klein

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大问题：建造一座巨大的图书馆成本极高

想象一下，你想建造一座拥有海量知识的巨大图书馆（即大语言模型，LLM），使其能够回答任何问题。传统上，为了建造这座图书馆，你必须雇佣数百万名员工（参数），给他们每人发一本空白笔记本，并让他们从头开始阅读整个互联网。这需要耗费大量的时间、金钱和电力。

虽然“小语言模型”（SLMs）因为规模较小而构建成本较低，但对于大多数普通研究人员或小型公司来说，从零开始构建它们仍然过于昂贵。他们需要一个捷径。

解决方案：“Whittle”框架

该论文的作者提出了一种构建这些小型图书馆的新方法。他们将这个框架称为 Whittle。与其雇佣新员工并从空白笔记本开始，不如从现有的、知识渊博的巨大图书馆（大型“教师”模型）中，尝试寻找一个能够完美胜任工作的更精简的小型团队。

他们通过三个主要技巧来实现这一点：

1. 寻找“黄金子团队”（子网络选择）

想象一下，这座巨大的图书馆是一个拥有数千名音乐家的庞大管弦乐队。你并不需要整个乐队来演奏特定的曲目，你只需要正确的乐器组。

旧方法： 通常，如果你想要一个小模型，你只需随机挑选一些音乐家，然后祈祷他们能演奏好。
新方法： 作者使用了一种智能搜索工具（称为进化搜索）在庞大的管弦乐队中进行搜寻，寻找那些已经完美掌握了那首曲子的特定音乐家组合。他们不只是随机挑选人，而是寻找已经在承担核心任务的特定“子网络”神经元。
结果： 他们发现，这些经过预先筛选的“子团队”比同等规模的随机组合学习能力更强。

2. “聪明模仿者”法（知识蒸馏）

一旦有了这个小型子团队，他们并不会让这群人仅从互联网中学习。相反，他们会将这群人放在一个由原始巨大图书馆（教师模型）组成的教室里。

运作方式： 教师不仅会说“答案是 A”。教师还会说：“答案是 A，但它非常有可能是 A，稍微有点可能是 B，且不太可能是 C。” 这给了小型团队对世界更丰富的理解。
类比： 这就像一位名厨在教徒弟。大师不仅展示最终的菜肴，还会解释细微的味道和技巧。徒弟的学习速度更快，也犯更少的错误。

3. “搜索空间”（尝试不同的组合）

作者意识到，并非所有的“子团队”都是平等的。有时你需要减少层数（比如移除建筑中的楼层），有时你需要减少宽度（比如移除柱子）。

他们测试了四种不同的切割模型的方法：
- 粗粒度 (Coarse)： 切割大块内容（如移除整个楼层）。
- 细粒度 (Fine-grained)： 切割微小的部分（如移除特定的砖块）。
- 均匀化 (Uniform)： 让整个模型均匀地变小。
- 逐层化 (Layer-wise)： 以不同的方式让模型的不同部分变小。
发现： 他们发现，对于非常小的模型，采用“细粒度”（切割微小部分）效果最好。但对于较大的小模型，采用“粗粒度”（切割大块内容）实际上表现更好。

结果：事半功倍

论文声称，通过使用这种方法，他们可以构建出性能与从头开始训练的模型相当的小型模型，但能大幅降低成本：

速度： 在特定的训练规模下，他们最好的模型达到了标准小型模型相同的智能水平，但所需的计算量（FLOPs）减少了 5.16 倍。
效率： 这就像建造一座既像豪宅一样坚固，却只用了 20% 的砖块和劳动力建造的房子。
开源： 不同于一些大科技公司使用的秘密方法，作者发布了所有的代码和工具（“Whittle”库），以便任何人都能使用这种方法。

总结

可以将这篇论文看作是一份关于如何回收利用一个巨大且昂贵的 AI 模型，将其转化为一个更小、更便宜且高效的模型的指南。他们不是从头造一辆新车，而是取用一辆豪华轿车，仔细拆除不需要的部分，并调整剩余的引擎，使其在消耗更少燃料的同时依然能完美运行，同时利用原豪华车的知识来教授它。

技术摘要：从何处开始：通过子网络选择与蒸馏实现高效预训练

问题陈述

尽管大语言模型（LLMs）通过大规模参数量实现了最先进的性能，但其训练和部署成本对于许多研究团队及资源受限的环境而言过于高昂。小语言模型（SLMs）提供了一种更易获取的替代方案，但其预训练仍需要大量的计算资源。现有的降本策略，如从开源权重 LLM 进行知识蒸馏或进行剪枝，往往依赖于闭源实现（例如 Llama 3.2 变体、Minitron），或者缺乏一个系统化、可复现的框架来发现最优的子网络初始化方式。此外，教师模型提升 SLM 预训练效率的具体机制目前仍不明确。

方法论

作者提出了一个由三个组件构成的框架，旨在显著提高 SLM 预训练的效率：

子网络初始化： 不同于随机初始化，该框架从较大的预训练教师模型中提取结构化稀疏的子网络。作者定义了四种不同的搜索空间，以探索如何选择权重：
- 粗粒度（Coarse）与细粒度（Fine-grained）： 粗粒度选择选取组件的前 $n$ 个条目（例如前 $n$ 个注意力头），而细粒度选择则采样任意索引。
- 统一（Uniform）与逐层（Layer-wise）： 统一配置在所有层应用相同的配置，而逐层配置允许每一层拥有独立的配置。
- 这产生了四个搜索空间：粗粒度统一、粗粒度逐层、细粒度统一以及细粒度逐层。
约束进化搜索： 为了在特定的参数预算内识别高质量的子网络，作者采用了进化搜索算法。
- 参数空间被划分为不相交的“箱体”（bins）（例如 385M–426M, 961M–1.06B），以确保覆盖范围的平衡。
- 在每个箱体内，通过变异和交叉算子对候选架构进行进化。
- 约束强制执行： 使用拒绝采样来确保候选架构严格符合目标箱体的参数数量。
- 适应度函数： 候选架构基于在小规模 Token 预算（2B tokens）上的验证困惑度（validation perplexity）进行评估，从而引导搜索向最优初始化方向进行。
知识蒸馏： 所选出的子网络使用来自其自身教师模型的知识蒸馏进行预训练。
- 损失函数结合了标准的交叉熵与教师和学生 Logit 分布之间的前向 KL 散度。
- 为了提高计算效率，作者使用了 Top-k Logit 蒸馏，将教师分布截断为最显著的 $k$ 个输出（例如 $k=1024$ ）。

核心贡献

系统化的开源框架： 作者推出了 Whittle，这是一个完全开源的库，能够实现从 Hugging Face 模型中提取子网络，支持灵活的搜索空间设计，并能促进从搜索到预训练及蒸馏的整个流程。
子网络初始化策略： 他们证明了从预训练教师模型中提取子网络比随机初始化能提供更好的起点。例如，较小的变体（约 410M 参数）达到与同规模从头训练的 Pythia 模型相同的验证困惑度所需的 FLOPs 少了 1.71 倍。
全面分析： 本文首次系统地比较了知识蒸馏下的子网络初始化与标准交叉熵训练，分析了教师模型大小、搜索空间粒度以及蒸馏超参数的影响。
实证验证： 在 Pythia 系列（6.9B 和 12B 教师模型）上的实验表明，通过进化搜索发现并利用教师权重初始化的最佳模型，在达到同规模 Pythia SLM 的性能的同时，显著减少了浮点运算次数（FLOPs）。

结果

效率增益：
- 对于 410M 参数的模型（Bin 1），该方法实现的验证困惑度与标准 Pythia-410M 预训练相比，减少了 1.71× 的 FLOPs。
- 对于 1B 参数的模型（Bin 2），节省程度增加至 1.75×。
- 对于 2.8B 参数的模型（Bin 3），该方法在匹配与使用 10B token 预算训练的同规模 Pythia-2.8B 模型达到相同验证困惑度时，减少了 5.16× 的 FLOPs。
- 在 100B token 预算下，2.8B 变体仍实现了 1.26× 的 FLOPs 削减，同时获得了比基准模型更低的最终验证困惑度。
搜索空间发现：
- 粒度： 最优搜索空间取决于模型大小。对于较大的变体（Bin 3），粒度最低的空间（粗粒度统一）效果最好。对于较小的变体（Bin 1），更细粒度的空间（如细粒度统一）表现更好。
- 蒸馏： 蒸馏在降低困惑度方面始终优于标准预训练。使用完整的 Logit 分布通常比 Top-k 截断产生更低的困惑度，尽管在主要实验中为了效率使用了 Top-k。
- 搜索指标： 在搜索阶段直接针对困惑度进行优化，优于权重量级或基于激活重要性得分的代理指标。
下游性能： 从超网络（supernet）初始化的子网络在常识和问答基准测试（COPA, MMLU, HellaSwag 等）中，始终优于同规模的随机初始化版本以及原始的 Pythia 模型。值得注意的是，从较小的 6.9B 教师模型中提取的子网络性能往往优于从 12B 教师模型中提取的子网络。

重要性与主张

本文声称是首个关于通过从更大规模的教师模型进行热启动（warm-starting）来预训练学生模型的系统性开源研究。其意义在于为开发高成本效益的 SLM 提供了一条实用且可复现的路径。通过结合用于子网络发现的进化搜索与知识蒸馏，作者证明了 SLM 的预训练速度可以提升高达 9.2×（如果不计入搜索成本）或 5.16×（如果计入搜索成本），对比基准 SLM。

作者强调，他们的方法通过利用现有的开源权重 LLM，避开了传统神经架构搜索（NAS）方法所需的昂贵的超网络预训练过程。他们总结道，虽然该方法非常有效，但未来仍需研究这些初始化策略的缩放法则（scaling laws），并调查教师模型选择（如多语言 vs 单语言）对学生模型性能的影响。作者明确指出，其研究结果为不同规模下的计算最优 SLM 预训练提供了指导方针。