Image generation models speed up by cutting low-value image tokens early
ToDoPro: Faster and Better Reference-Based Generation: Early Token Reduction via Downsampled Self-Probing
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Generating images using multiple reference images and text slows down because the model has to process many pieces of information all at once. The researchers propose a method called ToDoPro that quickly guesses which parts of the input are most important using a simpler, faster step. Then, it keeps only the most useful pieces for the main image generation, making the process faster and using less computer memory. This also helps the model focus better, improving image quality.
What this means in practice
- •For computer vision engineers: Accelerate multi-reference image generation models by early filtering of unimportant tokens to reduce memory and speed up inference.
- •For mobile app developers: Enable faster image generation on limited hardware by lowering computation through intelligent token pruning before heavy processing.
Authors
Caiyi Sun, Yujing Sun, Siu-Ming Yiu
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