[논문 리뷰] Module-wise Adaptive Distillation for Multimodality Foundation Models

Liang Chen, Jiahui Yu외 6명|arXiv (Cornell University)|2023. 10. 06.

Multimodal Machine Learning Applications인용 수 4

한 줄 요약

이 논문은 다중모态 기초 모델을 위한 모듈 단위 적응형 distillation 방법인 OPTIMA를 제안한다. 이 방법은 각 훈련 단계에서 가장 기여도가 높은 모듈(예: 이미지, 텍스트, 다중모달 인코더 등)을 동적으로 선택하기 위해 다중 팔 랜드미어(MAB) 프레임워크를 사용한다. 손실 감소량을 각 모듈별로 추적하고 수정된 톰슨 샘플링 알고리즘을 적용함으로써, 균일하거나 고정된 모듈 distillation보다 적은 distillation 단계로도 뛰어난 성능을 달성하며, 이미지 캡션 및 시각적 추론 벤치마크에서 최신 기술 수준의 성능을 입증한다.

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.

연구 동기 및 목표

• 대규모 다중모달 기초 모델에서 균일한 계층별 distillation이 정보가 풍부한 모듈을 우선시하지 못하는 비효율적 distillation 문제를 해결하기 위해.
• 목표 작업의 성능에 가장 크게 기여하는 모듈을 식별하고 우선순위를 정하기 위해.
• 고정된 훈련 예산 내에서 기여도 탐색과 고성능 모듈 활용 간의 균형을 맞추는 동적이고 적응적인 distillation 전략을 개발하기 위해.
• 하위 작업에서 손실 감소량 추정을 통한 모듈 수준 기여도 추적 기법이 고정 또는 동일 가중치 distillation 전략보다 더 효과적인 distillation을 가능하게 함을 입증하기 위해.

제안 방법

• 모듈 distillation을 다중 팔 랜드미어(MAB) 문제로 모델링하며, 각 모듈은 팔이며 보상은 고정된 수의 단계 동안 해당 모듈을 distillation한 후의 손실 감소량이다.
• 훈련 중에 기여도가 변화하는 비정적 모듈 기여도를 다루기 위해 수정된 톰슨 샘플링 알고리즘인 OPTIMA를 도입한다.
• OPTIMA는 최근의 보상 기록을 사용해 각 모듈의 변화하는 기여도를 추정함으로써, 각 훈련 라운드에서 가장 유익한 모듈을 적응적으로 선택할 수 있다.
• 모듈의 기여도는 고정된 수의 distillation 단계(P단계/라운드) 후의 손실 감소 비율을 측정함으로써 평가된다.
• 각 모듈에 대해 기대 보상(손실 감소량)의 사후 분포를 유지하고, 이러한 분포에서 샘플링된 값을 기반으로 팔을 선택함으로써 탐색과 이용의 균형을 이룬다.
• 이 방법은 CoCa-Large에 적용되며, 각 라운드에 대해 일부 모듈만 선택적으로 distillation함으로써 계산 비용을 줄이면서도 성능을 유지하거나 향상시킨다.

실험 결과

연구 질문

• RQ1 균일하거나 고정된 모듈 distillation과 비교해 볼 때, distillation에서 기여도 기반의 동적 모듈 선택 전략이 학생 모델의 성능 향상에 기여할 수 있는가?
• RQ2 훈련 중에 모듈의 distillation 손실 기여도는 어떻게 변화하는가? 이러한 변화는 효과적으로 모델링할 수 있는가?
• RQ3 톰슨 샘플링과 같은 적응형 샘플링 전략을 갖춘 다중 팔 랜드미어 프레임워크가 모듈 단위 distillation에서 탐색과 이용의 균형을 효과적으로 유지할 수 있는가?
• RQ4 모듈 별 손실 감소량을 추적하는 것이 하위 작업에서 가장 유익한 모듈을 식별하는 신뢰할 수 있는 신호를 제공하는가?

주요 결과

• 모든 모듈을 균일하게 distillation하는 것보다 이미지 인코더나 다중모달 디코더를 별도로만 distillation하는 것이 훨씬 뛰어난 성능을 보이며, 이는 일부 모듈이 하위 작업 정확도에 더 크게 기여한다는 것을 시사한다.
• OPTIMA의 모듈 팔에 대한 보상 분포는 시간이 지남에 따라 변화하며, 일부 모듈(예: COCO에서의 텍스트 인코더)은 기여도가 증가하는 경향을 보이고, 다른 모듈(예: 이미지 인코더)은 감소하는 경향을 보인다.
• 모듈 팔의 보상 분포 평균과 그 모듈만을 독립적으로 distillation했을 때의 최종 테스트 정확도 사이에 강한 정적 상관관계가 존재함을 확인하여, 기여도 추적 메커니즘의 신뢰성을 입증한다.
• OPTIMA는 'Img+Txt+Multi'가 가장 자주 선택되는 것으로 나타나 누적 기여도가 가장 높음을 확인함과 동시에 NLVR2 및 Microsoft COCO Caption 작업에서 베이스라인 distillation 방법보다 뛰어난 성능을 달성한다.
• 이 방법은 표준 계층별 distillation보다 추가적인 계산 비용이나 메모리 오버헤드 없이도 높은 효율성을 유지한다. 각 라운드에서 업데이트되는 모듈은 오직 선택된 모듈들 뿐이기 때문이다.