Variational Inference via Entropic Transport Descent

2026-06-23

AI Summary

This work addresses the challenges of variance collapse and mode collapse in particle-based variational inference under high-dimensional and multimodal settings by introducing a novel approach grounded in entropy-regularized optimal transport. The method models particle updates as a JKO proximal step in a coupling space and leverages a relaxed form of the KL chain rule together with Sinkhorn iterations to achieve score-free, globally coordinated particle evolution. To the best of our knowledge, this is the first integration of entropy-regularized optimal transport into particle variational inference, effectively mitigating collapse phenomena while preserving multimodal structure. Empirical results demonstrate significant improvements over SVGD, AGF-SVGD, and SGLD across tasks including high-dimensional Bayesian logistic regression, neural network posterior inference, and molecular Boltzmann distribution modeling, with particularly pronounced gains in multimodal and high-dimensional regimes.

Abstract

Particle-based variational inference (ParVI) methods approximate an intractable target distribution by evolving an ensemble of interacting samples. Existing approaches rely predominantly on kernel-based repulsion (e.g., SVGD), which suffers from variance collapse in high dimensions and mode collapse on multimodal targets -- pathologies caused by the absence of global transport structure. We introduce entropic transport descent (ETD), a ParVI family that frames each particle update as an entropy-regularized optimal transport problem. Derived from the JKO proximal scheme by lifting to the space of couplings and relaxing via the KL chain rule, each ETD iteration reduces to a Sinkhorn computation. The resulting transport plan provides global coordination, guiding each particle to nearby high-density proposals and naturally preserving multimodal structure. ETD can operate entirely score-free, requiring only pointwise evaluations of the unnormalized target density. Experiments on variance-collapse diagnostics, Bayesian logistic regression, neural networks, and molecular Boltzmann distributions show that ETD matches or outperforms SVGD, AGF-SVGD, and SGLD, with the largest gains in high-dimensional and multimodal settings.

Problem

• variance collapse
• mode collapse
• high-dimensional inference
• multimodal distributions
• particle-based variational inference

Innovation

• Entropic Transport Descent
• Particle-based Variational Inference
• Optimal Transport
• Sinkhorn Algorithm
• Multimodal Inference