Stick-Breaking Mixture Normalizing Flows with Component-Wise Tail Adaptation for Variational Inference

2025-10-09

AI Summary

Standard Gaussian normalizing flows struggle to model multimodal and heavy-tailed posterior distributions in Bayesian inference. To address this, we propose Stick-Breaking Mixture Normalizing Flows (SB-MNF), which innovatively integrates the modal expressiveness of Stick-Breaking mixture bases with learnable, component-wise tail transformations for anisotropic heavy-tail modeling. We introduce a weighted-component Evidence Lower Bound (ELBO) objective to mitigate mode collapse induced by reverse KL divergence. Furthermore, SB-MNF employs a shared backbone network coupled with component-specific tail parameterizations, ensuring both tractable density evaluation and training stability while enabling exact density estimation. Experiments demonstrate that SB-MNF significantly improves multimodal coverage and tail recovery accuracy on synthetic benchmarks, and achieves superior posterior inference performance over existing flow-based methods on real-world tasks.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Normalizing flows with a Gaussian base provide a computationally efficient way to approximate posterior distributions in Bayesian inference, but they often struggle to capture complex posteriors with multimodality and heavy tails. We propose a stick-breaking mixture base with component-wise tail adaptation (StiCTAF) for posterior approximation. The method first learns a flexible mixture base to mitigate the mode-seeking bias of reverse KL divergence through a weighted average of component-wise ELBOs. It then estimates local tail indices of unnormalized densities and finally refines each mixture component using a shared backbone combined with component-specific tail transforms calibrated by the estimated indices. This design enables accurate mode coverage and anisotropic tail modeling while retaining exact density evaluation and stable optimization. Experiments on synthetic posteriors demonstrate improved tail recovery and better coverage of multiple modes compared to benchmark models. We also present a real-data analysis illustrating the practical benefits of our approach for posterior inference.

Problem

• Capturing multimodal and heavy-tailed posteriors in variational inference
• Overcoming mode-seeking bias through flexible mixture base distributions
• Enabling anisotropic tail modeling while maintaining exact density evaluation

Innovation

• Stick-breaking mixture base with tail adaptation
• Component-wise ELBOs mitigate mode-seeking bias
• Shared backbone with calibrated tail transforms