（様式 20）成果報告書

TSUBAME 共同利用 令和年度 学術利用 成果報告書

利用課題名 22IAJ - Improving Sequential Bayesian Inference
利用課題責任者
Mohammad Emtiyaz Khan所属
Approximate Bayesian Inference Team, RIKEN Center for Advanced Intelligence Project
http://team-approx-bayes.github.io
Bayesian principles offer a mathematically rigorous way to fundamentally improve and rethinkcontinual learning and adaptive AI systems. However, performing approximate Bayesian inference
for deep learning and large neural networks is still a challenging open question and requires large
amounts of compute resources. In this project, we employed TSUBAME3.0 to develop newstate-of-the-art Bayesian deep learning methods. These lay the foundation for our future work on
developing robust practical AI systems which can continually learn and update their beliefs byperforming sequential Bayesian inference.

Keywords: Approximate Bayesian Inference, Continual Learning, Deep Learning, Optimization,
Numerical Algorithms

背景と目的

Current machine-learning and AI systems arevery rigid and struggle to adapt to newsituations. Sequential Bayesian inference
offers a principled way to update the current
beliefs based on newly observed information.
However, performing exact Bayesian inference
for large neural network models iscomputationally intractable. The goal of thisproject was to develop new state-of-the-art
methods for approximate Bayesian inference.
These algorithms are expected to enable AIsystems
tobecome
morerobust
toperturbations and continually adapt likehumans and animals.
概要

The report is structured in three sections. Eachsection corresponds to an algorithm forBayesian deep learning we researched anddeveloped using TSUBAME3.0.
1. Second-Order Optimization via Bayes(SOBA)
2. The Lie-Group Bayesian Learning Rule(LieGroup-BLR)
3. Sharpness-Aware Minimization as an

Most of modern deep learning uses optimizers
to train deep neural networks. The following
equations show the update-rule of the popular
RMSprop / Adam deep learning optimizer:

To enable deep learning systems to adapt to
new situations and continually learn, it isessential for these models to know what theyknow. However, these models should also
understand what they don’t know yet andshould not be overconfident in such situations.
Starting from a rigorous mathematical theory,
we have developed a modification of RMSpropwhich can estimate uncertainty in deeplearning by computing an approximate
Bayesian posterior. The algorithm is shown in
the following equations, where the changes toRMSprop/Adam are highlighted in red.

Optimal Relaxation of Bayes (bSAM)
結果および考察

1. Second-Order(SOBA)

Optimization

via

Bayes

This method allows the deep learning model tobe uncertain about its decisions in regions

（様式 20）成果報告書

where there is no data, as shown in thefollowing Figure 1.

Figure 1. When training neural networks, using
our SOBA optimizer instead of RMSprop allows
the model to be uncertain in situations wherethere is little observed data. This uncertainty isessential to quickly adapt to new beliefs once
more information is observed.
The SOBA optimizer offers state-of-the-art
performance for Bayesian deep learning. A
previous version was part of our winning solution
to the NeurIPS 2021 Challenge on Approximate
Inference
inBayesian
DeepLearning(https://izmailovpavel.github.io/neurips_bdl_competition/).
TSUBAME3.0 was essential in developing andbenchmarking this method. Large computational
resources are required when training neuralnetworks on big datasets such as ImageNet or
simply to train modern neural networkarchitectures. In future projects on sequential
Bayesian inference for large neural networks,
SOBA will be one of the key ingredients to obtainaccurate models with well-calibrated uncertainty.

2. The Lie-Group Bayesian Learning Rule
(LieGroup-BLR)
Due to its mathematical formulation, the SOBA
optimizer appears to be restricted to perturb
the weights with Gaussian noise (see theAlgorithm before Figure 1). In a follow-up work,
we found a way to generalize SOBA fromGaussian noise to arbitrary noise distributions.
This is achieved by using the mathematicalframework of Lie groups (see the following
Figure 2).

Figure 2. In approximate Bayesian inference, aprobability distribution (called the Bayesianposterior) is estimated. SOBA is restricted toGaussian distributions, but using the Lie-groupformalism, we were able to generalize to other
distributions. The figure shows Uniform,
Rayleigh and Laplace distributions.
In Figure 3, we compare the LieGroup-BLR to
SGD and the SOBA optimizer (here called“iVON”).
We observe that the additional
flexibility to choose a non-Gaussian noise
distribution can sometimes improve the results.

Figure3.
Thedeveloped
Lie-Group
optimization methods improve the performanceover traditional SGD optimization in terms ofaccuracy and uncertainty on image classification
datasets with neural network models.

3. Sharpness-Aware Minimization
Optimal Relaxation of BayesasanSharpness-Aware Minimization (SAM) is a
recent method proposed by Google researchers to
improve generalization and robustness in deeplearning. The method has been very influential
over the last years, accumulating over 500
citations in a short timespan.

（様式 20）成果報告書

まとめ、今後の課題

Figure 4. Sharpness-Aware Minimizationconsiders the worst-case perturbation in aneighborhood, whereas Bayesian approaches
average over many random perturbations.
In our recent work we discovered a connectionof SAM to Bayesian approaches, as illustrated in
the above Figure 4. This connection enabled us topropose a Bayesian-SAM method (called bSAM)
which extends SAM to obtain uncertainty
estimates.
We illustrate the uncertainty estimation inFigure 5. The decision boundary obtained bybSAM is more “blurred out”, especially in regions
where there is no data. In these regions, themodel is uncertain about the correct prediction.

Figure 5. bSAM extends SAM to obtain
uncertainty estimates without computationaloverhead.

Figure 6. The proposed bSAM method improves
accuracy and uncertainty estimation whencompared to SAM.
Using TSUBAME3.0 was essential indeveloping and researching our bSAM method onlarge deep learning models.

In this project, we have developed newstate-of-the-art methods for approximate
Bayesian inference (SOBA, LieGroup-BLR andbSAM). Our next step is to use thesealgorithms to improve sequential and continual
learning of AI systems, making them moreadaptive and robust. Furthermore, future effort
will be spent on applying the approximate
Bayesian inference algorithms developed inthis project on larger problems and machinelearning models. For instance, we want to showthat our methods can be useful in large-scalesettings like large language models.