p-MoD: Costruzione di Mixture-of-Depths MLLMs tramite Decadimento Progressivo del Rapporto

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

December 5, 2024

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

cs.AI

Abstract

Nonostante le prestazioni straordinarie dei modelli di linguaggio multimodali di grandi dimensioni (MLLMs) su diverse attività, i notevoli costi di addestramento e inferenza ne ostacolano l'avanzamento. La maggior parte della computazione deriva dall'enorme volume di token visivi elaborati dal decoder del trasformatore. In questo articolo, proponiamo di costruire MLLMs efficienti sfruttando il meccanismo Mixture-of-Depths (MoD), in cui ciascuno strato del decoder del trasformatore seleziona i token visivi essenziali da elaborare, evitando quelli ridondanti. Tuttavia, integrare MoD nei MLLMs è complesso. Per affrontare le sfide legate alla stabilità dell'addestramento e dell'inferenza, nonché alla limitata quantità di dati di addestramento, adattiamo il modulo MoD con due nuovi design: normalizzazione dei pesi con gate tangente (TanhNorm) e ripesatura simmetrica dei token (STRing). Inoltre, osserviamo che i token visivi mostrano una maggiore ridondanza nei livelli più profondi e progettiamo quindi una strategia di decadimento progressivo del rapporto (PRD), che riduce gradualmente il rapporto di ritenzione dei token strato per strato, utilizzando un programma cosinusoidale spostato. Questo design cruciale libera appieno il potenziale di MoD, migliorando significativamente l'efficienza e le prestazioni dei nostri modelli. Per convalidare l'efficacia del nostro approccio, conduciamo ampi esperimenti con due modelli di base su 14 benchmark. Il nostro modello, p-MoD, eguaglia o addirittura supera le prestazioni dei modelli di base, con soli il 55,6% dei TFLOPs e il 53,8% della memoria cache KV durante l'inferenza, e il 77,7% delle ore di GPU durante l'addestramento.

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.