Radial and Directional Posteriors for Bayesian Neural Networks

Changyong Oh; Kamil Adamczewski; Mijung Park

Source-linked AI summary

Radial and Directional Posteriors for Bayesian Neural Networks

Changyong Oh, Kamil Adamczewski, Mijung Park

arXiv:1902.02603v2Feb 7, 2019stat.MLcs.LG

TL;DR

The paper addresses Bayesian neural networks whose convenient mean-field posteriors neglect dependencies among weights and whose over-parameterization complicates compression. It proposes a posterior split into directional and radial components, using directional densities for correlations and radial densities for neuron strength and pruning. Experiments report improved predictive performance together with compressed networks, although practical computation remains constrained by modified Bessel-function evaluation and implementation support.

Problem
Mean-field Bayesian neural-network posteriors neglect statistical correlations among weights, while over-parameterized networks motivate structured compression.
Source: Our contributions · p. 2, §2 · p. 3
Method
RDP decomposes the variational posterior into directional and radial densities, using a von Mises-Fisher directional component and a Half-Cauchy/log-Normal radial component.
Source: Our contributions · p. 2
Results
RDP improves predictive performance and yields compressed networks simultaneously across the reported experiments.
Source: §6 · p. 8, §5.2 · p. 8
Takeaways & Limitations
Modeling weight dependencies and neuron strength can support Bayesian neural networks that target predictive quality and network compression together.
Source: §6 · p. 8
Efficient implementation is constrained by unstable modified Bessel-function ratios and limited support for these functions in TensorFlow and PyTorch versions discussed.
Source: §3.2 · p. 5, §4 · p. 13

Abstract

from arXiv · show

We propose a new variational family for Bayesian neural networks. We decompose the variational posterior into two components, where the radial component captures the strength of each neuron in terms of its magnitude; while the directional component captures the statistical dependencies among the weight parameters. The dependencies learned via the directional density provide better modeling performance compared to the widely-used Gaussian mean-field-type variational family. In addition, the strength of input and output neurons learned via the radial density provides a structured way to compress neural networks. Indeed, experiments show that our variational family improves predictive performance and yields compressed networks simultaneously.

1. Introduction

The paper introduces a Bayesian neural-network variational family that models weight dependencies while supporting structured compression. It combines directional modeling of correlations with radial modeling of neuron strength, alongside computational mechanisms for practical variational learning.

2. Variational Bayesian learning for neural networks

Mean-field variational families are computationally convenient but factorize the posterior and therefore ignore statistical correlations among weight parameters.The paper motivates modeling these dependencies as a route to improved predictive performance.Source: §1 · p. 1, §2 · p. 3

Our contributions

RDP decomposes the posterior into radial and directional densities to capture neuron strength and correlations among weight parameters while addressing compression.The radial component targets sparsification, whereas the directional component targets statistical dependencies.Source: Our contributions · p. 2
The directional density uses a von Mises-Fisher distribution to model weight correlations, with a polynomial approximation enabling stable gradients in dimensions of several thousands.The approximation avoids directly computing modified Bessel functions and supports GPU acceleration without extra CUDA implementations.Source: Our contributions · p. 2
The radial density combines a Half-Cauchy prior with a log-Normal posterior to bias weight norms toward zero and provide a closed-form KL divergence.This construction supplies a structured basis for pruning neurons according to learned radial strength.Source: Our contributions · p. 2
RDP evaluates models using both test likelihood and resulting model size, including remaining parameters and FLOPs.This reflects the paper’s joint goal of finding weight dependencies and reducing network parameters.Source: Our contributions · p. 2

2. Variational Bayesian learning for neural networks

Effective variational learning requires variance-reducing reparameterized Monte Carlo gradients and preferably a closed-form KL divergence.The reparameterization trick separates distribution parameters from random sources, reducing gradient-estimator variance.Source: §2 · pp. 2–3

3. Method

RDP factorizes each weight vector into radial and directional densities: the radial component supports structured sparsity, while the directional component models correlations among weights. The method also introduces a stable high-dimensional vMF optimization and grouping schemes for pruning input, output, or both neuron types.

3.1. Radial and Directional Posterior

The vMF directional component captures within-row weight correlations and shared layer-wise concentration induces structured correlations across rows.A shared concentration parameter reduces variational and prior parameters while allowing correlations to depend on the proximity of row mean directions.Source: §3.1 · pp. 3–4
The radial component uses products of Half-Cauchy variables with log-Normal posteriors to control local row-wise and global layer-wise sparsity.The Half-Cauchy prior biases weight norms toward zero, and the log-Normal posterior provides closed-form KL terms through Gamma and inverse-Gamma representations.Source: §3.1 · p. 4
RDP factorizes each weight row into a density over its L2 norm and a density over the direction of its normalized weight vector.The directional density uses a von Mises–Fisher distribution, while the radial density is designed for sparsity.Source: §3.1 · pp. 3, 5

3.2. Optimizing evidence lower bound for RDP

The ELBO combines expected log-likelihood and KL terms, but vMF gradients require stabilizing ratios of modified Bessel functions during optimization.The closed-form KL term still produces numerically unstable Bessel-function ratios, especially as dimensionality increases.Source: §3.2 · pp. 5–6
Approximating the Bessel-function ratio by the average of upper and lower bounds yields numerically stable gradients in dimensions of several thousands.The approximation also avoids direct Bessel-function computation and enables GPU acceleration without additional CUDA implementations.Source: §3.2 · p. 6, Figure 5 · p. 6

3.3. Structured compression using the radial density

The radial posterior supports structured pruning by grouping rows, columns, or both according to learned neuron contributions.Row grouping prunes output neurons, column grouping prunes input neurons, and double grouping prunes both; convolutional filters are grouped by output or input channels.Source: §3.3 · p. 6, Figure 6 · p. 6
Pruning uses the log posterior mode of combined local log-Normal radius factors as a cutoff for retaining or removing neurons.Global factors scale all rows or columns equally, so only the two local distributions determine row- or column-level pruning decisions.Source: §3.3 · pp. 6–7

4. Related Work

RDP captures dependencies between weights through a directional component rather than a matrix-normal covariance structure, while remaining connected to structured compression through uncertainty over weight magnitudes.

4. Related Work

RDP captures dependencies between weights through a directional component on a unit sphere, with coupled concentration parameters modeling correlations across rows and columns.This differs from matrix-normal approaches that use structured covariance matrices.Source: §4 · p. 7
The vMF reparametrization approximation makes dependency-aware variational inference scalable to several thousands of dimensions with corrected, stable gradient estimates.It avoids direct high-order Bessel-function computation and supports GPU acceleration without extra CUDA implementations.Source: §4 · p. 7
RDP is evaluated against mean-field and dependency-aware BNNs on UCI regression tasks, with better test likelihoods on five of nine datasets.The comparison uses average test log-likelihood across nine named UCI datasets.Source: Table 1 · p. 7
RDP prunes neurons using probability densities over weight magnitudes, extending Bayesian compression beyond methods based on scale parameters, mixture priors, or grouped Horseshoe priors.The radial density supplies uncertainty estimates for deciding which neurons to remove.Source: §4 · p. 7

5. Experiments

Experiments evaluate RDP for predictive performance on UCI regression and structured compression on MNIST. RDP improves comparison with dependency-aware baselines on regression and achieves competitive, balanced compression results for LeNet.

5.1. Regression using UCI data

RDP achieves better test log-likelihood than other methods on five of nine UCI regression datasets.The experiments use 90% training and 10% test splits and compare against mean-field BNNs, Dropout, and VMG.Source: §5.1 · pp. 7–8
RDP improves regression performance despite applying double grouping only to the first layer of an ninput−50−1 architecture.The second layer uses a fully factorized Gaussian; RDP outperforms VI, PBP, and Dropout, and beats VMG on five of nine datasets.Source: §5.1 · p. 8

5.2. Compression on MNIST Classification

RDP compresses LeNet53 on MNIST by pruning clearly separated clusters in the log-of-mode radius posteriors after selecting the model with the best training cross-entropy during the last 10 epochs.LeNet53 consists of two convolutional layers and two fully connected layers.Source: §5.2 · p. 8
RDP achieves competitive structured-pruning results across parameter count, FLOPs, and accuracy loss, with particularly strong convolutional-layer compression.It is only slightly more computationally heavy than FDOO and ranks second to BC-GHS in parameter number while using two-thirds as many parameters as Bayesian Compression-Group Normal Jeffrey.Source: §5.2 · p. 8, Table 2 · p. 8
Figure 7 summarizes pruning by plotting log mode against the number of rows or columns, with a red dotted pruning threshold and ratios of unpruned weights or filters.The top panel concerns rows and the bottom panel concerns columns.Source: Figure 7 · p. 8

6. Conclusion

The paper concludes that RDP captures weight dependencies while providing a natural structure for compression. It also introduces practical approximations for high-dimensional vMF inference, while leaving directional-based compression for future work.

6. Conclusion

RDP captures dependencies between weight parameters and scores better than other methods on multiple compression measures.Its empirical support comes from regression predictive performance and structured compression results.Source: §6 · p. 8
The practical method combines a Bessel-function-ratio approximation with a reparametrization trick for vMF distributions.The authors currently use only the radius density for pruning and identify the directional component as a future compression avenue.Source: §6 · p. 8

Supplementary Material Radial and Directional Posteriors for Bayesian Neural Networks

The supplementary material describes ELBO optimization for a vMF variational posterior using rejection-sampling reparametrization. It identifies the correction term needed for gradients through the sampled directional variable.

1. Gradient of ELBO with von-Mises Fisher(vMF) distribution

Variational inference optimizes the negative evidence lower bound, whose terms include reconstruction error and the prior-posterior KL divergence.The latent variable is W, with q(W) as the variational distribution, p(W) as the prior, and p(D|W) as the likelihood.Source: §1 · p. 11
The vMF posterior over the unit sphere is reparametrized through rejection sampling to enable gradient-based variational inference.The sample W is generated using a proposal-distribution sample ε.Source: §1 · p. 11

1. Gradient of ELBO with von-Mases Fisher(vMF) distribution

The gradient of reconstruction error includes a correction term associated with reparametrization through rejection sampling.The correction is attributed to the reparametrization procedure described by Naesseth et al. (2016) and Davidson et al. (2018).Source: §1 · p. 11

1. Gradient of ELBO with von-Mises Fisher(vMF) distribution

The vMF density uses a modified Bessel function of the first kind, whose order is determined by the distribution dimension.The supplementary material identifies Iν(·) as the modified Bessel function of the first kind.Source: §1 · p. 11

2. Instability of the calculation of the ratio of Bessel function Iν(z)

Direct evaluation of Iν(z) and especially consecutive-order ratios becomes numerically unstable as ν grows because both terms can become infinitesimally small. Bounds are therefore introduced to control Iν(z) over large orders and intervals.

Evaluating Iν(z) and Iν(z)/Iν−1(z) is expensive and unstable for large ν.The ratio can divide two infinitesimally small values, making numerical instability inevitable.Source: §2 · pp. 11–12
Upper bounds for Iν(z) are developed by combining Proposition 1 and Proposition 2 with inequalities describing its behavior for large ν, yielding a stated upper bound.The bound includes ϱ = −(1/e − log(2)) = 0.32526773938850295 > 0.Source: §2 · pp. 11–12
A larger order ν makes Iν(z) arbitrarily small over progressively longer intervals, explaining the numerical difficulty of direct computation.Source: §2 · pp. 11–12
Exponentially scaled Bessel functions do not resolve ratio instability because scaling produces division between two infinitesimally small numbers.For z > 0, e−z < 1, so scaling can exacerbate the numerical issue.Source: §2 · p. 12

3. Bound on the ratio of modified Bessel function of the first kind Iν(z)

The paper replaces unstable direct computation of consecutive-order Bessel ratios with tight algebraic bounds. The bounds are compared empirically with scipy.special.ive and remain reliable where that implementation becomes unstable.

Tight algebraic bounds on consecutive-order ratios avoid directly evaluating unstable Bessel functions for large ν.The bounds are derived from recurrence relations and provide a stable alternative to direct ratio computation.Source: §3 · pp. 12–13
Theorem 2 and Theorem 3 provide lower and upper bounds for Iν(z)/Iν−1(z) when ν ≥ 1/2.Source: §3 · p. 12
An Amos-type inequality supplies another algebraic bound for the ratio with smaller order in the numerator.Source: §3 · pp. 12–13
In Figure 8, the bounds closely match the ratio across z and ν, while scipy.special.ive fails for ν ≥ 500 and is unstable for ν = 100 when z < 100.Source: §3 · p. 12, Figure 8 · p. 13

4. Estimation of gradient of ELBO with vMF using Bound on the ratio of Bessel function Iν(z)

The bounded Bessel-ratio formulation addresses practical difficulties in computing vMF ELBO gradients. It avoids direct Bessel evaluation and supports stable gradient estimation despite framework and high-dimensionality constraints.

Computing vMF gradients is difficult because Bessel-ratio evaluation is unstable at large order and high-dimensional implementations lack broad framework support.The paper specifically notes missing TensorFlow and PyTorch support and unavailable CUDA implementations for the modified Bessel function.Source: §4 · p. 13
Tight rational-function bounds let the method avoid calculating Iν(z) directly and obtain stable gradient estimates.Source: §4 · pp. 13–14
The ratio approximation is used for the reconstruction-error gradient and for the KL-divergence of the location-agnostic directional prior.Source: §4 · p. 14

5. Approximate logarithm of modified Bessel function of the first kind log(Iν(z))

The paper approximates log(Iν(z)) using asymptotic bounds and a log-ratio series. These methods provide tight approximations, with the series remaining stable but slower for large ν.

Two approximation strategies are proposed for log(Iν(z)): an asymptotic-bound approach and a log-ratio series.Source: §5 · p. 14, §5.2 · p. 16

5.1. Asymptotic bound

The asymptotic-bound method uses different approximations for z < ν and z ≥ ν, with a tighter lower bound supplied by Proposition 3.Source: §5.1 · pp. 14–15
The lower bound from Proposition 3 is tighter than the earlier proposition, and algebraic simplification approximates log(cosh(z)) using z − log(2).Source: §5.1 · p. 14

5.2. Log Ratio Series

Figure 9 shows that the Theorem 3 approximation is tight but takes twice as long as scipy.special.ive.The approximation can support forward-pass KL estimates and ELBO monitoring during training.Source: Figure 9 · p. 15, §5.2 · p. 16
The log-ratio series exploits tight ratio bounds and remains stable for small and large ν, although the large-ν computation is slower.Source: §5.1 · p. 15, §5.2 · p. 16

6. Kullback-Leibler Divergence

Section 6 specifies KL-divergence components for directional and radial distributions, alongside gradients and special-function definitions used in their evaluation.

6.1. von-Mises Fisher q ∼vMF(µq, κq) and p ∼vMF(µp, κp)

The von-Mises Fisher subsection gives the gradient with respect to µq and κq for q and p parameterized by directional means and concentrations.Source: §6.1 · p. 16

6.2. q ∼Γ(αq, βq) and p ∼Γ(αp, βp) (Bauckhage, 2014)

The Gamma-distribution subsection defines Γ as the gamma function and Ψ as the digamma function.Source: §6.2 · p. 16

6.3. q ∼W(λq, kq) and p ∼W(λp, kp) (Bauckhage, 2013)

The Weibull-distribution subsection identifies γ as the Euler–Mascheroni constant and Γ as the gamma function.Source: §6.3 · p. 16

6.4. q ∼LN(µ, σ2) and p ∼Γ(α, β) (Louizos et al., 2017)

The remaining listed alternatives pair a log-normal variational distribution with Gamma or inverse-Gamma priors, with Γ identified as the gamma function.Source: §6.4 · pp. 16–17