PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Variational Bayesian Inference for a Nonlinear Forward Model

Michael A. Chappell, Martin S. Craig, Mark W. Woolrich

arXiv:2007.01675v1eess.SPstat.APstat.ML

TL;DR

Bayesian inference for nonlinear models is hindered by intractable posterior calculations and restrictive analytical approximations. The paper develops stochastic variational Bayes using sampling, reparameterization, and mini-batches, and finds comparable results to analytical VB with practical convergence and computational benefits. Its trade-offs include noisy optimization and sensitivity to algorithmic parameters.

  • Problem

    Posterior calculations for noisy nonlinear models are often intractable, while analytical VB can impose conjugacy and local-linearization restrictions.

  • Method

    The paper develops stochastic variational Bayes that optimizes free energy with sampled gradients, flexible priors, reparameterization, and mini-batches.

  • Results

    The sVB algorithm achieved very similar results to the existing aVB solution, with robust convergence over explored parameter ranges and reduced time to convergence using samples and batches.

  • Takeaways & Limitations

    Removing the analytical-integral and conjugacy requirements allows sVB to address a wider range of nonlinear model-inference problems, including different noise models and priors.

  • Takeaways & Limitations

    Noisy samples can hamper or prevent reaching the global optimum, and sampling makes convergence harder to monitor because free energy continues to vary.

Abstract

from arXiv · show

Variational Bayes (VB) has been used to facilitate the calculation of the posterior distribution in the context of Bayesian inference of the parameters of nonlinear models from data. Previously an analytical formulation of VB has been derived for nonlinear model inference on data with additive gaussian noise as an alternative to nonlinear least squares. Here a stochastic solution is derived that avoids some of the approximations required of the analytical formulation, offering a solution that can be more flexibly deployed for nonlinear model inference problems. The stochastic VB solution was used for inference on a biexponential toy case and the algorithmic parameter space explored, before being deployed on real data from a magnetic resonance imaging study of perfusion. The new method was found to achieve comparable parameter recovery to the analytic solution and be competitive in terms of computational speed despite being reliant on sampling.

I. INTRODUCTION

Bayesian inference incorporates prior information and quantifies uncertainty, but posterior calculations are often intractable. This motivates approximate methods for nonlinear parameter inference from noisy, large-scale datasets.

  • Bayesian inference combines prior information with uncertainty estimates for parameters inferred from noisy data.
  • Posterior distributions and their summary statistics are often difficult to compute because the required integrals are intractable.
  • Variational Bayesian methods approximate the true posterior with a tractable distribution optimized using KL divergence.
  • Nonlinear model fitting is especially relevant to functional MRI, where poor-SNR data must be analyzed across tens to hundreds of thousands of voxels.
  • The paper explores stochastic variational Bayes to increase flexibility in prior specification while achieving solutions on realistic timescales.

B. Variational Inference

Variational inference optimizes a tractable approximate posterior, while stochastic gradient estimates make direct optimization feasible for nonlinear models. Reparameterization and mini-batches address gradient variability and computational cost.

  • Variational Bayes selects a tractable approximate posterior q(θ) and minimizes its KL divergence from the true posterior.
  • Parameterized posteriors permit update equations for hyperparameters, but integration imposes constraints such as conjugate prior distributions.
  • Direct gradient descent on free energy uses Monte Carlo samples to approximate otherwise intractable integrals.
  • Noisy stochastic gradients can slow convergence but may broaden parameter-space exploration and help avoid local minima in nonlinear fitting.
  • The reparameterization trick draws samples independently of hyperparameters and transforms them to reduce stochastic-gradient variability.
  • Mini-batches reduce computational cost by calculating each stochastic-gradient update from only a subset of data points.

C. Inference for a Non-Linear Forward Model

The nonlinear forward-model problem represents measured data as model predictions plus noise. Under a white Gaussian-noise assumption, this formulation supplies the likelihood for parameter inference.

  • The forward model expresses the measured vector as y = g(θ) + e, with g nonlinear in parameters θ.
  • The data vector and noise vector each contain N samples, while the model maps θ to predicted observations.
  • Assuming white noise e~N(0, ϕI) specifies the noise covariance used to construct the log-likelihood.

D. Analytical solution for VB inference

The analytical VB approach factorizes model and noise posteriors into conjugate distributions and linearizes the nonlinear model locally. Its iterative optimization monitors free energy, but global convergence is not guaranteed.

  • The analytical VB scheme factorizes the approximate posterior into separate model-parameter and noise distributions.
  • It uses a multivariate normal distribution for model parameters and a gamma distribution for noise parameters and their priors.
  • Calculus of variations produces update equations for the posterior hyperparameters μ, Λ, s, and c.
  • A first-order Taylor expansion locally linearizes the nonlinear forward model to make the analytical scheme tractable.
  • Global convergence is not guaranteed for nonlinear forward models, so the method monitors free energy and uses trial iterations after reversals.

III. METHODS

The stochastic VB method uses a multivariate normal approximate posterior and reparameterization-based sampling, avoiding the need for conjugacy. Its free-energy optimization combines stochastic and analytically tractable terms, with implementation choices differing from the analytic VB baseline.

  • Stochastic VB formulation: The approximate posterior is multivariate normal over all model and noise parameters, allowing covariance between noise and model parameters.The noise parameter is represented on a log scale to constrain it to positive values.
  • Stochastic VB formulation: Samples are generated by reparameterizing the posterior as Θ*=m+Sε*, where ε follows a standard multivariate normal distribution.S is the Cholesky decomposition of covariance C.
  • Prior specification: The stochastic formulation does not require conjugacy between the prior and approximate posterior.For this work, the prior was nevertheless chosen as a multivariate normal distribution over all parameters.
  • Objective function: The KL divergence between approximate posterior and prior can be written analytically for the chosen multivariate normal prior.This reduces reliance on stochastic approximations for all terms in the free-energy objective.
  • Implementation: The sVB implementation minimizes negative free energy using TensorFlow and Adam, whereas aVB uses the C++ fabber implementation.The loss is defined as -F, so optimization seeks to minimize the non-negative free energy.

B. Assessing Convergence with Simulated Data

Convergence was assessed for stochastic VB on simulated biexponential data by varying posterior sample count, batch size, and learning rate. Experiments compared full and diagonal posterior covariance formulations across noisy realizations.

  • Algorithmic parameters: Three algorithmic parameters govern sVB convergence: posterior sample count L, batch size B, and optimizer learning rate α.These parameters were explored using a biexponential model fitted to simulated data corrupted by white noise.
  • Simulated data: Equivalent biexponential solutions were normalized by assigning A₁ and R₁ to the slower estimated rate.This resolves the label-switching symmetry caused by exchanging the two amplitude–rate pairs.
  • Optimization: Optimization was performed globally across all 1000 realizations using their mean free energy as the loss.The convergence analysis therefore evaluated a shared optimization problem rather than independently optimizing each realization.
  • Posterior covariance: Two posterior-covariance settings were compared: a full covariance matrix and a diagonal covariance containing only parameter variances.The diagonal setting provides a simpler alternative to inferring all covariance terms.
  • Parameter search: Learning rates, sample sizes, and batch sizes were varied to identify settings associated with low converged free energy and rapid convergence.The tested learning rates ranged from 0.005 to 0.5, and sample sizes ranged from 1 to 200.

C. Initialization

Initialization experiments examined how prior specification and initial posterior choice affect sVB convergence. Comparisons used converged mean free energy across 1000 realizations and runtime to convergence.

  • Evaluation: Initialization effects were evaluated using converged mean free energy over 1000 realizations and runtime to convergence.Both informative and noninformative prior settings were included in the comparison.
  • Initial posterior distribution: Four initial posterior settings were tested: True, Data, Wrong, and Uninformed.These settings differed in whether their means matched the prior, data-derived amplitudes, deliberately incorrect values, or a generic high-uncertainty distribution.
  • Initial posterior distribution: The True initialization matched the prior and therefore used the true parameter means when the informative prior was selected.Its standard deviations matched those of the prior.
  • Initial posterior distribution: The Data initialization estimated amplitudes from half the data’s maximum amplitude while using the true decay rates.Its standard deviations also matched those of the prior.
  • Initial posterior distribution: The Wrong initialization deliberately started far from the true solution, with decay-rate means of 1.0 and amplitude means of 100.The Uninformed initialization matched the noninformative prior and assigned all parameters generic high uncertainty.

D. Application

The methods were applied to perfusion estimation from real PCASL MRI data, where varying post-label delays produces voxelwise time series for nonlinear kinetic model fitting. Convergence was compared between sVB and aVB across brain voxels.

  • Application: The application used kinetic model fitting to estimate perfusion from real arterial spin labeling MRI brain data.The data came from a single individual and used pseudo-continuous labeling with label duration τ=1.8 s.
  • Application: PCASL labels arterial blood water before image acquisition after a post-label delay, and subtracting control images yields inflow measurements.Repeating this process across post-label delays generates voxelwise time-series data for estimating perfusion and other haemodynamic parameters.
  • Kinetic model: The kinetic model defines perfusion f, arterial transit time Δt, relaxation constants, label duration τ, and time since labeling t as parameters or timing quantities.The equilibrium arterial-blood magnetization M₀ᵃ is obtained from a separate calibration image.
  • Data acquisition: The dataset contained 96 measurements from 48 labeled–control pairs across six post-label delays, each repeated eight times.The post-label delays were 0.25, 0.5, 0.75, 1.0, 1.25, and 1.5 s.
  • Data acquisition: Slice-dependent delay differences were incorporated into the model because later slices were acquired 45.2 ms later than preceding slices.This adjustment accounted for longer effective post-label delays in later slices.
  • Convergence assessment: sVB and aVB convergence was assessed using minimum free energy and time to convergence across voxels within an existing brain mask.For the ASL data, batch sizes B=5, 6, 9, 12, 18, 24, and 48 were tested.

IV. RESULTS

The simulated biexponential experiments identified algorithm settings that reliably reached low free energy, while batching and posterior sampling affected convergence speed. Under selected settings, stochastic and analytic VB produced comparable parameter estimates, particularly with larger datasets.

  • Convergence settings: A learning rate of 0.05 most consistently achieved the minimum free energy across tested posterior sample sizes.
  • Convergence settings: Posterior sample sizes above 10 generally matched the minimum free energy achieved with larger sample sizes.Small sample sizes and larger learning rates were less reliable near the minimum.
  • Convergence settings: L=10–20 supported faster convergence for N=50–100, whereas L=20–50 benefited datasets with N<20.Larger sample sizes generally slowed convergence, except that very small values below 5 sometimes reversed this trend.
  • Batching: Batching reduced convergence time, with B=10 generally performing best for the displayed N=50 and 100 cases.
  • Batching: The same converged free energy was attainable across batch sizes, although smaller batches favored smaller learning rates.These results were obtained with posterior sample size 20.
  • Comparison of sVB and aVB: aVB and sVB produced comparable parameter estimates, especially for N=50 or 100; with smaller N, both tended to estimate a single exponential component.The sVB settings were α=0.05, L=20, and B=10.

C. Choice of prior and initial posterior distributions on simulated data

On simulated biexponential data, prior choice produced a small difference in minimum free energy, while initial posterior choice was generally unimportant unless non-informed. In the ASL-MRI application, comparable convergence was obtained with learning rate 0.05, posterior sample size 5, and batch sizes 12–18.

  • Batch sizes around 10 consistently produced the shortest convergence runtime in simulated biexponential fitting.The comparison used posterior sample size 20 and learning rate 0.05.
  • A small but significant difference in minimum free energy was observed between informative and noninformative priors.
  • Initial posterior choice was generally unimportant when the initial posterior was not non-informed.A distant initial mean mainly increased the number of epochs required for convergence.
  • Learning rate 0.05, posterior sample size 5, and batch sizes 12–18 produced the best convergence properties on real ASL-MRI data.Values around these settings showed similar performance.
  • Very small posterior sample sizes, including L=2, were generally acceptable compared with the best results.

V. DISCUSSION

The discussion presents sVB as a flexible stochastic alternative to analytic VB for nonlinear model fitting, while emphasizing sampling-related convergence trade-offs and algorithmic tuning. Across the explored models, sVB achieved similar results to aVB, and batching with limited posterior samples reduced convergence time.

  • sVB achieved very similar results to the existing aVB solution while avoiding analytic integrals and Taylor approximation of the nonlinear function.This removes the original conjugacy restriction and supports broader prior and noise-model choices.
  • sVB convergence depends on learning rate, posterior sample size, and batch size, although robust convergence was found across a parameter range.Large learning rates can make the search too coarse, while small sample sizes can prevent reaching the global minimum.
  • Batching combined with limited posterior sample size substantially reduced convergence time for large-scale ASL kinetic model fitting.This is particularly relevant when many series must be analyzed across a brain volume.
  • The sVB framework resembles a variational autoencoder, replacing its decoder with the nonlinear model and directly inferring interpretable latent parameters.
Loading 2007.01675v1…