TLDR: ReToM (Local Representative Token Guided Merging) is a new token merging strategy for text-to-image generation models like Stable Diffusion. It improves image quality and maintains computational efficiency by using adaptive window sizes for merging, selecting the most ‘representative’ tokens to preserve essential local features, and caching similarity computations to reduce overhead. This approach leads to better visual fidelity and text-image alignment without increasing generation time, making it practical for real-world applications.
Text-to-image generation models, such as Stable Diffusion, have revolutionized how we create digital art and imagery from simple text prompts. These models are incredibly powerful, capable of synthesizing high-quality images. However, their sophisticated internal workings, particularly the attention operations, demand significant computational resources, leading to slower generation times. This challenge makes it difficult to deploy these models in real-world applications where speed is crucial.
To address this, researchers have explored various methods to improve efficiency. One promising technique is ‘token merging,’ which aims to reduce the number of tokens processed during attention operations, thereby cutting down on computational overhead. While existing token merging methods have shown some success in speeding up models, they often overlook the unique characteristics of attention-based image generation models, limiting their overall effectiveness and sometimes compromising image quality.
A new research paper introduces a novel approach called Local Representative Token Guided Merging, or ReToM. This innovative strategy is designed to be applicable to any attention mechanism used in image generation, offering a significant leap forward in balancing visual quality with computational efficiency. You can read the full paper here.
How ReToM Works: Smart Merging for Better Images
ReToM tackles the limitations of previous token merging methods by introducing several key innovations:
Adaptive Window Sizes with Local Boundaries: Unlike older methods that use a fixed merging region, ReToM defines ‘local boundaries’ as windows within the attention inputs. Crucially, it adjusts the size of these windows. For instance, in parts of the model (like the U-Net’s downsampling and upsampling blocks) where preserving fine local details is important, smaller window sizes are used. In bottleneck layers, where capturing broader global context is more critical, larger window sizes are applied. This adaptive approach ensures that the model retains essential features while optimizing computational effort.
Local Token Merging with Representative Tokens: Traditional merging often involves randomly selecting tokens or using less effective matching methods, which can lead to information loss. ReToM introduces a ‘representative token’ concept. Within each adaptive window, it calculates the similarity between all tokens. The token that has the highest average similarity to all other tokens in that window is selected as the ‘representative token.’ This token is considered the most informative feature within that local area. Other less representative tokens are then merged into this chosen token, ensuring that the most salient local features are preserved, leading to higher quality and more stable image generation.
Similarity Computation Caching Strategy: In diffusion models, the latent representation of an image changes gradually across consecutive timesteps. ReToM leverages this consistency. Instead of recalculating token similarities at every single timestep, which is computationally expensive, ReToM computes and caches these similarities periodically. This cached information is then reused for subsequent timesteps, significantly reducing redundant computations without sacrificing performance, as the relative similarity between tokens remains largely consistent over short periods.
Also Read:
- Discrete Latent Codes: A New Approach to High-Fidelity and Creative Image Generation
- Advancing Diffusion Model Alignment for Enhanced Image Generation with Inversion-DPO
Impressive Results and Practical Applications
Experiments applying ReToM to the Stable Diffusion model demonstrate its effectiveness. Compared to the baseline Stable Diffusion model without token merging, ReToM achieves a notable improvement in image quality, as measured by the FID (Fréchet Inception Distance) score, which indicates better visual fidelity. It also achieves higher CLIP scores, signifying better semantic alignment between the generated images and their text prompts. Importantly, ReToM accomplishes these improvements while maintaining a comparable or even faster inference time.
Qualitative results show that images generated with ReToM are more natural, capturing both objects and complex backgrounds like leaves, grass, and sand with greater detail and consistency. Unlike some other methods that might distort shapes or express backgrounds unnaturally, ReToM maintains good detail without imbalance, even in intricate scenes.
The fact that ReToM does not require additional model training or fine-tuning means it can be directly integrated into existing attention-based models. This makes it a highly practical solution for real-world applications such as real-time image synthesis, generating high-resolution text-to-image content, and compressing models for use in environments with limited computational resources.
In conclusion, ReToM offers a sophisticated yet efficient token merging framework that significantly enhances the quality and computational performance of attention-based image generation models, paving the way for more accessible and powerful AI-driven creative tools.
