Approximate Bayesian Inference Team
Mohammad Emtiyaz Khan
近似ベイズ推論チーム

カーン

エムティヤーズ

Overview and Goals
Goal: AI that can continue to learn and improve throughout their lives, just like humans and animals. Currently, deeplearning (DL) requires a large amount of data which is costly and rigid (cannot quickly adapt). We aim to fix these issues
with a new learning paradigm based on Bayesian principles.
Summary of our research in the years 2020-2021:
A. Proposed Bayesian learning rule (BLR) yielding a wide-range of algorithms.
B. New BLR variants for DL, one of which won the NeurIPS-2021 Approximate Inference challenge.
C. Progress on adaptation and continual learning (FROMP, K-priors, Bayes-duality).
D. New theoretical results for online Bayes
E. Hyperparameter and architecture search using Bayesian methods.
F. A new paper on AI for social good in Nature communications.

Bayesian Learning Rule (BLR)

Standard Deep Learning

Deep Continual Learning

Continual Learning and Adaptation

Problem: Is there a common principle behind“successful” algorithms (e.g., those in DL)?

Problem: Approximate the expensive, exact Bayesianposterior (computed over several weeks on 512 TPUs)
but don’t exceed ~10x the cost of standard training.

Problem: Reduce catastrophic forgetting of the past. A
popular method is to use quadratic weight regularizers.

Solution: we propose the Bayesian Learning Rule [1]Solution: A BLR variant, called iVON [2], uses mixtureof-Gaussian posterior approximation. Won first prize!
Team consisted of Thomas Möllenhoff, Yuesong Shen,
Gian Maria Marconi, Peter Nickl, Emtiyaz Khan.

By choosing different approximations, we canderive a wide-variety of learning-algorithms. Better
approximations lead to better algorithms.

Table 1: A summary of learning algorithms derived from the BLR. Each algorithm is derived throughspecific approximations of the posterior and natural-gradient. New algorithms are marked with “(New)”.
Abbreviations: cov. ! covariance, STE ! Straight-Through-Estimator, VI ! Variational Inference,
VMP ! Variational Message Passing.

TeamMethod

Rank(LightTrack)

Rank
(Ext.
Track)

CIFAR
Agree

CIFAR
TVD

MedMNISTAgree

MedMNISTTVD

UCIGap
W2

🥇RIKEN TeamABI

Bayesian
Learning Rule

1

1.67

0.787

0.197

0.884

0.0994-0.094

Learning AlgorithmSec.

🥈ÉcolePolytechnique

MultiSWAG

2.5

2.5

0.777

0.218

0.89050.0983-0.166

🥉Universityof Liège

Seq AnchoredEnsembles

2.5

30.7730.210

0.87450.1066-0.115

Posterior Approx.

Natural-Gradient Approx.

Optimization Algorithms
Gradient Descent

Gaussian (fixed cov.)

Delta method

1.3

Newton’s method

Gaussian

—–“—–

1.3

Multimodal optimization (New)

Mixture of Gaussians

—–“—–

3.2

Deep-Learning Algorithms
Stochastic Gradient Descent

Gaussian (fixed cov.)

Delta method, stochastic approx.

4.1

RMSprop/Adam

Gaussian (diagonal cov.)

Delta method, stochastic approx.,
Hessian approx., square-root scaling, slow-moving scale vectors

4.2

Dropout

Mixture of Gaussians

Delta method, stochastic approx.,
responsibility approx.

4.3

STE

Bernoulli

Delta method, stochastic approx.

4.5

Online Gauss-Newton (OGN)

Gaussian (diagonal cov.)

Gauss-Newton Hessian approx. in
Adam & no square-root scaling

4.4

Variational OGN (New)—–“—–Remove delta method from OGN

4.4

BayesBiNN (New)

Bernoulli

Remove delta method from STE

4.5

(New)

Approximate Bayesian Inference AlgorithmsConjugate Bayes

Exp-family

Set learning rate ⇢t = 1

5.1

Laplace’s method

Gaussian

Delta method

4.4

Expectation-Maximization

Exp-Family + Gaussian

Delta method for the parameters

5.2

Stochastic VI (SVI)

Exp-family (mean-field)

Stochastic approx., local ⇢t = 1

5.3

VMP

—–“—–

⇢t = 1 for all nodes

5.3

Non-Conjugate VMP

—–“—–

—–“—–

5.3

Non-Conjugate VI (New)

Mixture of Exp-family

None

5.4

1. Khanandlearning
Rue, rule
TheasBayesian
Learning
Bayesiannatural-gradientdescentRule, arXiv, 2021

2.1¯
Given the objective L( ) = Eq [`(✓)+logq(✓)] in Eq. 2, the classical gradient-descent algorithm performs
the following update:L( ).Online Bayes(15)Theoretical Results⇢ rfort+1ttt

6

Problem: Theoretical analysis for online Bayesianlearning hold under restrictive conditions.
Solution: We propose to relax these conditions, byusing a generalize online Bayesian methods wherearbitrary divergences can be used (instead of KL) [8]More BLR variants:
o iVON [2] is proposed to ensure the steps of BLR
always lead to positive covariances.
o New generalizations in [3] for “structured”covariances allow low-rank and sparse structures (eg,recovering LBFGS/DFP style updates). This work uses
Lie-Group structures.
o BayesBiNN [4] is a BLR variant for Binary Neural
Networks which recovers the STE algorithm

8. Alquier, Non-exponentially Weighted Aggregation: RegretBounds for Unbounded Loss Functions, ICML 2021

In [6], we quantify “forgetting” in terms of past memory
represented via principal components analysis.
In [7], we present a generalization called K-priors tounify such adaptation methods. We show that thesemethods faithfully reconstruct the gradient of the past.

2. Lin, Schmidt, Khan, Handling the Positive-Definite Constraint
in the Batesian Learning Rule, ICML 2020
3. Lin, Nielsen, Khan, Schmidt, Tractable structured naturalgradient descent using local parameterizations, ICML 2021
4. Meng, Bachman, Khan, Training Binary Neural Networks
using the Bayesian Learning Rule, ICML 2020

5. Pan, Swaroop, Immer, Eschenhagen, Turner, Khan, Continual
Deep Learning by Functional Regularisation of Memorable
Past, NeurIPS 2020
6. Doan, Abbana Bennani, Mazoure, Rabusseau, Alquier, A
Theoretical Analysis of Catastrophic Forgetting through theNTK Overlap Matrix, AIStats 2021
7. Khan & Swaroop, Knowledge-Adaptation Priors, NeurIPS 2021

Architecture Selection for Deep Networks

A Summary of Other Works

Problem: Existing methods require validation data toselect architecture and hyperparameters.
Solution: A method based on marginal likelihood using
only training data. Uses Laplace approximation[9.10] withscalable Hessian approx (eg, KFAC).

Gaussian Process: Using BLR, we derive a fastalgorithm for state-space GP [11]. We also show that adual parameterization useful for sparse GPs [12]. Wederive a sparse representation using subset of data [13]Larger models, which give better test error, also
generally have higher marginal likelihoods.

11. Chang, Adam, Khan, Solin, Dual Parameterization of Sparse
Variational Gaussian Processes, ICML 2021
12. Chang, Wilkinson, Khan, Solin, Fast Variational Learning inState-Space Gaussian Process Models, MLSP, 2020
13. Jain, PK, Khan, Subset-of-Data Variational Inference for Deep
Gaussian-Process Regression, UAI 2021

Reinforcement Learning: We propose a replacementof “target networks” by functional regularization [14].
In [15], we propose imitation learning for diverse kindsof feedback, appropriately re-weighting them.

We derive an explicit formula for the updates
which we call generalized Bayes rule.

We prove a regret bound that holds for below theusual bounded setting (less restrictive).

Solution: We show that functional
regularization of “memorable past”(FROMP) [5] gives better results≈

1st Place in NeurIPS 2021 Challenge

14. Piche, Thomas, Marino, Marconi, Pal, Khan., Beyond Target
Networks: Improving Deep Q-learning with Functional
Regularization, arXiv 2021
15. Tangkaratt, Han, Khan, Sugiyama, VILD: Variational Imitation
Learning with Diverse-quality Demonstrations, ICML 2020
9. Immer, Bauer, Fortuin, Ratsch, Khan, Scalable marginal
likelihood for model selection in deep learning, ICML 2021
10. Immer, Korpeza, Bauer, Improving predictions of Bayesianneural networks via local linearization, Aistats 2021

AI for Social Good: We outline a few guidelines on how
to align AI systems for social good applications [16].
16. Tomasev et al., AI for Social Good: Unlocking the Opportunity
for Positive Impact, Nature communications 2020