Importance-weighted learning improves conditional distributions in flow models
Learning $ω$-Powered Conditional Distributions with Importance-Weighted Flow Matching
Machine LearningComputer Vision and Pattern Recognition
Summary
Generating data that matches specific conditions is important but tricky, especially when using popular guidance techniques that don't perfectly hit their target. The authors study a method to explicitly learn better versions of these guided distributions by reweighting examples during training using fixed models. Their approach works well on both simple test problems and complex image data, showing improved alignment to desired distributions. They also analyze how the model's rankings shift under this approach, revealing detailed class-specific effects.
What this means in practice
- •For machine learning engineers: Train conditional generative models that better approximate target powered distributions using importance-weighted flow matching.
- •For computer vision teams: Improve conditional image generation alignment on datasets like ImageNet by adding lightweight correction modules trained with importance weights.
Authors
Maciej Falkiewicz, Frantzeska Lavda, Jason Ramapuram, Alexandros Kalousis
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.