p-MoD: Construção de MLLMs de Misto de Profundidades via Decaimento Progressivo da Razão

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

Resumo

Apesar do desempenho notável dos grandes modelos de linguagem multimodais (MLLMs) em diversas tarefas, os substanciais custos de treinamento e inferência impedem seu avanço. A maioria da computação advém do volume avassalador de tokens de visão processados pelo decodificador transformer. Neste artigo, propomos construir MLLMs eficientes aproveitando o mecanismo de Mistura de Profundidades (MoD), onde cada camada do decodificador transformer seleciona tokens de visão essenciais para processamento, enquanto ignora os redundantes. No entanto, a integração do MoD nos MLLMs é complexa. Para enfrentar os desafios de estabilidade no treinamento e inferência, bem como a limitação de dados de treinamento, adaptamos o módulo MoD com dois projetos inovadores: normalização de peso com portas tanh (TanhNorm) e reponderação simétrica de tokens (STRing). Além disso, observamos que os tokens de visão apresentam maior redundância em camadas mais profundas e, portanto, projetamos uma estratégia de decaimento progressivo de razão (PRD), que reduz gradualmente a taxa de retenção de tokens camada por camada, empregando um cronograma cosseno deslocado. Este projeto crucial libera totalmente o potencial do MoD, impulsionando significativamente a eficiência e o desempenho de nossos modelos. Para validar a eficácia de nossa abordagem, realizamos experimentos extensivos com dois modelos de referência em 14 benchmarks. Nosso modelo, p-MoD, iguala ou até supera o desempenho dos modelos de referência, com apenas 55,6% de TFLOPs e 53,8% de armazenamento de cache KV durante a inferência, e 77,7% de horas de GPU durante o treinamento.

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.