[论文解读] One Student Knows All Experts Know: From Sparse to Dense

Fuzhao Xue, Xiaoxin He+3 more|arXiv (Cornell University)|Jan 26, 2022

COVID-19 diagnosis using AI被引用 5

一句话总结

该论文提出知识融合（knowledge integration），一种新颖的框架，通过两阶段训练将稀疏专家模型（MoE）的知识蒸馏到稠密的学生网络（OneS）中：首先通过SVD-KG或Top-KG进行知识收集，随后进行知识蒸馏。OneS仅使用1500万参数即实现78.4%的ImageNet top-1准确率，保留了61.7%的MoE优势，并实现3.7倍于MoE的推理速度，在NLP任务上相比基线模型最高提升51.7%。

ABSTRACT

Human education system trains one student by multiple experts. Mixture-of-experts (MoE) is a powerful sparse architecture including multiple experts. However, sparse MoE model is easy to overfit, hard to deploy, and not hardware-friendly for practitioners. In this work, inspired by the human education model, we propose a novel task, knowledge integration, to obtain a dense student model (OneS) as knowledgeable as one sparse MoE. We investigate this task by proposing a general training framework including knowledge gathering and knowledge distillation. Specifically, to gather key knowledge from different pre-trained experts, we first investigate four different possible knowledge gathering methods, \ie summation, averaging, Top-K Knowledge Gathering (Top-KG), and Singular Value Decomposition Knowledge Gathering (SVD-KG) proposed in this paper. We then refine the dense student model by knowledge distillation to offset the noise from gathering. On ImageNet, our OneS preserves $61.7\%$ benefits from MoE and achieves $78.4\%$ top-1 accuracy ImageNet with only $15$M parameters. On four natural language processing datasets, OneS obtains $88.2\%$ MoE benefits and outperforms the best baseline by $51.7\%$ using the same architecture and training data. In addition, compared with the MoE counterpart, OneS can achieve $3.7 imes$ inference speedup due to less computation and hardware-friendly architecture.

研究动机与目标

• 通过将稀疏MoE模型的知识迁移至稠密、可部署的学生模型，解决其在实际应用中的局限性，如过拟合、部署困难和硬件效率低下问题。
• 提出一项新任务——知识融合，即单一稠密模型从多个预训练专家中学习，模仿人类从多位教师学习的教育模式。
• 开发一种结合知识收集与知识蒸馏的通用训练框架，以有效初始化并优化来自预训练MoE的稠密学生模型。
• 证明所得到的稠密模型（OneS）可在性能上匹配或超越MoE，同时显著提升效率并更易于部署。

提出的方法

• 知识收集采用四种方法：求和、平均、Top-K知识收集（Top-KG）以及一种新颖的奇异值分解知识收集（SVD-KG），用于提取并融合MoE专家的关键知识。
• SVD-KG通过SVD分解专家权重，保留最重要的奇异分量，这些分量随后用于初始化稠密学生模型的前馈网络（FFN）层。
• Top-KG从各专家中选取激活最活跃的k个神经元或权重分量，用于初始化学生模型，聚焦于高影响力的知识。
• 知识蒸馏通过完整MoE教师模型生成的软标签来应用，以减少知识收集阶段引入的噪声，提升泛化能力。
• 该框架采用两阶段训练流程：首先通过知识收集初始化学生模型，随后使用软标签和交叉熵损失进行知识蒸馏微调。
• 学生模型架构与标准稠密模型相同，确保硬件兼容性并实现高效推理。

实验结果

研究问题

• RQ1 能否通过预训练MoE的知识有效初始化稠密学生模型，在保持其性能的同时提升可部署性？
• RQ2 在求和、平均、Top-KG和SVD-KG四种知识收集方法中，哪一种在从MoE专家初始化稠密学生模型时表现最佳？
• RQ3 知识蒸馏在知识收集后能在多大程度上提升学生模型的性能，特别是能否有效减少初始化阶段引入的噪声？
• RQ4 在准确率、参数效率和推理速度方面，最终的稠密学生模型与MoE相比表现如何？
• RQ5 知识融合的稠密模型是否在视觉和NLP基准测试中均优于现有的蒸馏方法和稀疏模型基线？

主要发现

• 在ImageNet上，OneS仅使用1500万参数即实现78.4%的top-1准确率，保留了61.7%的MoE性能优势。
• 在四个NLP基准测试中，OneS仅使用46%的参数量，便实现了MoE 88.2%的性能优势，且相比最佳基线模型最高提升51.7%。
• 由于计算量减少且采用硬件友好的稠密架构，OneS的推理速度达到MoE的3.7倍。
• 消融实验表明，知识收集与知识蒸馏均不可或缺：任一环节缺失均导致性能显著下降，其中知识收集对初始性能提升贡献最大。
• 在训练轮次相当的情况下，OneS优于基线模型，并且随着训练继续，性能持续提升，而MoE类模型（如WideNet）则趋于饱和。
• 在SQuAD1.1和SST-2等小规模数据集上，OneS甚至超越了MoE教师模型，表明其稠密、可泛化的学生结构有助于减少过拟合。