Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average
Authors: Benjamin Wiriyapong, Oktay Karakus, Can Eyupoglu, Kirill Sidorov
Organizations: School of Computer Science and Informatics, Cardiff University, Abacws, Cardiff, CF24 4AG, U.K. · Department of Computer Engineering, Turkish Air Force Academy, National Defence University, Istanbul, Türkiye.
Abstract
Normalising flows provide a powerful variational family for approximate inference, yet individual architectures often fail to generalise across heterogeneous posterior geometries. We revisit mixture-based flow formulations and introduce \emph{AMF\mbox{-}VI\mbox{-}sEMA}, a two-stage framework featuring a \emph{stable global weighting} mechanism based on a \emph{Simplex Exponential Moving Average} (sEMA) update. In Stage1, a heterogeneous set of experts (\textsc{RealNVP}, \textsc{MAF}, \textsc{RBIG}) are trained independently to specialise in distinct structural regimes. In Stage2, expert parameters are frozen and global mixture weights are learned through a temperature-controlled softmax of average log-likelihoods, followed by a smooth EMA update on the probability simplex. This design produces a tractable, data-agnostic gating mechanism (without per-sample gating or gradient backpropagation through weights) that adaptively reallocates capacity while avoiding component collapse. We evaluate the framework on ten posterior benchmarks: six canonical 2D synthetic families (Banana, X-Shaped, Bimodal, Multimodal, Two-moons, Rings) and four real/low-dimensional Bayesian targets (BLR, BPR, Weibull, Real-GMM2), with stronger baselines (\textsc{NICE}, \textsc{ResFlow}, and EM-Mixing). Comprehensive evaluation covers NLL, KL divergence, Wasserstein-2 distance, and MMD, together with diagnostics of mixture dynamics, hyperparameter sensitivity, and cross-seed robustness. Empirically, \emph{AMF\mbox{-}VI\mbox{-}sEMA} achieves consistent NLL improvements over its predecessor \emph{AMF\mbox{-}VI} and avoids the catastrophic transport failures of single-flow baselines, while maintaining stable weight trajectories (Neff>1.4 on all datasets) with minimal computational overhead.
Flow matching (FM) trains a time-dependent vector field that transports samples from a simple prior to a complex data distribution. However, for high-dimensional images, each training sample supervises only a single trajectory and intermediate point, yielding an extremely sparse and high-variance training signal. This under-constrained supervision can cause flow collapse, where the learned dynamics memorize specific source-target pairings, mapping diverse inputs to overly similar outputs, failing to generalize. We introduce Posterior-Augmented Flow Matching (PAFM), a theoretically grounded generalization of FM that replaces single-target supervision with an expectation over an approximate posterior of valid target completions for a given intermediate state and condition. PAFM factorizes this intractable posterior into (i) the likelihood of the intermediate under a hypothesized endpoint and (ii) the prior probability of that endpoint under the condition, and uses an importance sampling scheme to construct a mixture over multiple candidate targets. We prove that PAFM yields an unbiased estimator of the original FM objective while substantially reducing gradient variance during training by aggregating information from many plausible continuation trajectories per intermediate. Finally, we show that PAFM improves over FM by up to 3.4 FID50K across different model scales (SiT-B/2 and SiT-XL/2), different architectures (SiT and MMDiT), and in both class and text conditioned benchmarks (ImageNet and CC12M), with a negligible increase in the compute overhead. Code: https://github.com/gstoica27/PAFM.git.
One-step generative modeling has emerged as a leading approach for amortizing the inference cost of diffusion and flow-matching models. Among distillation-free methods, MeanFlow training is notoriously unstable, with non-decreasing loss and unbounded gradient variance. In this work, we establish a theory that attributes this pathology to a misuse of the conditional velocity field. We show that the conditional velocity plays two distinct statistical roles in the loss: both as an unbiased regression target and as a Monte Carlo control variate in a Jacobi-vector product, with the original MeanFlow loss assigning the wrong coefficient to the latter. We derive the optimal coefficient in closed form and show that a family of fixes in concurrent works corresponds to different practical realizations of the same optimum. A controlled sweep of this coefficient on two-dimensional benchmarks and on a latent Diffusion Transformer recovers the predicted bias-variance ordering. Our DiT experiment also reveals a quantitative FID-MSE landscape mismatch. Specifically, although the gradient-MSE is minimized at an interior coefficient value near β=0.94, the coefficient that minimizes FID prefers to use conditional velocity directly at the unbiased corner. Our analysis therefore explains why MeanFlow is unstable and unifies its concurrent remedies, and shows that the variance-optimal coefficient need not coincide with the quality-optimal one.
This paper builds bridges between two families of probabilistic algorithms: (hierarchical) variational inference (VI), which is typically used to model distributions over continuous spaces, and generative flow networks (GFlowNets), which have been used for distributions over discrete structures such as graphs. We demonstrate that, in certain cases, VI algorithms are equivalent to special cases of GFlowNets in the sense of equality of expected gradients of their learning objectives. We then point out the differences between the two families and show how these differences emerge experimentally. Notably, GFlowNets, which borrow ideas from reinforcement learning, are more amenable than VI to off-policy training without the cost of high gradient variance induced by importance sampling. We argue that this property of GFlowNets can provide advantages for capturing diversity in multimodal target distributions.
Esmeralda S. Whitammer, Salem Lahlou, Tristan Deleu +5