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

Zimeng Wu; Qijian Tian; Penghui Du; Changqian Yu; Jiaxin Chen; Yunhong Wang

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

Figure 1: Empirical observations on Bagel using the general attention metric. (a) Top-10% tokens at timesteps 1/31 across layers 9/18/27, with their union in the last column. (b) Top-10% tokens from the (Left) full-condition and (Middle) image-only branches, with (Right) their symmetric difference; red boxes indicate branch-specific tokens. (c) (Left) Reference images with manually selected synthesis-critical regions highlighted in red boxes. (Middle/Right) Heatmaps of the maximum attention from generated-image locations to these regions under dense and pooled attention, respectively.

Figure 2: Framework overview. (Top) UMM inference pipeline with CATP. It accelerates denoising by pruning visual KV cache tokens across layers at predefined timesteps. (Bottom) Illustration of the main modules. (a) GCP synchronizes pruning decisions from the full-condition to image-only branch, avoiding condition confusion from independent pruning. (b) TCAM estimates token importance using pooled-query attention and calibrates it with a theoretically grounded second-order term. (c) FSO efficiently realizes calibrated scoring via a two-pass tiled online-Softmax kernel.

Method	MacroBench					MICo-Bench
Method	1–3	4–5	6–7	8–10	Avg. $\uparrow$	Object	Human	HOI	De&Re	Overall $\uparrow$
Full Model	8.89	8.51	7.97	7.20	8.14	47.52	40.68	42.47	53.69	46.66
retain 25% tokens
Text-Aligned	8.21	8.22	8.07	7.44	7.98	52.11	35.77	46.95	52.92	47.64
VisionZip	8.18	8.07	8.00	6.97	7.81	49.65	22.99	38.79	56.40	43.38
SDTM	8.13	7.99	7.68	7.50	7.83	38.96	24.79	25.86	41.52	32.90

Table 1: Comparison of multi-image synthesis performance scores by various methods for Bagel, evaluated on the MacroBench and MICo-Bench datasets. The best results are highlighted in bold .

Method	Single		Multiple			Scene			Avg. $\uparrow$
Method	Char.	Obj.	Char.	Obj.	Char.+Obj.	Char.	Obj.	Char.+Obj.	Avg. $\uparrow$
Bagel (Latent Space)
Full Model	8.50	9.28	7.97	9.22	8.60	7.49	8.21	8.15	8.43
Text-Aligned	8.16	8.77	7.56	9.10	8.25	7.46	8.23	7.64	8.14
VisionZip	7.06	7.90	6.78	8.64	7.97	7.40	8.50	7.58	7.73
SDTM	7.16	8.32	6.44	7.10	7.56	6.32	7.45	6.97	7.16

Table 2: Comparison of performance scores by various methods for Bagel/SenseNova-U1 on the OmniContext dataset with 25% tokens retained. The best results are highlighted in bold .

Figure 3: Qualitative comparison of various methods on Bagel at a 25% token retention ratio. The failure/well-preserved regions are marked with red / green boxes.

Figure 4: (a) Attention core latency of the full model and CATP, along with the relative speedup of CATP, across varying numbers of reference images. (b) Visualization of (Left) window-wise query variance and (Right) the resulting aggregated calibration term over reference images in TCAM.

GCP	TCAM	Retain 25% Tokens					Retain 10% Tokens
GCP	TCAM	1–3	4–5	6–7	8–10	Avg. $\uparrow$	1–3	4–5	6–7	8–10	Avg. $\uparrow$
		8.32	8.19	8.14	7.62	8.07	6.65	6.87	7.25	6.85	6.91
$\checkmark$		8.75	8.65	8.36	7.45	8.30	8.23	8.28	8.39	7.38	8.07
$\checkmark$	$\checkmark$	8.76	8.69	8.51	7.87	8.46	8.40	8.33	8.31	7.54	8.15

Table 3: Ablation of GCP and TCAM modules for Bagel on MacroBench at two retention ratios.

Pooling	TCAM	FSO	Latency $\downarrow$	Avg. $\uparrow$
			102.97	8.48
$\checkmark$			7.24	8.30
$\checkmark$	$\checkmark$		30.12	8.46
$\checkmark$	$\checkmark$	$\checkmark$	5.39	8.46

Table 4: Effect of TCAM and FSO modules on scoring latency (ms) and MacroBench average score (Avg.).

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Single		Multiple			Scene			Avg. $\uparrow$
Method	Char.	Obj.	Char.	Obj.	Char.+Obj.	Char.	Obj.	Char.+Obj.	Avg. $\uparrow$
Bagel (Latent Space)
Full Model	8.50	9.28	7.97	9.22	8.60	7.49	8.21	8.15	8.43
Text-Aligned	7.06	7.36	5.17	8.33	7.33	6.13	7.95	7.26	7.07
VisionZip	4.38	4.62	3.67	7.50	4.01	4.27	7.11	6.06	5.20
SDTM	6.57	6.23	5.89	5.77	6.68	5.31	6.31	5.75	6.06

Appendix

Table A: Comparison of performance scores by various methods for Bagel/SenseNova-U1 on the OmniContext dataset with 10% tokens retained. The best results are highlighted in bold .

$\lambda$	1–3	4–5	6–7	8–10	Avg. $\uparrow$
0	8.75	8.65	8.36	7.45	8.30
0.25	8.73	8.70	8.34	7.73	8.38
0.50	8.76	8.69	8.51	7.87	8.46
0.75	8.79	8.76	8.49	7.77	8.45
1.00	8.81	8.67	8.37	7.74	8.40
1.25	8.83	8.65	8.35	7.68	8.38

Appendix

Table B: Effect of the correction strength

\lambda

on MacroBench score for Bagel at a 25% token retention ratio. The default setting is underlined .

$w$	1–3	4–5	6–7	8–10	Avg. $\uparrow$	Latency $\downarrow$
2	8.75	8.73	8.49	7.98	8.49	17.87
4	8.76	8.69	8.51	7.87	8.46	5.39
8	8.68	8.57	8.27	7.68	8.30	2.31

Appendix

Table C: Effect of the window size

w

on MacroBench score and scoring latency (ms) for Bagel at a 25% token retention ratio. The default setting is underlined .

Figure A: Qualitative ablation of GCP and TCAM modules on Bagel at a 25% token retention ratio. The failure/well-preserved regions are marked with red / green boxes.

Method	1–3	4–5	6–7	8–10	Avg. $\uparrow$	Latency $\downarrow$
Average Pooling	8.75	8.65	8.36	7.45	8.30	7.24
Max Pooling	8.59	8.56	8.29	7.50	8.24	7.23
Mixed Query	8.80	8.74	8.48	7.66	8.42	31.47
TCAM + FSO	8.76	8.69	8.51	7.87	8.46	5.39

Appendix

Table D: Effect of various design choices for the pooling-induced bias on MacroBench for Bagel at a 25% retention ratio. The best results are highlighted in bold and the default setting is underlined .

Method	Qwen2.5-VL-72B-Instruct					Gemini-3-Flash
Method	1–3	4–5	6–7	8–10	Avg. $\uparrow$	1–3	4–5	6–7	8–10	Avg. $\uparrow$
Full Model	8.43	7.65	6.45	5.42	6.99	9.01	8.48	6.93	4.45	7.21
retain 25% tokens
Text-Aligned	8.07	7.52	6.26	5.34	6.80	8.05	8.34	6.76	4.82	6.99
VisionZip	7.81	7.36	5.88	5.02	6.52	7.43	7.71	6.36	3.72	6.31
SDTM	7.68	6.75	5.28	3.88	5.90	7.84	7.69	6.42	4.59	6.63

Appendix

Table E: Comparison of multi-image synthesis performance scores by various methods for Bagel on MacroBench, evaluated by Qwen2.5-VL-72B-Instruct and Gemini-3-Flash. The best results are highlighted in bold .

Method	1–3	4–5	6–7	8–10	Avg. $\uparrow$
Full Model	8.810 $\pm$ 0.045	8.412 $\pm$ 0.083	8.048 $\pm$ 0.078	7.256 $\pm$ 0.099	8.132 $\pm$ 0.047
Text-Aligned	8.210 $\pm$ 0.041	8.164 $\pm$ 0.060	8.064 $\pm$ 0.048	7.406 $\pm$ 0.064	7.960 $\pm$ 0.027
VisionZip	8.128 $\pm$ 0.065	8.010 $\pm$ 0.091	7.966 $\pm$ 0.168	7.012 $\pm$ 0.153	7.780 $\pm$ 0.065
SDTM	8.110 $\pm$ 0.222	8.048 $\pm$ 0.120	7.876 $\pm$ 0.157	7.252 $\pm$ 0.165	7.822 $\pm$ 0.066
ToPi	8.570 $\pm$ 0.056	8.278 $\pm$ 0.073	8.240 $\pm$ 0.087	7.532 $\pm$ 0.113	8.152 $\pm$ 0.032
CATP (Ours)	8.814 $\pm$ 0.073	8.640 $\pm$ 0.104	8.540 $\pm$ 0.039	7.838 $\pm$ 0.058	8.458 $\pm$ 0.020

Appendix

Table F: Comparison of MacroBench scores (mean

\pm

std) by various methods on Bagel at a 25% token retention ratio. The best results are highlighted in bold .

Method	1–3	4–5	6–7	8–10	Avg. $\uparrow$	Speedup $\uparrow$
Full Model	8.89	8.51	7.97	7.20	8.14	1.00 $\times$
SDTM	8.95	8.79	8.41	7.65	8.45	1.04 $\times$
SDTM + CATP (Ours)	8.91	8.98	8.63	8.06	8.64	1.41 $\times$
TaylorSeer	8.76	8.48	8.21	7.70	8.29	2.39 $\times$
TaylorSeer + CATP (Ours)	8.78	8.95	8.65	7.81	8.55	3.06 $\times$

Appendix

Table G: MacroBench scores and overall denoising speedup of SDTM and TaylorSeer, individually and in combination with CATP, on Bagel. The best results are highlighted in bold .

Figure B: Qualitative generation results by SDTM and TaylorSeer, individually and combined with CATP, on Bagel. The failure/well-preserved regions are marked with red / green boxes.

Efficient Multimodal Inference

Efficient multimodal inference reduces the compute, memory, latency, and energy needed to process image, video, audio, text, and sensor inputs.

See all 231 papers

G $^2$ TR: Generation-Guided Visual Token Reduction for Separate-Encoder Unified Multimodal Models
May 12, 2026Junxian Li, Kai Liu, Zizhong Ding +4Efficient Multimodal Inference Unified Multimodal Models
Shanghai Jiao Tong University · Huawei Technologies Ltd
Rethinking UMM Visual Generation: Masked Modeling for Efficient Image-Only Pre-training
Mar 17, 2026Peng Sun, Jun Xie, Tao LinUnified Multimodal Models Multimodal Pretraining
Zhejiang University · Westlake University · Shanghai Innovation Institute
Reconstruction Alignment Improves Unified Multimodal Models
Sep 8, 2025Ji Xie, Trevor Darrell, Luke Zettlemoyer +1Unified Multimodal Models Cross-Modal Alignment
1UC Berkeley · University of Washington · 3Duke University

May 12, 2026J/K move · Enter open · S save

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

Junxian Li, Kai Liu, Zizhong Ding, Zhixin Wang+3

Shanghai Jiao Tong University · Huawei Technologies Ltd

Efficient Multimodal Inference Unified Multimodal Models Visual Token Merging Visual Token Compression Multimodal Understanding

Open paper AskPDF arXiv

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 G

^2

TR, 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. G

^2

TR 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 G

^2

TR 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.

Abstract

Efficient Multimodal Inference

Explore similar work

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

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

Reconstruction Alignment Improves Unified Multimodal Models

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