This paper introduces Bernstein Normalizing Flows (BNFs) for universal learning of stochastic dynamics to enable exact belief propagation. The core idea is to combine the expressiveness of normalizing flows for density estimation with the analytical tractability of Bernstein polynomials. This approach allows for the representation of probability density functions (PDFs) as polynomials, which are closed under multiplication and integration, thereby supporting exact analytical belief propagation without approximation. Empirical results demonstrate that BNFs outperform state-of-the-art data-driven methods for belief propagation, particularly in highly non-linear systems with non-additive, non-Gaussian noise.
Problem
The paper addresses several bottlenecks in probabilistic reasoning under uncertainty:
Intractability of Analytical Belief Propagation: Nonlinear dynamics or non-Gaussian noise often make analytical belief propagation intractable, necessitating approximate methods.
Lack of Universal Approximation and Analytical Tractability: When system models are unknown and learned from data, there is a challenge in finding a model that can both universally approximate general nonlinear stochastic dynamics and support analytical belief propagation.
Approximation Errors in Existing Methods: Current approximate methods for belief propagation, such as linearization-based methods (e.g., EKF), Monte Carlo methods, and GMM-based approaches, introduce significant errors, especially in nonlinear systems, or lack formal error bounds.
Computational Scaling of GMM-based Methods: Component splitting techniques for GMMs can lead to an exponential growth in components, making long-horizon prediction intractable.
Key Contributions
Introduction of Bernstein Normalizing Flows (BNFs) as the first class of general nonlinear Markov chain models that simultaneously support universal approximation of stochastic dynamics and exact analytical belief propagation.
An explicit-constraint-free training procedure for learning valid PDFs.
A method for enhancing model expressiveness without increasing the number of parameters.
Empirical validation demonstrating the effectiveness of BNFs for belief propagation compared to state-of-the-art data-driven approaches.
Method
The proposed method, Bernstein Normalizing Flows (BNFs), models probability density functions (PDFs) using Bernstein polynomials embedded within a normalizing flow architecture. This allows for universal density estimation while maintaining analytical tractability for belief propagation.
Method Flow:
Problem Formulation: The system dynamics are defined as Xk+1 = f(Xk, Vk) (1), where Vk ~ p(v) is stochastic process noise. The state-transition distribution is p(xk|xk-1) = f(xk-1,·)#p(v) (2). Belief propagation involves recursively computing p(xk) = ∫Xk-1 p(xk|xk-1)p(xk-1)dxk-1 (3) and probabilistic reachability P(XK ∈ R) = ∫R p(xK)dxK (4).
Polynomial Distribution Modeling: To enable polynomial representation of PDFs, the unbounded state space X is mapped to a bounded unit box Un := (0, 1)n using a diffeomorphism Ω. This transforms the density px(x) to pu(u) = pu(Ω(x)) |det JΩ(x)| (7a), which simplifies to pu(Ω(x)) Πni=1 |dΩi/dxi| (7b) for a diagonal Ω. This transformation preserves the integration property, allowing marginalization over hyper-rectangular regions to be equivalent to integration over their image in Un.
Bernstein Normalizing Flow (BNF): A normalizing flow g: Un → Un is used, where the latent density is uniform over Un. Each component gi is a multivariate Bernstein polynomial gi(u) = π(u<i) of degree d. The resulting density is p(u) = Πni=1 ∂π/∂ui (u<i) (9).
Constraints for Valid PDFs: To ensure g is a diffeomorphism and p(u) is a valid PDF, two conditions are enforced:
Monotonicity: ∂gi/∂ui (u<i) > 0 (10a)
Full Range Coverage: ∫01 ∂gi/∂ui (u<i)dui = 1 (10b)
These conditions are enforced by constraining the coefficients bj of the Bernstein polynomial representation of ∂π/∂ui (u<i). Specifically, bj ≥ 0 for (10a) and Σdi-1 ji=0 bji = di along the i-th axis for (10b).
Conditional Density Estimation: For conditional distributions p(u|w), a flow function h: Un × Un → Un is constructed, where each component hi is a polynomial hi(u, w) = πi(u<i, w). The conditional polynomial distribution is p(u|w) = Πni=1 ∂π/∂ui h(u<i, w) (12).
Exact Belief Propagation: By modeling p(u0) as a BNF and p(u'|u) as a conditional BNF, operations (3) and (4) can be performed exactly using tensor operations on the coefficients of each model. Both multiplication and integration are carried out in the Bernstein basis, improving numerical stability.
Technical Details
State Space Transformation: A diffeomorphism Ω: X → Un maps the original state space X to the n-dimensional open unit box Un = (0, 1)n. A valid choice for each component Ωi is a univariate cumulative distribution function (CDF) of a continuous distribution supported on R, e.g., the Gaussian CDF.
Polynomial Basis: Bernstein polynomials are used to model PDFs. A multivariate polynomial π(x) of degree d can be expressed as π(x) = Σj bjβj(x), where βj(x) = Πni=1 (di choose ji)xji(1 - xi)di-ji.
Loss/Objective: Constrained Log-Likelihood Optimization is used. For p(u0), the objective is arg maxb∈B Eu0~p*(u0) [log p(u0)] (13a), subject to bji ≥ 0 (13b) and Σdi-1 ji=0 bji = di (13c). A similar objective is used for p(u'|u).
Optimization: Stochastic Gradient Descent (SGD) is used. To avoid explicit constraints, a differentiable function Ψ: Θ → Bfeas is defined, where Θ is an unconstrained parameter space. Ψ is constructed by mapping each vector to be positive (e.g., softplus) and then normalizing coefficients to sum to di along dimension i: σ(θi) = θi / Σdi-1 ji=0 θji (14).
Expressiveness Enhancement: A procedure to tighten the Bernstein relaxation involves lifting a given Bernstein polynomial to a higher degree d+ > d. This is done by applying a linear transformation b+ = Md+b (15) to the original coefficients. During training, a soft-constraint violation penalty Σj max(0, b+j) is used, followed by iterative projection steps to ensure b+ ≥ 0.
Model Parameters: The number of parameters in BNF scales with the number of coefficients in a degree-d 2n-variate Bernstein polynomial, i.e., O(d2n).
Results
Datasets:
Two highly non-linear latent stochastic systems: Van der Pol with additive Gaussian noise and a stable oscillator with multiplicative non-Gaussian noise.
Training data: 1K initial state data points, 10K state-transition data points.
Test data: Distinct large test data set, consisting of Monte-Carlo samples from the true system.
Metrics: Average log-likelihood (13a).
Baselines:
First-order linearization around component means (EKF-style).
Whitened Spherical Average Second-Order Stretching (WSASOS).
Grid-based approach.
All baselines rely on Gaussian assumptions and use GP regression for p(x'|x) and expectation maximization on a GMM for p(x0).
Numeric Results:
Additive Gaussian Noise System: BNF performs comparably to the grid-based method and generates more accurate predictions than EKF. The grid method performs best. WSASOS (learned) times out after 6 time steps.
Multiplicative Non-Gaussian Noise System: BNF significantly outperforms all baselines. Baselines struggle to capture the stochastic nature of the system due to highly non-linear, non-Gaussian noise.
Polynomial Degree Effect: Performance of BNF improves with increasing polynomial degree (degrees 10, 20, and 30 were used).
Variance: All log-likelihood results were found to have a variance below 10-3 across 10 experiments with random training seeds.
Runtime/Training Cost: Not stated in the paper, except that WSASOS (learned) times out after 6 time steps.
Weaknesses / Limitations
Computational Scalability: The number of parameters in BNF scales with the number of coefficients in a degree-d 2n-variate Bernstein polynomial, i.e., O(d2n), which is memory intensive for large n. The paper recognizes sparse models as valuable future work for scalability.
Conservatism of Bernstein Relaxation: Constraining coefficients b ≥ 0 is only a sufficient (but not necessary) condition for π ≥ 0. This means the feasible set Bfeas inner-approximates the true set of all degree-d polynomial diffeomorphisms, which limits the expressiveness of the model.
Evidence
Method evidence: "In our approach, dubbed Bernstein Normalizing Flows (BNFs), the underlying PDFs are modeled using Bernstein polynomials, which offer both favorable analytical properties and universal approximation capabilities." (section: 1 Introduction)
Result evidence: "Overall, we observe that BNF performs comparably to the grid-based method in the additive Gaussian case and significantly outperforms all baselines under non-Gaussian noise. Additionally, performance improves with increasing polynomial degree." (section: 6 Evaluations)
Limitation evidence: "Note that number of parameters in BNF scales with the number of coefficients in a degree-d 2n-variate (for u and u') Bernstein polynomial, i.e. O(d2n), which is memory intensive for large n. For scalability purposes, we recognize sparse models as valuable future work." (section: 5 Training Procedure)
Additional Figures
Fig. 1 shows a visual comparison of the computed beliefs at time step k = 9 for the Van der Pol system (Fig. 1a–d) and the stable oscillator (Fig. 1e–h). Figs. 2–3 present the average log-likelihood results. Overall, we observe that BNF performs comparably to the grid-based method in the additive Gaussian