📊 statistics

Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average

This paper introduces AMF-VI-sEMA, a two-stage framework that enhances normalizing flow mixtures by training diverse experts independently and then applying a stable, data-agnostic global weighting mechanism via Simplex Exponential Moving Average to achieve robust posterior approximation across heterogeneous geometries without component collapse.

Original authors: Benjamin Wiriyapong, Oktay Karakus, Can Eyupoglu, Kirill Sidorov

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin Wiriyapong, Oktay Karakus, Can Eyupoglu, Kirill Sidorov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to draw a perfect map of a very strange, complex landscape. This landscape might have deep valleys, sharp peaks, winding rivers, or multiple islands. In the world of machine learning, this "landscape" is called a posterior distribution, and it represents all the possible answers to a problem, along with how likely each answer is.

The paper introduces a new method called AMF-VI-sEMA to help computers draw these maps more accurately. Here is how it works, explained through simple analogies.

The Problem: One Artist Can't Paint Everything

Traditionally, computers try to learn these maps using a single "artist" (a mathematical model called a Normalizing Flow).

Artist A (RealNVP) is great at drawing straight lines and blocky shapes but struggles with curves.
Artist B (MAF) is excellent at drawing complex, twisting paths but might miss the big picture.
Artist C (RBIG) is a master of handling weird, heavy-tailed shapes but might be slow.

If you hire just Artist A, they will fail if the map needs Artist B's skills. If you hire just Artist B, they will fail on Artist A's terrain. The old way of trying to make them work together was like forcing them to paint on the same canvas at the same time; they would argue, get confused, and the final picture would be a mess.

The Solution: A Two-Stage "Team of Experts"

The authors propose a new strategy that treats these artists like a specialized team of experts who work in two distinct phases.

Stage 1: Independent Training (The "Specialization" Phase)

First, the computer hires each artist and tells them to work alone.

Artist A practices on the data until they are a master of their specific style.
Artist B does the same.
Artist C does the same.
They don't talk to each other yet. They just get really good at their own specific jobs. This ensures that by the time they meet, they are all highly skilled specialists.

Stage 2: The "Smart Manager" (The sEMA Phase)

Once the artists are trained, the computer freezes their skills (so they don't change) and brings in a Manager. This manager's job is to decide how much of the final map should be painted by Artist A, Artist B, or Artist C.

Here is where the paper's secret sauce comes in: The Simplex Exponential Moving Average (sEMA).

Imagine the Manager is looking at a test map (data the artists haven't seen before) and asking, "Who did the best job here?"

If Artist A did a great job, the Manager gives them a bigger share of the work.
If Artist B did poorly, the Manager gives them less work.

The "Moving Average" Trick:
In the past, managers would change their minds instantly based on the last thing they saw. If Artist A had one bad day, the Manager might fire them immediately. This causes the team to be unstable and chaotic.

The sEMA method is like a wise, calm manager who doesn't panic. Instead of changing the team's composition instantly, they adjust the weights slowly and smoothly over time.

If Artist A is consistently good, their share of the work grows gradually.
If Artist B has a bad batch of data, the Manager doesn't fire them immediately; they wait to see if it was just a fluke.

This "smooth averaging" prevents the team from collapsing into just one artist (who might be wrong) and ensures the final map is a stable, reliable blend of everyone's best work.

Why This Matters

The paper tested this method on 10 different types of "landscapes," ranging from simple shapes to complex, multi-peak mountains.

No More Crashes: Single artists often fail completely on complex shapes (like trying to draw a circle with a ruler). The team approach never crashes; it always finds a good solution.
Stability: Because the Manager adjusts weights slowly (using the sEMA), the system doesn't wobble or oscillate. It finds a steady rhythm.
Smart Allocation: The system automatically figures out which artist is best for which part of the map. For example, on a "Ring" shape, it might let Artist A do 90% of the work, while on a "Two-Moons" shape, it splits the work evenly between Artist A and B.

The Bottom Line

The paper claims that AMF-VI-sEMA is a robust, stable, and efficient way to combine different machine learning models. Instead of forcing them to fight for dominance or training them together in a chaotic mess, it lets them specialize first, then uses a calm, smoothing algorithm to blend their strengths together.

The result is a system that is better at understanding complex, confusing data than any single model could be on its own, without requiring extra computing power or complex adjustments. It's like having a dream team where everyone knows their job, and a wise manager ensures they work together smoothly.

Technical Summary: Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average

Problem Statement
Normalising flows have significantly expanded the flexibility of variational inference (VI) by transforming simple base distributions into expressive probability densities. However, individual flow architectures (e.g., RealNVP, MAF, RBIG) often fail to generalise across heterogeneous posterior geometries. Different flows possess distinct inductive biases, optimisation dynamics, and numerical characteristics, leading to a situation where no single architecture performs consistently across diverse target distributions. While mixture-based approaches offer a natural strategy to combine complementary flows, existing frameworks suffer from significant limitations:

Joint Optimisation: Learning mixture weights and expert parameters simultaneously often leads to unstable training or "expert domination" due to the coupling of heterogeneous optimisation dynamics.
Fixed Global Weighting: Estimating fixed coefficients during training and holding them constant during inference lacks the flexibility to exploit complementary expert behaviour, often resulting in inefficient capacity allocation.

Consequently, there is a need for a stable, principled optimisation framework that can combine heterogeneous experts without requiring prior knowledge of the most appropriate architecture or incurring the instability of joint training.

Methodology: AMF-VI-sEMA
The paper proposes AMF-VI-sEMA (Adaptive Mixture of Flows - Variational Inference with Simplex Exponential Moving Average), a two-stage framework designed to decouple expert learning from mixture optimisation.

Stage 1: Independent Expert Specialisation
A heterogeneous set of experts (specifically RealNVP, MAF, and RBIG) are trained independently to specialise in distinct structural regimes. Each expert $q_k(z|\phi_k)$ is optimised to maximise its own log-likelihood on the training data. Once trained, the expert parameters $\phi_k$ are frozen. This stage ensures that each component converges under its own optimisation dynamics, avoiding the coordination pathologies of joint training.
Stage 2: Simplex EMA (sEMA) Global Weighting
With expert parameters frozen, the framework learns a global mixture weight vector $\pi$ on the probability simplex. This is achieved through a likelihood-driven, temperature-controlled mechanism:
1. Responsibility Calculation: For a batch of held-out validation data, responsibilities $r_k(z)$ are computed using a softmax of the average log-likelihoods of the experts, scaled by a temperature parameter $\tau$ .
2. Smoothing and Averaging: The responsibilities are averaged over the validation batch and optionally smoothed toward a uniform distribution to dampen transients.
3. Simplex EMA Update: The global weights $\pi$ are updated via an Exponential Moving Average (EMA) directly on the simplex: $\pi \leftarrow \alpha \pi + (1-\alpha)\bar{r}$ . A small floor $\epsilon$ is applied to prevent zeroing, followed by renormalisation.
This design creates a tractable, data-agnostic gating mechanism that reallocates capacity adaptively without per-sample gating or backpropagation through the weights. The framework monitors the effective number of experts ( $N_{eff} = \exp(H(\pi))$ ) to diagnose component collapse.

Key Contributions

Decoupled Optimisation Framework: The paper introduces a two-stage strategy that separates expert learning from mixture adaptation. By freezing experts before updating weights, it avoids the instability of joint responsibility-weighted optimisation while remaining architecture-agnostic.
Practical Adaptive Weighting Mechanism: The authors propose a computationally efficient, responsibility-based Simplex EMA update. This method estimates global mixture weights without requiring joint optimisation or per-sample gating, providing stable expert cooperation with minimal computational overhead.
Comprehensive Empirical Evaluation: The framework is evaluated on ten benchmark problems, including 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). The evaluation extends beyond standard log-likelihood to include Wasserstein-2 distance, Maximum Mean Discrepancy (MMD), calibration analysis, and robustness across random seeds.
Analysis of Optimisation Dynamics: The paper provides an analysis of expert utilisation, mixture-weight evolution, and hyperparameter sensitivity, offering insights into why heterogeneous flow mixtures succeed and how the sEMA mechanism stabilises training.

Results
Empirical evaluations demonstrate that AMF-VI-sEMA achieves consistent improvements over its predecessor (AMF-VI) and avoids the catastrophic transport failures observed in single-flow baselines (such as ResFlow and EM-Mix).

Performance: The framework achieves best or near-best Negative Log-Likelihood (NLL) across all synthetic and real datasets. It maintains robust transport behaviour (Wasserstein-2) and discrepancy metrics (MMD), avoiding the geometric failures seen in single architectures (e.g., NICE and ResFlow showing high transport errors on curved geometries).
Stability: The method maintains stable weight trajectories with an effective number of experts $N_{eff} > 1.4$ on all datasets, indicating that the mixture does not collapse to a single expert.
Robustness: Cross-seed evaluations show low variance in NLL, KL divergence, and Wasserstein distance, confirming that the framework is reproducible and robust to random initialisation.
Adaptivity: The learned global weights adapt to the target geometry. For instance, on regression-style posteriors (BLR, BPR), the mechanism suppresses RealNVP in favour of MAF and RBIG, while on ring-like structures (Rings), it heavily concentrates mass on RealNVP.

Significance and Claims
The paper claims that AMF-VI-sEMA provides a robust and reliable optimisation strategy for posterior approximation across heterogeneous landscapes without requiring prior knowledge of which flow architecture is most appropriate. The significance of the work lies in:

Stability: The sEMA mechanism introduces a temporal regularisation effect that stabilises gating dynamics and mitigates volatile expert switching.
Simplicity and Efficiency: The approach is architecture-agnostic, requires no modification to underlying flow models, and introduces negligible computational overhead compared to joint training.
Reliability: By avoiding the fragility of joint responsibility-weighted training, the framework offers a practical solution for modelling complex, non-convex, and multimodal posterior geometries.

The authors note that while specialised single-flow architectures may remain preferable for simple unimodal targets, AMF-VI-sEMA consistently outperforms them in robustness across diverse geometries. Future work is suggested to explore conditional or amortised gating strategies and extensions to higher-dimensional settings.