summaries: Inspired by asymptotic results from singular learning theory, the authors of the paper propose using a generalised gamma mean-field distribution with a normalising flow (that targets the "desingularization map") to perform variational inference. The authors additional build on prior work to derive a tighter bound on the log normalised evidence for variational inference using this approximate posterior combined with the "correct" desingularization map.   Note: while I am familiar with the Bayesian Neural Network literature, I am not familiar with singular learning theory. This is reflected in my confidence score. As such, I was unable to verify the theory in the paper, and assess the significance of the theoretic results included in the paper.

summaries: The paper analyzes the approximation of the log-ratio (called normalized evidence in the paper) (the log-probability under the model minus the log-probability under the true model). Since the model is singular, the Laplacian approximation of the normalized evidence needs a new method. The mean-field distribution and a resolution map allow us to estimate two important numbers (lambda and m) in the approximation.   The paper has made several contributions to the study of this problem in several aspects. First, the paper improves previous results by relaxing the assumption. Second, the paper actually uses a normalizing flow to actually learn the transformation.

summaries: The paper builds on the recent theoretic results of [Bhattacharya et al. (2020)], which shows that a mean-field variational approximation with carefully chosen approximation family leads to an ELBO $\Psi$ which is sharp up to a constant C(d) which only depends on the dimensionality d of the parameter space. For the proof, [Bhattacharya et al. (2020)] assumes that the approximation family is a Gamma distribution truncated to [0,1].  Therefore the authors of this paper conjecture that the Gamma distribution truncated to [0,1] is a good choice for the source distribution of a normalizing flow. Their experimental results suggest that this choice leads a higher ELBO than using as source distribution a normal distribution.

summaries: Working with Bayesian neural networks, this paper proposed a variational algorithm to approximate posterior distribution of the network weights. To overcome model singularities, the authors used the idea of normalizing flow by transforming the weights through an affine coupling network, and subsequently worked on the desingularized parameter space. In addition, they derived an asymptotic expression for the ELBO, and compared the Gaussian and generalized gamma approximating families in the experiments.

abstractText: Predicated on the premise that neural networks are best viewed as singular statistical models, we set out to propose a new variational approximation for Bayesian neural networks. The approximation relies on a central result from singular learning theory according to which the posterior distribution over the parameters of a singular model, following an algebraic-geometrical transformation known as a desingularization map, is asymptotically a mixture of standard forms. From here we proceed to demonstrate that a generalized gamma mean-field variational family, following desingularization, can recover the leading order term of the model evidence. Affine coupling layers are employed to learn the unknown desingularization map, effectively rendering the proposed methodology a normalizing flow with the generalized gamma as the source distribution.

authors:
Miki Aoyagi
Sumio Watanabe

title: Stochastic complexities of reduced rank regression in Bayesian estimation

venue: Neural Networks,

title: Resolution of Singularities and the Generalization Error with Bayesian Estimation for Layered Neural Network

venue: IEICE Trans,

authors:
Anirban Bhattacharya
Debdeep Pati
Sean Plummer

title: Evidence bounds in singular models: Probabilistic and variational perspectives

authors:
Charles Blundell
Julien Cornebise
Koray Kavukcuoglu
Daan Wierstra

title: Weight Uncertainty in Neural Network

venue: In International Conference on Machine Learning,

authors:
Tianqi Chen
Emily Fox
Carlos Guestrin

title: Stochastic gradient hamiltonian monte carlo

venue: Proceedings of the 31st International Conference on Machine Learning,

authors:
Andrew Y.K. Foong
David R. Burt
Yingzhen Li
Richard E. Turner

title: On the expressiveness of approximate inference in bayesian neural networks

venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems

authors:
Yarin Gal
Zoubin Ghahramani

title: Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning

venue: In Proceedings of The 33rd International Conference on Machine Learning},

title: Practical variational inference for neural networks

venue: Advances in Neural Information Processing Systems,

authors:
Chuan Guo
Geoff Pleiss
Yu Sun
Kilian Q. Weinberger

title: On calibration of modern neural networks

venue: In Proceedings of the 34th International Conference on Machine Learning - Volume 70,

title: Variational Approximation Error in Bayesian Non-negative Matrix Factorization

title: Well-Calibrated Bayesian Neural Networks

venue: PhD thesis, University of Cambridge,

authors:
Jose Hernandez-Lobato
Yingzhen Li
Mark Rowland
Thang Bui
Daniel Hernandez-Lobato
Richard Turner

title: Black-box alpha divergence minimization

venue: Proceedings of the 33rd International Conference on Machine Learning,

authors:
T. Hosino
K. Watanabe
S. Watanabe

title: Stochastic complexity of variational Bayesian hidden Markov models

venue: In Proceedings

authors:
Pavel Izmailov
Dmitrii Podoprikhin
T. Garipov
Dmitry P. Vetrov
Andrew Gordon Wilson

title: Averaging weights leads to wider optima and better generalization

venue: Uncertainty in Artificial Intelligence (UAI),

authors:
Mohammad Khan
Didrik Nielsen
Voot Tangkaratt
Wu Lin
Yarin Gal
Akash Srivastava

title: Fast and scalable bayesian deep learning by weight-perturbation in adam

authors:
Masahiro Kohjima
Sumio Watanabe

title: Phase Transition Structure of Variational Bayesian Nonnegative Matrix Factorization

venue: Lecture Notes in Computer Science,

authors:
Ranganath Krishnan
Omesh Tickoo

title: Improving model calibration with accuracy versus uncertainty optimization

authors:
Balaji Lakshminarayanan
Alexander Pritzel
Charles Blundell

title: Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles

authors:
Yingzhen Li
Richard E. Turner

title: Rényi divergence variational inference

venue: In Proceedings of the 30th International Conference on Neural Information Processing Systems,

title: Algebraic Methods for Evaluating Integrals in Bayesian Statistics

venue: PhD thesis, University of California Berkeley,

authors:
Christos Louizos
Max Welling

title: Multiplicative Normalizing Flows for Variational Bayesian Neural Networks

title: A practical Bayesian framework for backpropagation networks

venue: Neural Computation,

authors:
Wesley J Maddox
Pavel Izmailov
Timur Garipov
Dmitry P Vetrov
Andrew Gordon Wilson

title: A simple baseline for bayesian uncertainty in deep learning

authors:
Stephan Mandt
Matthew D. Hoffman
David M. Blei

title: Stochastic Gradient Descent as Approximate Bayesian Inference

authors:
Daniel Murfet
Susan Wei
Mingming Gong
Hui Li
Jesse Gell-Redman
Thomas Quella

title: Deep Learning is Singular, and That’s Good

authors:
Shinichi Nakajima
Sumio Watanabe

title: Variational Bayes Solution of Linear Neural Networks and Its Generalization Performance

title: Bayesian Learning for Neural Networks, volume 118 of Lecture Notes in Statistics

authors:
Hippolyt Ritter
Aleksandar Botev
David Barber

title: A scalable laplace approximation for neural networks

venue: In 6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings,

authors:
Shengyang Sun
Guodong Zhang
Jiaxin Shi
Roger Grosse

title: Functional Variational Bayesian Neural Networks

venue: In International Conference on Learning Representations,

authors:
Kazuho Watanabe
Sumio Watanabe

title: Stochastic Complexities of Gaussian Mixtures in Variational Bayesian Approximation

venue: The Journal of Machine Learning Research,

title: Algebraic analysis for non-regular learning machines

title: Algebraic Geometry and Statistical Learning Theory

title: Mathematical Theory of Bayesian Statistics

venue: Chapman and Hall/CRC,

authors:
M. Welling
Y.W. Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: In Proceedings of the 28th International Conference on Machine Learning,

authors:
Andrew Gordon Wilson
Pavel Izmailov

title: Bayesian Deep Learning and a Probabilistic Perspective of Generalization

authors:
Guodong Zhang
Shengyang Sun
David Duvenaud
Roger Grosse

title: Noisy Natural Gradient as Variational Inference

authors:
Ruqi Zhang
Chunyuan Li
Jianyi Zhang
Changyou Chen
Andrew Gordon Wilson

title: Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning

title: 2020) that the variational parameters λ and k (and

heading: 1 INTRODUCTION

text:
Singular statistical models are ubiquitous in modern machine learning. In contrast to their regular counterpart, singular models need not be identifiable nor possess a positive-definite Fisher information matrix. These departures from classic regularity conditions have considerable implications for both theory and practice. Since neural networks are (strictly) singular models, accounting for such differences may be critical to some of the most pressing challenges in deep learning theory including model selection and the generalization puzzle (Murfet et al., 2020).
In this work, singular learning theory (Watanabe, 2009) is brought to bear on the challenge of performing inference in Bayesian neural networks (MacKay, 1992; Neal, 1996). In particular, singular learning theory helps us understand the large-sample properties of the posterior distribution over neural network weights w ∈ Rd. By (Watanabe, 2018, Chapter 6), there exists a so-called desingularization map (also known as a resolution map), g : Rd → Rd, g(ξ) = w, such that when n is large, the posterior distribution in a local weight set is proportional to
exp(−nξ2k11 ξ 2k2 2 · · · ξ 2kd d )|ξ h1 1 · · · ξ hd d |b(ξ). (1)
The k’s and h’s are natural numbers, and b(·) is a real-valued C∞ function. We say the posterior distribution has been put into standard form if coordinates ξ have been found that allow the posterior to be locally written as in equation 1. Remarkably, it is possible under very general conditions to write the posterior distribution over the parameters of a singular model as just such a mixture of standard forms. Based on equation 1, we proceed to demonstrate that a certain mean-field variational family, following desingularization, can recover the leading order term of the log (normalized) evidence, up to constant terms that do not depend on sample size n or dimension d.
Let (x, y) denote the input-target pair modeled jointly as p(x, y|w) = p(y|x,w)p(x). Let us assume the parameter space W is a compact set in Rd and p0(x, y) = p0(y|x)p(x) is the true data-generating mechanism. Throughout we suppose there exists w0 ∈ W such that p0(y|x) = p(y|x,w0). In the parlance of singular learning theory, this condition is known as realizability. Let ϕ(w) be a compactly-supported prior. We shall refer to (p(·, ·), p0(·, ·), ϕ(·)) as a model-truth-prior triplet. Define K(w) to be the Kullback-Leibler divergence between the truth and the model, as a function of the model parameters w:
K(w) = KL(p0(x, y)||p(x, y|w)) = Ep0 log p0(y|x) p(y|x,w) .
Following Watanabe (2009), we say a model is regular if it is 1) identifiable i.e., {w : K(w) = 0} is a singleton, and 2) its Fisher information matrix I(w) is positive definite for arbitrary w ∈W . We call a model strictly singular if it is not regular. The term singular will refer to either regular or strictly singular models.
Henceforth let p(y|x,w) be a neural network model. We wish to approximate the intractable posterior distribution over neural network weights, p(w|Dn) = ∏n i=1 p(yi|xi,w)ϕ(w) Z(n) , whereDn = {(xi, yi)} n i=1
is a dataset of n input-output pairs. The normalizing constant, Z(n) = ∫ ∏n
i=1 p(yi|xi, w)ϕ(w) dw, is variously known as the (model) evidence and the marginal likelihood.
The asymptotic expansion of Z̄K(n), a variant of Z(n) defined in equation 5, will prove crucial to justifying the proposed variational approximation. For strictly singular models, quantities such as Z(n) and Z̄K(n) manifest as a singular integral, i.e., an integral of the form ∫ W e−nf(w)ϕ(w) dw where W ⊂ Rd is a compact semi-analytic subset, and f and ϕ are real analytic functions. The behavior of a singular integral critically depends on the zeros of f .

heading: 2 VARIATIONAL INFERENCE FOR SINGULAR MODELS

text:
Variational inference is an approximate inference technique in which a family of densities Q, often called the variational family, is first posited and a member of the variational family, some q∗ ∈ Q, is then found via optimization according to some criterion that measures closeness to the desired target density. To begin, let us write the posterior distribution p(w|Dn) in two alternate forms:
Ln(w) = − 1
n n∑ i=1 log p(yi|xi, w)
p(w|Dn) ∝ e−nLn(w)ϕ(w)
Z(n) = ∫ W e−nLn(w)ϕ(w) dw
Kn(w) = 1
n n∑ i=1 log p0(yi|xi) p(yi|xi, w)
p(w|Dn) ∝ e−nKn(w)ϕ(w)
Z̄(n) = ∫ W e−nKn(w)ϕ(w).
On the left, the posterior is written in terms of the average negative log likelihood Ln(w) and on the right in terms of the average log likelihood ratio Kn(w), which is the empirical counterpart to K(w). Since Z̄(n) = Z(n)/ ∏n i=1 p0(yi|xi), we refer to Z̄(n) as the normalized evidence.
Now, consider some general change of variables g : Rd → Rd, g(ξ) = w. In the new coordinate ξ, the KL divergence between a variational distribution q(ξ) and the desired posterior target,
p(ξ|Dn) ∝ e−nKn(g(ξ))ϕ(g(ξ))|g′(ξ)|, is given by
KL(q(ξ)||p(ξ|Dn)) = EqnKn(g(ξ)) + KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|) + log Z̄(n).
As long as the support of q is contained in the support of the posterior, we have KL(q(ξ)||p(ξ|Dn)) ≥ 0 and hence the following bound:
Ψ(q, g) := −EqnKn(g(ξ))−KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|) ≤ log Z̄(n). (2) Equality in equation 2 is achieved if and only if q(ξ) = p(ξ|Dn). It is easy to recognize that maximizing Ψ(q, g) is equivalent to maximizing the so-called evidence lower bound,
ELBO(q, g) :=Eq [ n∑ i=1 log p(yi|xi, g(ξ))− log q(ξ) + log(ϕ(g(ξ))|g′(ξ)|) ] , (3)
since Ψ(q, g) = ELBO(q, g) + nSn where Sn = − 1n ∑n i=1 log p0(yi|xi) is the empirical entropy. Indeed, just as Ψ(q, g) is a lower bound on the log (normalized) evidence log Z̄(n), so too is ELBO(q, g) a lower bound on the log (unnormalized) evidence logZ(n).
To facilitate theoretical analysis, we will work with the deterministic counterparts to Ψ(q, g) and Z̄(n), respectively given by
ΨK(q, g) := −EqnK(g(ξ))−KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|), (4)
and Z̄K(n) := ∫ W e−nK(w)ϕ(w) dw. (5)
It is plain to see that just as Ψ(q, g) ≤ log Z̄(n), we have ΨK(q, g) ≤ log Z̄K(n). Our theoretical results address the possibility that supq∈QΨK(q, g) ≈ log Z̄K(n) for some variational family Q and change-of-variables g.
By (Watanabe, 2009, Theorem 6.7), for the leading order term in the asymptotic expansion of Z̄K(n), we have Z̄K(n) ≈ Cn−λ(log n)m−1. (6) The rational number λ ∈ [0, d/2] is an important quantity in singular learning theory known as the real log canonical threshold (RLCT) and the integer m ≥ 1 is its associated multiplicity. These two quantities, to be defined in equation 11, are uniquely associated to a model-truth-prior triplet.
Equipped with equation 6, (Bhattacharya et al., 2020, Theorem 3.1) established that when g is a resolution map, the mean-field variational family Q(0,1] of equation 12 can achieve
sup q∈Q(0,1]
ΨK(q, g) ≥ −λ log n− constant independent of n. (7)
(In Appendix C, we prove this result under more general conditions than assumed in Theorem 3.1 of Bhattacharya et al. (2020).) Unfortunately, the bound in equation 7 provides little reassurance that Q(0,1] is a desirable variational family.
Contribution Following Lin (2011), let us call C in equation 6 the leading coefficient of Z̄K(n). We go beyond the analysis in Bhattacharya et al. (2020) by taking into account those terms in the leading coefficient that depend on the dimension d, call it C(d). Our main result, Theorem 5.1, shows that if g is a resolution map, then the same variational family Q(0,1] of equation 12 can achieve
sup q∈Q(0,1]
ΨK(q, g) = −λ log n+ logC(d). (8)
Next, rather than presuming the resolution map theoretically tractable as in Bhattacharya et al. (2020), we employ a normalizing flow to learn the unknown resolution map g at the same time as learning the variational parameters in q. We are aided by the fact that a resolution map can attain the optimal value of ΨK(q, g) and therefore justifies learning the resolution map via optimization of Ψ(q, g). Finally let us note that although the result in equation 8 is stronger than that in equation 7, it does come at the cost of additional assumptions as we will discuss at the end of Section 5. Remark. When the model is regular, we need not bother with singular learning theory and may obtain Z̄K(n) ≈ ϕ(w0) √ (2π)d
detH(w0) n−d/2 via the Laplace approximation. The Laplace approximation,
however, is egregiously inappropriate for singular models, in particular neural network models. Since λ = d/2 and m = 1 in regular models, equation 6 is a true generalization of the Laplace approximation, holding for both regular and strictly singular models.

heading: 3 SINGULAR LEARNING THEORY

text:
That the posterior distribution in singular models can be written, under quite general conditions, as a mixture of standard forms is predicated on the monomialization of K(w). The following theorem from Watanabe (2009), adapted for notational consistency, gives precise conditions for the existence of the resolution map, an algebraic geometrical transformation such that K(w) can be written as a monomial. The result is itself based on Hironaka’s resolution of singularities, a celebrated result in modern algebraic geometry. To prepare, let W = {w ∈ W : K(w) ≤ } for some small positive constant and W (R) be some real open set such that W ⊂W (R) . The theorem below will make use of the multi-index notation: for a given w = (w1, . . . , wd) ∈ Rd, define wk := wk11 · · ·w kd d where the multi-index k = (k1, . . . , kd) with each kj a nonnegative integer. Theorem 3.1 (Theorem 6.5 of Watanabe (2009)). Suppose the model-truth-prior triplet (p, p0, ϕ) satisfies Fundamental Conditions I and II with s = 2 in Watanabe (2009). We can find a real analytic manifold M (R) and a proper and real analytic map g : M (R) →W (R) such that
1. M = g−1(W ) is covered by a finite set M = ∪αMα where Mα = [0, b]d.
2. In each Mα, K(g(ξ)) = ξ2k = ξ2k11 · · · ξ 2kd d , (9)
where kj ∈ N such that not all kj are zero.
3. There exists C∞ function b(ξ) such that
ϕ(g(ξ))|g′(ξ)| = ξhb(ξ) = ξh11 · · · ξ hd d b(ξ), (10)
where hj ∈ N, |g′(ξ)| is the absolute value of the determinant of the Jacobian and b(ξ) > c > 0 for ξ ∈ [0, b]d.
If we “plug in" equation 9 and equation 10 into the transformed posterior p(ξ|Dn), we obtain the first display of the paper, equation 1. Theorem 6.5 of Watanabe (2009) holds for regular statistical models as well, e.g., by the transform w = g(ξ) = w0 + I(w0)1/2ξ, we can put a regular model-truth-prior triplet into the standard form. It is worth noting that neither the resolution map g nor the multi-indices k and h are unique for a given triplet (p, p0, ϕ).
A crucial quantity that appears in singular learning theory is the real log canonical threshold (RLCT). Let {Mα : α} be as in Theorem 3.1 and λj = hj+12kj , j = 1, . . . , d where hj and kj are the entries of the multi-indices h and k in a local coordinate Mα. (For brevity, the dependence on α has been suppressed.) When kj = 0, λj is taken to be infinity. Uniquely associated to a triplet (p, p0, ϕ) are its real log canonical threshold (RLCT) and its multiplicity defined, respectively, as
λ = min α min j∈1,...,d λj , m = max α
#{j : λj = λ}. (11)
Let {α∗} be the set of those local coordinates in which both the min and max in equation 11 are attained. Watanabe (2000) calls such a set the essential coordinates.
Assuming the prior is proper, the RLCT of a model-truth-prior triplet is at most d/2 (Watanabe, 2009, Theorem 7.2). When the model is regular, the RLCT is exactly equal to d/2 and the multiplicity m = 1 (Watanabe, 2009, Remark 1.15). Murfet et al. (2020) argues that (twice) the RLCT might be a most natural count of parameters in singular models. In Appendix A, we recall a simple toy example where the resolution map, the RLCT, and the multiplicity can be calculated explicitly. Remark. In Section 6, we will learn the resolution map using affine coupling layers. However as indicated by Theorem 3.1, the resolution map g(ξ) as well as the multi-indices k and h in fact depend on α. Despite this, it is unclear if learning multiple resolution maps gα would have any practical advantages since the RLCT is determined entirely by the essential local coordinates.

heading: 4 RELATED WORK

text:
Bayesian learning in neural networks well precede the advent of modern deep learning MacKay (1992); Neal (1996). The resurgence of interest in Bayesian learning for deep neural networks, sometimes called Bayesian deep learning Wilson & Izmailov (2020), has been prompted by concerns of overconfidence and miscalibration. Since exact inference for Bayesian neural networks is intractable, all methods proceed by approximate inference. A major class of approximate inference techniques is based on scaling classic MCMC to modern settings of large datasets and deep neural networks. Some scalable variants of MCMC suitable for deep neural networks include Welling & Teh (2011); Chen et al. (2014); Zhang et al. (2019). Another major approximate inference technique for Bayesian neural network is represented by variational inference which learns the target posterior via optimization. The various flavors of variational inference can be commonly characterized by two ingredients: 1) an approximating family, e.g., a class of distributions over the neural network weights and 2) a criterion for measuring closeness of two distributions. The most commonly employed approximating family is undoubtedly the mean-field family of fully factorized Gaussian distributions. Many variational inference techniques share this approximating family even if they use different criterion to measure closeness to the target (Graves, 2011; Blundell et al., 2015; Hernandez-Lobato et al., 2016; Li & Turner, 2016; Khan et al., 2018; Sun et al., 2019). The limitations of the mean-field Gaussian approximating family are well known however (MacKay, 1992). The desire to move beyond
mean-field Gaussian has motivated many recent methods to make use of more realistic covariance structures (Zhang et al., 2018) or more expressive approximating families, e.g., via normalizing flows (Louizos & Welling, 2017). Monte Carlo dropout (Gal & Ghahramani, 2016) is another popular approximate inference technique for Bayesian neural networks which, despite first appearances, can in fact be viewed as variational inference. Another strain of work is based on the idea of using stochastic gradient descent as a sampler for the underlying posterior distribution of interest. Works in this spirit include Mandt et al. (2018), Izmailov et al. (2018) and Maddox et al. (2019). Finally there are various approximate inference techniques for Bayesian neural networks that are not easily classifiable according to the distinctions above, e.g., temperature scaling (Guo et al., 2017) and deep ensembles (Lakshminarayanan et al., 2017).

heading: 5 THE GENERALIZED GAMMA MEAN-FIELD APPROXIMATION

text:
Consider the mean-field variational family proposed in Bhattacharya et al. (2020) Q(0,1] = {qλ,k,β = d∏ j=1 qλj ,kj ,βj (ξj) : λ = (λ1, . . . , λd) ∈ Rd>0,k = (k1, . . . , kd) ∈ Rd>0,β = (β1, . . . , βd) ∈ (0,∞)d}
qj(ξj) := qλj ,kj ,βj (ξj) ∝ ξ 2kjλj−1 j exp(−βjξ 2kj j )1(0,1](ξj), (12)
where each univariate density qj supported on (0, 1] is the density of a truncated generalized gamma random variable. Let j∗ ∈ {1, . . . , d} be the dimension that attains λ̃j∗ = λ̃ where λ̃ is the RLCT of the underlying (unknown) model-truth-prior triplet. It was established in Bhattacharya et al. (2020) that if g is a resolution map, then Q(0,1] satisfies equation 7 in particular by setting λ and k to their respective true values, and βj∗ = n and all other βj = 1.
In Theorem 5.1 below we will instead establish equation 8. In the proof we show this can be accomplished by setting λ and k to their true values as in Bhattacharya et al. (2020) though the optimal values of β will be different. Finally, although we assume below that m = 1, this is not necessary; as long as m d, we can set βj = n1/m for all j such that λ̃j = λ̃. Theorem 5.1. Suppose the model-truth-prior triplet is such that Theorem 3.1 holds with [0, b] = [0, 1], K(g(ξ)) = ξ2k̃ and ϕ(g(ξ))|g′(ξ)| = ξh̃. Let λ̃ denote the RLCT of the triplet and assume the multiplicity m = 1. Then supq∈Q(0,1] ΨK(q, g) = −λ̃ log n + logC(d) where C(d) is the term in the leading coefficient C that depends on d.
Proof. If ξ ∈ [0, 1]d and ϕ(g(ξ))|g′(ξ)| ∝ ξh̃, then necessarily b(ξ) = ∏d j=1 1/(h̃j + 1). Let
λ̃j = (h̃j + 1)/(2k̃j). Applying Corollary 5.9 in Lin (2011), we obtain
Z̄K(n) = ∫ [0,1]d exp(−nξ2k̃)ξh̃ dξ ≈ Cn−λ̃(log n)m−1 = Cn−λ̃ (13)
where
C = Γ(λ̃) (m− 1)! ∏d j=1(2k̃j) ∏d j=m+1(2k̃j)(λ̃j − λ̃) . (14)
If we denote by C(d) the terms in the leading coefficient that depend on d, then we have
logC(d) = − d∑ j=1 log(2k̃j)− d∑ j=m+1 log(2k̃j)− d∑ j=m+1 log λ̃j − d∑ j=m+1 log(1− λ̃/λ̃j).
Returning to our variational distribution, if the kj’s are well specified and we additionally set βj∗ = n where j∗ is the dimension that attains λ̃j∗ = λ̃, we then have EqnK(g(ξ)) = λj∗ ∏ j 6=j∗ G(λj , βj),
where G(λ, β) is as in B.1. Next, we have
KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|) = KL ( q ∥∥∥ ξh̃b(ξ)) = d∑
j=1
KL ( qj ∥∥∥ ξh̃jj /(h̃j + 1))
= d∑ j=1 Eqj log qj − h̃jEqj log ξj + log(h̃j + 1).
If in addition the λj’s are all well-specified, then
KL(q||ϕ(g(ξ))|g′(ξ)|) = d∑ j=1 [ −βjG(λ̃j , βj)− logB(k̃j , h̃j , βj) + log(2k̃j) + log λ̃j ] ,
where B(k, h, β) is as in B.1. Now let us make use of the fact that logB(k, h, β) −λ log β and G(λ, β) λ/β for large β. Since βj∗ = n is large, we get ΨK(q, g) = −λ̃ ∏ j 6=j∗ G(λ̃j , βj) + d∑ j=1 [ βjG(λ̃j , βj) + logB(k̃j , h̃j , βj)− log(2k̃j)− log λ̃j ]
−λ̃ log n+ λ̃(1− d∏
j=m+1
G(λ̃j , βj)) + d∑ j=m+1 [βjG(λ̃j , βj) + logB(k̃j , h̃j , βj)]− d∑ j=1 log(2k̃j)− d∑ j=1 log(λ̃j).
If βj are sufficiently large for j 6= j∗, we get ΨK(q, g) −λ̃ log n− d∑ j=1 log(2k̃j)− d∑ j=1 log(λ̃j) + d∑ j=m+1 λ̃j(1− log βj) + λ̃(1− d∏ j=m+1 λ̃j βj ).
(15)
Then there exist βj , j 6= j∗ so that equation 15 matches −λ̃ log n+ logC(d).
Theorem 5.1 critically assumes that b(ξ) ∝ 1. We do not expect this to hold in reality. For example, even for the simple one-hidden layer tanh network considered in Section 7, it does not appear that b(ξ) ∝ 1 (Watanabe, 2000). Currently, we are prevented from stating a more general version of Theorem 5.1 because there is no off-the-shelf derivation of the leading coefficient when the singular integral is of the general form ∫ [0,b]d
exp(−nξ2k̃)ξh̃b(ξ) dξ. We expect the generalization to be technically feasible but as its development requires advanced knowledge of algebraic geometry, it is beyond the scope of this paper and best left as separate investigation.
In the next section, we proceed to learn the resolution map rather than presume it is known. This is an improvement over Bhattacharya et al. (2020) which, due to the difficulty of deriving resolution maps, was limited in its single experiment to the toy neural network fw(x) = b tanh(ax) with weight w = (a, b) ∈ R2. Again, though Theorem 5.1 is stated in terms of the truncated generalized gamma mean-field family, we do not believe the assumption b(ξ) ∝ 1 is critical and proceed henceforth to work with the untruncated generalized gamma mean-field family, which we denote Q.

heading: 6 LEARNING TO DESINGULARIZE

text:
When the resolution map g is known, the preceding results suggest to 1) apply the change-of-variables g(ξ) = w and 2) maximize ELBO(q, g) over the untruncated generalized gamma mean-field family Q via e.g., stochastic gradient descent. However, theoretically, the resolution map g is notoriously difficult to derive. Computationally, the recursive blow-up procedure in algebraic geometry is entirely not scalable to high dimensions.
We propose to learn the resolution map via a normalizing flow, which is commonly used to model complex distributions as the push-forward of a simple source distribution through an invertible neural network G. An interesting direction of future work would be to exploit properties of the resolution
map that may aid in the design of the normalizing flow architecture. For now, we simply make use of a common type of invertible architecture consisting of affine coupling layers. With r denoting a binary mask, a so-called affine coupling layer acts as follow:
u, v ∈ Rd, u 7→ v = (1− r) u+ r (u exp(s(r u)) + t(r u)).
The binary mask r must alternate from one affine coupling layer to the next for otherwise there would be little expressive power in the resulting network. Let Gθ be a network consisting of alternating affine coupling layers, where θ denotes the collective parameters. Since the resolution map is to be learned, there is no loss in generality to assuming j∗ = 1. We shall need the gradient with respect to the variational parameters and the normalizing flow weights as part of employing stochastic gradient descent, i.e. ∇λ,k,β,θELBO(qλ,k,β, Gθ). Note that here we abuse the notation slightly as we do not need to update β1 which should be set to the sample size n according to the proof of Theorem 5.1.
Although the source distribution in a normalizing flow can have its own trainable parameters, it is common practice to adopt a parameter-less source distribution. In particular if the source distribution is itself reparametrizable, then the learning of the associated parameters can be absorbed into the invertible transformation. The generalized gamma distribution is not easily reparametrizable for general values of λ,k,β. However for certain settings of λ, we can in fact avail ourselves to the reparametrization trick, at least approximately. Let Vj be a gamma random variable with shape λj and rate βj , then V 1/(2kj) j ∼ qj(ξj) := ξ 2kjλj−1 j exp(−βjξ 2kj j )1(0,∞)(ξj). We will mostly be interested in settings where λj is large, in which case Vj is approximately Gaussian with mean λj/βj and variance λj/β2j . Letting T ( ) := (F −1 1 ( 1) 1/(2k1), . . . , F−1d ( d) 1/(2kd)), where F−1j ( ) ≈
(λj + √ λj )/βj , the reparametrization trick then leads to the objective function
E ∼N(0,I) [ n∑ i=1 log p(yi|xi, Gθ(T ( ))) + logϕ(Gθ(T ( ))) + log |G′θ(T ( ))| ] − Eq log q. (16)
Note that equation 16 is not the same as ELBO(qλ,k,β, Gθ) because we have made use of the Gaussian approximation for Vj in the case of large λj .
The entropy component of equation 16, −Eqλ,k,β log qλ,k,β, can be derived analytically, see Appendix B.2. Next, the specific architecture of Gθ has rendered the log Jacobian term, log |G′θ(·)|, computationally tractable. The final piece is to replace E ∼N(0,I) with an empirical average over M samples. In the experiments that follow we will consider either learning the variational parameters λ,k,β or fixing them at some initial value since other values can be learned through a transformation that gets absorbed into G.

heading: 7 EXPERIMENTS

text:
In this section, we compare the effect of two source distributions for a normalizing flow given by the affine coupling network Gθ consisting of 2 pairs of alternating couplings with scaling and translation networks each consisting of 2 hidden layers with 16 hidden units, see Appendix D for exact details of the architecture. We denote by nf_gamma_λ0_k0_β0_flag the variational family that results from pushing forward the untruncated generalized gamma distribution with variational parameters initialized at λ0 = (1, λ0, . . . , λ0),k0 = (k0, . . . , k0), and β0 = (n, β0, . . . , β0), and a boolean flag indicating whether the variational parameters are subject to updating. For the other approximation resulting from pushing forward a Gaussian source distribution N(µ0, v0), we write analogously nf_gaussian_µ0_v0 where µ0 and v0 are the fixed mean and variance. For each combination of λ0 and β0 considered, we set µ0 = λ0/β0 and v0 = λ0/β20 .
We employ the widely adopted variant of (minibatch) stochastic gradient descent known as Adam. The number of epochs and batch size were set to 2000 and n/10, respectively. Different constant learning rates were employed according to the parameter type: 1e−3 for affine coupling layer weights and k, and 1e−1 for λ and β. Expectations in the objective function that are not analytically tractable are replaced with an average over M = 5 samples. At the end of training, we evaluate Ψ(q, g) in nf_gamma without recourse to the (approximate) reparametrization and by using the analytic expression for the entropy component of Ψ and 100 samples from q to approximate other components in Ψ under expectation.
To date, there is a very small collection of model-truth-prior triplets where the RLCT λ̃ and multiplicity mare known. We shall limit our experiments to two such triplets. This will allow us to compare the achieved Ψ(q̂, Ĝ) following training to −λ log n+ (m− 1) log log n. Note that even in such triplets where the RLCT and multiplicity are known, the exact value of the leading coefficient C is still usually unknown. Now, to compare nf_gamma and nf_gaussian to each other, we can simply see which achieves higher Ψ(q̂, Ĝ) after training.
Looking ahead to downstream tasks, it is natural to ask whether the variational posterior predictive distribution, pvb(y|x,Dn) = 〈p(y|x,Gθ̂(ξ))〉q̂(ξ), inherits the desirable properties of the Bayes posterior predictive distribution, p(y|x,Dn) = 〈p(y|x,w)〉p(w|Dn). The answer turns out to depend on the relationship between the variational real log canonical threshold λvb and the RLCT of the model-truth-prior triplet. It may very well be that a variational family which is closer to the true posterior (in the KL sense) than another variational family may induce a worse approximation of the true Bayes posterior predictive distribution. We provide an in-depth account of this phenomenon through the lens of singular learning theory in Appendix E.
Our first example concerns the tanh model-truth-prior triplet. Consider input x ∈ R following the uniform distribution on [−1, 1], and response variable y ∈ R modeled as p(y|x,w) =
1√ 2π exp(− 12 (y − f(x,w)) 2), where fw(x) = ∑H h=1 bh tanh(ahx) is a tanh network with H hidden units and w is the collection of neural network weights {(ah, bh)}Hh=1. If the true distribution is given by p0(y|x) = p(y|x, 0) = 1√2π exp(− 1 2y 2) and the prior ϕ is a C∞ function of w with compact support, satisfying ϕ(0) > 0, Aoyagi & Watanabe (2006) showed that λ = H+i 2+i
4i+2 and m = 2 if i2 = H , and m = 1 if i2 < H where i is the maximum integer satisfying i2 ≤ H . In contrast, were this a regular statistical model, we would have λ = H . We consider two settings for p0(y|x) = p(y|x,w0) in the experiment: 1) w0 = 0 and 2) w0 = 5. Note for the latter, the corresponding RLCT and multiplicity are unknown.
Our next example is the reduced rank model-truth-prior triplet. Consider input x ∈ RM and response variable y ∈ RN modeled as p(y|x,w) = (2π)−N/2 exp{− 12 ||y−BAx||
2}where {w = (A,B)|A ∈ RH×M , B ∈ RN×H}. This model is readily seen to be a special case of a neural network with hidden units H and identity activation function. We shall set M = H + 3 and N = H . In the realizable case, i.e., p0(y|x) = p(y|x,A0, B0), if the prior ϕ is a C∞ function with compact support satisfying ϕ(A0, B0) > 0, the RLCT was derived in Aoyagi & Watanabe (2005) for various values of N,M,H and r = rank(B0A0). Below we set B0 = IN×N and A0 = [IH×H ; JH×3]. The rank r for B0A0 equals H . This then falls under Case (3) in Aoyagi & Watanabe (2005) since N + H < M + r, leading to λ = (NH −Hr+Mr)/2,m = 1. Note that were this a regular model, we would instead have λ = (MH +NH)/2.
Table 1 provides a summary of the values ofH considered in each of the triplets and the corresponding RLCT and dimension. In the experiments, we tested various combinations of values of H , settings for w0 and priors. Throughout, the priors were designed to be mis-specified. The full results of these experiments can be found in Tables 2–4. Note that the comparison between different variational approximations have been made subject to the value of H , the prior, and the relationship λ0/β0 = µ0. We should also point out that not all runs reached convergence as can be seen from the count column in Tables 2–4. This seems largely to be a result of the constant learning rate we used for all training.
The first setting (Table 2) is the realizable tanh network with w0 = 0. In this case, the true RLCT is known as discussed above. As Table 2 demonstrates, there is no clearly discernible difference between nf_gamma and nf_gaussian where the source distribution for both starts near zero. However when the source distribution has λ0/β0 = µ0 = 5, i.e. very far from w0, nf_gamma is more robust than nf_gaussian. The second setting (Table 3) is analogous to the previous except the true generating mechanism is given by w0 = 5. In this setting, the untruncated generalized gamma source distribution is seen to provide a better variational approximation than the Gaussian source distribution. This appears to hold quite independent of the initial values of λ,k,β as well as the gradient flag. Interestingly, in contrast to the previous setting, there is a clear benefit to learning λ,k,β. The third setting (Table 4) is the realizable reduced rank regression. In this experiment, we do not see a discernible difference between the source distributions until we look at λ0/β0 = µ0 = 5, which is far from w0 = (A0, B0). As in the first setting, we see nf_gamma is more robust than nf_gaussian.
Finally, though we have cautioned against judging the quality of a variational approximation according to the approximate posterior predictive distribution induced, it is nonetheless informative to visualize downstream uncertainty quantification. In Figure 1, we choose one line from the tanh w0 = 0 experiment (Table 2) to further examine. For the two methods illustrated in in Figure 1, their performance in terms of Ψ is quite close as indicated by Table 2. It is then comforting that the respective confidence bands are quite similar. Next, in Figure 2, we choose one line from the tanh w0 = 5 experiment (Table 3) to further examine. This time, nf_gamma achieves significantly better Ψ than nf_gaussian. We see a corresponding relationship between their respective posterior predictive distributions. Namely, the confidence band resulting from nf_gaussian are far too conservative while that from nf_gamma is less so.

heading: 8 CONCLUSION

text:
In this work we propose a variational approximation for Bayesian neural networks by leveraging insights from singular learning theory. Namely, for large n, the posterior distribution over neural network weights is not Gaussian but rather can be put into a mixture of standard forms. From this, we demonstrate that the generalized gamma mean-field family, following desingularization, can in theory achieve the leading order term of the log normalized evidence. Because we choose to learn the desingularization map using affine coupling layers, the proposed work can be cast as a normalizing flow with an unconventional source distribution. Interestingly, for large values of the variational parameter λ, the source distribution in each dimension is approximately N(λ, λ/β)1/(2k), which is reparametrizable in terms of the conventional source distribution N(0, 1). Though learning the variational parameters λ,k,β in the source distribution goes against conventional wisdom in normalizing flows, our experiments suggest some performance may be gained by learning the optimal untruncated generalized gamma source distribution at the same time as learning the resolution map.
ETHICS STATEMENT
Bayesian learning for neural networks is often touted as a panacea to the challenges of uncertainty quantification in deep learning. In this work, we propose a variational approximation that performs well in terms of the ELBO achieved. However, we make no claim that our variational approximation is superior in downstream uncertainty quantification. We explain extensively in the appendix that a variational approximation which achieves higher ELBO may induce a worse approximate posterior predictive distribution than a variational approximation which achieves lower ELBO. According to singular learning theory, the performance of a variational approximation in terms of the posterior predictive distribution it induces critically depends on the relationship between the variational RLCT and the underlying RLCT of the model-truth-prior triplet.
REPRODUCIBILITY STATEMENT
A zip file of the source code has been submitted as supplementary materials. For theoretical results, explanations of assumptions and proof of the claims are included in the main text. No datasets are used; the experiments are solely based on simulated data.

heading: A TOY EXAMPLE OF RLCT CALCULATION

text:
We recall Example 27 from Watanabe (2018) to illustrate the concepts of resolution map, RLCT and multiplicity for a simple model-truth-prior triplet. For univariate input x ∈ [0, 1] and univariate output y ∈ R, consider the model with parameter w = (a, b) ∈ [0, 1]2 given by
p(x, y|w) = 1√ 2π exp(−1 2 (y − a tanh(bx))2) (17)
Suppose the prior is uniform, i.e., ϕ(w) = 1 and the truth is given by p0(x, y) = p(x, y|0, 0). Then we can easily see that
K(w) = b2a2 1
2 K0(w),
where
K0(w) = ∫ 1 0 ( tanh(bx) b )2 dx
The following desingularization map puts the triplet in standard form:
ξ1 =
√ K0(w)
2 a
ξ2 = b
Furthermore we have ϕ(g(ξ)) = ξh where h = (0, 0) and b(ξ) = |g′(ξ)|. Since (k1, k2) = (1, 1) and (h1, h2) = (0, 0) we have (λ1, λ2) = (1/2, 1/2). Therefore for this particular model-truth-prior triplet, the RLCT is 1/2 with multiplicity 2.

heading: B BASIC CALCULATIONS FOR THE GENERALIZED GAMMA DISTRIBUTION

text:
B.1 TRUNCATED
We summarize some basic calculations for the [0, 1] truncated generalized gamma distribution from Bhattacharya et al. (2020). Let qj be the univariate truncated generalized gamma density given in equation 12.
• The normalizing constant of qj is given byB(λj , kj , βj) whereB(λ, k, β) = β−λΓ(λ)γ(λ,β)
2k and γ(a, x) = 1Γ(a) ∫ x 0 ta−1e−t dt is the (regularized) lower incomplete gamma function.
• The quantity Eqjξ 2kj = G(λj , βj) where G(λ, β) = λβ γ(λ+1,β) γ(λ,β) .
B.2 UNTRUNCATED
Consider the univariate density
q(ξ) ∝ ξh exp(−βξ2k)
supported on (0,∞) where h, k, β > 0. Let λ = (h+ 1)/(2k). Elementary calculations give
• Eqξ2k̃ = β− k̃ k
Γ( k̃k+λ)
Γ(λ)
• Eq log ξ = 12k (ψ(λ)− log β) where ψ is the digamma function.
• Eq log q = h2k (ψ(λ)− log β)− λ− logZ where, Z = β−λΓ(λ)
2k is the normalizing constant of q.

heading: C GENERALIZATION OF (BHATTACHARYA ET AL., 2020, THEOREM 3.1)

text:
Lemma C.1 below states that any variational distribution qλ,k,β in the untruncated generalized gamma mean field family Q with well-specified k where we additionally set βj∗ = n is capable of achieving ΨK(q, g) ≥ −λ̃ log n + A where A is some constant that does not depend on n. Lemma C.1 is a straightforward extension of Theorem 3.1 of Bhattacharya et al. (2020)from the [0, 1]d truncated generalized gamma mean field variational family to the untruncated case. We also relax the condition in Theorem 3.1 of Bhattacharya et al. (2020) that the variational parameters λ and k (and hence h) are well-specified; here we only require that k be well-specified.
Lemma C.1. Suppose the model-truth-prior triplet is such that Theorem 3.1 holds with K(g(ξ)) = ξ2k̃ and ϕ(g(ξ))|g′(ξ)| = ξh̃. Consider the untruncated generalized gamma mean-field family Q, we have supq∈QΨK(q, g) ≥ −λ̃ log n+A, where A is a constant that does not depend on n.
Proof. If kj = k̃j for all j = 1, . . . , d, we have
EqnK(g(ξ)) = n d∏ j=1 β − k̃j kj j Γ( k̃j kj + λj) Γ(λj) = n 1− k̃j∗ kj∗ Γ( k̃j∗ kj∗ + λj∗) Γ(λj∗) ∏ j 6=j∗ β − k̃j kj j Γ( k̃j kj + λj) Γ(λj) = λj∗ ∏ j 6=j∗ λj βj .
Next, we have
KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|)
= d∑ j=1 Eqj [log qj − h̃j log ξj ]− Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2kj (ψ(λj)− log βj)− λj + λj log βj + log(2kj)− log Γ(λj)
] − Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2kj ψ(λj)− λj +
h̃j + 1
2kj log βj + log(2kj)− log Γ(λj)
] − Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2k̃j ψ(λj)− λj + λ̃j log βj + log(2k̃j)− log Γ(λj)
] − Eq log b(ξ)
= λ̃ log n+ ∑ j 6=j∗ λ̃j log βj + d∑ j=1
[ hj − h̃j
2k̃j ψ(λj)− λj + log(2k̃j)− log Γ(λj)
] − Eq log b(ξ).
The last line follows from setting βj∗ = n where j∗ ∈ {1, . . . , d} is such that λ̃j∗ = λ̃, breaking ties arbitrarily. Finally, the desired inequality follows from the fact that Eq log b(ξ) is bounded below by some constant.

heading: D EXPERIMENTS

text:
In the experiments, we train an affine coupling network with two pairs of alternating couplings. The translation t is a feedforward (leaky) ReLU neural network with tanh output activation function. The scaling t is another feedforward (leaky) ReLU neural network with identity output activation function. The models implemented are given below.
Translation network
Sequential( (0): Linear(in_features=dim, out_features=16, bias=True) (1): LeakyReLU(negative_slope=0.01) (2): Linear(in_features=16, out_features=16, bias=True) (3): LeakyReLU(negative_slope=0.01) (4): Linear(in_features=16, out_features=16, bias=True) (5): LeakyReLU(negative_slope=0.01) (6): Linear(in_features=16, out_features=dim, bias=True) )
Scaling network
Sequential( (0): Linear(in_features=dim, out_features=16, bias=True) (1): LeakyReLU(negative_slope=0.01) (2): Linear(in_features=16, out_features=16, bias=True) (3): LeakyReLU(negative_slope=0.01) (4): Linear(in_features=16, out_features=16, bias=True) (5): LeakyReLU(negative_slope=0.01) (6): Linear(in_features=16, out_features=dim, bias=True) (7): Tanh() )

heading: E APPROXIMATE POSTERIOR PREDICTIVE DISTRIBUTION

text:
Bayesian prediction proceeds by marginalization, i.e., averaging over all possible values of the model parameter. This results in the posterior predictive distribution1,
p(y|x,Dn) := ∫ w p(y|x,w)p(w|Dn) dw = 〈p(y|x,w)〉p(w|Dn). (18)
According to singular learning theory, Bayesian prediction via equation 18 is superior to MAP and MLE, in the sense that the expected generalization error of the posterior predictive distribution is smaller than that of the MAP or MLE. Specifically, let
Gn(p̂n(y|x)) := KL(p0(y|x)p(x)||p̂n(y|x)p(x))
be the generalization error of a predictive distribution p̂n(y|x). According to Theorems 1.2 and 7.2 in Watanabe (2009), we have
EGn(p(y|x,Dn)) = λ/n+ o(1/n) (19)
where the expectation is taken with respect to Dn and λ is the RLCT of the model-truth-prior triplet. On the other hand, Theorem 6.4 of Watanabe (2009) shows that the expected generalization error of MLE (and similarly of MAP) is
EGn(p(y|x, ŵmle)) = C/n+ o(1/n) (20)
where C, the maximum of a Gaussian process, can be much larger than λ. Such a distinction cannot be made in regular models in which the difference between the three estimators becomes negligible in the large n regime.
Now, consider the variational approximation to the posterior predictive distribution given by
pvb(y|x,Dn) = 〈p(y|x,w)〉q∗(w),
where q∗ is the optimal variational distribution in some variational family Q. It may be tempting to compare different variational approximations according to how well their respective variational posterior predictive distributions approximate the true Bayes posterior predictive distribution, e.g., Blundell et al. (2015); Louizos & Welling (2017). This turns out to be a thorny issue, as documented in various works on Bayesian neural networks Heek (2018); Krishnan & Tickoo (2020); Foong et al. (2020). In particular, pvb(y|x,Dn) does not necessarily inherit the desirable properties of the Bayes posterior predictive distribution p(y|x,Dn). Here, we offer insights from singular learning theory to account for this.
Consider the (normalized) variational free energy
F̄vb(n) = EqnKn(w) + KL(q(w)||ϕ(w))
where q is some variational distribution. If the minimum (normalized) variational free energy F̄ ∗vb(n) = minq F̄vb(n) admits an asymptotic expansion 2, then it would be of the form
F̄ ∗vb(n) = λvb log n+ (mvb − 1) log log n+Rn.
Note that λvb ≥ λ necessarily. Now, if the generalization error of pvb(y|x,Dn) would admit an asymptotic expansion3, it would be of the form
EGn(pvb(y|x,Dn)) = λ̃vb/n+ o(1/n). 1Though this may seem fundamentally distinct from the maximum likelihood estimator (MLE) or maximum a posterior (MAP) solution commonly employed in training deep networks, both MLE and MAP may in fact be regarded as impoverished estimates of equation 18 whereby p(w|Dn) is approximated with δ(w = ŵ), a point mass at ŵ.
2We should disclose that general conditions for such an asymptotic expansion of the minimum variational free energy is still an open problem. The issue has so far been addressed on a case-by-case basis, e.g., reduced rank regression Nakajima & Watanabe (2007), nonnegative matrix factorization Kohjima & Watanabe (2017); Hayashi (2020), normal mixture model Watanabe & Watanabe (2006), hidden Markov model Hosino et al. (2005).
3Similar to the minimum variational free energy, an expansion of EGn(pvb(y|x,Dn)) has not been established in full generality at this point.
Importantly, λ̃vb 6= λvb in general. Examples can be found where sometimes one is bigger, sometimes the other Nakajima & Watanabe (2007).
This development is to be contrasted with the Bayes free energy and the Bayes generalization error. Recall the normalized Bayes free energy, F̄n := − log Z̄(n), admits the following asymptotic expansion F̄n = λ log n+ (m− 1) log log n+Rn. (21) The very same λ coefficient appears in the asymptotic expansion of the Bayes generalization error in equation 19. This means that minimizing the Bayes free energy is equivalent to minimizing the Bayes generalization error. In contrast, minimizing the variational free energy (equivalent to maximizing the ELBO) does not necessarily lead to a lower generalization error since a variational family with higher variational free energy (higher λvb) may have lower variational generalization error (lower λ̃vb).


summaries: This paper considers the problem of domain generalization (DG), wherein predictors are trained on a related set of training domains and evaluated on an unseen test domain.  The authors first present a learning-theoretic bound on the performance of an average-case formulation for DG, and then present a set of experiments that consider the trade-off between complexity and out-of-distribution (OOD) performance.  The experiments indicate that such trade-offs exist for linear models, and there is also evidence that the trade-off persists for shallow neural networks.

summaries: In domain generalization, recent work has shown that ERM has comparable out-of-domain accuracy to state-of-the-art DG methods. The paper aims to explain this result through model complexity. The main argument is that ood generalization requires a smaller model complexity. The upshot is that DG methods can utilize cross-domain validation to obtain better generalizing models.

summaries: This paper presents a theory of domain generalization based on statistical learning theory (Rademacher complexity) and demonstrates a trade-off between training loss and model complexity. They show that existing methods are actually controlling the model complexity. Based on the analysis, the authors argue that proper model selection is critical for complexity control. Instead of hyper-parameter search, they propose to use domain-wise cross-validation as the model selection strategy. Experiments on the DomainBed benchmark show the effectiveness of the proposed method.

summaries: This paper has several findings taking a bias-variance trade-off of DG performance, that explains existing DG algorithm performance variability. (ii) The complexity control strategy used to determine bias-variance trade-off is crucial in practice, with peak DG performance achieved when optimizing model complexity based on domain-wise validation. (iii) Regularisation required for optimal DG is greater than for conventional optimization for within-domain performance.

abstractText: The domain generalization (DG) problem setting challenges a model trained on multiple known data distributions to generalise well on unseen data distributions. Due to its practical importance, a large number of methods have been proposed to address this challenge. However most of the work in general purpose DG is heuristically motivated, as the DG problem is hard to model formally; and recent evaluations have cast doubt on existing methods’ practical efficacy – in particular compared to a well chosen empirical risk minimisation baseline. We present a novel learning-theoretic generalisation bound for DG that bounds novel domain generalisation performance in terms of the model’s Rademacher complexity. Based on this, we conjecture that existing methods’ efficacy or lack thereof is a variant of the standard empirical risk-predictor complexity trade-off, and demonstrate that their performance variability can be explained in these terms. Algorithmically, this analysis suggests that domain generalisation should be achieved by simply performing regularised ERM with a leave-one-domain-out cross-validation objective. Empirical results on the DomainBed benchmark corroborate this.

authors:
Kartik Ahuja
Jun Wang
Amit Dhurandhar
Karthikeyan Shanmugam
Kush R. Varshney

title: Empirical or invariant risk minimization? a sample complexity perspective

venue: In ICLR,

authors:
Isabela Albuquerque
João Monteiro
Mohammad Darvishi
Tiago H Falk
Ioannis Mitliagkas

title: Generalizing to unseen domains via distribution matching

venue: arXiv preprint arXiv:1911.00804,

authors:
Yogesh Balaji
Swami Sankaranarayanan
Rama Chellappa

title: Metareg: Towards domain generalization using meta-regularization

venue: In NeurIPS,

authors:
Peter L Bartlett
Shahar Mendelson

title: Rademacher and gaussian complexities: Risk bounds and structural results

authors:
Peter L Bartlett
Dylan J Foster
Matus J Telgarsky

title: Spectrally-normalized margin bounds for neural networks

authors:
Sara Beery
Grant Van Horn
Pietro Perona

title: Recognition in terra incognita

venue: In ECCV,

authors:
Mikhail Belkin
Daniel J Hsu
Partha Mitra

title: Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate

venue: NeurIPS,

authors:
Gilles Blanchard
Gyemin Lee
Clayton Scott

title: Generalizing from several related classification tasks to a new unlabeled sample

venue: In NIPS,

authors:
Gilles Blanchard
Aniket Anand Deshmukh
Ürün Dogan
Gyemin Lee
Clayton Scott

title: Domain generalization by marginal transfer learning

authors:
Mathilde Caron
Hugo Touvron
Ishan Misra
Hervé Jégou
Julien Mairal
Piotr Bojanowski
Armand Joulin

title: Emerging properties in self-supervised vision transformers

venue: arXiv preprint arXiv:2104.14294,

title: Domain Adaptation in Computer Vision

authors:
Steve Dias Da Cruz
Bertram Taetz
Thomas Stifter
Didier Stricker

title: Illumination normalization by partially impossible encoder-decoder cost function

venue: In WACV,

authors:
Chen Fang
Ye Xu
Daniel N Rockmore

title: Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias

venue: In CVPR,

authors:
Stuart Geman
Elie Bienenstock
René Doursat

title: Neural networks and the bias/variance dilemma

authors:
Muhammad Ghifary
W. Bastiaan Kleijn
Mengjie Zhang
David Balduzzi

title: Domain generalization for object recognition with multi-task autoencoders

venue: In ICCV,

authors:
Henry Gouk
Timothy Hospedales

title: Distance-based regularisation of deep networks for finetuning

authors:
Ishaan Gulrajani
David Lopez-Paz

title: In search of lost domain generalization

authors:
Moritz Hardt
Ben Recht
Yoram Singer

title: Train faster, generalize better: Stability of stochastic gradient descent

authors:
Trevor Hastie
Robert Tibshirani
Jerome Friedman

title: The elements of statistical learning: data mining, inference, and prediction

venue: Springer Science & Business Media,

authors:
Shoubo Hu
Kun Zhang
Zhitang Chen
Laiwan Chan

title: Domain generalization via multidomain discriminant analysis

venue: In UAI,

authors:
Zeyi Huang
Haohan Wang
Eric P Xing

title: Self-challenging improves cross-domain generalization

title: Causal regularization

authors:
Pang Wei Koh
Shiori Sagawa
Henrik Marklund
Sang Michael Xie
Marvin Zhang
Akshay Balsubramani
Weihua Hu
Michihiro Yasunaga
Richard Lanas Phillips
Irena Gao
Tony Lee
Etienne David
Ian Stavness
Wei Guo
Berton Earnshaw
Imran Haque
Sara M Beery
Jure Leskovec
Anshul Kundaje
Emma Pierson
Sergey Levine
Chelsea Finn
Percy Liang

title: Wilds: A benchmark of in-the-wild distribution shifts

venue: In ICML,

authors:
David Krueger
Ethan Caballero
Jörn-Henrik Jacobsen
Amy Zhang
Jonathan Binas
Rémi Le Priol
Aaron C. Courville

title: Out-of-distribution generalization via risk extrapolation (rex)

authors:
Da Li
Yongxin Yang
Yi-Zhe Song
Timothy Hospedales

title: Deeper, broader and artier domain generalization

authors:
Haoliang Li
Sinno Jialin Pan
Shiqi Wang
Alex C. Kot

title: Domain generalization with adversarial feature learning

authors:
Yiying Li
Yongxin Yang
Wei Zhou
Timothy M. Hospedales

title: Feature-critic networks for heterogeneous domain generalization

authors:
Jonathan Lorraine
Paul Vicol
David Duvenaud

title: Optimizing millions of hyperparameters by implicit differentiation

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Andreas Maurer
Massimiliano Pontil

title: Empirical bernstein bounds and sample variance penalization

venue: In COLT,

authors:
Mehryar Mohri
Afshin Rostamizadeh
Ameet Talwalkar

title: Foundations of machine learning

venue: MIT press,

authors:
Krikamol Muandet
David Balduzzi
Bernhard Schölkopf

title: Domain generalization via invariant feature representation

authors:
Behnam Neyshabur
Zhiyuan Li
Srinadh Bhojanapalli
Yann LeCun
Nathan Srebro

title: Towards understanding the role of over-parametrization in generalization of neural networks

authors:
Mohammad Pezeshki
Sékou-Oumar Kaba
Yoshua Bengio
Aaron Courville
Doina Precup
Guillaume Lajoie

title: Gradient starvation: A learning proclivity in neural networks

venue: arXiv preprint arXiv:2011.09468,

authors:
Massimiliano Pontil
Andreas Maurer

title: Excess risk bounds for multitask learning with trace norm regularization

title: Early stopping-but when

venue: In Neural Networks: Tricks of the trade,

authors:
Elan Rosenfeld
Pradeep Ravikumar
Andrej Risteski

title: An online learning approach to interpolation and extrapolation in domain generalization

venue: arXiv preprint arXiv:2102.13128,

authors:
Shai Shalev-Shwartz
Shai Ben-David

title: Understanding machine learning: From theory to algorithms

venue: Cambridge university press,

authors:
Shirish K Shevade
S Sathiya Keerthi
Chiranjib Bhattacharyya
Karaturi Radha Krishna Murthy

title: Improvements to the smo algorithm for svm regression

venue: IEEE transactions on neural networks,

authors:
Baochen Sun
Kate Saenko

title: Deep coral: Correlation alignment for deep domain adaptation

authors:
Hemanth Venkateswara
Jose Eusebio
Shayok Chakraborty
Sethuraman Panchanathan

title: Deep hashing network for unsupervised domain adaptation

authors:
Yufei Wang
Haoliang Li
Alex C Kot

title: Heterogeneous domain generalization via domain mixup

venue: In ICASSP,

authors:
Haotian Ye
Chuanlong Xie
Tianle Cai
Ruichen Li
Zhenguo Li
Liwei Wang

title: Towards a theoretical framework of out-of-distribution generalization

venue: arXiv preprint arXiv:2106.04496,

heading: 1 INTRODUCTION

text:
Machine learning systems have shown exceptional performance on numerous tasks in computer vision and beyond. However performance drops rapidly when the standard assumption of i.i.d. training and testing data is violated. This domain-shift phenomenon occurs widely in many applications of machine learning (Csurka, 2017; Zhou et al., 2021; Koh et al., 2021), and often leads to disappointing results in practical machine learning deployments, since data ‘in the wild’ is almost inevitably different from reference training sets.
Given the practical significance of this issue, a large number of methods have been proposed that aim to improve models’ robustness to deployment under a different distribution than used for training (Zhou et al., 2021), a problem setting known as domain generalisation (DG). These methods span diverse approaches such as specialised neural architectures, data augmentation strategies, and regularisers. Nevertheless, the DG problem setting is difficult to model formally for principled derivation and theoretical analysis of algorithms, the target domain of interest is unobservable during training, and cannot be directly approximated by the training domains due to unknown distribution shift. Therefore the majority of these existing approaches are based on poorly understood empirical heuristics.
To make matters worse, a recent study by Gulrajani & Lopez-Paz (2021) assessed the state of DG research with a carefully conducted comparative evaluation of algorithms on a large benchmark suite under a common platform. They found that published methods were not as effective as claimed, and in particular reported that ‘no existing method reliably beats a well tuned empirical risk minimization (ERM) baseline’. We argue that this negative result highlights the need for better theory in this area in order to understand why existing algorithms have such erratic performance, and to guide the development of principled algorithms that are more effective and reliable.
To this end, our first contribution is to present an intuitive learning-theoretic bound for DG performance. Intuitively, while the held-out domain of interest is indeed unobservable during training, we can bound its performance using learning theoretic tools similar to the standard ones used to bound the performance on (unobserved) testing data given (observed) training data. In particular we show that the performance on a held out target domain is bounded by the performance on known source
domains, plus two additional model complexity terms, that describe how much a model can possibly have overfitted to the training domains. This theoretical contribution leads to several insights.
Firstly, our theory suggests that DG performance is governed by a trade-off between empirical risk and model complexity that is analogous to the corresponding and widely understood trade-off that explains generalisation in standard i.i.d. learning as an overfitting-underfitting trade-off (Geman et al., 1992). Based on this, we hypothesise that performance variability is determined by implicit or explicit regularisation. That is, the plethora of different strategies available (Zhou et al., 2021) – from data-augmentation to specialised optimisers – actually affect DG performance by explicitly or implicitly choosing different fit-complexity trade-offs. We corroborate this hypothesis by evaluating a number of models in the DomainBed suite in terms of complexity, and showing that their apparently erratic performance in Gulrajani & Lopez-Paz (2021) is actually consistent with an explanation in terms of implied complexity.
Practically, our analysis suggests that the model selection strategy (Hastie et al., 2009) is a factor in DG performance that is at least as important as the actual mechanism of model complexity control (i.e., Tuning of regularisation strength vs specific parametric design of regulariser). In particular, regularisation should be stronger when optimizing for future DG performance than when optimizing for performance on seen domains. Unfortunately, model complexity is hard to carefully control in deep learning due to the large number of relevant factors (architecture, regularisers, implicit regularisation from optimiser, etc). Gulrajani & Lopez-Paz (2021) attempted to address this by hyper-parameter search in the DomainBed benchmark, but are hampered by the computational infeasibility of accurate hyper-parameter search. In this paper, we use linear models and off-the-shelf self-supervised features to demonstrate much more clearly how cross-domain performance depends on complexity. Specifically, our theoretical and empirical results show that, contrary to the conclusion of Gulrajani & Lopez-Paz (2021), simple domain-wise cross-validation is a better objective to drive DG model selection.
In summary, based on our new generalisation bound, and associated empirical analysis, our takehome messages are: (i) Model fit vs complexity trade-off is a key determinant of DG performance, that explains existing DG algorithm performance variability. (ii) The complexity control strategy used to determine bias-variance trade-off is crucial in practice, with peak DG performance achieved when optimizing model complexity based on domain-wise validation. (iii) Regularisation required for optimal DG is greater than for conventional optimization for within-domain performance.

heading: 2 RELATED WORK

text:
Theoretical Analysis of the DG Setting and Algorithms The DG problem setting was first analysed in Blanchard et al. (2011). Since then there have been some attempts to analyse DG algorithms from a generalisation bound perspective (Muandet et al., 2013; Blanchard et al., 2021; Hu et al., 2020; Albuquerque et al., 2020; Rosenfeld et al., 2021). However these studies have theoretical results that are either restricted to specific model classes, such as kernel machines, or make strong assumptions about how the domains seen during training will resemble those seen at test time—e.g., that all domains are convex combinations of a finite pre-determined set of prototypical domains. In contrast, our Rademacher complexity approach can be applied to a broad range of model classes (including neural networks), and makes comparatively milder assumptions about the relationship between domains—i.e., they are i.i.d. samples from another arbitrary distribution over domains.
The majority of the existing work investigating the theoretical foundations of DG follow the initial formalisation of the domain generalisation problem put forth by Blanchard et al. (2011), where the goal is to minimise the expected error over unseen domains. However, several recent works have also explored the idea of bounding the error on a single unseen domain with the most pathological distribution shift (Janzing, 2019). This type of analysis is typically rooted in methods from causal inference, rather than statistical learning theory. As a consequence, they are able to make stronger claims for the problems they address, but the scope of their analysis is necessarily limited to the scenarios where their assumptions about the underlying causal structures are valid. For example, Janzing (2019) provides bounds that assume problems conform to a specific class of structural equation models, and the analysis is performed under the assumption that infinite training data is available within each of the observed training domains. Throughout the work we address the stan-
dard DG formalisation given by Blanchard et al. (2011), where one is concerned with the expected performance of a model on domains sampled from some distribution over domains.
Others rely on trying to link between domain adaptation objectives (where target domains are observable for alignment to source domains) and domain generalisation (where target domains are not observable and thus cannot correctly be used in a learning objective). Albuquerque et al. (2020) proceed by making assumptions on the structure of the distribution over possible domains (i.e., that it has support determined by the convex hull of a finite set of prototypical domains), which allows them to upper bound the domain alignment metric. Ye et al. (2021) provide a bound that depends on an unobservable domain distance quantity, which they then approximate in experiments using kernel density estimates.
Rosenfeld et al. (2021) is another piece of work that theoretically investigates the generalisation of ERM in a DG setting. They deal with online DG, where each time-step corresponds to observing a new domain, and the learner must produce a new model capable of generalising to novel domains. Another point of difference between their work and the standard DG problem setting of Blanchard et al. (2011) is that the domain at each time-step is chosen by an adversary. They analyse this game for a finite number of time-steps, but they assume each domain has an infinite amount of data. They also put some limitations on the adversary: e.g., it must choose a domain that is a convex combination of a finite number of pre-determined domains. In contrast, our theoretical analysis is in the more realistic setting where one has a finite amount of data per domain, and the domains we consider are not limited to convex combinations of a set of prototypical domains. Possibly the most similar work to our theoretical contributions is due to Ahuja et al. (2021), who also provide learningtheoretic generalisation bounds for DG. However, their analysis only applies to finite hypothesis classes (which does not include, e.g., linear models or neural networks), whereas ours can be applied to any class amenable to analysis with Rademacher complexity.
Empirical Analysis The main existing empirical analysis on DG is Gulrajani & Lopez-Paz (2021), who compared several state of the art DG methods under a common evaluation and hyper-parameter tuning protocol called DomainBed. They ultimately defend Empirical Risk Minimization (ERM) over more sophisticated alternatives on the grounds that no competitor consistently beats it across the benchmark suite. We also broadly defend ERM, and build on the same benchmark, but differently we provide a much deeper analysis into when and why ERM works. More specifically: (i) We provide a new theoretical analysis of ERM’s generalisation quality unlike the prior purely empirical evaluation, (ii) We re-use the DomainBed benchmark to directly corroborate this theory under controlled conditions using linear models where model complexity can be tractably and accurately tuned. (iii) We use our complexity-based analysis to explain the previously erratic results of prior DomainBed competitors in terms of model complexity. (iv) We identify, and empirically validate, the preferred model selection criterion for DG, a point which was inconclusive in Gulrajani & Lopez-Paz (2021).

heading: 3 BOUNDING RISK FOR DOMAIN GENERALIZATION

text:
I.i.d. learning is concerned with learning a mapping from some input space X , to a label space Y , given data drawn from a distribution on X × Y . One aims to find a model, f∗ ∈ F , that minimises the expected loss (also called risk) on unseen data,
f∗ = arg min f∈F Lp(f), Lp(f) = E(~x,y)∼p[`(f(~x), y)], (1)
where p is the data distribution and `(·, ·) is the loss function. In practice, we only have access to a finite set of data, S = {(~xi, yi)}mi=1, sampled i.i.d. from this distribution, so must minimise an empirical risk estimate,
f̂ = arg min f∈F
L̂p(f), L̂p(f) = 1
m m∑ i=1 `(f(~xi), yi). (2)
wherem is the number of training examples. One of the central focuses of statistical learning theory is to bound the difference between these two types of risk. For example, in the standard single domain setting this can be done via
Lp(f) ≤ L̂p(f) + 2Rm(F) +O (√ ln(1/δ) m ) , (3)
which holds with probability at least 1− δ, andR(F) is known as the empirical Rademacher complexity of the hypothesis class, F . This complexity term is defined as
Rm(F) = E~σ [ sup f∈F 1 m m∑ i=1 σi`(f(~xi), yi) ] . (4)
For a hypothesis class consisting of norm-constrained linear classifiers, one can achieve the following upper bound on the empirical Rademacher complexity (Shalev-Shwartz & Ben-David, 2014),
Rm(F) ≤ XB√ m , F = {~x 7→ ~x · ~w : ‖~w‖2 ≤ B},
where we assume that each input, ~x, has a Euclidean norm of at mostX . There are a variety of ways to define hypothesis classes for neural networks, but most recent approaches take the view of fixing a particular architecture and specifying constraints on the norms of weights or distances they can move from their initialisations (Bartlett et al., 2017; Neyshabur et al., 2019; Gouk et al., 2021).
Domain generalisation. While standard i.i.d. learning assumes all data come from the same distribution, the DG problem setting assumes the existence of an environment E , of distributions p (Blanchard et al., 2011). Note that we do not restrict what types of differences one could see between different domains sampled from E : for any two distributions in p, qsupp(E) it could be the case that either p(~x) 6= q(~x), or p(y|~x) 6= q(y|~x), or even that both types of distribution shift have occurred. The conceptually simplest—and often implicitly assumed—goal of DG methods is to minimise the expected risk across different distributions that could be sampled from the environment,
LE(f) = Ep∼E [Lp(f)]. (5)
This object is also the most commonly analysed idealised objective in the learning theory literature (Blanchard et al., 2011; Muandet et al., 2013; Blanchard et al., 2021; Rosenfeld et al., 2021), but other formulations also exist (Arjovsky et al., 2019). As with the single domain learning problem, we only have access to an empirical estimate of the risk,
L̂E(f) = 1
n n∑ j=1 L̂pj (f), (6)
where we assume for ease of exposition that all n domains have the same number of examples.

heading: 3.1 BOUNDING THE GENERALISATION GAP

text:
We next bound the generalisation gap between the observed empirical risk, L̂E , on the source domains and the expected risk, LE , on unseen domains that holds uniformly for all hypotheses in F . Theorem 1. For a 1-Lipschitz loss, `(·, ·), taking values in [0, 1], with confidence at least 1− 2δ for all f ∈ F we have that
LE(f) ≤ L̂E(f) + 2Rmn(F) + 2Rn(F) + 3 √
ln(2/δ) 2mn + 3
√ ln(2/δ)
2n ,
where n is the number of training domains andm is the number of training examples in each training domain.
The proof can be found in Appendix A.
Discussion Theorem 1 tells us that expected risk on unseen domains is bounded by the empirical risk (training loss) on seen domains, plus Rademacher complexity terms and sampling error terms that decay with the number of domains and training instances per domain. As with typical singledomain bounds, the latter sampling error terms are not under control of the model designer. The former Rademacher terms describe the complexity of the chosen model class and govern how much it could possibly overfit to the seen domains while minimising empirical risk. In the case of linear models, these terms depend on weight norms, while in the case of deep models they further depend on properties of the chosen network architecture. Mirroring conventional generalisation in standard i.i.d. learning, a very simple model may minimise the Rademacher terms, R, while producing high
empirical risk, L̂E , and vice-versa. Thus good generalisation critically depends on a carefully chosen empirical risk vs model complexity trade-off. The difference between Theorem 1 and single domain bounds (e.g., Equation 3) is the additional dependence on the number of domains n and additional Rademacher complexity term. This implies that when the goal is to generalise to new domains, the risk of overfitting is higher. Therefore we expect a lower complexity model to me optimal for held out domain performance compared to seen domain performance in standard i.i.d. learning.

heading: 3.2 BOUNDING THE EXCESS RISK OF ERM

text:
As a second theoretical contribution, we next bound the excess expected risk between the ERM solution f̂ and the best possible model f∗ within the function class F . Note that bounding excess risk, as opposed to the generalisation gap, overcomes some of the known issues with theoretically analysing the generalisation properties of deep networks (Belkin et al., 2018).
Corollary 1. With probability as least 1 − 2δ, the excess risk of the empirical risk minimiser in F is bounded as
E[LE(f̂)− LE(f∗)] ≤ 2Rmn(F) + 2Rn(F) + 2 √
ln(2/δ) 2mn + 2
√ ln(2/δ)
2n ,
The proof (in Appendix B) uses the same technique as the usual proof for excess risk (see, e.g., Pontil & Maurer (2013)), and some of the intermediate results required for Theorem 1.
Discussion Corollary 1 tells us that the gap between ERM and the best possible predictor in the function class depends on the same complexity terms observed in Theorem 1. In particular, for any typical hypothesis class, ERM converges to the optimal classifier at a rate of O(1/ √ mn) + O(1/ √ n). To justify itself theoretically, any future DG method that claims to be better than ERM should either: (1) demonstrate a faster convergence rate than this—at least by an improved constant factor; or (2) formally show that the chosen hypothesis class is composed of models that can extrapolate to new domains without additional data. The latter would likely involve making specific assumptions about the underlying data generation process, coupled with analysis of a specific hypothesis class using methods from causal inference. An example of such analysis is given by Janzing (2019). Methods based on causal inference have the potential to give bounds with much better convergence rates than the rate given above. However, because one must make assumptions about the underlying family of structural equation models, the applicability of such bounds is much more restricted than our Rademacher complexity technique, which does not require these assumptions.

heading: 4 EXPERIMENTS

text: Based on our previous theoretical analysis, we conduct experiments on DomainBed (Gulrajani & Lopez-Paz, 2021). In particular we aim to address the following questions: (1) Theorem 1 shows that novel domain generalisation performance is governed by empirical-risk complexity trade-off. Can we directly demonstrate the dependence of generalisation on complexity by controlling for model complexity? (2) Can the erratic performance of state of the art methods previously observed on DomainBed be explained in terms of model complexity? (3) Given that model complexity is a crucial determinant of generalisation, what is the best objective for tuning regularisation strength?

heading: 4.1 DOMAIN GENERALISATION PERFORMANCE DEPENDS ON MODEL COMPLEXITY

text:
To directly investigate the impact of model complexity on domain generalisation performance, we first work with Linear SVM applied to pre-computed deep features1. In this function class model complexity can be directly controlled by a scalar hyper-parameter, the objective is convex so confounding factors in deep network training (optimiser choice, early stopping, etc) disappear, and training is fast enough that we can densely and exhaustively evaluate a wide range of complexities.
1Note that when a fixed feature extractor, or any other preprocessing is used, it does not impact the model complexity or associated generalisation bound.
Setup We use DINO (Caron et al., 2021) pre-trained models, such as DINO-ViTB8 and DINOViTS8 to extract features, and then train LinearSVC classifiers. We experiment on six different DG benchmarks, including RotatedMNIST (Ghifary et al., 2015), PACS (Li et al., 2017), VLCS (Fang et al., 2013), OfficeHome (Venkateswara et al., 2017), SVIRO (Dias Da Cruz et al., 2021) and Terra Incognita (Beery et al., 2018). Different trials have different numbers of training instances in practice, so we used a linear SVM objective of the form Cn ∑n i `(fw(xi), yi)+‖w‖2, with parameter C controlling the complexity of model w through loss-regularisation trade-off. We searched a wide range of log2 C in {−10, . . . , 10}. We then conduct two experiments, holding training and test split size constant across both: (i) Training on the train splits of all domains, and testing on the test splits of all domains (i.e., standard i.i.d. learning). (ii) Training on the train splits of three domains, and testing on the test splits of a held out domain (i.e., DG performance).
Results The results in Fig 1 average over 5 random seeds for dataset splitting, and all choices of hold-out target domain. From these we can see that: (i) all experiments exhibit the classic trade-off between fitting the data well and constraining hypothesis class complexity appropriately. We observe underfitting for high regularisation (small C), and overfitting at low regularisation (large C). (ii) Comparing the within- and across-domain evaluation: across-domain leads to lower performance— which is to be expected, due to the distribution shift. A more noteworthy observation is that, the optimal regularisation for novel-domain performance is stronger than for seen-domain performance (red vertical lines left of blue). This corroborates our theoretical result that the ideal model complexity is lower for DG than for conventional i.i.d. learning. Note that these results are using an exhaustive oracle to explore mode complexity, and do not constitute a complete algorithm, which would need to chose a specific complexity. We discuss an algorithm in Sec 4.3.

heading: 4.2 EXPLAINING DOMAINBED DG PERFORMANCE IN TERMS OF COMPLEXITY

text:
Our experiment in Section 4.1 showed that linear models exhibit a very clean trade-off between complexity and generalization. However there are many subtle factors that influence effective model complexity in neural networks besides architecture: learning rate, learning rate schedule, weight decay, early stopping, etc; and the repeatably of learning is hampered by non-convexity, stochastic initialization, and stochastic mini-batches, etc. The previous evaluation study of Gulrajani & Lopez-Paz (2021) attempted some degree of control over these factors for state of the art neural DG methods by applying a common hyperparameter search strategy that aimed to find the best tuning (cf: Fig 1) for each competitor. They nevertheless found existing methods to perform unreliably. Our conjecture is that the number of seed samples (3) and hyper-parameter search trials (20) used was far too few to reliably tune the deep networks evaluated, and therefore the tuning of each model
was inaccurate and performance estimates unreliable. Since dense and accurate hyper-parameter tuning for neural networks is more challenging than our linear models in Section 4.1, we take a different approach. Rather than trying to control complexity directly, we take trained DomainBed methods and measure their complexity retrospectively.
Setup To ensure that complexity can be measured accurately, we work with 2-layer MLPs. Specifically, we take fixed ImageNet pre-trained ResNet-18 features, and feed them to MLPs, which are then trained using the DomainBed framework. 2-layer MLPs are sufficient to instantiate many of the state of the art DG algorithms in DomainBed. We train and compare ERM, CORAL (Sun & Saenko, 2016), Mixup (Wang et al., 2020), MMD (Li et al., 2018), RSC (Huang et al., 2020), SD (Pezeshki et al., 2020), VRex (Krueger et al., 2020), and IRM (Arjovsky et al., 2019) with two hyperparameter choices. We also report the results of an ERM model, checkpointed at a range of training iterations.
Measuring Model Complexity Because we have limited our attention to 2-layer MLPs, we can take advantage of a model capacity metric that is specialised for this class of models. We retrospectively determine the complexity of a trained network using the measure proposed by Neyshabur et al. (2019), who also use this measure to bound the Rademacher complexity of a 2-layer MLP hypothesis classes. More concretely, the expression used for computing complexity is
C(fU,V ) = ‖V ‖F (‖U − U0‖F + ‖U0‖2), (7)
where U is the weight matrix of the first layer, U0 is its random initialisation, and V is the weight matrix of the second layer. We use ‖·‖F to denote Frobenius norm, and ‖·‖2 to indicate the spectral norm. Note that for simplicity we have omitted constant factors that depend only on the architecture and problem setting, and not the learned weights, as we use the same architecture for all methods we investigate.
Results The results in Figure 2 summarise the trade-off between measured model complexity (xaxis), and held-out domain test accuracy (y-axis), averaged over all choices of held out domain for each dataset (See Supplementary for full breakdown across all held-out domains). The top plot compares several of the published neural DG methods implemented in DomainBed. The main message of the top plot is that results are consistent with the hypothesis that the resulting complexity of the models trained by various DG methods is a key factor in determining domain generalisation accuracy (compare Fig. 2 with the clean result for linear where models we are able to intervene and control complexity directly in Fig. 1)—thus explaining the previously erratic behaviour observed in domain-bed.
To provide a different view of the same issue, the bottom plot reports a vanilla untuned ERM model check-pointed every 300 training iterations between (up to a total of 15,000) but expressed in the same complexity units as above. Because neural models can gain complexity with iterations (see, e.g., Prechelt (1998); Hardt et al. (2016)) we also see typical overfitting-underfitting trade-off curves. This shows that, as expected, proper choice of early stopping criterion is important. It also shows that over a similar dynamic range of complexity (0.5-3), ERM (below) and alternative models (above) span a similar dynamic range of accuracy (e.g., 6% for PACS). This suggests that complexity, as controlled by whatever mechanism, is a key determinant of DG performance.
How Do DG Models Control Complexity? By introducing different modifications to standard ERM, DG models explicitly or implicitly modify the bias and variance of the function class to be learned. Gulrajani & Lopez-Paz (2021) highlight neural models as being dependent on learning rate, batch size, dropout rate, and weight decay; with other factors being choice of optimiser and weight initialiser, etc. For example, higher dropout rate and weight decay tend to reduce complexity, while some other factors influence complexity in a less transparent way. RSC introduce more sophisticated dropout-like mechanisms, which would be expected to reduce complexity. Meanwhile alignment-based methods like CORAL and MMD effectively add auxiliary losses, which will implicitly affect complexity, but are hard to explicitly link to it. Consistency based methods like IRM and VRex penalise loss variance, which also tends to reduce the generalisation gap in the single task case (Maurer & Pontil, 2009). The specific setting of all the corresponding hyper-parameters (e.g., regulariser strength, dropout rates) influence final model complexity, which we argue is the key determinant of performance.

heading: 4.3 EVALUATING LINEAR MODELS ON DOMAINBED

text:
Having demonstrated control of generalisation with linear model regularisation (Sec. 4.1), and explained DomainBed performance post-hoc via measuring complexity of neural models (Sec. 4.2), we next evaluate algorithms for practical complexity tuning. From Fig. 1 we saw that within-domain and cross-domain evaluation have different optimal tuning. If we consider automating the search for a good regularisation parameter, these two final evaluation criteria correspond respectively to hyperparameter selection based on a validation set from the seen source domains, vs based on a validation set drawn from a held out source domain. The latter validation criterion corresponds to a unbiased estimator of expected cross-domain generalisation performance RE(f), which our theorem bounds.
Setup We performed DG evaluation on DomainBed using ERM with linear models using the same DINO features as Section 4.1. For each held out domain, we performed hyperparameter tuning with either instance-wise cross-validation (where validation folds are drawn from the same source domains used for training) or domain-wise cross-validation (where validation folds are drawn from held-out domains not seen during training).
Results The results in Table 1 report the average accuracy across held out domains, and the average selected logC regularisation strength (See Supplementary for the full breakdown of results for each held-out domain). We can see that the domain-wise cross-validation objective leads to similar or better accuracy, and similar or smaller C value selection (i.e., stronger regularisation). This outcome is opposite to the observation made by Gulrajani & Lopez-Paz (2021). Given the theoretical support for our domain-wise objective, and our clear empirical outcome when freed of confounding factors such as stochasticity in neural network training, we consider our result to be decisive and attribute the different observation Gulrajani & Lopez-Paz (2021) to the inaccuracy and stochasticity of hyperparameter tuning in neural networks.

heading: 4.4 DISCUSSION

text:
At the start of this section we set out to address three questions, which can be answered as: (1) The dependence of cross-domain generalisation on complexity can be directly and precisely when using linear models. (2) The erratic performance of state of the art methods in DomainBed can be largely explained in terms of implied model complexity. (3) If the goal is to optimise for domain generalisation, then then domain-wise validation is preferred to instance-wise validation as a modelselection objective.
Our analysis suggests that several existing methods that tried to tune learnable DG hyper-parameters by performance on a held-out domain (Balaji et al., 2018; Li et al., 2019) were broadly on the right track, and concurs with Gulrajani & Lopez-Paz (2021) that those methods with underspecified hyperparameter and model selection procedures are unhelpful. However, given that most neural methods have many more complexity hyperparameters than the single hyperparameter that we were able to carefully control for linear models, obtaining accurate tuning and reliable performance evaluation is likely to be a challenge. Gradient-based hyper-parameter estimation methods, as initially attempted in Balaji et al. (2018); Li et al. (2019), together with efficient methods for long-inner loop hypergradient calculation (Lorraine et al., 2020), may benefit the former problem by making search in larger numbers of hyperparameters more feasible. Alternatively, using general purpose pre-trained features (Caron et al., 2021) as we did here, and focusing on learning shallow models that can be accurately tuned for DG may be another promising avenue in practice. Although achieving state of the art performance is not our focus, we note that our results in Table 1 are quite competitive with end-to-end trained state of the art (Gulrajani & Lopez-Paz, 2021), despite using fixed features and shallow models.

heading: 5 CONCLUSION

text: In this paper we explained the performance of domain generalisation methods in terms of model complexity (bias-variance trade-off) from both theoretical and empirical perspectives. Both perspectives show that complexity impacts cross-domain generalisation in a way that closely mirrors the bias-variance trade-off in conventional i.i.d. learning – but where stronger regularisation is required if optimising for cross-domain generalisation than if optimising for conventional within-domain generalisation. We clarified the preferred model selection criterion in each case. This analysis demystifies the problem being posed by the DG problem setting, why existing algorithms succeed or fail to work, and sets the bar for future theoretical theoretical studies to surpass in terms of convergence rates, as well as for future empirical studies to surpass in terms of strong baselines.

heading: A PROOF OF THEOREM 1

text:
We use a slightly modified version of the standard empirical Rademacher complexity bound on generalisation error, as stated by Mohri et al. (2018) but originally shown by Bartlett & Mendelson (2002), where the modification is to weaken the i.i.d. assumption to be only an independence assumption—i.e., allow instances to be drawn from different distributions. The proof is exactly the same (because McDiarmid’s inequality requires only independence, and not identical distributions), but we re-state it here for completeness. Theorem 2. For p1, ..., pn independent samples from E , and 1-Lipschitz loss `(·, ·) taking values in [0, 1], the following holds with confidence at least 1− δ,
1
n n∑ j=1 Lpj (f) ≤ 1 n n∑ j=1 L̂pj (f) + 2Rmn(F) + 3 √ ln(2/δ) 2mn , (8)
where L̂pj (f) is measured on Spj = {(~xij , yij)}mi=1, a collection of m i.i.d. samples from pj .
Proof. Let S = Sp1 ⋃ ... ⋃ Spn and define
Φ(S) = sup f∈F
1
n n∑ j=1 (Lpj (f)− L̂pj (f)). (9)
Note that Φ(S) satisfies the bounded differences property required by McDiarmid’s inequality: i.e., if we construct S ′ by replacing any one of the (xij , yij) in S with another random variable also drawn from pj , then |Φ(S) − Φ(S ′)| ≤ 1mn . Therefore McDiarmid’s inequality implies that with confidence at least 1− δ2
Φ(S) ≤ ESp1:n∼p1:n [Φ(S)] + √ ln(2/δ) 2mn . (10)
We continue by bounding the expected value of Φ(S),
ESp1:n∼p1:n [Φ(S)] (11)
= ESp1:n∼p1:n [ sup f∈F 1 n n∑ j=1 (Lpj (f)− L̂pj (f)) ] (12)
= ESp1:n∼p1:n [ sup f∈F 1 n n∑ j=1 ( ES′pj∼pj [ 1 m m∑ i=1 `(f(~x′ij , y ′ ij) ] − 1 m m∑ i=1 `(f(~xij , yij) )] (13)
≤ ESp1:n∼p1:nES′p1:n∼p1:n [ sup f∈F 1 n n∑ j=1 1 m m∑ i=1 (`(f(~x′ij), y ′ ij)− `(f(~xij), yij)) ] (14)
= ESp1:n∼p1:nES′p1:n∼p1:nE~σ [ sup f∈F 1 n n∑ j=1 1 m m∑ i=1 σij(`(f(~x ′ ij), y ′ ij)− `(f(~xij), yij)) ] (15)
≤ ES′p1:n∼p1:nE~σ [ sup f∈F 1 n n∑ j=1 1 m m∑ i=1 σij`(f(~x ′ ij), y ′ ij) ] (16)
+ ESp1:n∼p1:nE~σ [ sup f∈F 1 n n∑ j=1 1 m m∑ i=1 −σij`(f(~xij), yij) ] (17)
= 2ESp1:n∼p1:nE~σ [ sup f∈F 1 n n∑ j=1 1 m m∑ i=1 σij`(f(~xij), yij) ] (18)
= 2ESp1:n∼p1:n [Rmn(F)], (19)
where the first inequality is from moving the supremum inside the expectation and the second is from subadditivity of suprema. Observing that the absolute difference of computing Rmn(F) on S and
S ′ cannot exceed 1mn , another application of McDiarmid’s inequality tells us that with confidence at least 1− δ2
2ESp1:n∼p1:n [Rmn(F)] ≤ 2Rmn(F) + 2 √ ln(2/δ) 2mn . (20)
Combining Equations 10 and 20 with the union bound concludes the proof.
We now prove Theorem 1.
Proof. Theorem 2 tell us that with confidence at least 1− δ,
1
n n∑ j=1 Lpj (f) ≤ L̂E(f) + 2Rmn(F) + 3 √ ln(2/δ) 2mn . (21)
Thus, we must provide a (high confidence) upper bound on
LE(f)− 1 n n∑ j=1 Lpj (f) (22)
that holds uniformly for all f ∈ F . We can use the same idea the proof for Theorem 2 to show how Rademacher complexity controls generalisation to novel domains, rather than novel instances within the same domain. Begin by letting P = {p1, ..., pn} be an i.i.d. sample of n domains from E , and define
Φ(P ) = sup f∈F LE(f)− 1 n n∑ j=1 Lpj (f). (23)
If we construct P ′ by replacing any pj ∈ P with p′j ∼ E , then we have |Φ(P ) − Φ(P ′)| ≤ 1n , so McDiarmid’s inequality tells us that with confidence at least 1− δ2
Φ(P ) ≤ Ep1:n∼E [Φ(P )] + √
ln(2/δ) 2n . (24)
We proceed by bounding the expected value of Φ(P ),
Ep1:n∼E [ sup f∈F ( Eq∼E [Lq(f)]− 1 n n∑ j=1 Lpj (f) )] (25)
= Ep1:n∼E [ sup f∈F Eq1:n∼E [ 1 n n∑ j=1 ( Lqj (f)− Lpj (f) )]] (26)
≤ Ep1:n,q1:n∼E [ sup f∈F 1 n n∑ j=1 ( Lqj (f)− Lpj (f) )] (27)
= Ep1:n,q1:n∼EE~σ [ sup f∈F 1 n n∑ j=1 σj ( Lqj (f)− Lpj (f) )] (28)
≤ 2Ep1:n∼EE~σ [ sup f∈F 1 n n∑ j=1 σjLpj (f) ] (29)
= 2Ep1:n∼EE~σ [ sup f∈F 1 n n∑ j=1 σjE(~x,y)∼pj [`(f(~x), y)] ] (30)
≤ 2Ep1:n∼EE(~xj ,yj)∼pjE~σ [ sup h∈F 1 n n∑ j=1 σj`(f(~xj), yj) ] (31)
= 2Ep1:n∼EE(~xj ,yj)∼pj [Rn(F)], (32)
where the first inequality comes from moving the supremum inside the expectation, ~σ is a vector of Rademacher random variables, the second inequality is due to the subadditivity of suprema, and the final inequality comes from moving the expectation outside of the supremum. Noting that replacing one of the (~xj , yj) pairs in the final equality will result in Rn(F) changing by at most 1n , McDiarmid’s inequality can be used to say with confidence 1− δ2 that
2Ep1:n∼EE(~xj ,yj)∼pj [Rn(F)] ≤ 2Rn(F) + 2 √
ln(2/δ) 2n . (33)
Combining Equations 21, 24, and 33 using the union bound completes the proof.

heading: B PROOF OF COROLLARY 1

text:
Proof. We begin with the expected excess risk, which can be bounded from above by
ES [LE(f̂)− LE(f∗)] = ES [LE(f̂)− L̂E(f̂)] + ES [L̂E(f̂)− L̂E(f∗)] + ES [L̂E(f∗)− LE(f∗)]
(34)
≤ ES [LE(f̂)− L̂E(f̂)] + ES [L̂E(f∗)− LE(f∗)] (35) = ES [LE(f̂)− L̂E(f̂)] (36)
= ES [ LE(f̂)− 1
n n∑ j=1 Lpj (f̂)
] + ES [ 1
n n∑ j=1 Lpj (f̂)− L̂E(f̂)
] , (37)
where the inequality arises because, by definition, the empirical risk of f̂ must be less than or equal to the empirical risk of the optimal model. The second equality comes from the fact that f∗ is determined independently of the particular training set that we sample, so ES [L̂E(f∗)] = LE(f∗). For the final equality: the first term can be bounded from above (with confidence 1 − δ) using the bound for the expected value of Φ(S) derived in the proof of Theorem 2; and the second term can be bounded from above (with confidence 1 − δ) using the bound for the expected value of Φ(P ) derived in the proof for Theorem 1. Combining these two high confidence bounds using the union bound yields the result.

heading: C DETAILED RESULTS

text:
Complexity Measurement for DomainBed Competitors In Figure 2 we illustrated the performance of DomainBed competitors in terms of their resulting model complexity by averaging over held-out domains for each dataset. In Figure 3, we present detailed results broken down by each held out dataset.
Detailed results on each DG benchmark In Table 1 we reported results of our linear model competitors on DomainBed summarised over choice of held out domain. In Tables 2-7, we report the detailed results of accuracy and choice of C paramater for each held out domain.