[Paper Review] Scalable Training of Inference Networks for Gaussian-Process Models

Jiaxin Shi, Mohammad Emtiyaz Khan+1 more|arXiv (Cornell University)|May 27, 2019

Gaussian Processes and Bayesian Inference53 references4 citations

TL;DR

This paper proposes a scalable, minibatch-trainable inference network for Gaussian process models using a stochastic functional mirror-descent algorithm that tracks a Bayesian filter in function space. The method enables flexible, high-capacity posterior approximations—outperforming sparse variational GP methods in regression and classification tasks—by efficiently maintaining output correlations across data batches without inducing point limitations.

ABSTRACT

Inference in Gaussian process (GP) models is computationally challenging for large data, and often difficult to approximate with a small number of inducing points. We explore an alternative approximation that employs stochastic inference networks for a flexible inference. Unfortunately, for such networks, minibatch training is difficult to be able to learn meaningful correlations over function outputs for a large dataset. We propose an algorithm that enables such training by tracking a stochastic, functional mirror-descent algorithm. At each iteration, this only requires considering a finite number of input locations, resulting in a scalable and easy-to-implement algorithm. Empirical results show comparable and, sometimes, superior performance to existing sparse variational GP methods.

Motivation & Objective

• To address the challenge of training flexible inference networks for Gaussian process models in a minibatch setting, where maintaining function-space correlations is difficult.
• To overcome the limitations of sparse variational GP methods, which rely on a small number of inducing points and often underfit large datasets.
• To enable scalable, stable, and high-performance training of deep neural networks as posterior approximators for GP models using stochastic gradients.
• To provide a more flexible alternative to inducing-point-based GP inference that supports complex, structured kernels like those from deep neural networks.

Proposed method

• The method uses a stochastic functional mirror-descent algorithm to track a Bayesian filter in the function space, which is computationally intractable but serves as a target for approximation.
• At each training step, the algorithm uses a minibatch of data to compute stochastic gradients that update the inference network, ensuring meaningful correlations across function outputs.
• The inference network is trained by bootstrapping from its own current approximation, enabling iterative refinement without requiring full data passes.
• The approach supports various network architectures, including random feature expansions and deep ConvNets, and is compatible with non-conjugate likelihoods such as softmax for classification.
• The method is interpreted in a teacher-student framework, where the 'teacher' is a mirror-descent update derived from the current network, making it efficient and scalable.
• It enables training of deep ConvNets for GP inference in infinite-width Bayesian ConvNets, overcoming numerical instability issues seen in standard SVGP with such kernels.

Experimental results

Research questions

• RQ1 Can we train flexible inference networks for GP models using minibatch stochastic gradients while preserving meaningful correlations across function outputs?
• RQ2 Does a functional mirror-descent-based training algorithm outperform existing sparse variational GP methods in terms of predictive accuracy and generalization?
• RQ3 Can the proposed method scale to complex, structured kernels such as those derived from deep neural networks, especially in classification tasks?
• RQ4 Is it possible to train deep ConvNets as inference networks for GP models without relying on inducing points, and does this lead to better performance?

Key findings

• GPNet achieves the best RMSE among all methods on the Bikes and Kegg datasets, with a significant performance gap over SVGP when using a small number of inducing points.
• The RMSE gap between M=100 and M=500 for SVGP is larger than for GPNet, indicating that GPNet's performance is not limited by the number of inducing points.
• On MNIST and CIFAR10, GPNet achieves a test error rate of 1.12% and 24.63%, respectively, outperforming SVGP and other sparse GP baselines while using a flexible, deep ConvNet inference network.
• The method successfully trains a deep ConvNet for a CNN-GP prior, achieving state-of-the-art performance in classification despite the numerical instability typically encountered with SVGP on such kernels.
• The approach avoids the overfitting and underfitting issues seen in previous minibatch training of inference networks, particularly in the FBNN model where performance degrades with larger M.
• Empirical results show that GPNet provides a more flexible and scalable alternative to sparse GP methods, especially for complex, structured kernels and high-dimensional data.