Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

2026-06-22

AI Summary

Existing reference-guided diffusion models suffer from high computational overhead and low inference efficiency under multi-reference conditions. This work proposes Sparse Context, a method that constructs sparse reference representations by introducing a random token dropping strategy during training and enabling task-aware selection of critical tokens at inference time, thereby decoupling token selection from model training. The approach achieves 2× and 4× inference speedups on single-reference and multi-reference generation tasks, respectively, while preserving strong spatial alignment capabilities and high-quality subject-driven generation performance.

Abstract

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.

Problem

• reference-based generation
• computational efficiency
• token redundancy
• diffusion models
• inference speed

Innovation

• token dropping
• reference-based generation
• sparse representation
• diffusion models
• inference efficiency