Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping

Source: arXiv:2606.23682 · Published 2026-06-22 · By Rishubh Parihar, Ayush Raina, R. Venkatesh Babu, Or Patashnik

TL;DR

This paper addresses the computational inefficiency challenge of reference-conditioned image diffusion models, whose runtime and memory costs grow significantly with the number of input reference images due to dense token representations. The authors observe that many reference tokens are redundant and can be dropped with only minor degradation in generation quality. Building upon this, they propose Sparse Context, a method that fine-tunes the diffusion model with random token dropping of reference tokens during training to improve robustness to sparse conditioning at inference. At inference, task-aware token selection strategies prioritize informative reference tokens (e.g., edges for editing, salient objects for personalization) to flexibly control token budgets and maximize efficiency. Extensive experiments on high-quality editing and personalization datasets demonstrate up to 4× speedup in multi-reference settings and up to 2× in single-reference generation without compromising visual quality or prompt adherence. The approach is also complementary to other efficiency methods like token merging and KV caching, enabling compound speedups.

Key findings

• Dropping up to 80% of reference tokens without model fine-tuning preserves coarse scene layout but loses fine details (Naïve Context Pruning).
• Fine-tuning with random token dropping (keep fractions 5%-25%) significantly improves robustness to sparse reference tokens, closing the quality gap with full-token models.
• On instruction-driven image editing, Sparse Context achieves a 1.49× speedup at 20% token fraction, maintaining LPIPS of 0.576 (versus 0.421 baseline) and higher DINO (0.738 vs. 0.761).
• For personalization, Sparse Context retains strong identity preservation (CLIP-I of 0.754 vs. 0.708 baseline) and text-image alignment (CLIP T-I ~0.274) at 20% tokens.
• In multi-reference generation, speedups scale with number of references; at 6 reference images and 10% tokens, 4× speedup is achieved.
• Task-aware token dropping using edge maps for editing and saliency maps for personalization outperforms random token dropping on structural consistency and identity preservation.
• Integration with token merging shows improved inference time reductions and quality restoration compared to token merging without fine-tuning tokens dropping.
• Combining Sparse Context with KV caching yields up to 6.72× speedup on multi-reference generation with 6 references.

Threat model

The adversary considered is a user or external agent who can submit reference images to guide generation, potentially multiple references. However, the adversary does not attempt to attack the model by adversarial inputs or manipulate the token-dropping mechanism. The focus is on improving efficiency and robustness to partial reference inputs rather than security against malicious manipulation.

Methodology — deep read

The paper addresses the inefficiency in transformer-based reference-conditioned diffusion models, where concatenating dense reference image tokens to noisy latents causes quadratic scaling in self-attention. The threat model involves inference-time generation with one or multiple reference images; the focus is on improving inference cost without compromising generation quality or controllability.

The starting point uses the FLUX-2-Klein-9B DiT model pretrained on standard text- and image-conditioned generation tasks. Reference images are encoded by the VAE encoder into latent token grids representing spatial patches.

Initially, the authors perform Naïve Context Pruning: at inference, random dropping of large fractions (up to 80%) of reference tokens without any retraining. This reveals that even aggressively pruning tokens preserves coarse scene layout but reduces fine semantic details, indicating large redundancies in reference tokens.

To address train-test mismatch, the core methodology is Sparse Context: the authors fine-tune the model via lightweight LoRA adaptation on a curated dataset covering two tasks — instruction-driven image editing and personalization. During training, random reference token dropping with varying keep fractions (f between 0.05 and 0.25) is applied, so the model learns to denoise target latents conditioned on sparse partial references. This encourages robustness to incomplete references. The lightweight fine-tuning lasts for 15K iterations on two NVIDIA A100 GPUs, with curriculum learning starting from single-reference to multi-reference samples.

At inference time, rather than random dropping, task-aware token selection strategies are applied: for image editing, tokens are sampled preferentially from Canny edge regions to ensure structural consistency; for personalization, sampling is weighted by saliency maps that approximate subject locations, improving fine identity preservation. These strategies adapt the token budget flexibly to task demands and user compute constraints.

The evaluation protocol includes multiple benchmarks: PIE-Bench (700 image editing examples), Subject-200K subset (200 personalization single-ref examples), and CustomDiffusion-105 dataset (500 multi-reference compositional scenes). Quality metrics include LPIPS, CLIP-Image (CLIP-I), DINO for reference-image similarity, and CLIP text-image alignment (CLIP T-I). Additional evaluation of Kernel Inception Distance (KID) supplements metrics. Inference speedups are measured on a single NVIDIA A100 GPU averaged over 100 runs, comparing the baseline full-token model to sparse token variants.

Integration experiments combine Sparse Context with existing inference-time speedups like Token Merging and KV Caching. Token Merging merges similar reference tokens before passing them to the model, while KV Caching reuses stored attention key-value pairs across diffusion steps. Results show Sparse Context fine-tuning yields additive speedups and quality preservation when combined with these methods.

One concrete training example: a 512×512 HQ-Edit image and its editing instruction are converted to latent tokens; during a training iteration, 5%-25% of the reference tokens are randomly dropped, and the denoising DiT model is trained with LoRA to predict noise given the sparse tokens concatenated with noisy latents. The model learns to leverage the sparse conditioning and generalizes better to token dropping at inference.

Reproducibility: code and pretrained weights are not explicitly stated as released; the datasets used include public subsets of HQ-Edit, Subject-200K, and CustomDiffusion-105 detailed in the appendix. The lightweight LoRA fine-tuning strategy facilitates efficient replication on standard hardware.

Technical innovations

• Introduction of Sparse Context, a fine-tuning strategy using random token dropping of reference tokens during training to enable robustness to sparse reference conditioning at inference.
• Task-aware token dropping at inference that prioritizes spatially informative tokens (edges for editing, saliency for personalization) instead of random dropping.
• Demonstration that aggressive token pruning (up to 80%) is feasible without major quality degradation when coupled with robust training and task-specific sampling.
• Showcasing compatibility and synergy of Sparse Context with existing inference-time acceleration methods, such as Token Merging and KV Caching, for compounded speedups.

Datasets

• HQ-Edit (subset) — 30K images with editing instructions and reconstructions generated using Flux-Klien2 [Labs, 2025]
• Subject-200K (subset) — 20K subject-based reference images and compositions [Tan et al., 2024]
• CustomDiffusion-105 — 13K multi-object compositional scenes with 2-6 references [Kumari et al., 2023]
• PIE-Bench — 700 image editing examples with source images, edit instructions, and target images [Ju et al., 2023]

Baselines vs proposed

• Baseline (full tokens, f=1.0): LPIPS = 0.421 (editing), CLIP-I = 0.826, DINO = 0.761, Speedup = 1x vs Sparse Context (f=0.2): LPIPS = 0.576, CLIP-I = 0.817, DINO = 0.738, Speedup = 1.49x
• Naïve pruning (random drop tokens) f=0.20: LPIPS = 0.687, CLIP-I = 0.762, DINO = 0.631, Speedup = 1.49x
• Sparse Context (fine-tuned) f=0.10: LPIPS = 0.628 vs Naïve f=0.10 at 0.719 on editing task
• Sparse Context + Token Merging f=0.10: LPIPS = 0.570 vs Token Merging only 0.690, Speedup both ~1.76x
• Sparse Context + KV Caching on 6 reference images: combined speedup up to 6.72x
• For personalization, Sparse Context at f=0.20: CLIP-I = 0.754 vs baseline 0.708, Speedup 1.49x

Figures from the paper

Figures are reproduced from the source paper for academic discussion. Original copyright: the paper authors. See arXiv:2606.23682.

Fig 1: Sparse Context enhances the efficiency of reference-conditioned generation by replacing

Fig 2: Redundancy in reference tokens. We

Fig 3 (page 1).

Fig 4 (page 1).

Fig 5 (page 1).

Fig 6 (page 1).

Fig 7 (page 1).

Fig 8 (page 1).

Limitations

• Quality metrics decline as token fraction drops below 5%, indicating a lower bound on token pruning.
• The method requires fine-tuning on curated, paired reference-conditioned datasets, which may not be available for all domains.
• Evaluation focuses on 512×512 resolution; scalability to higher resolutions is untested.
• Multi-reference identity preservation is mostly evaluated qualitatively; quantitative metrics are limited.
• No explicit adversarial robustness or adaptive attacker evaluation—token dropping strategies may be vulnerable if attacker controls token selections.
• Code and pretrained models are not yet publicly released, limiting immediate reproducibility.

Open questions / follow-ons

• How does Sparse Context perform under distribution shifts, such as unseen reference image domains or resolutions?
• Can task-aware token dropping be further improved via learned token importance scores rather than heuristic edge or saliency maps?
• Does the method generalize to real-time interactive editing scenarios where tokens dropped dynamically during user interaction?
• What are the effects of Sparse Context on model robustness to adversarial or corrupted references?

Why it matters for bot defense

Reference-conditioned diffusion models often require processing multiple high-dimensional token streams from input images, which leads to increased computational cost and latency in real-time or interactive generation scenarios. Sparse Context’s approach of robustly training the model to function with partial sparse reference tokens enables substantial inference speedups with minimal quality loss. For bot-defense engineers and CAPTCHA practitioners interested in leveraging image-based human verification through generative challenges or reference-conditioned image tasks, this work offers a practical strategy to reduce latency and computational overhead. Furthermore, task-aware token selection can guide focus to the most informative regions of the input, improving quality-efficiency tradeoffs. Being compatible with other acceleration techniques like token merging and KV caching makes Sparse Context an appealing efficiency layer to adopt in reference-conditioned generation frameworks, especially for scaling multi-reference or high-resolution scenarios relevant to CAPTCHA challenges with visual puzzles or compositional elements.

Cite

@article{arxiv2606_23682,
  title={ Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping },
  author={ Rishubh Parihar and Ayush Raina and R. Venkatesh Babu and Or Patashnik },
  journal={arXiv preprint arXiv:2606.23682},
  year={ 2026 },
  url={https://arxiv.org/abs/2606.23682}
}

Read the full paper

• arXiv:2606.23682 (abstract)
• arXiv:2606.23682 (PDF)