[Paper Review] Function Space Particle Optimization for Bayesian Neural Networks

Ziyu Wang, Tongzheng Ren+2 more|arXiv (Cornell University)|Feb 26, 2019

Adversarial Robustness in Machine Learning14 citations

TL;DR

This paper proposes Function Space Particle Optimization (f-SVGD), a novel variational inference method for Bayesian neural networks that performs particle optimization directly in the space of regression functions rather than weight space. By avoiding gradient mixing and particle collapse in over-parameterized models, f-SVGD achieves superior performance in prediction, adversarial robustness, and reinforcement learning, outperforming strong baselines on multiple benchmarks.

ABSTRACT

While Bayesian neural networks (BNNs) have drawn increasing attention, their posterior inference remains challenging, due to the high-dimensional and over-parameterized nature. To address this issue, several highly flexible and scalable variational inference procedures based on the idea of particle optimization have been proposed. These methods directly optimize a set of particles to approximate the target posterior. However, their application to BNNs often yields sub-optimal performance, as such methods have a particular failure mode on over-parameterized models. In this paper, we propose to solve this issue by performing particle optimization directly in the space of regression functions. We demonstrate through extensive experiments that our method successfully overcomes this issue, and outperforms strong baselines in a variety of tasks including prediction, defense against adversarial examples, and reinforcement learning.

Motivation & Objective

• Address the failure mode of weight-space particle optimization in over-parameterized Bayesian neural networks, where particles collapse to a single mode.
• Overcome the curse of dimensionality and mode collapse in high-dimensional, multi-modal weight-space posteriors.
• Improve posterior approximation quality by shifting optimization from weight space to function space, where particles represent regression functions.
• Enable scalable, parallelizable inference for large models by exchanging only function-space error signals between particles.
• Demonstrate consistent performance gains across diverse tasks, including regression, adversarial defense, and reinforcement learning.

Proposed method

• Perform particle optimization in the function space by representing each particle as a regression function, parameterized by neural network weights.
• Use a function-value kernel to compute particle interactions based on functional similarity rather than weight distance.
• Derive a particle update rule that avoids gradient mixing by applying gradients from particle j to particle i through the Jacobian of f^j with respect to θ^j.
• Implement a mini-batch stochastic approximation of the update rule to scale to large datasets and models.
• Leverage model parallelism: only function-space error signals are communicated between particles, reducing communication overhead.
• Use a reparameterization trick to compute gradients of the kernel function with respect to particle weights, enabling efficient optimization.

Experimental results

Research questions

• RQ1 Can particle optimization in function space mitigate the mode collapse and gradient mixing issues seen in weight-space particle methods for over-parameterized BNNs?
• RQ2 Does optimizing in function space lead to better generalization and uncertainty calibration compared to weight-space methods?
• RQ3 How does f-SVGD perform in comparison to strong baselines on downstream tasks like regression, adversarial robustness, and reinforcement learning?
• RQ4 What is the impact of function-space kernel design on convergence and approximation quality in BNNs?
• RQ5 Can f-SVGD be efficiently scaled using model parallelism with minimal communication overhead?

Key findings

• f-SVGD successfully avoids the particle collapse problem that plagues weight-space particle optimization in over-parameterized models.
• On regression tasks, f-SVGD achieves lower test error and better uncertainty calibration than SVGD and other baselines.
• In adversarial defense, f-SVGD significantly improves robust accuracy on MNIST and CIFAR-10, outperforming standard BNNs and SVGD.
• In deep reinforcement learning, f-SVGD enables better exploration and higher cumulative returns in MuJoCo environments compared to baseline methods.
• The method scales efficiently: with hundreds of particles fitting in a single GPU, and communication overhead remaining low due to function-space signal exchange.
• Ablation studies confirm that gradient mixing in weight-space methods is a key cause of failure, and f-SVGD's design effectively eliminates this issue.