💻 computer science

Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

The paper introduces Sparse Context, a method that significantly accelerates reference-based diffusion models by fine-tuning them to be robust against random token dropping and then applying task-aware token selection at inference to reduce computational costs by up to 4x without compromising visual quality.

Original authors: Rishubh Parihar, Ayush Raina, R. Venkatesh Babu, Or Patashnik

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rishubh Parihar, Ayush Raina, R. Venkatesh Babu, Or Patashnik

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to recreate a complex dish based on a photo of the original meal and a list of instructions.

In the current state of AI image generation, the "chef" (the AI model) looks at the reference photo and tries to study every single pixel of it, one by one, before starting to cook. If you give the chef a high-resolution photo, that's like handing them a library of millions of tiny notes describing every crumb, shadow, and texture. Even if the chef only needs to know the general shape of the plate or the color of the sauce, they are forced to read every single note. This makes the cooking process incredibly slow and requires a massive kitchen (computer memory), especially if you ask them to look at five or six different photos at once.

The paper "Keep The Essentials" proposes a smarter way to do this. Here is the breakdown of their idea:

1. The Problem: Too Much Noise

The authors noticed that reference photos are full of redundancy.

The Analogy: Imagine you are describing a cat sitting on a rug to a friend. You don't need to describe the texture of every single fiber of the rug or the dust motes in the air. You just need to say, "There's a cat in the middle, and a rug underneath."
The Reality: Current AI models treat the reference image like a dense grid of millions of "tokens" (tiny data chunks). The paper found that if you randomly throw away 80% of these tokens without changing the model, the AI can still guess the general layout of the scene. It's like trying to read a book where 80% of the letters are missing, but you can still figure out the story.

2. The Solution: "Sparse Context"

The team created a method called Sparse Context. Think of it as teaching the AI to be a "skimmer" rather than a "deep reader."

Step 1: Training the Skimmer (The Fine-Tuning)
If you just start throwing away tokens during the actual cooking (inference), the AI gets confused because it was trained to read every single note. It's like asking a student who studied every page of a textbook to take a test with 80% of the pages ripped out—they might fail.
- The Fix: The authors retrained the AI model by intentionally ripping out random pages (dropping tokens) during its training sessions. They taught the model to be robust, meaning it learned to understand the "gist" of the image even when parts of the reference were missing.
Step 2: Smart Selection (The Chef's Intuition)
Once the model is trained to handle missing pieces, the team didn't just throw away random pieces. They taught the AI to be smart about what to keep, depending on the task:
- For Image Editing: If you want to change the color of a car but keep the shape, the AI focuses on the edges (the outline of the car). It ignores the smooth, empty parts of the sky or the background. It's like tracing a drawing with a highlighter only on the lines.
- For Personalization (Putting a specific person/object in a new scene): If you want to put a specific dog in a new park, the AI focuses on the salient (most important) parts—the dog itself—and ignores the background of the original photo. It's like cutting out the dog from a magazine and ignoring the rest of the page.

3. The Results: Speed Without Sacrifice

By using this method, the AI doesn't have to process millions of data points. It only processes the "essential" ones.

The Speed Boost:
- For a single reference image, the AI runs 2 times faster.
- For multiple reference images (like 5 or 6 photos), it runs 4 times faster.
The Quality: Despite throwing away most of the data, the final image looks just as good as if the AI had read every single note. The details, the style, and the identity of the objects are preserved.

4. It Plays Well with Others

The paper also notes that this method is like a "universal adapter." It can be combined with other existing speed-up tricks (like "KV-caching" or "token merging") to make the AI even faster, like stacking two different types of fuel to get a car to go even quicker.

Summary

In short, the paper says: "Stop reading the whole book to understand the plot."

Reference-based AI models are currently wasting time and energy reading every single detail of a reference image. This new method teaches the AI to ignore the boring, repetitive parts and focus only on the critical details needed to do the job. The result is an AI that is much faster and cheaper to run, without losing any of its ability to create beautiful, accurate images.

Technical Summary: Sparse Context – Efficient Reference Conditioned Generation via Token Dropping

Problem Statement

Reference-based diffusion models have become pivotal for controllable image generation, enabling tasks such as image editing and subject-driven synthesis by conditioning generation on input images. However, these models face significant computational bottlenecks. In modern transformer-based diffusion models (DiTs), reference images are typically encoded into dense spatial grids of tokens via a VAE and concatenated with the noisy input tokens. Since self-attention operations scale quadratically with the number of tokens, the inclusion of one or more high-resolution reference images drastically increases inference time and memory usage. This cost becomes prohibitive in real-world applications involving multiple references. While efficiency techniques exist for prompt-driven generation, they remain under-explored for reference-conditioned models, which present unique challenges due to the dense representation of visual inputs.

Methodology

The paper introduces Sparse Context, a framework designed to construct sparse reference representations by retaining only a reduced subset of reference tokens. The methodology addresses the problem through three core components:

1. Observation of Redundancy

The authors first demonstrate that reference tokens contain significant redundancy. Through a "Naïve Context Pruning" experiment on a pretrained FLUX.2-Klien model, they show that randomly dropping up to 80% of reference tokens at inference time largely preserves the coarse layout of the scene, despite a loss in fine details. This suggests that reference-conditioned models are inherently robust to sparse inputs, but the performance gap between full and partial references is exacerbated by a train-test distribution mismatch.

2. Training Strategy: Random Token Dropping

To bridge the distribution gap and enable robust sparse generation, the authors fine-tune a state-of-the-art reference-conditioned generator using random token dropping during training.

Process: During training, a random fraction of reference tokens is dropped (keeping a fraction $f \in [0.05, 0.25]$ ) before concatenation with the noisy target tokens.
Goal: This encourages the model to learn robust representations that function effectively with partial reference information, decoupling the model from any specific token selection rule.
Implementation: A lightweight LoRA (Low-Rank Adaptation) is applied to the FLUX-2-Klein-9B base model on a curated dataset covering instruction-driven editing and personalization.

3. Inference Strategy: Task-Aware Token Selection

While training uses random dropping to ensure robustness, inference employs task-aware token selection strategies to maximize quality within a fixed compute budget:

Instruction-Driven Image Editing: To preserve spatial structure, the method uses an edge-aware sampling strategy. A Canny edge map is extracted from the reference image, and tokens are sampled based on this probability distribution, prioritizing structural boundaries over uniform regions.
Personalization: To preserve subject identity, the method uses saliency-based sampling. A saliency map approximates the subject mask, allowing the model to concentrate tokens on the most informative regions of the reference image.

Key Contributions

Sparse Context Framework: A novel approach that constructs sparse reference representations by dropping tokens, significantly reducing the quadratic cost of self-attention without modifying the underlying model architecture.
Robust Fine-Tuning Strategy: A training regimen using random token dropping that enables the model to generalize across varying levels of sparsity, allowing flexible control over the inference-time compute budget.
Task-Aware Inference: The proposal of specific token selection heuristics (edge-based for editing, saliency-based for personalization) that outperform random dropping by prioritizing task-critical information.
Orthogonality to Existing Methods: Demonstration that Sparse Context is compatible with and complementary to other efficiency techniques, such as Token Merging (ToMe) and KV-caching, allowing for compounded speedup gains.

Experimental Results

The method was evaluated on instruction-driven image editing (PIE-Bench) and reference-based personalization (Subject-200K, CustomDiffusion-105) using the FLUX-2-Klein-9B model.

Efficiency:
- Single Reference: Achieves up to 2× inference speedup (e.g., 1.83× at 10% token retention).
- Multi-Reference: Achieves up to 4× inference speedup for scenarios with multiple references (e.g., 6 references), as reference tokens dominate the attention computation.
- Combined Efficiency: When integrated with KV-caching for multi-reference generation, the method achieves speedups of up to 6.72×.
Quality:
- Image Editing: The method preserves scene structure and identity significantly better than the Naïve Context Pruning baseline. Quantitative metrics (LPIPS, DINO, CLIP-I) show that Sparse Context maintains high fidelity even at low token fractions (e.g., $f=0.1$ ), whereas the baseline suffers from identity corruption.
- Personalization: Saliency-based sampling successfully preserves fine-grained subject details (e.g., specific colors, patterns) that are lost with random dropping.
Integration: Combining Sparse Context with Token Merging further improves performance, demonstrating that the fine-tuned model effectively leverages additional token reduction techniques.

Significance and Claims

The paper claims that Sparse Context unlocks the untapped efficiency potential inherent in reference-based diffusion models. By recognizing that dense token grids are often unnecessary, the authors demonstrate that reference-based generation can be made significantly more efficient without compromising controllability or visual fidelity.

The authors emphasize that their approach is simple and modular, requiring only lightweight fine-tuning and offering flexible inference-time control. They assert that this method addresses a critical bottleneck in the deployment of high-resolution, multi-reference generative models, making them more practical for real-world applications. The work highlights that the redundancy in reference representations is a fundamental property that can be exploited to reduce computational costs, a finding that extends beyond the specific models tested to the broader class of DiT-based image and video generation systems.