[论文解读] Module-wise Adaptive Distillation for Multimodality Foundation Models

Liang Chen, Jiahui Yu+6 more|arXiv (Cornell University)|Oct 6, 2023

Multimodal Machine Learning Applications被引用 4

一句话总结

该论文提出OPTIMA，一种用于多模态基础模型的模块级自适应蒸馏方法，采用多臂赌博机框架，在每个训练步骤中动态选择最具贡献的模块（如图像、文本或多模态编码器）进行蒸馏。通过追踪每个模块的损失减少量，并应用改进的汤普森采样算法，OPTIMA在较少蒸馏步骤下实现了更优性能，相较于均匀或固定模块蒸馏方法，在图像字幕生成和视觉推理基准测试中展现出最先进结果。

ABSTRACT

Pre-trained multimodal foundation models have demonstrated remarkable generalizability but pose challenges for deployment due to their large sizes. One effective approach to reducing their sizes is layerwise distillation, wherein small student models are trained to match the hidden representations of large teacher models at each layer. Motivated by our observation that certain architecture components, referred to as modules, contribute more significantly to the student's performance than others, we propose to track the contributions of individual modules by recording the loss decrement after distillation each module and choose the module with a greater contribution to distill more frequently. Such an approach can be naturally formulated as a multi-armed bandit (MAB) problem, where modules and loss decrements are considered as arms and rewards, respectively. We then develop a modified-Thompson sampling algorithm named OPTIMA to address the nonstationarity of module contributions resulting from model updating. Specifically, we leverage the observed contributions in recent history to estimate the changing contribution of each module and select modules based on these estimations to maximize the cumulative contribution. We evaluate the effectiveness of OPTIMA through distillation experiments on various multimodal understanding and image captioning tasks, using the CoCa-Large model (Yu et al., 2022) as the teacher model.

研究动机与目标

• 为解决大型多模态基础模型中蒸馏效率低下的问题，其中均匀的层级蒸馏无法优先考虑最具信息量的模块。
• 识别并优先选择对目标任务学生模型性能贡献最大的模块。
• 开发一种动态自适应蒸馏策略，在固定训练预算下平衡对模块贡献的探索与对高性能模块的利用。
• 证明通过损失减少量估计的模块级贡献追踪，相较于固定或等权重蒸馏策略，能实现更有效的蒸馏。

提出的方法

• 该方法将模块蒸馏建模为多臂赌博机（MAB）问题，其中每个模块为一个臂，奖励为在固定步数内蒸馏该模块后损失的减少量。
• 提出一种改进的汤普森采样算法，称为OPTIMA，以处理训练过程中动态变化的模块贡献。
• OPTIMA通过近期奖励历史估计每个模块的动态贡献，实现在每轮训练中自适应选择最具效益的模块。
• 通过测量在固定数量蒸馏步骤（每轮P步）后损失减少的相对变化来评估每个模块的贡献。
• 维护每个模块预期奖励（损失减少量）的后验分布，并基于从这些分布中采样的值选择臂，以平衡探索与利用。
• 该方法应用于CoCa-Large模型，每轮仅蒸馏部分模块，降低计算成本的同时保持或提升性能。

实验结果

研究问题

• RQ1 与均匀或固定模块蒸馏相比，蒸馏过程中是否可通过动态的、基于贡献的模块选择提升学生模型性能？
• RQ2 模块对蒸馏损失的贡献在训练过程中如何变化？这些变化是否可被有效建模？
• RQ3 采用自适应采样策略（如汤普森采样）的多臂赌博机框架能否在模块级蒸馏中有效平衡探索与利用？
• RQ4 追踪每个模块的损失减少量是否能提供可靠信号，以识别在下游任务中最具效益的蒸馏模块？

主要发现

• 仅蒸馏图像编码器或多模态解码器即可显著优于均匀蒸馏所有模块，表明某些模块对下游任务准确率的贡献更大。
• OPTIMA中模块臂的奖励分布随时间演变，部分模块（如COCO中的文本编码器）贡献呈上升趋势，而另一些模块（如图像编码器）则呈下降趋势。
• 模块臂奖励分布的均值与该模块被单独蒸馏时达到的最终测试准确率之间存在显著正相关，验证了贡献追踪机制的可靠性。
• 与基线蒸馏方法相比，OPTIMA在NLVR2和Microsoft COCO Caption任务上实现了更优性能，其中'Img+Txt+Multi'被最频繁选择，表明其持续提供最高累积收益。
• 该方法保持高效率，除标准层级蒸馏外无额外计算或内存开销，因为每轮仅更新选定模块。