Learning $ω$-Powered Conditional Distributions with Importance-Weighted Flow Matching

AI summary

Overview

Research area: Generative modeling — specifically guidance in diffusion and flow-matching models, and the theory of "powered" conditional distributions.

Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion ODEs, classifier-free guidance, importance sampling, and density estimation via the instantaneous change-of-variables formula.

Scope: The paper proposes Importance-Weighted Conditional Flow Matching (IW-CFM), a method that explicitly learns the ω-powered conditional distribution that classifier-free guidance (CFG) is commonly assumed to sample from, and evaluates it on synthetic 2D distributions, ImageNet 256×256 generation, and an exploratory CLIP-based analysis of powered-density rankings of real ImageNet images.

What This Paper Is About

Classifier-Free Guidance is usually justified by saying that a CFG model with guidance scale γ samples from the powered distribution p_ω(x|y) ∝ p(x|y)^ω p(x)^(1−ω) when ω = γ, but the authors note this identification does not hold in general — the CFG-guided backward vector field lacks a path-consistent forward process, so its endpoint distribution need not equal p_ω. The paper asks whether that powered target can instead be learned explicitly, by reweighting conditional flow-matching training examples with density ratios computed from frozen conditional and unconditional flow models.

Key Contributions

1. A conditional instantiation of energy-weighted flow matching. Building on Energy-Weighted Flow Matching (Zhang et al., 2025), the authors define endpoint weighting for model-defined ω-powered conditional targets using density ratios from frozen conditional and unconditional flow models, and show that with exact normalized weights the population objective has the same gradient as flow matching on the powered target.

2. A hierarchical self-normalized estimator. For heterogeneous (y, ω) pairs, weights are normalized only within samples sharing the same (y_j, ω_j), preserving the intended outer distribution p(y)p(ω) rather than reweighting pairs by their partition functions. The finite-sample procedure is explicitly distinguished from the exact population objective.

3. A lightweight, amortized corrector plus a high-dimensional ImageNet construction. ω is treated as a conditioning variable for a single corrector (about 5% additional parameters over the baseline SiT), and the ImageNet implementation uses regularized precomputed weights, proposal sampling, and interpolation across powering parameters.

4. Empirical validation and an exploratory CLIP study. Recovery of analytical powered targets on synthetic distributions, alignment with regularized empirical references on ImageNet, and an exploratory characterization of powered-density rankings on real images through selected CLIP contrasts, which show heterogeneous and nonmonotonic classwise responses.

Main Findings

• Synthetic recovery succeeds; matched CFG does not. On two-dimensional synthetic datasets (a three-mode Gaussian mixture and a fractal-inspired distribution), CFG with numerically matched guidance γ = ω departs visibly from the analytical target contours, whereas IW-CFM samples at ω = 4 cover the region enclosed by the target's 99.7% probability-mass contour.

• Better alignment with regularized ImageNet references. Under the guided practical configuration, IW-CFM shows substantially lower Fréchet distances (FID, FD-DINOv2) to the "weighted" and "filtered" references than fixed-guidance CFG at γ = 1.5 or matched-guidance CFG at γ = ω, across the evaluated target values of ω.

• PRDC scores favor IW-CFM. Against the filtered regularized references, IW-CFM achieves the best or near-best Precision, Recall, Density, and Coverage scores across ω = 4, 8, 16, 32, 64, 128. For example, at ω = 128: CFG (γ = 1.5) gives Precision 0.61, Recall 0.57, Density 0.94, Coverage 0.86; CFG (γ = 128.0) gives 0.10, 0.00, 0.08, 0.01; IW-CFM (ω = 128.0) gives 0.68, 0.63, 1.06, 0.96.

• Matched-guidance CFG degrades at high guidance. In the displayed qualitative ImageNet example, matched-guidance CFG loses coherent image structure for γ ≥ 32, while IW-CFM remains visually coherent at the corresponding target values with inference-time guidance fixed at γ = 4. The PRDC table shows the same pattern, with CFG at γ = 64.0 dropping to Precision 0.04 and Coverage 0.05.

• Interpolation across ω works. IW-CFM evaluated at non-integer values of ω has FID and FD-DINOv2 consistent with the neighboring integer targets, indicating the learned conditioning interpolates effectively.

• Density-estimation precision interacts with powering. In an ablation on a fixed subset of 100 ImageNet categories, results are relatively close at ω ∈ {4, 8}, begin to separate at ω = 16, and diverge substantially for ω ≥ 32 — consistent with the paper's analysis that an additive error in the estimated log-density ratio is multiplied by (ω − 1) before exponentiation.

• Regularization parameters behave qualitatively as expected. Across 20 configurations of temperature T ∈ {1, 1.2, 1.5, 2.5, 5} and clipped upper fraction τ ∈ {0.5%, 1%, 2.5%, 5%}, increasing τ consistently reduces target-relative distances because it moves the operational target toward the unweighted class-conditional empirical distribution. The default (T, τ) = (1.2, 1%) is described as comparatively mild and stable.

• No claim of superiority over related methods. REG and FKC were evaluated on the exact conditional and unconditional vector fields of the three-Gaussian problem; under the tested configurations neither closely matched the analytical powered contours, so no image-scale comparison was conducted and the ImageNet results make no claim of superiority over REG or FKC.

• Real-data rankings are heterogeneous. In the exploratory analysis, rankings of ImageNet instances by the powered-density score change with ω. In the displayed "hammerhead shark" examples, top-ranked images at lower ω have relatively simple compositions while those at higher ω include more complex scenes; one anomalous image (a great white shark with digitally added saber teeth) is bottom-ranked at ω = 32 but does not appear among the extremes at lower ω. Aggregate shifts across selected prompt-specific contrasts coexist with heterogeneous classwise responses, including nonmonotonic trajectories.

• Memorization check. Because distributional metrics do not rule out memorization, a separate test for training-data retrieval is reported in the paper (Section D.5) and reports low rates under its specified test.

Methodology in Plain English

Flow matching trains a model to move samples from a simple prior (a Gaussian) to data by learning a velocity field along an interpolating path. CFG steers such a model toward a condition y by mixing conditional and unconditional velocity predictions, with a guidance scale γ. The paper's key idea is that instead of relying on CFG to approximate the powered distribution, you can train directly on it by weighting each training example.

The weights come from the ratio between the conditional and unconditional densities of the training example: samples that the conditional model considers much more likely than the unconditional model get up-weighted, and the strength of that up-weighting depends on the powering parameter ω. Because the exact normalizing constant of the powered distribution is unknown, the authors use self-normalized importance sampling within a batch. Crucially, when training across many (y, ω) pairs at once, they use hierarchical batches and normalize only within samples sharing the same (y, ω); normalizing globally would accidentally reweight the pairs by their partition functions.

For density estimation, they use the instantaneous change-of-variables formula and approximate the divergence with the Hutchinson trace estimator, with the base conditional and unconditional models kept frozen so that log-densities can be precomputed.

On ImageNet, exact densities are unavailable, so the authors substitute densities induced by frozen pretrained SiT vector fields. They note that powering amplifies density-estimation errors multiplicatively by exp((ω − 1)ε), so unregularized weights become extremely concentrated. They therefore apply temperature scaling with T > 1 and per-class clipping at an upper fraction τ, then interpolate these regularized weights across a grid of integer powering parameters to define a continuous operational family. Training samples instances using the nearest-integer proposal.

The corrector itself is a parameter-efficient adapter inserted into the frozen Transformer blocks of a pretrained SiT-XL/2, conditioned on ω, adding about 5% of parameters over the baseline. Training uses ω ∼ U[4, 128], and the class token is replaced with a null token during 10% of iterations so that CFG can optionally be applied at inference. Evaluation uses FID, FD-DINOv2, and Inception-based Precision, Recall, Density, and Coverage against weighted and filtered references derived from the regularized weights.

Why This Matters

This work sharpens a long-standing but often glossed-over gap between the motivation for CFG and its actual behavior. Earlier analyses (Bradley and Nakkiran, 2024; Chidambaram et al., 2024) diagnosed that CFG does not generally sample the powered distribution; IW-CFM offers a practical route to approximate it, moving the heavy computation into density precomputation and corrector training so that sampling still uses a single learned vector field.

• Controllable image generation: Providing a principled knob (ω) for trading off condition adherence against diversity, rather than a heuristic guidance scale that saturates or produces artifacts at high values.
• Content moderation and safety filtering: Studying distributions that emphasize features discriminative of a condition may help characterize what a model considers "representative" of a class, with the caveat the authors note that model likelihood and perceptual quality are not the same thing (citing Karczewski et al., 2025).
• Data curation and dataset auditing: Powered-density rankings on real data can surface anomalous or mislabeled examples — as in the reported shark example — without retraining a generative model.
• Guidance distillation: The amortized, ω-conditioned corrector is related to guidance-distillation architectures, suggesting a route to deploying a single network that spans a family of guidance strengths.

Industry relevance. The approach builds directly on widely used components (Stable Diffusion VAE latents, pretrained SiT backbones) and adds only a small parameter overhead, so it is architecturally compatible with existing diffusion infrastructure. The cost, however, is density precomputation over the training set, which the authors describe as the main computational shift and which requires care with numerical precision as ω grows.

Future Directions

• Extending exact target recovery to high dimensions. The paper's exact gradient-equivalence claim applies to the population formulation and controlled synthetic experiments; on ImageNet the target is an explicitly defined regularized empirical reference. Closing that gap, or building reliable exact densities at scale, is a natural next step.
• Improving density estimation under strong powering. The ablation shows performance divergence for ω ≥ 32, tied to the (ω − 1) amplification of log-density error. Better estimators or precision-adaptive schemes would extend the usable range of ω.
• Principled selection of regularization parameters. The authors note that lower distances under more aggressive clipping do not imply better target fidelity, and suggest τ should be the smallest value that prevents pathological weight concentration. A more principled criterion for choosing (T, τ) remains open.
• Broader comparison with inference-time alternatives. The paper explicitly makes no claim of superiority over REG or FKC at image scale, leaving an image-scale comparison as an open question. The exploratory CLIP analysis of powered-density rankings also invites a fuller study of classwise and nonmonotonic behavior.

Target Audience

Researchers and practitioners working on diffusion and flow-matching generative models, particularly those interested in guidance mechanisms, controllability, and the theory connecting CFG to powered distributions. It is also relevant to readers interested in importance sampling and density-ratio estimation applied to high-dimensional generative models, and to engineers who need to control the trade-off between sample fidelity and condition adherence in deployed image generation systems.

Authors’ abstract

Classifier-Free Guidance (CFG) with guidance scale $γ$ is commonly motivated by the $ω$-powered conditional distribution obtained by setting $ω=γ$, but CFG does not, in general, sample from this distribution exactly. We study whether the powered target can instead be learned explicitly. We introduce Importance-Weighted Conditional Flow Matching (IW-CFM), which reweights conditional flow-matching examples using density ratios obtained from frozen conditional and unconditional flow models. With exact normalized weights, the resulting objective has the same gradient as flow matching on the powered target; in practice, we use self-normalized estimation and a regularized empirical approximation for high-dimensional data. On synthetic distributions, IW-CFM recovers analytical powered targets for which numerically matched CFG, $γ=ω$, exhibits a visible mismatch. On ImageNet, an amortized lightweight corrector achieves better alignment than fixed- and matched-guidance CFG with the corresponding weighted and filtered reference distributions. We additionally present an exploratory CLIP analysis of powered-density rankings on real ImageNet data. Aggregate shifts across selected prompt-specific contrasts coexist with heterogeneous classwise responses, including nonmonotonic trajectories. These results provide a practical route for studying and approximating model-defined powered conditional distributions.