Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Token Pruning for LLaDA-V

Zhewen Wan, Tianchen Song, Chen Lin, Zhiyong Zhao, Xianpeng Lang

TL;DR

This work tackles the high computational cost of diffusion-based LLaDA-V by revealing that cross-modal semantic alignment concentrates in middle-to-late layers, unlike autoregressive models. It introduces LLaDA-FastV, a spatiotemporal token pruning method that prunes visual tokens at designated middle-to-late layers during the first denoising step and locks them for subsequent steps, yielding substantial FLOPs reductions while preserving task performance. Empirical results across AI2D, MMMU, RealWorldQA, and ChartQA show up to about a 65% reduction in computation with an average of roughly 95% of the original task performance, validating both the pruning strategy and its scalability with more denoising steps. The findings provide a practical acceleration mechanism for diffusion-based multimodal models and motivate future exploration of pre-inference token pruning to further enhance efficiency without compromising semantic fidelity.

Abstract

Diffusion-based large multimodal models, such as LLaDA-V, have demonstrated impressive capabilities in vision-language understanding and generation. However, their bidirectional attention mechanism and diffusion-style iterative denoising paradigm introduce significant computational overhead, as visual tokens are repeatedly processed across all layers and denoising steps. In this work, we conduct an in-depth attention analysis and reveal that, unlike autoregressive decoders, LLaDA-V aggregates cross-modal information predominantly in middle-to-late layers, leading to delayed semantic alignment. Motivated by this observation, we propose a structured token pruning strategy inspired by FastV, selectively removing a proportion of visual tokens at designated layers to reduce FLOPs while preserving critical semantic information. To the best of our knowledge, this is the first work to investigate structured token pruning in diffusion-based large multimodal models. Unlike FastV, which focuses on shallow-layer pruning, our method targets the middle-to-late layers of the first denoising step to align with LLaDA-V's delayed attention aggregation to maintain output quality, and the first-step pruning strategy reduces the computation across all subsequent steps. Our framework provides an empirical basis for efficient LLaDA-V inference and highlights the potential of vision-aware pruning in diffusion-based multimodal models. Across multiple benchmarks, our best configuration reduces computational cost by up to 65% while preserving an average of 95% task performance.

Efficient Token Pruning for LLaDA-V

TL;DR

This work tackles the high computational cost of diffusion-based LLaDA-V by revealing that cross-modal semantic alignment concentrates in middle-to-late layers, unlike autoregressive models. It introduces LLaDA-FastV, a spatiotemporal token pruning method that prunes visual tokens at designated middle-to-late layers during the first denoising step and locks them for subsequent steps, yielding substantial FLOPs reductions while preserving task performance. Empirical results across AI2D, MMMU, RealWorldQA, and ChartQA show up to about a 65% reduction in computation with an average of roughly 95% of the original task performance, validating both the pruning strategy and its scalability with more denoising steps. The findings provide a practical acceleration mechanism for diffusion-based multimodal models and motivate future exploration of pre-inference token pruning to further enhance efficiency without compromising semantic fidelity.

Abstract

Diffusion-based large multimodal models, such as LLaDA-V, have demonstrated impressive capabilities in vision-language understanding and generation. However, their bidirectional attention mechanism and diffusion-style iterative denoising paradigm introduce significant computational overhead, as visual tokens are repeatedly processed across all layers and denoising steps. In this work, we conduct an in-depth attention analysis and reveal that, unlike autoregressive decoders, LLaDA-V aggregates cross-modal information predominantly in middle-to-late layers, leading to delayed semantic alignment. Motivated by this observation, we propose a structured token pruning strategy inspired by FastV, selectively removing a proportion of visual tokens at designated layers to reduce FLOPs while preserving critical semantic information. To the best of our knowledge, this is the first work to investigate structured token pruning in diffusion-based large multimodal models. Unlike FastV, which focuses on shallow-layer pruning, our method targets the middle-to-late layers of the first denoising step to align with LLaDA-V's delayed attention aggregation to maintain output quality, and the first-step pruning strategy reduces the computation across all subsequent steps. Our framework provides an empirical basis for efficient LLaDA-V inference and highlights the potential of vision-aware pruning in diffusion-based multimodal models. Across multiple benchmarks, our best configuration reduces computational cost by up to 65% while preserving an average of 95% task performance.
Paper Structure (8 sections, 4 equations, 4 figures, 1 table)

This paper contains 8 sections, 4 equations, 4 figures, 1 table.

Table of Contents

  1. Introduction
  2. Inefficient Attention in LLaDA-V
  3. Method
  4. Experiments
  5. Experimental Setup
  6. Main Results Analysis
  7. Ablation Study: Impact of Pruning Layer and Strategy
  8. Conclusion

Figures (4)

  • Figure 1: Overview of LLaDA-FastV. We propose a token pruning framework tailored for iterative generative models. Right: By applying FastV at Layer $K$ during the initial step (Step 1), redundant visual tokens (grey blocks) are permanently discarded. Subsequent generation steps (e.g., Step $t$) operate on the reduced sequence, significantly lowering computational costs. Left: Performance comparison across different pruning layers. While aggressive early pruning ($K=7$) leads to hallucinations (highlighted in red), pruning at a deeper layer ($K=15$) maintains semantic fidelity comparable to the baseline (highlighted in green) while reducing FLOPs by $\sim 54\%$.
  • Figure 2: Illustration of inefficient visual attention phenomena.
  • Figure 3: Comparison of Attention Aggregation Patterns.Top Row (LLaDA-V): Attention maps remain sparse in early layers (Layers 1--4), with bright vertical bands (semantic aggregation) emerging only after Layer 16. Bottom Row (LLaVAfastv): In contrast, the autoregressive baseline exhibits immediate, high-intensity attention aggregation starting from Layer 4, indicating early semantic saturation.
  • Figure 4: Ablation study on ChartQA. We compare the impact of pruning strategies (Attention vs. Random) and depth (Layer 3 vs. Layer 15). Layer 3 (Raw-FastV) suffers from immature feature selection, while Random pruning fails to preserve critical details. Our method (Layer 15) effectively balances efficiency and accuracy.