p-MoD: Building Mixture-of-Depths MLLMs via Progressive Ratio Decay (2412.04449v1)

Abstract: Despite the remarkable performance of multimodal LLMs (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.

• The paper presents p-MoD, which integrates TanhNorm, STRing, and PRD to stabilize training and optimize token processing in transformer decoders.
• The approach uses a cosine schedule to progressively reduce vision token retention in deeper layers, yielding about 44.4% TFLOPs and 46.2% KV cache savings.
• Empirical results show that p-MoD matches or surpasses baseline models on 14 benchmarks, highlighting its potential for resource-constrained multimodal AI applications.

Building Efficient Multimodal LLMs with p-MoD

The paper "p-MoD: Building Mixture-of-Depths MLLMs via Progressive Ratio Decay" presents a novel approach for enhancing the efficiency of Multimodal LLMs (MLLMs). The research tackles a significant challenge in the field: the computational and memory costs associated with processing extensive vision tokens in transformer decoders. The authors propose a strategic use of the Mixture-of-Depths (MoD) mechanism, previously explored in LLMs, to alleviate these bottlenecks in MLLMs.

Key Contributions

1. Tanh-gated Weight Normalization (TanhNorm): This technique is introduced to stabilize training when MoD modules are inserted into a pre-trained LLM. By normalizing token weights to be around zero and controlling variance through a hyperparameter, this normalization ensures stable convergence and alleviates numerical instability during inference.
2. Symmetric Token Reweighting (STRing): To fully exploit the constrained amount of multimodal training data, the researchers implement STRing, which applies token reweighting to both selected and skipped tokens. This innovation leverages linguistic supervision signals more effectively, enabling the learning model to better assess token importance.
3. Progressive Ratio Decay (PRD): Reflecting the observation that vision tokens are more redundant in deeper layers of the transformer, PRD progressively reduces the token retention ratio across layers following a cosine schedule. This approach enhances efficiency by prioritizing the processing of the most informative tokens in deeper layers, optimizing the computation load without sacrificing performance.

Implications

The proposed model, referred to as p-MoD (progressive Mixture-of-Depths), integrates these advancements and has been tested against current models like LLaVA-1.5 and LLaVA-NeXT. The empirical results are promising, showcasing that p-MoD matches or even surpasses baseline models in performance across 14 benchmarks while requiring significantly reduced computational resources. Specifically, during inference, p-MoD demonstrates a reduction in TFLOPs and KV cache storage by approximately 44.4% and 46.2%, respectively.

These findings suggest a transformative potential for practical applications where resource constraints may limit the deployment of traditional MLLMs. By decreasing both the computational demand and memory footprint, this approach could foster the deployment of advanced MLLMs in more resource-constrained environments, such as mobile or edge computing devices. Practically, it enables large-scale vision-language tasks to be executed more efficiently, broadening the applicability of multimodal AI.

Future Directions

The paper leaves room for further exploration in several directions:

• Generalization Across Modalities: While the current paper focuses on image data, extending this approach to other modalities, such as video or audio, could provide further efficiency gains.
• Integration with Other Model Architectures: Experimenting with p-MoD in conjunction with other state-of-the-art architectures might yield insights into its versatility and limitations.
• Trade-offs Between Efficiency and Accuracy: Further explorations can optimize the balance between computational efficiency and model accuracy, especially in applications where precision is critical.

Overall, this paper contributes significantly to the field of efficient multimodal AI, providing a foundation for future innovations that could unlock new possibilities in AI-mediated understanding of complex environments.