FY2022/2022年度

Approximate Bayesian Inference Team
Mohammad Emtiyaz Khan
近似ベイズ推論チーム
Overview and Goals
Goal: AI systems that can learn and improve continually throughout their lives, just like humans and animals. Currently, deeplearning (DL) requires a large amount of data which is costly and rigid, leading to a system that is unable to quickly adapt. We
aim to fix this with a new learning paradigm based on Bayesian principles.

Summary of our research highlights in the year 2022 (paper number shown in green boxes)-1 We show a robust deep-learning method is related to Bayes and add uncertainty to it (top 75/5000 papers at ICLR2023)
2- We extend the Bayesian learning rule using Lie-groups, simplifying gradient computations and eliminating retractions.
3- We design a generally unimprovable procedure for relaxation time estimation in non-reversible Markov chains.
-4 We simplify momentum-based Riemannian optimization over positive-semi-definite matrix submanifold.
5- We use low-rank matrix completion techniques to reconstruct partially-observed high-dimensional time series.
-6-10 Analysis of MMD estimation, Neural processes, A Dataset for African Languages, Deviation inequalities, and more …!
imal Relaxation of Bayes
SAM as an Optimal Relaxation of Bayes

SAM as an Optimal Relaxation of Bayes= 0.5
Variance =Variance0.5=2Variance = Variance
2Standard Deep Learning

Deep Continual Learning

The Lie-Group Bayesian Learning Rule

Variance = 4Variance = 4

Expected-Loss
Expected-Loss
Biconjugate
Biconjugate
Loss
Loss

Problem: Many popular algorithms can be derived from the Bayesian learning rule
Problem: Sharpness-Aware Minimization (SAM) by Foret et al. improve significantlySAM
SAM
0minimizer minimizer
0
SAM
SAM
of
Khan
and
Rue
(2021)
but
the
rule
can
be
difficulttoapply,
e.g.,
gradients
are
over
SGD
but
the
reasons
behindits
successareunclear.
minimizer minimizer
cy "
NLL #
ECE #
AUROC "SAM
SAM
Accuracy
"
NLL
#
ECE #
AUROC "
difficult to compute, and steps can lead to invalid distributions (e.g., –ve variances).
etter) Method(lower is better)(lower is better)(higher is better) minimizer minimizer
Solution: We show
SAM is(lower
equivalentrelaxation
ofBayes
Bayes obtained(higherthat
is better)is better)(lowerisbetter) Relaxation(higheris better)
SAM
as
an to
OptimalRelaxationof BayesSAM
asan
an optimal
Optimalof
0.044(0.0012)
0.00280.979(0.002)-1
Fenchelbiconjugate(left
figure).
SAM
canbe seen as0.979“smoothing” the-1(0.06)by using(0.0007)
Solution:We
extend
the
rule
by
using
Lie-groups
which
solvestheabove
problems:
SGD
98.63(0.06) 0.044(0.0012) 0.0028(0.0007)(0.002)0.039
0.0031
0.978(0.03)objective(0.0005)(0.0002)(0.003)using a 98.82“posteriorvariance”
andgradients can always be obtained by reparameterizations, and steps always stay on
SAM-SGD
0.039
0.0031upper-bounding
0.978Bayes(0.03)(0.0005) always(0.0002)(0.003) (right).
0.044(0.0002) 0.0028(0.0005)
0.976(0.003)(0.04)
-5
0
5 -5
-5 0
0 5 -5
5 -50
05
5
-5
0
5
SWAG
98.65
0.044
0.00280.976
Our0.046
paper [1] is among
top-5%
ofall
accepted
papers
(75
out
of
5000
submissions)(0.04)(0.0002)(0.0005)(0.003)themanifold.
Fisher computation is also simplified and only need to be done once.
0.00330.980(0.03)(0.0007)(0.002)
VOGN (0.0008) 98.54(0.03)
0.046(0.0008)
0.0033(0.0007)
0.980(0.002)
translation
0.050(0.0012) 0.0036
0.979
Conceptual
algorithm:(0.06) Adam
(0.0007)(0.002)⇡(g)q (✓) = q (✓ g)q
g98.41(0.06) 0.050(0.0012) 0.0036
0.979(0.0007)(0.002)
Variance= 0.5
Variance=2
Variance = 4

Variance = 0.5
Variance = 2
Variance = 4
sup ( sup+0.980) ( +(0.002)
0.046
0.0044SAM:(0.0009)(0.0005))
SAM:
SAM-Adam
98.58(0.03)| |< 0.046(0.0009) 0.0044
0.980(0.002)
g
g exp( ↵Y )(0.0005)
Expected-Loss
qg
0.038(ours)
0.982(0.001)| |<Expected-Loss(0.06) bSAM
(0.0011) 0.0024(0.0004)0.038
98.78(0.06)
0.0024
0.982
Biconjugate(0.0011)(0.0004)(0.001)
scaling
Biconjugate
Loss
g
SAM
Loss
0.32
0.04290.897
SAM(0.27) SGD(0.007)(0.0015)(0.004)q0
91.37(0.27)
0.32(0.007) 0 0.0429
0.897
minimizer(0.0015)(0.004)
1✓
SAM
qg
minimizer
0
⇡(g)q
(✓)
=
q
(
)
`(✓)
0
0
SAM
Fenchel0.22
0.00620.917
gg(0.17) SAM-SGD(0.004)(0.0007)(0.002)minimizer
Fenchel
91.95(0.17)
0.22
0.917
SAM(0.0007)(0.004) 0.0062(0.002)minimizer
SAM
Biconjugate
0.31
0.0397
0.901
minimizer(0.27) SWAG (0.005)(0.0026) 0.31
(0.005)
Biconjugate
91.38(0.27)
0.03970.901(0.005)0minimizer(0.005)(0.0026)translation & scaling
0.24(0.004) 0.0071
0.916(0.25) VOGN
(0.0012) 0.24
(0.002)91.24(0.25)
0.00710.916(0.004)(0.0012)(0.002)-1⇡(g)q0 (✓) = A1 q0 ( ✓A b )
q
g
=
(A,
b)
0.33
0.04500.897
0(0.25) Adam
(0.005)(0.0008) 0.33(0.005)(0.005)-1
qg
91.14Bayes:
0.04500.897
CC(0.25)(0.0008)(0.005)A
A
Bayes:
0.25(0.004) 0.0225
0.913(0.19) SAM-Adam(0.0026)
91.66(0.19)
0.25
0.0225(0.0026) B 0.913(0.002) B[(+)] (0.002)
2(0.004)(0, )2)[ ( + )]
-5
0
5 -5
(0, 0.920
0
5 -5
0
5
0.22
0.0066(0.26) bSAM
(0.005)(0.0022)(0.002)(ours)
92.10(0.26)
0.22(0.005)
0.00660.920-5 (0.0022)0
5 -5(0.002) 0
5 -50
5(a) 0.0839
SAM vs Bayes(b) Higher variances give smoother objectives
Result: Reduces to the update [3] when applied to exp. fam. and linearized.
0.56
0.878
(0.35) SGD(0.014)(0.003)(0.005)0.0839(0.003)
(a)
SAM
vs
Bayes
(b)
Highervariancesgivesmootherobjectives
86.55
0.56
0.878
(0.35)(0.014)(0.005)Contribution#2:
Our
Bayesian-SAM
improves
0.49(0.019) 87.49
0.07100.891(0.26) SAM-SGD(0.003) 0.49
(0.004)0.0710(0.003)
0.891(0.004) SAM is(0.26)(0.019)Affineaccuracy
over1:SGD(by
8%),
SAM-Adam(by
10%),
Additive
Multiplicative
FigurePanel(a)highlightsthemaindifference
betweenSAM
andBayes.
SAM
uses
max-loss
0.53
0.07740.880(0.10) SWAG
(0.017)(0.001)(0.006)86.80
0.53(0.017)
0.07740.880SAM(0.10)(a) highlights(0.001) between(0.006)andoverconfident
Figure
1:
Panel
the
main
difference
Bayes.
SAM
uses
2 max-loss
(the
red
dot),
while
Bayes
uses
expected-loss(the blue
dot shows`(✓(0.003)+ ✏) for an ✏ ⇠ N (✏ | 0, )).
0.38
0.0315
0.890
Adam
(by
22%),
but
also
improves
AUROC
2(0.24) and(0.004)(0.003)(0.003)0.0315
RN
VOGN
87.30
0.38
0.890
(0.24)(0.004)(0.003)(narrow
white
(the
red
dot),
while
Bayes
uses
expected-loss
(the
blue
dot
shows
`(✓
+
✏)
for
an
✏
⇠
N
(✏
|
0,
)).
Our
main
resultin
Theorem2
connects
the
two
by
using
the
optimalconvexrelaxation(Fenchel
0.83(0.063) using
0.13170.820(0.011) with
(0.013)0.1317(0.011)
0(0.98) (for
Adam
80.85
0.83
0.820
CIFAR-100
ResNet-20
270K(0.98)(0.063)(0.013)Our
mainresultin
Theorem2
connects
the
two
by
using
the
optimalconvexrelaxation(Fenchelregion)biconjugate)oftheexpected
loss.
It
shows
that
the
role
of
⇢
and
are
exactly
the
same,
and
a
0.46(0.007) 85.26
0.02280.874
0.874(0.15) SAM-Adam(0.002) 0.46
(0.004)0.0228(0.002)(0.15)(0.007)(0.004)params).
biconjugate)oftheexpected
loss.
It
shows
that
the
role
of
⇢theandare exactly
the .same,
and a
SAM
minimizer
obtained
for
a
fixed
⇢
can
alwaysberecoveredfromrelaxation
for
some
An
0.34
0.01630.903
Accuracy(ours)
88.72
0.01630.903(0.003)(0.24) bSAM
(0.005)(0.002) 0.34
(0.003)(0.24)(0.005)(0.002) be recovered
AUROC
SAM
minimizer
obtained
for
a
fixed
⇢
can
alwaysthe relaxation. An
example is shown in Panel (b) for a loss (gray line) with 3 local minimafrom
indicatedby A, B,for
andsomeC.
Multiplicative Learning (nll=0.058)
exampleis shown
in Panel(b)line)fora lossline) whichwith
3upperlocal
minimait,indicatedby A, B, and C.
0.1695
0.811
+8%
55.82
1.91
0.16950.811(0.97) SGD1.91(0.025)(0.005)(0.004)(0.97)(0.025)(0.005)(0.004)The
expectedloss
is smoother
(blue
but
the (grayrelaxation,boundsis even smoother
1.60
0.09890.827
The expectedlosssmoother
(blue
line)
but(0.005)
therelaxation,whichboundsit, is even
58.58
1.60
0.09890.827(0.59) SAM-SGD(0.022)(0.005)(0.003)(0.59)(0.022)(0.003)(redline).
Highergiveissmootherobjectives
where
sharp
minimaA
andupper
CBayesian-SAM
slowlydisappear.
Thesmoother
1.86
0.1604(redline).
Highergivesmootherobjectives
where
sharp
minimaA
andthisC slowly disappear. The
56.53
1.86
0.1604
0.814
fixes(0.40) SWAG
(0.018)(0.004)with(0.004)(0.40)(0.018)(0.004)(0.004)SAM
minimizer
isshown
a0.814
red
star
which
matchesthe minimizerof therelaxation.
1.44(0.019)
0.0756
0.830
RN
59.83
1.44
0.0756
SAM minimizer
is shown
witha red star
which
matches0.830
the minimizerof therelaxation. Additive Learning (nll=0.083)
blurry(0.75) VOGN
(0.005)(0.002)(0.75)(0.019)(0.005)(0.002) (more
0(0.97) Adam
39.73(0.97)
0.0295(0.018)
0.775(0.004) away from the2.29(0.045) 0.0295
0.775+10%2.29(0.045)(0.018)(0.004)2 (0.035)
SAM 53.25
AS
AN
O PTIMALR
ELAXATION
OF BAYES
1.71
0.04010.818(0.005) data)
1.71
0.04010.818(0.80)(0.035)(0.005)(0.80) SAM-Adam(0.005)(0.005)+22%
2 SAM
AS AN O1.32
PTIMALR
ELAXATION
OF0.841B
AYES
bSAM(ours)
62.64
0.03111.32
0.03110.841(0.33)(0.001)(0.003)(0.004)(0.33)(0.001)(0.003)(0.004)
In SAM, the training loss `(✓) is replaced by the maximum loss within a ball of radius ⇢ > 0 aroundP loss `(✓) is replaced by the maximum loss within a ball of radius ⇢ > 0 around
In
SAM,
the
trainingtheparameter✓2⇥⇢
R, as shown
below foran `2 -regularized
problem,
1. Moellenhoff, Khan, SAM as an Optimal
RelaxationofBayes,
ICLR
2023(oral,
notable top-5%)
2. Kiral, Moellenhoff, Khan, The Lie-Group Bayesian Learning Rule, AISTATS 2023.
P
the parameter ✓ 2 ⇥ ⇢ R , as shown below for an `2 -regularized problem,rison
without
data
results areaveragedover 5 random seeds,
withion. The
results
areaugmentation.
averaged overThe
5 randomseeds,
with2
ESAM (✓; ⇢, ) = sup `(✓ + ✏) + k✓k ,
(1)
on
shown
in the brackets.
The inbest
performance
is ⇢,
shown
bold`(✓using
he bestperformanceis shownbold
using
statistical
k✏k⇢
ESAM
(✓;
) = insup+2 ✏)statistical+ k✓k2 ,(1)
Estimation
Time in Non-Reversible Markov Chains
AM (gray row)
givesperformance
either comparable
or better
performancethanoftheRelaxation
rest. 2k✏k⇢comparableor
better
than
the
rest.
with > 0 as the regularization parameter. The use of ‘maximum’ above differs from Bayesianwiththat>use0 ‘expectation’,as the regularization
The useformulationof ‘maximum’
above(1988),
differs from
Bayesianstrategies
for ofexample,the variational
by Zellner
Contributions:
Problem:
Estimate
the relaxation
time
anparameter.
unknown,
ergodic,
non-reversible
Markovstrategies that use ‘expectation’,the variational
formulation by Zellner (1988),
Ø (2)Prove there exists an estimatorforsuch that for any ergodic chain, whenL(q) = E for[ example,`(✓)] D(q(✓) k p(✓)),
00

0(0.03)

1

2

3

✓⇠q
chain from a single trajectory of length 𝒎started fromKL
an arbitrary state.

L(q) = E(Zhang,
DKL (q(✓)
k p(✓)),(2)✓⇠q [ `(✓)]
where q(✓) is the generalized posterior
1999; Catoni,2007),p(✓) is the prior, and
where.
DKL (·k ·) isAccuracy
theECE
Kullback-Leibler
divergence(KLD).
For anisotropic
Gaussian
=
where
q(✓)
is "#the generalized
posterior(Zhang,1999;
Catoni,2007),posteriorp(✓) is q(✓)
the prior,and is the minimum stationary probability, it holds w.h.p. that
NLL
#
ECE
#
AUROC
"
acy "Application:
NLL
#
AUROC
"
This
is
2 useful in reinforcement learning, Markov Chain Monte Carlo
MethodN is(✓better)
|Dm,
I)
prior
p(✓)
=
N
(✓is |better)0, I/), theis objective
to the Gaussian
mean m becomesk(lower·)andtheKullback-Leibler
divergence(KLD).with
Forrespectanisotropic
posteriorq(✓) bound=(higher
isisbetter)(loweris better)(lowerbetter)(higheris better)
KL (·
Upper
matches known lower bounds, thus generally unimprovable.
better)(lowerisbetter)(higherdiagnostics, anddata. to the mean m becomesN (✓deriving| m, 2 I)generalizationand prior p(✓) bounds= N (✓ | with
0, I/ Markov-dependent), the objective with respectØDesignalgorithm
with
fully
empirical
confidence
intervals
that
decay
in
.
0
2
SGD
91.68
0.29
0.03970.915(0.26)(0.008)(0.002)(0.002)
EBayes (m; 0.915, )=E✏0 ⇠N (0, 2 I) [`(m + ✏ )] + kmk ,(3)
8(0.26)0.29(0.008) 0.0397(0.002)(0.002)
20
RN /3. Wolfer
SAM-SGD
92.29
0.25
0.02660.920
2(0.39)(0.004)(0.003)(0.003)90.25
0.02660.920
02
EBayes (m;(0.003), Relaxation) = E✏ ⇠NTime
[`(m
+
✏
)]
+
kmk
,Chains, arXiv 2022. (3)and(0.004)
Kontorovich, Improvedofin
Non-reversible
Markov(0.39)(0.003) Estimation(0, I)
07
Adam
0.06100.900bywhich closely89.97
resembles
1,0.900butwith themaximuman expectation.
The above2(0.27) Eq. 0.41
(0.021)(0.003) replaced(0.003)0.41
0.0610(0.27)(0.021)(0.003)(0.003)
91.57
0.26
objective
is obtainedby simply
plugging
in0.0329
the posterior
and 0.918
prior,
then collecting
the terms that
(0.21)(0.004)(0.002)(0.001)whichclosely
resembles
Eq.
1,
but
with
the
maximum
replacedbyanexpectation.
The above
7(0.21) SAM-Adam
0.26
0.03290.918(0.004)(0.002)(0.001)0bSAM(ours)
92.16
0.23
0.00570.925
dependonm,andfinally
rewriting✓=m+✏. in(0.002)(0.16)(0.003)(0.001)
objective
is obtainedbysimplypluggingtheposterior
andprior, then collecting
the terms that
Other Works6(0.16)0.23(0.003)
0.00570.925(0.002)(0.001)0dependonm,andfinally
rewriting
✓
=
m
+
✏
.
SGD
66.48
1.20
0.0524
0.846
There
are
some
resemblances
betweenthetwoobjectives
which
indicate(0.10)(0.007)(0.004)(0.002) that they might be con8(0.10)
1.20
0.05240.846(0.007)(0.004)(0.002)
Improving
Neural
Process:We
improvetest-time
inferencefor
Neural
Processes
bySimplify
momentum-based
Riemannian
optimization:We
considerthespecificcase
RN/
SAM-SGD
67.27
1.19
0.04810.848
nected.
For
example,
both
use
local
neighborhoods:
theisotropic
Gaussian
can
be
seen
as
a
‘softer’(0.22)(0.011)(0.001)(0.002)indicate that they might be conThere aresomeresemblances
betweenthe two
objectives which7(0.22)
1.19(0.011)
0.04810.848(0.001)(0.002)
incorporatingand exploiting graphical-model structure among context points.
submanifold
containing
positive-definite
matrices.
The
methoduses
a0 of aAdam
61.76
1.66
0.1582
0.826
versionof
the‘hard-constrained’
ball
used
in
SAM.
The
size
of
the
neighborhood
isdecidedbytheir(0.67)symmetric(0.049)(0.006)(0.003)nected.
For
example,
both
use
local
neighborhoods:
theisotropic
Gaussian
can
be
seen
as
a
‘softer’6(0.67)
1.66(0.049) 0.1582(0.006)
0.826(0.003)
SAM-Adam
65.34
1.23
0.01660.847
respective
parameters
(
or
⇢
in
Fig.
1(a)).
But,
apart
from
these
resemblances,
are there any
deeper(0.32)(0.012)(0.003)(0.002)
7. Tailor,
generalizedversionof
local
coordinates
which“trivializes”
theFishermatrix.
versionof
the‘hard-constrained’
ball
used
in
SAM.
The
size
of
the
neighborhood
is decided
by Khan,their Nalisnick, Exploiting Inference Structure in Neural Processes, UAI 2022 Workshop on
4(0.32)
1.23(0.012) 0.0166(0.003)
0.847(0.002)
bSAMconnections?
(ours)
68.22
1.10
0.0258
0.857
Here,
we
answerthisquestion.
(0.44)(0.013)(0.003)(0.004)
Tractable
Probabilistic Modelingrespective
parameters
(
or
⇢
in
Fig.
1(a)).
But,
apart
from
these
resemblances,
are
there
any
deeper2(0.44)1.10
0.02580.857(0.013)(0.003)(0.004)
4. Lin, Duruisseaux, Leok, Nielsen, Khan, Schmidt, Practical Structured Riemannian Optimization
connections?
Here,
weanswerthisquestion.
SGD We will
52.01
1.98(0.007)
0.0330(0.002)
0.832
Name-Entity Recognition (NER) dataset for Sub-Saharan African languages: We create
show
that
Eq.
1
is
equivalenttooptimizing
withrespect
toWorkshopman optimal
convex relaxation(0.36)(0.002)
with Momentum
Generalized
NormalCoordinates,
NeurIPS2022
on
1(0.36)
1.98(0.007)by using
0.03300.832(0.002)(0.002)
RN / SAM-SGD
52.25
1.97See0.01550.827(0.005)is drawn between the(0.26)(0.013)(0.002) an equivalence
of
Eq.
3
while
keeping
fixed.
Theorem2
whereandlargest human-annotated dataset called MasakhaNER 2.0, and analyze features
5
1.97
0.01550.827
Symmetry
andGeometryin
NeuralRepresentation
We
will
show
that
Eq.
1
is
equivalenttooptimizing
withrespecttoman
optimalconvexrelaxation(0.26)(0.013)(0.002)(0.005)
Net
Adam⇢, revealing similarities
49.04(0.38) between
2.14the0.0502used
0.820(0.024)(0.004)(0.004)mechanisms
in
SAM
and
Bayes to favor
flatter regions.
4(0.38) SAM-Adam
2.14(0.024)
0.05020.820
of Eq. 351.17
while(0.004)
keeping2.02
fixed.
See
Theorem2
whereanequivalenceisdrawn
betweenandthat contributeto cross-lingual transfer, giving large gains for 0-shot learning.
(0.004)0.0460
0.828(0.45)(0.014)(0.004)(0.004)
Fig. 1(b) shows
anseries
illustration
where
the Werelaxationis shownmatrix
to upperbound the expected loss. We
High-dimensional
timecompletion:use
low-rank
completion
7(0.45)2.02
0.04600.828⇢,
revealingsimilarities
betweenthemechanisms
used
in
SAM
and Bayes to favor flatter regions.
(0.014)(0.004)(0.004)bSAMwill(ours)
52.90
1.94
0.01990.831
8. Buzaaba with many others, MasakhaNER 2.0: Africa-centric Transfer Learning for Named Entity(0.35)(0.009)(0.003)(0.001)now
describe
theconstructionoftherelaxation
basedonFenchelbiconjugates.
techniques
toreconstructpartiallyobserved
high-dimensional
timeseriesandshow
0(0.35)
1.94(0.009)
0.01990.831where
Fig. 1(b)shows(0.003)
an illustration(0.001)the relaxation is shown to upper bound the expected loss. We
Recognition, EMNLP 2022
will
now
describe
theconstructionoftherelaxation
basedonFenchelbiconjugates.
that periodicity or smoothness can even lead to faster rates than in the independent
2.1 F ROM EXPECTED - LOSS TO MAX - LOSS VIA F ENCHEL BICONJUGATESDeviation inequalities for stochastic approximation by averaging: We establish

setting.
arisonwith data augmentation. Similar to Table 1, the shaded row show that bSAM
deviation
inequalities for separately Lipschitz functions of Markov chains belonging to
2.1
F
ROM
EXPECTEDLOSSTOMAX -expected-loss
LOSS VIA F ENCHEL
BICONJUGATESn.
Similar
to
Table
1,
the
shaded
row
show
that
bSAM
We
use
results
fromconvexdualitytoconnecttomax-loss.
Our
mainresultrelieson
provesover theMarie,
baselines
andis thebest
method.
5. Alquier,
Rosier,Tight
riskoverallbound forhighdimensional time series completion, EJS 2022
certain class we define, and which includes models of stochastic approximation byis the overallbest method.
Fenchelconjugates (Hiriart-Urruty & Lemaréchal, 2001), also known as convex conjugates. aConWe use results from convex duality to connect expected-loss to max-loss. Our main result
relies on
averagingand non-averaging.
Finite sampleFenchel
properties
ofparametricMMD
estimation:Wetackletheproblemofconjugates (Hiriart-Urruty & Lemaréchal, 2001), also known as convex conjugates. Con8minimum distance estimator based on Maximum Mean
9. Fan, Alquier, Doukhan, Deviation inequalities for stochastic approximation by averaging,SPA 2022
universal
estimation
using
a
2
8
Discrepancy, and we show its robustness to both dependence and presence ofPrioritization of minibatches: We give empirical support for the hypothesis that2
outliers.
improving calibration can help in prioritizing minibatches during training.
6. Chérief-Abdellatif, Alquier, Finite Sample Properties of Parametric MMD Estimation:
Robustness to Misspecification and Dependence, Bernoulli 2022

10. Tata, Gudur, Chennupati, Khan, Can calibration improve sample prioritization, NeurIPS 2022
Workshop on Has It Trained Yet