TokenRepel: Generating Diverse Image Sets Without Compromising Quality

AI summary

Overview

Research area: Computer Vision / generative modeling — specifically inference-time diversity enhancement for text-to-image (T2I) diffusion and flow-matching models.

Technical level: Advanced. The method relies on Riemannian gradient descent on a von Mises–Fisher mixture over unit hyperspheres, plus score/velocity decompositions from flow-matching theory. The core intuition is explainable, but the formal machinery is dense.

Scope: A training-free, model-agnostic algorithm that increases the visual diversity of image batches produced from a single prompt while keeping average image fidelity close to the original model.

What This Paper Is About

Modern text-to-image models produce nearly identical images for the same prompt regardless of the random starting noise — a failure mode the authors call "mode gravity," where the model keeps landing on only the highest-confidence samples. The goal is to make a batch of images generated from one prompt genuinely varied (different poses, layouts, objects) without the usual side effect of the diversity tricks: blurry, artifact-ridden, or prompt-misaligned outputs. TokenRepel attacks this at the level of individual feature tokens inside the transformer, pushing semantically similar tokens in different images apart, then correcting that push so it does not drag the model out of the region of latent space it considers plausible.

Key Contributions

1. Token repulsion via Riemannian gradient descent. Normalized hidden tokens from other images in the batch are modeled as an isotropic, equal-weight von Mises–Fisher mixture on the unit hypersphere. Each token takes a negative Riemannian gradient step on this density, which in practice reduces to rotating it away from its most similar cross-image token by an angle δ when the angular separation is below a threshold θ̄. Working on the sphere means the update changes semantic direction, not feature magnitude.

2. Token density preservation. The repulsion-induced change in the model's prediction is projected onto a half-space aligned with the model's own score direction, approximated as the velocity difference between the positive and negative prompts. If the repulsion pushes against increasing density, the offending component is removed; otherwise it is left alone. This is the explicit fidelity-preserving mechanism the authors argue prior token-level work (NegToMe) lacks.

3. Layerwise extension of the fidelity constraint. The same projection is applied at the output of each intermediate transformer layer, not just at the final velocity prediction, using three parallel streams (positive prompt, negative prompt, positive prompt with repulsion). This gives more frequent, finer-grained corrections and counters error accumulation across layers.

4. A geometry-based, interpretable trade-off knob. The repulsion strength is specified directly as the token rotation angle δ (with the constraint θ̄ + δ ≤ 90° to prevent overshooting orthogonality), which the authors argue is easier to set consistently across models than the less interpretable hyperparameters of competing methods.

Main Findings

• Best trade-off on FLUX.2-Klein-9B, second-best on FLUX.1-dev and SDXL. Trade-off curves use DreamSim (diversity) against HPSv3 (quality); the authors describe TokenRepel's performance as the most stable across the three models.

• Large DreamSim gains at the selected operating points (Table 1, HPSv3 thresholds 7.4 / 8.8 / 5.5). FLUX.1-dev: 0.2561 → 0.3314. FLUX.2-Klein-9B: 0.1964 → 0.2643. SDXL: 0.3212 → 0.3700.

• Quality stays close to baseline. HPSv3: FLUX.1-dev 7.5266 → 7.4321; FLUX.2-Klein-9B 9.0595 → 8.9825; SDXL 5.9234 → 5.6233. Vendi Score and LPIPS rise alongside DreamSim (e.g. FLUX.1-dev Vendi 1.8747 → 2.0090, LPIPS 0.5821 → 0.6519).

• Moderate overhead. Latency relative to baseline: 1.15× (FLUX.1-dev), 1.32× (FLUX.2-Klein-9B), 1.05× (SDXL) — far below Group Inference at 10.38× on SDXL and Contextual Repulsion at 3.52× on FLUX.2-Klein-9B.

• Competitors trade off differently. Contextual Repulsion reaches higher HPSv3 than TokenRepel on FLUX.1-dev (7.6795, above the 7.5266 baseline) but drops CLIP (0.3533 vs 0.3613) and GenEval (0.6404 vs 0.6653). Group Inference achieves the highest SDXL DreamSim (0.3969) but at 10.38× latency and with much smaller gains on FLUX.2-Klein-9B.

• Layerwise preservation helps the FLUX models. At comparable diversity, layerwise correction generally yields higher HPSv3 than no preservation or final-block-only preservation; the effect is less pronounced on SDXL. Correcting more layers improves fidelity but reduces diversity, and correcting only early layers saves computation since the two auxiliary branches can stop early.

• Repulsion saturates early in sampling. Extending the repulsion window from 3 to 24 of FLUX.1-dev's 25 steps improves diversity quickly but saturates while HPSv3 declines. The measured fraction of tokens above the threshold (θ̄ = 45°) falls from about 86% at the first step to about 10% near the end — early steps set coarse structure and contain more similar cross-image patches.

• Both δ and θ̄ increase diversity at a quality cost. Sweeps on FLUX.2-Klein-9B cover δ ∈ {0°, 10°, 20°, 30°, 40°} at θ̄ = 60°, and θ̄ ∈ {0°, 20°, 40°, 60°, 80°} at δ = 30°; δ = 0 or θ̄ = 0 recovers the baseline. Higher κ also raises diversity with a slight quality drop, and the zero-temperature (argmax) setting is used for simplicity.

• Batch size behaves statistically, not monotonically. Over batch sizes 2 to 12 on FLUX.2-Klein-9B, the HPSv3 gap stays relatively stable, DreamSim decreases slightly, and Vendi Score rises — the authors attribute the DreamSim decrease to more pairwise comparisons after a one-step repulsion, not to worse diversity.

• It transfers to image editing. On KontextBench (1,026 image–instruction pairs from 108 source images) with FLUX Kontext: DreamSim 0.098 → 0.209, Vendi 1.425 → 1.699, LPIPS 0.320 → 0.502, with HPSv3 9.384 → 9.301, CLIP 0.364 → 0.357, and ImageReward 0.582 → 0.509. Gains are limited for deterministic edits like object removal and stronger for open-ended edits.

Methodology in Plain English

Think of each image as a cloud of feature tokens — one per patch or concept — living on the surface of a high-dimensional sphere (3072 dimensions for FLUX.1-dev). The angle between two tokens says how semantically alike they are, which is why the method rotates tokens rather than stretching them.

Step one — repulsion. For every token in every image, the method finds the most similar token belonging to a different image in the batch. If the angle between them is below a threshold θ̄, the token is rotated away along the great circle connecting them, by an angle δ. A token for a dog's paw pushes away from a similar paw in another image but is left alone by a sky token. This is presented as a negative gradient step on a mixture of von Mises–Fisher distributions; because the density is extremely concentrated at these dimensionalities, the softmax-weighted mean collapses to the single nearest token, so the authors use the simpler argmax form. A guardrail ensures θ̄ + δ never exceeds 90°, so a token can't be rotated past orthogonality, and coincident tokens get a random tangent direction.

Step two — damage control. A pure repulsion would drift the images into territory the model never learned, producing artifacts. Since the model can't hand over a score function directly, the authors use a proportional stand-in: the difference between the model's velocity prediction with and without repulsion gives the repulsion-induced change, and the difference between predictions under the positive prompt and a negative prompt gives a direction pointing toward higher-density, on-prompt content. If the repulsion change points against that high-density direction, its opposing component is subtracted off — a Euclidean projection onto a half-space. The corrected velocity drives the next denoising step.

Step three — do it everywhere. Rather than checking only at the final output, the same projection is applied after each transformer layer, running three parallel streams through the network (positive prompt, negative prompt, positive prompt plus repulsion). In practice the authors apply repulsion in all layers but restrict density preservation to the early layers to balance cost and benefit.

Setup. Inference steps are 25 for FLUX.1-dev, 4 for FLUX.2-Klein-9B and 50 for SDXL, with repulsion active during the first 3, 1 and 5 steps respectively. Four images per prompt are generated at resolution 512 with shared initial noise across all compared methods, and everything is evaluated on GenEval. Diversity is measured with DreamSim (ensemble of CLIP, DINO and OpenCLIP, all ViT-B/16), Vendi Score (InceptionV3 features), and LPIPS (VGG); quality with HPSv3 (Qwen-7B-based), ImageReward (BLIP-based), CLIP (ViT-H/14), PickScore, and GenEval's detection-based compositionality score.

Why This Matters

Impact on research. The paper makes a case that fidelity preservation should be an explicit, principled part of any diversity-promoting intervention rather than an afterthought or a hyperparameter tuned by hand. Framing token repulsion as gradient ascent/descent on a spherical density, and the correction as a half-space projection using the model's own score direction, gives the diversity–quality trade-off a geometric interpretation — and gives practitioners a knob (δ, an angle in degrees) that transfers across architectures. It also demonstrates that training-free interventions can work on aggressively distilled few-step models like FLUX.2-Klein-9B, where there is very little sampling trajectory to intervene in.

Real-world applications:

• Design and creative ideation — giving art directors a spread of genuinely different concepts from one brief instead of near-duplicates.
• Synthetic data augmentation — producing varied training data for downstream vision models from fixed captions.
• Test-time scaling — supplying a wider candidate pool for best-of-N selection or verifier-guided generation.
• Representation and inclusion — increasing the range of ethnicities and genders the model surfaces for a given prompt, which the paper cites as a motivation.

Industry relevance. The method requires no retraining, no external models, and no extra decoding passes; overhead is 1.05×–1.32× at the reported settings, and repulsion is confined to the first few steps so the extra branches can terminate early. That profile is compatible with production inference budgets in a way that Group Inference's 10.38× does not appear to be.

Future Directions

• Raising the diversity ceiling. The authors state that achievable diversity is fundamentally bounded by the support of the learned training distribution, and note the method is notably more effective on the FLUX models than on SDXL. Closing that model-dependent gap is open.

• Extending to autoregressive generators. The density-preservation mechanism as formulated does not apply to autoregressive image models, whose predictions are probabilities rather than gradient-like score estimates — a substitute constraint would be needed.

• Better quality preservation on the metrics that lag. Gains are clear on HPSv3 but weaker on GenEval and PickScore, and the editing experiments show a visible ImageReward and PickScore drop (0.582 → 0.509 and 21.777 → 21.591). Tightening the projection or adapting it per metric is a natural next step.

• Adaptive repulsion scheduling. Since the eligible-token fraction falls from roughly 86% to about 10% over sampling, a schedule that varies δ and θ̄ per timestep or per layer — rather than a fixed early window — could extract more diversity per unit of quality lost.

Target Audience

Graduate students and research engineers working on diffusion and flow-matching generative models, especially those interested in inference-time guidance, sampling diversity, and the diversity–fidelity trade-off. Readers need comfort with score-based generative modeling, transformer internals, and differential geometry on manifolds to follow the derivations, though the algorithm itself is short enough (two pseudocode blocks) to reimplement from the main text. Practitioners deploying T2I systems who want a cheap diversity boost without retraining will find the empirical tables and latency numbers directly useful.

Authors’ abstract

We present an inference-time approach that enhances the visual diversity of image sets generated by text-to-image models, without compromising the average image quality. Our method leverages the intrinsic geometry of hidden representations in pretrained models to diversify images generated from the same prompt. First, we introduce a token repulsion mechanism that increases the angle between feature tokens across samples. It takes a negative Riemannian gradient step on a von Mises-Fisher mixture distribution of tokens on the unit hypersphere. Second, to preserve the quality of each output image, we propose a token density preservation mechanism that projects the repulsion-induced prediction change onto a half-space aligned with the model's score direction. This encourages the updates to remain within the high-density regions learned by the generative model. We further extend this mechanism to the output of each intermediate layer, enabling finer-grained iterative guidance. We evaluate our method on three types of pretrained text-to-image models. Experiments show substantial gains in visual diversity under identical prompts, while maintaining a comparable level of image fidelity.