PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational Inference with Tail-adaptive f-Divergence

Dilin Wang, Hao Liu, Qiang Liu

arXiv:1810.11943v3cs.LGstat.ML

TL;DR

KL-based variational inference can underestimate variance and miss posterior modes, while α-divergences can be difficult to estimate safely when importance-weight tails are heavy and their tail index changes. The paper introduces tail-adaptive f-divergences that adapt to those tails while retaining mass coverage. Across Bayesian neural networks and reinforcement learning, the approach outperforms standard α-divergence methods, with especially strong results in high-dimensional, multimodal settings.

  • Problem

    KL-based variational inference can underestimate variance and miss posterior modes, while a fixed α may be unsafe when the unknown importance-weight tail index changes during optimization.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    The paper develops f-divergences that adapt to the density-ratio tail, using a rank-based empirical ratio function to support stable estimation and mass coverage.

    Source: §1 · p. 2, §4 · pp. 3–4, §5 · p. 5

  • Results

    Across Bayesian neural networks and reinforcement learning, the approach outperforms standard α-divergence, especially for high-dimensional multimodal distributions.

    Source: §7 · p. 9

  • Takeaways & Limitations

    In reinforcement learning, the method improves soft actor-critic and can optimize multimodal loss functions more efficiently.

    Source: §1 · p. 2

    The stated gradient setup assumes reparameterized samples x = g_θ(ξ), a twice-differentiable f, and an existing gradient of log(p(x)/q_θ(x)).

    Source: §4 · p. 4

Abstract

from arXiv · show

Variational inference with {\alpha}-divergences has been widely used in modern probabilistic machine learning. Compared to Kullback-Leibler (KL) divergence, a major advantage of using {\alpha}-divergences (with positive {\alpha} values) is their mass-covering property. However, estimating and optimizing {\alpha}-divergences require to use importance sampling, which could have extremely large or infinite variances due to heavy tails of importance weights. In this paper, we propose a new class of tail-adaptive f-divergences that adaptively change the convex function f with the tail of the importance weights, in a way that theoretically guarantees finite moments, while simultaneously achieving mass-covering properties. We test our methods on Bayesian neural networks, as well as deep reinforcement learning in which our method is applied to improve a recent soft actor-critic (SAC) algorithm. Our results show that our approach yields significant advantages compared with existing methods based on classical KL and {\alpha}-divergences.

1 Introduction

The introduction motivates alternatives to KL because it can underestimate posterior variance and miss local modes, while positive α-divergences offer mass coverage but can be unstable to estimate. It presents a tail-adaptive f-divergence and reports benefits in variational inference and reinforcement learning.

  • Variational inference approximates an intractable distribution by optimizing over a family of tractable distributions, a practical alternative to MCMC's slow convergence that can yield biased results.Source: §1 · p. 1
  • KL-based variational inference can underestimate variance and miss important local modes of the true posterior.Source: §1 · p. 1
  • Positive α-divergences encourage mass coverage, but estimating them can have high or infinite variance when importance weights are heavy-tailed.Source: §1 · p. 2
  • Tail-adaptive f-divergence changes f according to the density-ratio tail to combine stable empirical estimation with mass coverage.The method uses a rank-based function of the empirical ratio at each gradient-descent step.Source: §1 · p. 2
  • The method better recovers multiple modes in variational inference and improves soft actor-critic, optimizing multimodal reinforcement-learning losses more efficiently.Source: §1 · p. 2

2 f-Divergence and Friends

This section formulates variational inference as minimizing an f-divergence between a target distribution and its approximation, then situates KL and α-divergences among its cases.

  • Variational inference minimizes the f-divergence between the target distribution p and a simpler distribution q_θ from a parameterized family.Source: §2 · p. 2
  • When f is strictly convex at 1, zero f-divergence implies the two distributions are equal; in practice, the expectation is approximated with samples from q_θ during stochastic optimization.Source: §2 · p. 2
  • The f-divergence family includes both KL directions, with KL(q || p) the typical variational-inference objective.Source: §2 · p. 2
  • α-divergence is an f-divergence family whose limits include KL(q || p) as α → 0 and KL(p || q) as α → 1.Source: §2 · p. 3
  • At α = 2, χ²-divergence equals the variance of the importance weight, so minimizing it corresponds to finding an optimal importance-sampling proposal.Source: §2 · p. 3

3 α-Divergence and Fat Tails

Positive α-divergences encourage mass coverage, but large α can make empirical objectives unstable or undefined when importance weights have heavy tails. The unknown tail index can change as the approximation is updated, making a fixed α difficult to choose safely.

  • For positive α, α-divergence penalizes large p(x)/q(x) ratios, encouraging q to account for p's probability mass and local modes.Source: §3 · p. 3
  • For α ≤ 0, the divergence has the opposite support condition: p(x) = 0 must imply q(x) = 0 for the divergence to be finite.Source: §3 · p. 3
  • Large α strengthens mass coverage but can increase estimator variance or make the expectation nonexistent when the density ratio is fat-tailed.Source: §3 · p. 3
  • A fat-tailed ratio with tail index α* has a finite α-th moment only when α < α*.Source: §3 · p. 3
  • The ratio p(x)/q(x) can be fat-tailed when q is more peaked than p, while it is bounded when the distributions' standard deviations satisfy σp ≤ σq.Source: §3 · p. 3
  • Because α* is unknown and can change as q updates, fixing α can be suboptimal; estimating α* may require many samples and can destabilize optimization if overestimated.Source: §3 · p. 3

4 Hessian-based Representation of f-Divergence

The section represents convex f-divergences through their curvature and derives gradient forms that support adaptive divergences without explicitly constructing f. This representation underlies a method that aims to keep gradient expectations finite while retaining mass-covering strength tied to the tail index.

4 Hessian-based Representation of f-Divergence

  • The adaptive construction guarantees finite gradient expectations while theoretically matching the strongest mass-covering property associated with α = α∗.The divergence is specified through gradient functions rather than requiring explicit construction of f.Source: §4 · pp. 3–5
  • Any twice-differentiable convex f can be decomposed using nonnegative curvature h = f′′ and linear terms set by a = f′(0) and b = f(0).Conversely, any nonnegative h and real a, b specify a convex function.Source: §4 · p. 4
  • The f-divergence also admits a representation through f′′, so its curvature determines the divergence apart from linear terms that do not change it.The passage cites Feldman & Osterreicher (1989), Osterreicher (2003), Liese & Vajda (2006), Reid & Williamson (2011), and Sason (2018) for related integral representations.Source: §4 · p. 4
  • The gradient of Df(p || qθ) is directly related to f′′ and depends on f through ρf or γf.The section gives score-function and reparameterization gradient forms; the latter uses ∇x log(p(x)/qθ(x)).Source: §4 · pp. 4–5
  • Any increasing ρ or nonnegative γ can specify a convex f, allowing the gradient functions to be chosen without explicitly defining f.The stated conditions include ρf increasing and γf nonnegative on R+.Source: §4 · p. 5
  • If ρf is strictly increasing at 1 or γf(1) is positive, then Df(p || q) = 0 implies p = q.Source: §4 · p. 5

5 Safe f-Divergence with Inverse Tail Probability

The section defines adaptive divergences through powers of the importance-weight tail probability, aiming for finite gradient expectations while retaining mass-covering strength. It estimates that tail probability from samples and uses the resulting weights in stochastic gradient updates.

5 Safe f-Divergence with Inverse Tail Probability

  • For β < 0, the tail-probability power is nonnegative and increasing, with asymptotic scale t^-βα∗; the relevant expectation is finite only when β > -1.Source: §5 · p. 5
  • Inverse tail-probability powers define implicit adaptive f-divergences whose gradients are computable without deriving f or the divergence itself.The construction targets finite gradient expectations while keeping the function large enough for mass coverage.Source: §5 · p. 5
  • The algorithm estimates the weight tail probability empirically and uses each sample’s estimated tail probability to set its gradient weight.The empirical estimate counts sampled weights at least as large as a threshold; the resulting update uses stochastic approximation.Source: §5 · p. 6
  • With -1 < β < 0, rank-based weighting penalizes high density ratios without being overly aggressive, and β = -1 almost always gives the best empirical performance despite the theoretical requirement β > -1.Source: §5 · pp. 5–6
  • When α is too large, classical α-divergence weights can let one data point dominate and make the gradient estimate unstable.Source: §5 · p. 6

6 Experiments

Experiments evaluate tail-adaptive f-divergence on Gaussian mixtures, Bayesian neural network regression, and MuJoCo reinforcement-learning benchmarks. The reported results favor the adaptive method most clearly in higher-dimensional or more complex settings.

6.1 Gaussian Mixture Toy Example

  • At dimension 10, tail-adaptive f-divergence performs best on Gaussian mixtures, while all tested algorithms perform similarly at dimension 2.Figure 1 reports mode-shift distance and mean and variance MSE, averaged over 10 random trials.Source: §6.1 · pp. 6–7, Figure 1 · p. 7
  • For 10-dimensional mixtures, adaptive weights significantly outperform the baselines when the target becomes highly non-Gaussian; all algorithms perform well when it is close to Gaussian.Figure 2 tracks mode-shift distance and mean and variance MSE as non-Gaussianity increases.Source: Figure 2 · p. 7, §6.1 · p. 7

6.2 Bayesian Neural Network

  • In Bayesian neural regression, the adaptive method achieves the best test RMSE and test log-likelihood in most cases against α = 0 and α = 0.5.The reported comparison summarizes results over the datasets in Table 1.Source: §6.2 · p. 7, Table 1 · p. 8

6.3 Application in Reinforcement Learning

  • Across MuJoCo benchmarks, tail-adaptive f-divergence outperforms the tested α-divergences in final performance and learns faster in most environments.The reported gains are especially significant on the high-dimensional, complex Ant and Humanoid environments.Source: Figure 3 · p. 9, §6.3 · p. 9

7 Conclusion

The paper presents tail-adaptive f-divergence for variational inference and reinforcement learning. It reports finite density-ratio moments and improved empirical performance over standard α-divergences, particularly for high-dimensional multimodal distributions.

7 Conclusion

  • Tail-adaptive f-divergence adapts the divergence function to the importance-weight tail, and guarantees finite moments of the density ratio.The conclusion also reports more stable importance weights and gradient estimates than classic α-divergence.Source: §7 · p. 9
  • On Bayesian neural networks and reinforcement learning, the approach outperforms standard α-divergence in the reported empirical results.Source: §7 · p. 9
  • The reported advantages are especially pronounced for high-dimensional multimodal distributions.Source: §7 · p. 9

A Proof of Proposition 4.1

The proof relates the convexity of f to its second derivative h and derives expressions for f using h and boundary terms.

A Proof of Proposition 4.1

B Proof of Proposition 4.2

The proof gives equivalent score-function and reparameterization expressions for an expectation gradient. The appendix also summarizes a score-function implementation of tail-adaptive variational inference.

B Proof of Proposition 4.2

  • The score-function and reparameterization tricks express the same expectation gradient using either a log-density gradient or a reparameterized sample.The reparameterization assumes x = gθ(ξ) with ξ sampled from q0.Source: B Proof of Proposition 4.2 · p. 12

C Tail-adaptive f-divergence with Score-Function Gradient Estimator

  • Algorithm 2 updates the variational parameter with a score-function gradient and weights samples using an empirical tail estimate raised to β.The algorithm describes β as an index and uses a stepsize for the parameter update.Source: C Tail-adaptive f-divergence with… · p. 13

D More Results for Bayesian Neural Network

The appendix reports that the adaptive approach performs best in most Bayesian neural network cases and describes the reinforcement-learning experiments and environments.

D More Results for Bayesian Neural Network

E Reinforcement Learning

  • The reinforcement-learning appendix provides additional experiment information and results, including comparisons using score-function gradient estimators.Source: E Reinforcement Learning · p. 14

E.1 MuJoCo Environments

  • Experiments cover six MuJoCo environments, with action-space dimensions ranging from 2 for Swimmer(rllab) to 21 for Humanoid(rllab).Figure 4 shows the environments in the order HalfCheetah, Hopper, Swimmer(rllab), Humanoid(rllab), Walker, and Ant.Source: §E.1 · p. 14, Figure 4 · p. 14

E.2 Different Choices of α

On Hopper and Walker, α=0.5 and α=+∞ outperform KL; score-function experiments also show better performance tends to hold across all tested environments.

E.2 Different Choices of α

  • The α-divergence comparison evaluates average rewards on Hopper and Walker using both score-function and reparameterization gradient estimators.Source: §E.2 · p. 14, Figure 5 · p. 15
  • α=0.5 and α=+∞ perform consistently better than standard KL divergence (α=0) on Hopper and Walker.Figure 5 presents average rewards for different α choices with score-function and reparameterization gradient estimators.Source: §E.2 · p. 14, Figure 5 · p. 15

E.3 Tail-adaptive f-divergence with score function estimation

  • With score-function gradient estimation, tail-adaptive f-divergence tends to yield better performance across all tested environments.Figure 6 reports average rewards.Source: Figure 6 · p. 15, §E.3 · p. 15
Loading 1810.11943v3…