Don't Settle at the Mode! Mitigating Diversity Collapse in Pretrained Flow Models via Feature Self-Guidance

Paper: 2606.27371 Authors: Pradhaan S Bhat, Rishubh Parihar, Abhijnya Bhat, R. Venkatesh Babu Categories: cs.CV

The Gap

Existing flow models (e.g., Rectified Flow, Flow Matching) can generate stunning images, but when asked to produce multiple samples under the same text/image conditioning, they collapse to near-identical outputs. Two families of fixes exist: (1) latent guidance—manipulating the initial noise or intermediate latent codes—which is cheap but only pushes diversity slightly; (2) sample selection—generating many candidates and picking a diverse subset via an external reward model (e.g., CLIP-based diversity score)—which is expensive at inference time and requires a separately trained judge. Neither solves the root cause: the flow model’s internal features themselves are insufficiently disentangled across the batch. This paper asks: can we directly perturb the model’s own feature activations during generation to force diversity, while keeping those perturbations on the natural image manifold?

[Problem] --> Diversity Collapse under same conditioning
       |
       v
[Assumption] --> Internal features can be dispersed for diversity,
                 but must stay on-manifold to avoid quality loss
       |
       v
[Method] --> Feature Self-Guidance (disperse) + Manifold Regularization (project)
       |
       v
[Evidence] --> Higher diversity metrics (LPIPS, Vendi) with minimal FID/CLIP loss
       |
       v
[Conclusion] --> Training-free, plug-and-play, marginal cost

The Increment

One sentence: Before this paper, you could either accept limited diversity (latent guidance) or pay heavy inference overhead (sample selection); after this paper, you can get high diversity for free by nudging internal features during batch generation.

Core Mechanism

The method requires no training or fine-tuning. It acts on a batch of conditioning inputs (e.g., same text prompt, or same depth map) that share the same flow model. At each denoising step, the model produces intermediate feature maps (e.g., from U-Net blocks). The key idea is to disperse these features across the batch so that each sample’s features differ from others, then regularize them back toward the natural image manifold to prevent artifacts.

Concretely, for each feature position (e.g., spatial location in a feature map), the method computes the mean feature over the entire batch. Then it pushes each sample’s feature away from that mean (repulsion). This creates diversity. However, unconstrained repulsion can push features off the manifold, leading to unnatural outputs. So a regularization step projects the dispersed features onto the low-dimensional subspace spanned by the batch’s principal components of features (i.e., the local linear approximation of the manifold). This ensures that after dispersion, features still lie close to the data manifold.

+------------------+      +-----------------------------+
| Input batch      |      | Batch of features at step t |
| (same condition) |----->| [z_1, z_2, ..., z_B]       |
+------------------+      +-----------------------------+
                                     |
                                     v
                     +-------------------------------+
                     | Compute batch mean: mu_z       |
                     | For each i: z_i' = z_i +       |
                     |   alpha * (z_i - mu_z)         |
                     |   / ||z_i - mu_z||             |
                     +-------------------------------+
                                     |  Feature Self-Guidance
                                     v
                     +-------------------------------+
                     | Compute batch covariance     |
                     | PCA on z_i'                   |
                     | Project each z_i' onto        |
                     | top-k principal components    |
                     +-------------------------------+
                                     |  Manifold Regularization
                                     v
                      +-------------------------------+
                      | z_i'' = regularized features |
                      | (used for next denoising step)|
                      +-------------------------------+
                                     |
                                     v
                            +--------+--------+
                            | Decoder/Output  |
                            +-----------------+

Structural metaphor: a choir singing in unison, then given individual improv rules.
Think of a choir (the batch) given the same melody (conditioning). Without intervention, every singer produces almost the same sound (diversity collapse).

• Feature Self-Guidance is like whispering a unique “twist” to each singer: “slightly raise your third note” or “add a vibrato on the fifth”. These instructions are different per singer, so the overall sound diversifies.
• But if the twist is too extreme, the singer might sound like a cat. Manifold Regularization is the choir director who listens to the modified sound and gently nudges it back onto a “natural singing” subspace (a set of notes that still fit the song’s style). The director doesn’t tell each singer exactly what to sing—just ensures the twist stays within the acceptable range.
  The whole process repeats at every section of the piece (every denoising step), and no script (training) is needed—just the director’s ear (PCA).

Key Concepts

• Diversity Collapse: When a generative model produces multiple outputs from the same condition, the outputs are often nearly identical. This is especially severe in flow models because they map a simple prior to the data distribution via a deterministic ODE; if the initial noise is the same, the trajectory collapses. Intuitively, the model “settles at the mode” of the conditional distribution instead of exploring the full manifold.

• Feature Self-Guidance: Instead of using an external reward model to judge diversity after generation, this method directly modifies the intermediate features *during generation. It computes a batch-level repulsion force on the features. Concretely, if two samples have similar features, they are pushed apart in feature space. This is done feature-wise (spatially) and can be applied at any layer. The strength is controlled by a hyperparameter alpha.

• Manifold Regularization: The repulsion step alone can produce out-of-distribution features. The method projects the dispersed features back onto the local linear manifold defined by the batch’s principal components (top-k eigenvectors of the feature covariance). This is a cheap, batch-level operation that prevents quality degradation without requiring an explicit density model. It ensures that the diversity injection stays “on the manifold” of plausible images.

Framework Shift

Before (mainstream approach):        After (this paper):

  External Reward Model                Internal Feature Manipulation
  +---------------------------+        +---------------------------+
  | Generate N candidates    |        | Generate batch together   |
  | Score each for diversity |        | At each step, spread      |
  | Pick top-k              |        | features via repulsion    |
  | (cost: Nx compute +     |        | + manifold projection     |
  |  reward model inference)|        | (cost: marginal extra per |
  +---------------------------+        |  step, no external model)|
                               |
  Diversity limited by reward         Diversity controlled by alpha
  model's own biases                  + PCA rank; no bias

One sentence: From external model-based selection to internal feature-based guidance, the core shift is moving diversity control inside the generation loop, making it both cheaper and more direct.

Expert Assessment

Problem choice: Real gap—diversity collapse in flow models is a known frustration for practitioners. The paper correctly identifies that existing remedies are either weak or expensive. However, it’s a niche within generative modeling; not a field-defining problem.

Method maturity: Clever insight—using batch-level repulsion on intermediate features is simple and elegant. The manifold regularization via PCA is a bit brute-force (linear approximation of a nonlinear manifold), but it works empirically. Could there be a simpler approach? Possibly just using dropout-like noise, but the paper shows this method outperforms that.

Experimental integrity: Baselines are reasonable: they compare against vanilla flow, latent guidance, and sample selection with CLIP. Tables show consistent gains in diversity (LPIPS, Vendi) with negligible loss in FID and CLIP score. One red flag: all experiments are on a single flow model architecture (Stable Diffusion v2.1-based); generalizability to other flow backbones is not strongly demonstrated. Also, the hyperparameter alpha is tuned per task—may require calibration.

Writing quality: Clear structure, but the “Feature Self-Guidance” section could use a toy diagram (included here). The paper’s theory section is sparse; more intuition about why repulsion works and why PCA projection doesn’t kill diversity would help.

Verdict: weak accept — A practical, training-free technique that addresses a real pain point with a clever yet simple mechanism. Not a breakthrough, but a solid contribution that many practitioners will adopt.

Takeaways

• Steal the idea of batch-level repulsion on intermediate features: This can be applied to any iterative generative model (diffusion, flow, score-based) to increase diversity. Just compute batch mean and push apart.
• Steal the manifold regularization trick: When you add diversity to latent features, always project back to the batch’s principal subspace to avoid quality collapse. This is a cheap post-hoc fix.
• Steal the plug-and-play design: The method requires no training, no extra reward model. It can be injected into existing inference pipelines with a few lines of code. If you run batch generation, you can adopt this immediately.