cs.CVOct 29, 2026

Calibrating Attention-Based Token Pruning for Efficient Multi-Image Synthesis in UMMs

Authors: Zimeng Wu, Qijian Tian, Penghui Du, Changqian Yu, Jiaxin Chen, Yunhong Wang

Organizations: Beihang University · KlingAI Research · Shanghai Jiao Tong University

Abstract

Unified Multimodal Models (UMMs) enable powerful multi-image synthesis but incur substantial computation from long visual contexts. While token pruning offers a direct route to accelerating inference, existing methods are primarily adapted from understanding tasks, suffering from insufficient reference image exploration, confused conditional guidance, and pooling-induced ambiguity that degrades generation quality. To address these issues, we propose CATP, the first training-free context-compression method tailored to UMM generation, which Calibrates the discrepancies of general Attention-based Token Pruning for better quality-efficiency trade-off. Concretely, building on independent layer-wise pruning during denoising, we introduce Guidance-Consistent Pruning (GCP) to synchronize pruning decisions across conditioning branches, and a Taylor-Calibrated Attention Metric (TCAM) that derives a theoretically grounded second-order correction for pooling-induced bias. In addition, a Fast Scoring Operator (FSO) fuses the calibration into tiled online Softmax to reduce pruning cost. Extensive experiments across representative UMMs and benchmarks demonstrate that CATP consistently outperforms state-of-the-art token-reduction methods, achieving up to 2.09×\times attention-core speedup.

Figures & tables

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

May 12, 2026

G2^2TR: Generation-Guided Visual Token Reduction for Separate-Encoder Unified Multimodal Models

The development of separate-encoder Unified multimodal models (UMMs) comes with a rapidly growing inference cost due to dense visual token processing. In this paper, we focus on understanding-side visual token reduction for improving the efficiency of separate-encoder UMMs. While this topic has been widely studied for MLLMs, existing methods typically rely on attention scores, text-image similarity and so on, implicitly assuming that the final objective is discriminative reasoning. This assumption does not hold for UMMs, where understanding-side visual tokens must also preserve the model's capabilities for editing images. We propose G2^2TR, a generation-guided visual token reduction framework for separate-encoder UMMs. Our key insight is that the generation branch provides a task-agnostic signal for identifying understanding-side visual tokens that are not only semantically relevant but also important for latent-space image reconstruction and generation. G2^2TR estimates token importance from consistency with VAE latent, performs balanced token selection, and merges redundant tokens into retained representatives to reduce information loss. The method is training-free, plug-and-play, and applied only after the understanding encoding stage, making it compatible with existing UMM inference pipelines. Experiments on image understanding and editing benchmarks show that G2^2TR substantially reduces visual tokens and prefill computation by 1.94x while maintaining both reasoning accuracy and editing quality, outperforming baselines on almost all benchmarks. Code is at: https://github.com/lijunxian111/G2TR.
Mar 17, 2026

Rethinking UMM Visual Generation: Masked Modeling for Efficient Image-Only Pre-training

Unified Multimodal Models (UMMs) are often constrained by the pre-training of their visual generation components\textbf{visual generation components}, which typically relies on inefficient paradigms and scarce, high-quality text-image paired data. In this paper, we systematically analyze pre-training recipes for UMM visual generation\textbf{UMM visual generation} and identify these two issues as the major bottlenecks. To address them, we propose Image-Only Training for UMMs (IOMM)\textbf{Image-Only Training for UMMs (IOMM)}, a data-efficient two-stage training framework. The first stage pre-trains the visual generative component exclusively\textbf{exclusively} using abundant unlabeled image-only data, thereby removing the dependency on paired data for this costly phase\textbf{for this costly phase}. The second stage fine-tunes the model using a mixture of unlabeled images and a small curated set of text-image pairs, leading to improved instruction alignment and generative quality. Extensive experiments show that IOMM not only improves training efficiency but also achieves state-of-the-art (SOTA) performance. For example, our IOMM-B (3.6B) model was trained from scratch using only ∼1050\sim \textbf{1050} H800 GPU hours (with the vast majority, 1000\textbf{1000} hours, dedicated to the efficient image-only pre-training stage\textbf{image-only pre-training stage}). It achieves 0.89\textbf{0.89} on GenEval and 0.55\textbf{0.55} on WISE--surpassing strong baselines such as BAGEL-7B (0.82 & 0.55) and BLIP3-o-4B (0.84 & 0.50). Code is available https://github.com/LINs-lab/IOMM.
Sep 8, 2025

Reconstruction Alignment Improves Unified Multimodal Models

Unified multimodal models (UMMs) unify visual understanding and generation within a single architecture. However, conventional training relies on image-text pairs (or sequences) whose captions are typically sparse and miss fine-grained visual details, even when they use hundreds of words to describe a simple image. We introduce Reconstruction Alignment (RECA), a resource-efficient post-training method that leverages visual understanding encoder embeddings as dense "text prompts", providing rich supervision without captions. Concretely, RECA conditions a UMM on its own visual understanding embeddings and optimizes it to reconstruct the input image with a self-supervised reconstruction loss, thereby realigning understanding and generation. Despite its simplicity, RECA is broadly applicable: across autoregressive, masked-autoregressive, and diffusion-based UMMs, it consistently improves generation and editing fidelity. With only 27 GPU hours, post-training with RECA substantially improves image generation performance on GenEval (0.73 →\rightarrow 0.90) and DPGBench (80.93 →\rightarrow 88.15), while also boosting editing benchmarks (ImgEdit 3.38 →\rightarrow 3.75, GEdit 6.94 →\rightarrow 7.27). Notably, RECA surpasses much larger open-source models and applies broadly across diverse UMM architectures, establishing it as an efficient and general post-training alignment strategy for UMMs.