Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

Paper: 2606.23682 Authors: Rishubh Parihar, Ayush Raina, R. Venkatesh Babu, Or Patashnik Categories: cs.CV

The Gap

Reference-based diffusion models (e.g. IP-Adapter, ReferenceNet) allow highly controllable image generation by injecting tokens from a reference image into the diffusion process. However, they treat reference features as dense token grids (e.g. 16x16 = 256 tokens per image) and feed all tokens into cross-attention layers. This makes inference cost scale linearly with the number of references—a 2-reference generation costs twice as much as a single-reference one. While prior work has extensively studied efficiency for text-conditioned diffusion (e.g. pruning, distillation), the unique setting of reference-based generation has been largely ignored. The key oversight is that most reference tokens carry redundant or irrelevant information for the target generation, especially in tasks like subject-driven generation where only the object region matters. This paper identifies that dropping a large fraction of reference tokens (up to 75%) barely hurts quality, and builds a practical method around this observation.

+-----------------------+       +----------------------------+
| Problem: Linear cost  | ----> | Assumption: Most ref       |
| growth with number of |       | tokens are redundant;      |
| reference images in   |       | dropping them preserves    |
| cross-attention       |       | generation capability      |
+-----------------------+       +----------------------------+
          |                              |
          v                              v
+-----------------------+       +----------------------------+
| Method: Sparse Context| ----> | Evidence: 4x speedup for   |
| (random drop training |       | multi-ref generation, 2x   |
| + task-aware selection|       | for single ref, no quality |
| at inference)         |       | degradation in metrics     |
+-----------------------+       +----------------------------+
          |
          v
+-----------------------+
| Conclusion: Reference |
| token redundancy can  |
| be exploited to make  |
| generation efficiency |
| scale sub-linearly    |
+-----------------------+

The Increment

One sentence: Before this paper, every generation with reference images paid the full token cost; after this paper, we can routinely drop 75% of reference tokens with a simple fine-tuning step, maintaining quality while cutting inference time by 2-4x, and adapt the token budget to the task on the fly.

Core Mechanism

Sparse Context operates in two phases: training and inference. During training, the entire diffusion model (including cross-attention layers) is fine-tuned with random token dropping at different ratios. For each training step, a random subset of reference tokens (e.g. 25%–100% of the original grid) is retained, and the model learns to predict the noise conditioned on this partial context. This forces the cross-attention mechanism to become robust to missing tokens—it cannot rely on any specific token always being present.

At inference time, the user can choose a token budget (e.g., keep only 64 out of 256 tokens) and a selection strategy. The selection strategy is task-aware: for spatially-aligned editing (e.g., face swapping), tokens are chosen based on their spatial location in the reference image (e.g., keep tokens inside a mask); for subject-driven generation (e.g., “a dog wearing that sweater”), tokens are selected by the highest cross-attention scores from the original full-token run (a cheap pre-pass). These selected tokens are then used as the sole context for cross-attention in the denoising U-Net, drastically reducing the key/value sequence length.

Reference Image
       |
       v
Image Encoder -> Dense Token Grid (256 x D)
       |
       v
+-----------------------------------------+
| Token Dropping Module                   |
|    Training: random mask (ratio r)      |
|    Inference: task-aware selection      |
+-----------------------------------------+
       |
       v (only kept tokens, e.g. 64)
       |
       v
+-----------------------------------------+
| Cross-Attention Layers                  |
| Q from UNet features                    |
| K, V from sparse reference tokens      |
| (sequence length 64 vs 256)            |
+-----------------------------------------+
       |
       v
Denoising U-Net (rest of diffusion steps)
       |
       v
Generated Image

Think of it like packing for a trip. The reference image is a closet full of clothes (all tokens). Normally you dump the entire closet into your suitcase—that’s the dense grid approach. It’s heavy, slow, and most of it won’t be used. Sparse Context is like a smart packing algorithm. Before the trip, you practice packing randomly (training): sometimes you forget your coat, sometimes your socks. You learn to survive with whatever you grab. On the actual trip (inference), you pick only what you need for that specific destination: for a beach vacation (subject-driven), you take swimwear; for a business meeting (spatial editing), you take formal wear. The key insight is that the item-level granularity (each token) is fine enough that you never need the full closet. The model (the traveller) has been trained to make do with less, and the task-aware selection (the destination guide) ensures the right subset is chosen.

Key Concepts

• Token Dropping: In neural networks, input features are often arranged as a grid (e.g., patches from an image). Token dropping means deliberately removing some of these grid entries before feeding them into subsequent layers. For a 16×16=256 token grid, dropping 75% leaves only 64 tokens. The remaining tokens still carry enough information because visual features are spatially redundant—neighbouring patches look similar, and many patches (e.g., background) are irrelevant. Training the model with random dropping makes it robust to missing neighbours, so generation quality stays high.

• Task-aware Token Selection: Random dropping works well when the model has been fine-tuned, but not all tokens are equally informative for a given task. Task-aware selection uses task-specific heuristics to pick the best tokens. For example, in image editing with a mask, only tokens inside the mask region are kept. In subject-driven generation (e.g., “a corgi wearing a red hat”), you can run a quick forward pass with all tokens (cheap because it’s one step) and record which tokens get the highest attention from the text prompt; then keep only those. This gives better quality than random dropping at the same token budget.

• Sparse Context: The core idea of representing the reference condition as a sparse set of tokens rather than a dense grid. The sparsity is not fixed—it can vary per reference image and per task. Sparse Context reduces the cross-attention complexity from O(UNet_tokens ** reference_tokens) to O(UNet_tokens * sparse_tokens). Since cross-attention is a major bottleneck in diffusion, slashing the reference tokens by 4× yields a near-4× speedup in those layers, and in practice the overall inference time improves by 2–4× depending on the ratio of time spent in cross-attention vs other components.

Framework Shift

Before (mainstream approach):        After (this paper):
Reference images                     Reference images
       |                                  |
       v                                  v
+--------------------+              +--------------------+
| Dense Token Grid   |              | Task-aware Selector|
| (all tokens kept)  |              | (keep only ~25%)   |
+--------------------+              +--------------------+
       |                                  |
       v                                  v
+--------------------+              +--------------------+
| Cross-Attention    |              | Cross-Attention    |
| K,V len = 256*tok  |              | K,V len = 64*tok   |
| (slow)             |              | (fast)             |
+--------------------+              +--------------------+

From “treat all reference tokens as equally important” to “selectively keep only the most informative tokens,” the core shift is recognizing that redundancy is a feature to exploit, not a bug to ignore. The training strategy (random dropping) makes the model compatible with sparse inputs, transforming an efficiency gain from a hack into a principled approach.

Expert Assessment

Problem choice: Real gap. The efficiency of reference-based generation was genuinely underexplored, and the observation that tokens can be dropped so aggressively without quality loss is surprising and useful. However, the problem is more engineering than science—it’s about optimization, not new modeling.

Method maturity: Clever combination of known tricks (random dropping, attention-based selection) plus fine-tuning. Not a breakthrough algorithm, but the systematic study and practical recipe are solid. Simpler alternatives (e.g., just using a fixed low-resolution feature map) might also work but are not compared.

Experimental integrity: Baselines are fair (IP-Adapter, etc.) and metrics (FID, CLIP score, user study) support claims. One red flag: the 4x speedup for multi-reference generation is only achieved when using a very high drop ratio (75%), and the paper should report the trade-off curve more granularly. Also, the task-aware selection overhead (pre-pass) is not fully accounted for in the speedup numbers.

Writing quality: Clear and well-structured, but the experimental section feels rushed—ablation studies are there, but discussion of failure cases or limitations (e.g., when token selection fails) is missing. A dedicated “Limitations” subsection would elevate the paper.

Verdict: weak accept — The contribution is practical and the speedups are impressive, but the scientific novelty is modest and the analysis could be deeper.

Takeaways

• Random dropping as a regularizer: If you have a model that consumes variable-length conditioning tokens (e.g., multiple references, video frames, etc.), fine-tuning with random token masking can make it robust to token loss at inference, enabling straightforward efficiency gains.

• Task-aware selection beats generic pruning: The paper shows that the same token budget yields better results when tokens are selected based on task-specific cues (mask, attention) rather than uniformly. This principle transfers to any token-based conditioning (e.g., text tokens, audio tokens).

• Measure redundancy before optimizing: The key enabler was first verifying that dropping 75% of tokens barely hurts quality. Before applying aggressive pruning, check your own data/modality for similar redundancy; it’s often there.