[Paper Review] Module-wise Adaptive Distillation for Multimodality Foundation Models

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

Multimodal Machine Learning Applications4 citations

TL;DR

This paper proposes OPTIMA, a module-wise adaptive distillation method for multimodal foundation models that uses a multi-armed bandit framework to dynamically select the most contributive modules (e.g., image, text, or multimodal encoders) for distillation at each training step. By tracking loss decrement per module and applying a modified Thompson sampling algorithm, OPTIMA achieves superior performance with fewer distillation steps compared to uniform or fixed module distillation, demonstrating state-of-the-art results on image captioning and visual reasoning benchmarks.

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.

Motivation & Objective

• To address the challenge of inefficient distillation in large multimodal foundation models, where uniform layerwise distillation fails to prioritize the most informative modules.
• To identify and prioritize modules that contribute most significantly to student model performance on target tasks.
• To develop a dynamic, adaptive distillation strategy that balances exploration of module contributions and exploitation of high-performing modules under a fixed training budget.
• To demonstrate that module-level contribution tracking via loss decrement estimation enables more effective distillation than fixed or equal-weighted distillation strategies.

Proposed method

• The method models module distillation as a multi-armed bandit (MAB) problem, where each module is an arm and the reward is the loss decrement after distilling that module for a fixed number of steps.
• It introduces a modified Thompson sampling algorithm called OPTIMA to handle nonstationary module contributions that evolve during training.
• OPTIMA estimates the changing contribution of each module using recent reward history, enabling adaptive selection of the most beneficial module at each training round.
• The algorithm evaluates each module’s contribution by measuring the relative decrease in distillation loss after a fixed number of distillation steps (P steps per round).
• It maintains posterior distributions over the expected reward (loss decrement) for each module and selects arms based on sampled values from these distributions to balance exploration and exploitation.
• The method is applied to CoCa-Large, distilling only a subset of modules per round, reducing computational cost while maintaining or improving performance.

Experimental results

Research questions

• RQ1 Can dynamic, contribution-based module selection in distillation improve student model performance compared to uniform or fixed module distillation?
• RQ2 How do module contributions to distillation loss change during training, and can these changes be effectively modeled?
• RQ3 Can a multi-armed bandit framework with adaptive sampling strategies like Thompson sampling effectively balance exploration and exploitation in module-wise distillation?
• RQ4 Does tracking loss decrement per module provide a reliable signal for identifying the most beneficial modules for distillation on downstream tasks?

Key findings

• Distilling only the image encoder or multimodal decoder alone yields significantly better performance than distilling all modules equally, indicating that certain modules contribute more to downstream task accuracy.
• The reward distributions of module arms in OPTIMA evolve over time, with some modules (e.g., text encoder in COCO) showing increasing contribution, while others (e.g., image encoder) show decreasing trends.
• A strong positive correlation exists between the mean of the reward distribution of a module arm and the final test accuracy achieved when that module is distilled exclusively, validating the reliability of the contribution tracking mechanism.
• OPTIMA achieves superior performance on NLVR2 and Microsoft COCO Caption tasks compared to baseline distillation methods, with 'Img+Txt+Multi' being the most frequently selected arm, indicating it consistently offers the highest cumulative benefit.
• The method maintains high efficiency, incurring no additional computational or memory overhead beyond standard layerwise distillation, as only selected modules are updated per round.