p-MoD: Construcción de MLLMs de Mezcla de Profundidades mediante Decaimiento de Proporción Progresiva

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

December 5, 2024

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

cs.AI

Resumen

A pesar del notable rendimiento de los modelos de lenguaje multimodales grandes (MLLMs) en diversas tareas, los significativos costos de entrenamiento e inferencia obstaculizan su avance. La mayoría de la computación proviene del abrumador volumen de tokens visuales procesados por el decodificador del transformador. En este artículo, proponemos construir MLLMs eficientes aprovechando el mecanismo de Mezcla de Profundidades (MoD), donde cada capa del decodificador del transformador selecciona tokens visuales esenciales para procesar mientras omite los redundantes. Sin embargo, integrar MoD en MLLMs es complejo. Para abordar los desafíos de estabilidad en el entrenamiento e inferencia, así como la limitada cantidad de datos de entrenamiento, adaptamos el módulo MoD con dos diseños novedosos: normalización de pesos con compuerta tanh (TanhNorm) y reponderación simétrica de tokens (STRing). Además, observamos que los tokens visuales muestran una mayor redundancia en capas más profundas y, por lo tanto, diseñamos una estrategia de decaimiento progresivo de ratio (PRD), que reduce gradualmente la proporción de retención de tokens capa por capa, utilizando un horario de coseno desplazado. Este diseño crucial libera completamente el potencial de MoD, mejorando significativamente la eficiencia y el rendimiento de nuestros modelos. Para validar la efectividad de nuestro enfoque, realizamos experimentos extensos con dos modelos de referencia en 14 pruebas. Nuestro modelo, p-MoD, iguala o incluso supera el rendimiento de los modelos de referencia, con solo un 55.6% de TFLOPs y un 53.8% de almacenamiento de caché KV durante la inferencia, y un 77.7% de horas de GPU durante el entrenamiento.

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.