ToDoPro: Faster and Better Reference-Based Generation: Early Token Reduction via Downsampled Self-Probing

Caiyi Sun; Yujing Sun; Siu-Ming Yiu

cs.CVOct 26, 2026

ToDoPro: Faster and Better Reference-Based Generation: Early Token Reduction via Downsampled Self-Probing

Authors: Caiyi Sun, Yujing Sun, Siu-Ming Yiu

Organizations: The University of Hong Kong · Nanyang Technological University

Abstract

Recent generic reference-based image generation models natively support multiple conditions by concatenating all reference image and text tokens into a joint sequence processed by a single Diffusion Transformer (DiT). However, this versatile paradigm comes at an efficiency cost: accommodating multiple references leads to a token explosion, incurring massive $\mathcal{O}(N^2)$ computational complexity. Most acceleration methods operate inside the DiT, showing limited capability in reducing spatial redundancy and failing to alleviate peak VRAM. While some pre-model token reduction alternatives exist, they inevitably introduce external models. This raises a question: Is it possible to achieve acceleration and memory savings while preserving compatibility with in-model techniques and requiring no additional parameters? In this work, we realize this goal by proposing the early Token reduction via Downsampled self-Probing (ToDoPro). By executing a few denoising steps on spatially $2\times$ downsampled reference images, this cheap probing stage utilizes the DiT model itself to effectively and efficiently extract reference token importance. Using the probed importance, we apply a hierarchical reduction before the main denoising stage: we completely discard low-importance tokens, substitute mid-importance tokens with their downsampled counterparts to cheaply preserve global context, and keep high-importance tokens. Experiments across three base models and two benchmarks demonstrate that our early reduction accelerates inference significantly and mitigates peak GPU memory. Furthermore, we can improve generation quality by safely filtering out noisy, uninformative reference tokens, enabling the model to concentrate its limited attention capacity on genuinely useful tokens. The code will be released at https://github.com/scy639/ToDoPro

Figures & tables

Figure 1: Compared to base model FLUX.2-klein-4B, our method reduces inference latency while improving generation quality.

Figure 2: (a) Most existing acceleration methods operate inside the DiT during the main denoising stage. (b) ToDoPro reduces reference tokens before the main denoising stage begins. A probing stage on downsampled references cheaply extracts token importance, guiding token preservation / low-resolution substitution / discarding. The compressed reference images then enter the main denoising stage, slashing end-to-end latency and activation VRAM while remaining compatible with in-model methods in (a).

Figure 3: Alignment of reference utilization scores with pseudo ground-truth importance across different query sets (

\Omega

). The set (

\Omega=G\cup P

) consistently yields the best overall alignment, demonstrating that querying reference tokens via both generative and prompt tokens optimally proxies the importance.

Panel of Figure 4 — Figure 4: Analysis of layer-wise attention dynamics. (a) Layer-wise visualizations of the reference utilization score show stochastic anomalous values, and a single layer is usually biased because different layers may have different functions and focuses. (b) Computing the mean utilization score across layers smooths out the bias and instantaneous noise, achieving better alignment with the oracle token importance.

Figure 5: The illustration of our token reduction by a spatial-temporal downsampled self-probing. By executing a downsampled time step schedule denoising loop on spatially $2\times$ downsampled input images, this probing stage cheaply extracts the reference importance. The extracted importance scores then guide the reduction strategy outside the main denoising run. The visual example is generated using FLUX.2-Klein-4B (4-step).

Figure 6: Average quality versus speedup across three base models and two benchmarks.

Method	MultiBanana									MACRO
	Quality $\uparrow$						Efficiency			Quality $\uparrow$					Efficiency
	Inst. [-3pt]Align	Ref. [-3pt]Cons.	B-S [-3pt]Match.	Phys. [-3pt]Real.	Visu. [-3pt]Qual.	Avg	Latency -3pt $\downarrow$	Speed $\uparrow$	Act. VRAM -3pt $\downarrow$	Cust.	Illus.	Spat.	Temp.	Avg	Latency -3pt $\downarrow$	Speed $\uparrow$	Act. VRAM -3pt $\downarrow$
Baseline	4.57	4.86	3.96	4.40	5.72	4.70	5.74	1.00 $\times$	15.3	7.49	6.17	3.83	7.23	6.18	7.93	1.00 $\times$	24.7
Fast
ToMeSD	4.45	4.51	3.72	4.17	5.55	4.48	2.62	2.19 $\times$	11.2	7.10	5.82	3.45	7.10	5.87	3.23	2.46 $\times$	17.4
VisPruner	4.54	4.53	3.72	4.24	5.62	4.53	2.55	2.25 $\times$	8.7	7.27	5.99	3.87	7.20	6.08	3.39	2.34 $\times$	13.7

Table 1: Comparison of efficiency and quality on MultiBanana and MACRO (Flux.2-Klein-4B).

Figure 7: Visual comparison of generation quality across diverse tasks from MultiBanana and Macro datasets (all methods are under the Slow level in Table 1 with comparable speedups). Best viewed when zoomed in.

Method	MultiBanana									MACRO
	Quality $\uparrow$						Efficiency			Quality $\uparrow$					Efficiency
	Inst. [-3pt]Align	Ref. [-3pt]Cons.	B-S [-3pt]Match.	Phys. [-3pt]Real.	Visu. [-3pt]Qual.	Avg	Latency -3pt $\downarrow$	Speed $\uparrow$	Act. VRAM -3pt $\downarrow$	Cust.	Illus.	Spat.	Temp.	Avg	Latency -3pt $\downarrow$	Speed $\uparrow$	Act. VRAM -3pt $\downarrow$
Baseline	4.57	4.86	3.96	4.40	5.72	4.70	5.7	1.00 $\times$	15.3	7.49	6.17	3.83	7.23	6.18	7.9	1.00 $\times$	24.7
Ours	4.63	4.85	3.97	4.38	5.77	4.72	2.5	2.30 $\times$	7.3	7.50	6.15	3.88	7.27	6.20	3.2	2.48 $\times$	11.5
RepreShift	4.51	4.69	3.93	4.44	5.78	4.67	3.5	1.63 $\times$	11.3	7.45	6.12	3.83	7.11	6.13	4.1	1.93 $\times$	22.1
RepreShift + Ours	4.61	4.66	3.96	4.40	5.72	4.67	2.4	2.38 $\times$	8.3	7.44	6.10	3.85	7.11	6.12	3.0	2.65 $\times$	15.2

Table 2: Compatibility of Ours(fast) with existing in-model acceleration methods on FLUX.2-Klein-4B, yielding compounded inference speedups.

Method	Instruction [-3pt]Alignment $\uparrow$	Reference [-3pt]Consistency $\uparrow$	BG-Subject [-3pt]Match $\uparrow$	Physical [-3pt]Realism $\uparrow$	Visual [-3pt]Quality $\uparrow$	Overall $\uparrow$
Only Low-Res Tier (i.e., infer baseline on 2 $\times$ downsampled images)	4.25	4.10	3.53	4.12	5.63	4.33
No Low-Res Tier	4.56	4.67	3.88	4.37	5.72	4.64
Random Reduction	4.25	4.08	3.62	4.22	5.60	4.35
Ours	4.63	4.85	3.97	4.38	5.77	4.72

Table 3: Ablation study on Ours(fast) on MultiBanana (controlled to the same reduction ratio). Base model: Flux.2-Klein-4B.

Figure 8: Qualitative evolution of generated images. At 0% reduction (unpruned baseline), the model mistakenly retains visual distractors (e.g., an extra floating glove). Within a broad window (30% to 94%), our method successfully filters out these distractors, boosting quality. At extreme reduction levels ( $>95$ %), the irreversible loss of informative reference tokens causes quality degradation.

Layers used for $I_{r}$ Extraction	Effective Attention Ratio $\uparrow$	Instruction [-3pt]Alignment $\uparrow$	Reference [-3pt]Consistency $\uparrow$	BG-Subject [-3pt]Match $\uparrow$	Physical [-3pt]Realism $\uparrow$	Visual [-3pt]Quality $\uparrow$	Overall $\uparrow$
Baseline (No Reduction)	43.6	4.57	4.86	3.96	4.40	5.72	4.70
Top-1 layer	51.7	4.81	5.08	3.92	4.42	5.85	4.82
Top $20\%$ of layers	52.9	4.89	5.13	4.01	4.47	5.93	4.89
All layers	52.7	4.86	5.06	4.03	4.54	5.92	4.88

Table 4: Effective Attention Ratio (EAR) on all samples from the object composition task of MultiBanana. Baseline: Flux.2-Klein-4B. All variants use the same reduction ratio (Slow level).

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Average score and latency of Ours (slow) (Flux.2-Klein) on MACRO with 2–10 reference images.

Image Generation

Image generation studies models that create or edit images from text, examples, layouts, and other structured controls.

See all 474 papers

Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping
Jun 22, 2026Rishubh Parihar, Ayush Raina, R. Venkatesh Babu +1Geometry-Conditioned Image Generation Conditional Image Generation
1IISc Bangalore · 2Tel Aviv University
Rethinking Token Reduction for Diffusion Models via Output-Similarity-Awareness
May 21, 2026Hangyeol Lee, Hyojeong Lee, Joo-Young KimDiffusion Transformer Diffusion-Based Image Generation
KAIST, Daejeon, Republic of Korea
MoDDiT: Simple Token Routing in Diffusion Transformers
Oct 20, 2026Jarek Liesen, Matthew T. Jackson, Benjamin Ellis +2Diffusion Transformer Adaptive Computation for Diffusion Models
British Open-ended Learning and Discovery

Jun 22, 2026J/K move · Enter open · S save

Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

Rishubh Parihar, Ayush Raina, R. Venkatesh Babu, Or Patashnik

1IISc Bangalore · 2Tel Aviv University

Geometry-Conditioned Image Generation Conditional Image Generation Diffusion Model Fine-Tuning Token Pruning Diffusion-Based Image Editing Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Reference-based diffusion models enable highly controllable image generation by leveraging elements from input images to guide prompt-driven synthesis. However, these models are computationally expensive in runtime, and their cost scales severely with the number of input references. While the efficiency of diffusion models has been extensively studied in the context of prompt-driven generation, it remains largely under-explored in the realm of reference-based models. This setting presents unique challenges not addressed by methods focusing solely on generation. In particular, the wasteful representation of references as dense token grids offers significant opportunities for improvement. In this work, we present Sparse Context, a method for constructing sparse reference representations by retaining only a reduced subset of reference tokens. We observe that even without modifying the model, dropping a significant portion of reference tokens at inference time largely preserves its generation capabilities. To fully realize this potential, we fine-tune the model with random token dropping at varying ratios, encouraging robustness to partial reference representations. Crucially, this training strategy decouples the model from any specific token selection rule, allowing flexible control at inference time. At inference time, instead of random dropping, we apply task-aware token selection strategies that prioritize the most informative regions of the reference images, adapting the token budget to the input and task requirements. Extensive experiments show our method achieves a 4x increase in inference speed for multi-reference generation and an 2x for single reference generation. Importantly, this efficiency is achieved without compromising visual quality across both spatially-aligned editing and subject-driven generation.

Abstract

Image Generation

Explore similar work

Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

Rethinking Token Reduction for Diffusion Models via Output-Similarity-Awareness

MoDDiT: Simple Token Routing in Diffusion Transformers