p-MoD : Construction de MLLM à Mélange de Profondeurs via Décroissance de Ratio Progressive

p-MoD: Building Mixture-of-Depths MLLMs via Progressive Ratio Decay

December 5, 2024

Auteurs: Jun Zhang, Desen Meng, Ji Qi, Zhenpeng Huang, Tao Wu, Limin Wang

cs.AI

Résumé

Malgré les performances remarquables des grands modèles de langage multimodaux (MLLMs) dans diverses tâches, les coûts substantiels d'entraînement et d'inférence freinent leur avancée. La majorité des calculs provient du volume écrasant de jetons visuels traités par le décodeur transformateur. Dans cet article, nous proposons de construire des MLLMs efficaces en exploitant le mécanisme de Mélange de Profondeurs (MoD), où chaque couche de décodeur transformateur sélectionne des jetons visuels essentiels à traiter tout en sautant ceux redondants. Cependant, l'intégration de MoD dans les MLLMs est complexe. Pour relever les défis de la stabilité de l'entraînement et de l'inférence ainsi que du manque de données d'entraînement, nous adaptons le module MoD avec deux conceptions novatrices : la normalisation des poids à portes tanh (TanhNorm) et la repondération symétrique des jetons (STRing). De plus, nous observons que les jetons visuels présentent une redondance plus élevée dans les couches plus profondes et concevons donc une stratégie de décroissance progressive du ratio (PRD), qui réduit progressivement le ratio de rétention des jetons couche par couche, en utilisant un calendrier cosinus décalé. Cette conception cruciale libère pleinement le potentiel de MoD, améliorant significativement l'efficacité et les performances de nos modèles. Pour valider l'efficacité de notre approche, nous menons des expériences approfondies avec deux modèles de référence sur 14 bancs d'essai. Notre modèle, p-MoD, égale voire dépasse les performances des modèles de référence, avec seulement 55,6 % de TFLOPs et 53,8 % de stockage cache KV lors de l'inférence, et 77,7 % d'heures GPU lors de l'entraînement.

English

Despite the remarkable performance of multimodal large language models (MLLMs) across diverse tasks, the substantial training and inference costs impede their advancement. The majority of computation stems from the overwhelming volume of vision tokens processed by the transformer decoder. In this paper, we propose to build efficient MLLMs by leveraging the Mixture-of-Depths (MoD) mechanism, where each transformer decoder layer selects essential vision tokens to process while skipping redundant ones. However, integrating MoD into MLLMs is non-trivial. To address the challenges of training and inference stability as well as limited training data, we adapt the MoD module with two novel designs: tanh-gated weight normalization (TanhNorm) and symmetric token reweighting (STRing). Moreover, we observe that vision tokens exhibit higher redundancy in deeper layer and thus design a progressive ratio decay (PRD) strategy, which gradually reduces the token retention ratio layer by layer, employing a shifted cosine schedule. This crucial design fully unleashes the potential of MoD, significantly boosting the efficiency and performance of our models. To validate the effectiveness of our approach, we conduct extensive experiments with two baseline models across 14 benchmarks. Our model, p-MoD, matches or even surpasses the performance of the baseline models, with only 55.6% TFLOPs and 53.8% KV cache storage during inference, and 77.7% GPU hours during training.