A theory of representation learning in neural networks gives a deep generalisation of kernel methods
Authors: Adam X. Yang
Maxime Robeyns
Edward Milsom
Nandi Schoots
Laurence Aitchison
Venue: Submitted to ICLR 2023 (2023)
Keywords: Gaussian process
infinite-width neural networks

Abstract: The successes of modern deep neural networks (DNNs) are founded on their ability to transform inputs across multiple layers to build good high-level representations. It is therefore critical to understand this process of representation learning. However, standard theoretical approaches (formally NNGPs) involving infinite width limits eliminate representation learning. We therefore develop a new infinite width limit, the representation learning limit, that exhibits representation learning mirroring that in finite-width networks, yet at the same time, remains extremely theoretically tractable. In particular, we derive an elegant objective that describes how each network layer learns representations that interpolate between input and output, and we confirm that the modes of the objective match the behaviour of finite but wide networks. Moreover, we use this limit and objective to develop a flexible, deep generalisation of kernel methods, that we call deep kernel machines (DKMs). We show that DKMs can be scaled to large datasets using inducing point methods from the Gaussian process literature, and we show that DKMs exhibit superior performance to other kernel-based approaches.

A THEORY OF REPRESENTATION LEARNING IN NEURAL NETWORKS GIVES A DEEP GENERALISATION OF KER- NEL METHODS
Anonymous authors Paper under double-bling blind review
ABSTRACT
The successes of modern deep machine learning methods are founded on their ability to transform inputs across multiple layers to build good high-level rep- resentations. It therefore critical to understand this process of representation learning. However, standard thcoretical approaches (formally NNGPs) involving infinite width limits climinate representation learning. We therefore develop a new infinite width limit, the Bayesian representation learning limit, that exhibits representation learning mirroring that in finite-width models, yet at the same time, retains some of the simplicity of standard infinite-width limits. In particular, we show that Deep Gaussian processes (DGPs) in the Bayesian representation learn- ing limit have exactly multivariate Gaussian posteriors, and the posterior covari- ances can be obtained by optimizing an interpretable objective combining a log- likelihood improve performance with a series of KL-divergences which keep the posteriors close to the prior. We confirm these results experimentally in wide but finite DGPs. Next, we introduce the possibility of using this limit and ohjec- tive as a flexible, deep generalisation of kernel methods, that call deep kernel machines (DKMs). Like most naive kernel methods, DKMs scale cubically in the number of datapoints. We therefore use methods from the Gaussian process induc- ing point literature to develop a sparse DKM that scales linearly in the number datapoints. Finally, we extend these approaches NNs (which have non non- Gaussinn posteriors) in the Appendices.
INTRODUCTION
The successes of modern machine learning methods from neural networks (NNs) deep Gaussian processes (DGPs Damianou & Lawrence 2013) Salimbeni & Deisenroth] 2017) is based on their ability to depth to transform the input into high-level representations that are good for solving difficult tasks (Bengio et al.] 2013] LeCun et al. 2015). However, theoretical approaches using infinite limits to understand deep models struggle to capture representation learning. In particular, there are two broad families of infinite limit, and while they both use derel--matrxx like objects they ultimately very different. The neural network Gaussian process (NNGP Neal 1996: Lee al.
2017 Matthews et al. 2018) applies to Bayesian models like Bayesian neural networks (BNNs) and DGPss and describes the representations each layer (formally, the NNGP kernel is raw second moment of the activities). In contrast, the neural tangent kernel (NTK Jacot et al.) (018)) is a very different quantity that involves gradients, and describes how predictions at all datapoints change we do gradient update a single datapoint. As such, the NNGP and NTK are suited to asking very different theoretical questions. For instance, the NNGP is better suited to understanding the transformation of representations across layers, while the NTK is better suited to understanding how predictions change through NN training.
While challenges surrounding representation learning have recently been addressed in the NTK setting Yang & Hu (2020), we are the first to address this challenge in the NNGP setting.
At the same time, kernel methods (Smola & Schölkopf, 1998 Shawe-Tary Taylor & Cristianini 2004 Hofmann et al. (2088) were a leading machine learning approach prior to the deep learning revolu- tion Krizhevsky et al. 2002). However, kernel methods were eclipsed by deep NNs because depth

---

gives NNs the flexibility to learn a good top-layer representation (Aitchison 2020). In contrast, in standard kernel method, the kernel (or equivalently the representation) is highly inflexible - there usually few tunable hyperparameters, but nothing that approaches the enormous flexibility of the top-layer representation in a deep model. There is therefore need to develop flexible, deep gen- eralisations of kernel method. Remarkably, our advances in understanding representation learning in DGPs give such a flexible, deep kernel method.
2 CONTRIBUTIONS
We present new infinite width limit, the Bayesian representation learning limit, that re- tains representation learning in deep Bayesian models including DGPs. The key insight is that as the width goes to infinity, the prior becomes stronger, and eventually overwhelms the likelihood. We can fix this by rescaling the likelihood to match the prior. This rescaling
can be understood a Bayesian context as copying the labels (Sec. 4.3).
We show that in the Bayesian representation learning limit, DGP posteriors are exactly zero-mean multivariate Gaussian, P (f5 y) = N (fş: G.) where fy, is the activation of the Ath feature in layer for all inputs (Sec. 4.4 and Appendix
We show that the posterior covariances can be obtained by optimizing the "deep kernel machine objective".
L(G,,..,,Gi) - - Po DE (V (0, GDIlW (o, (G+-))))
where G, are the posterior covariances, K ((G+1) are the kernel matrices, and Vt accounts for any differences in layer width (Sec. 4.3
We give an interpretation of this objective, with log P (Y(),) encouraging improved
regulariser, keeping posteriors, N' (0, Gy), close to the prior, (0, (GG,1)) (Sec. 4.5).
performance, while the KL-divergence terms act
We introduce sparse DKM, which takes inspiration GP induuiing point literature obtain a practical, scalable method that is linear in the number of datapoints. In contrast, naively computing/optimizing the DKM objective cubic in the number of datapoints (as with most other naive kernel methods: Sec. 4.7)
We extend these results to BNNs (which have non- non-Gaussian posteriors) in Appendix A
3 RELATED WORK
Our work is focused on DGPs and gives new results such as the extremely simple multivariate Gaussian form for DGP true posteriors. As such, our work is very different from previous work
on NNs, where such results are not available. There are at least three families of such work. First,
there is recent work on representation learning in the very different NTK setting (Jacot et al. 2018
Yang, 2019 Yang & Hu 2020) (see Sec. [T. In contrast, here we focus NNGPs NNal 1996
2017, Matthews et al.) 2018 Novak et al.) 2018 Garriga-Alonso et al. 2018 Jacot et al. 2018), where the challenge of representation learning has yet be addressed.
Williams| 1996] Lee et al.
Second, there is a body of work using methods from physics to understand representation learning
in neural networks (Antognini 2019 Dyer & Gur-Ari 2019 Hanin & Nica 2019. Aitchison 2020)
Li & Sompolinsky Yaida Yaida 2020 Naveh et al. [2000] Zavatone- Veth et al. Zavatone Veth & Pehlevan [202]] Roberts al. [202] [2021] Naveh & Ringel |aleesnn et al.) 2021). This work
focuses on perturbational, rather than variational methods. Third, there is a body of theoretical work including (Mei et al., 22088] Nguyen| 2019, SSrignaao Spiliopoulos 2020a by Nguyen & Pham 2020) which establishes properties such as convergence to the global optimum. This work is focused on two-layer (or one-hidden layer network) networks, and like the NTK, considers learning under SGD rather than Bayesian posteriors.
Another related line of work uses kernels to give a moosed-form expression for the weights of a neural network, based greedy, layerwise objective (Wu et al. 2022). This work differs in that it uses the HSIC objective, and therefore does not have a link to DGP' or Bayesian neural networks, and in that it uses greedy- layerwiie objective, rather than end-to-end gradient descent.
2

---

Layer
Layer
Layer 3
X
F3
Y
X
Go
F1
G1
E2
22
F3
G3
Y
X Go
G1
G2
G3
Y
The graphical model structure for each of our generative models for L 3. Top. The standard model (Eq. D. written purely in terms of features, Fe. Middle. The standard model, including Gram matrices as random variables (Eq.(5) Bottom. Integrating out the activations, Ft.
Figure 1:
4
RESULTS
We start by defining a DGP, which contains Bayesian NN (BNNs) as a special case (Appendix [A]. This model maps from inputs, X R to outputs, Y Rxxut+ where P is the number of input points, Vn is the number of input features, and U141 is the number of output features. The model has L intermediate layers, indexed & E {1,..., L,,, and at each intermediate layer there Ne features, F, E RP Both F; and Y can be written as stack of vectors,
F2 F==(f; (ff F
()))
Y (y1 Y2
Notti):
where f E IR gives the value of one feature and yx E R. gives the value of one output for all P input points. The features, f,,..., FL,, and (for regression) the outputs, Y, are sampled from Gaussian process (GP) with covariance which depends on the previous layer features (Fig. top),
(F. , (,, ((G(F))))))
(la)
P (Y W (yx:O, +
(lb)
Note only the regression likelihood to give concrete example; could equally use alternative likelihood e.g. for classification (Appendix(B) The distinction between DGPs and BNNs arises through the choice of K(-) and G(-). For BNNs, Appendix A For DGPs, G(-), which takes the features and computes the corresponding x Gram matrix, is
J.Ex G(F (
No FE-IFT--
Now, we introduce random variables representing the Gram matrices, G G(F where Ge-1 is random variable representing the Gram matrix layer é 1, whereas G(.) is deter- ministic function that takes features and computes the corresponding Gram matrix using Eq. (2). Finally, K(-), transforms the Gram matrices, G to the final kernel. Many kernels of interest are isotropic, meaning they depend only on the normalized squared distance between datapoints, Rig.
Kisotropies (G = Aisocropic (Rag (G
(3)
Importantly, we can compute this squared distance from Gę-1, without needing F(-1-
R.y > &E F2x)² N ((Px) 2Fx (Fxx)
= G# - 2G + Gyj.
(4)
where À indexes features, i and index datapoints and we have omitted the layer index for simplicity.
Importantly, we not restricted to isotropic kernels: other kernels that depend only the Gram matrix, such as the arccos kernels from the infinite NN literature (Cho & Saul 2009) can also be used (for further details, see Aitchison et al. 2020).
4.1 BNN AND DGP PRIORS CAN BE WRITTEN PURELY IN TERMS OF GRAM MATRICES
Notice that F. depends on F Fe-1 only through G = G(Fz-1). and Y depends on Fz only through G, G(F,) (Eq. I. We can therefore write the graphical model in terms of those Gram matrices (Fig. middle).
P (F = I (f): ((G-I))
(5a)
(G = 8 G(F.))
(5b)
PY|)r,) (xx oo'))
(5c)
3

---

where 8 is the Dirac-delta, and Go depends X (e.g. Go 1 XX"). Again, for concreteness have used regression likelihood, but other likelihoods could also be used.
vn
Now, we can integrate Fz out the model, in which case, we get an equivalent generative model written solely in terms of Gram matrices (Fig. [bottom), with
f /dF P(GGE))((/GG...
P P
and with the usual likelihood (e.g. Eq. Scp. This looks intractable (and indeed, in general it is intractable). However for DGPs, an analytic form is available. In particular, note the Gram matrix (Eq. [2p is the outer product of IID Gaussian distributed vectors (Eq. Tap. This matches the definition of the Wishart distribution (Gupta & Nagar 2018), so we have,
P (Ge/ (Ge| GG-1)
Wishart G5 K ((G-))), No
log P (G.|G,-)) Nr-P-1 log Gyl No log (G; (G- + const.
This distribution over Gram matrices is valid for DGPs of any width (though we need to be careful in the low-rank setting where N₂- < P). We are going leverage these Wishart distributions to understand the behaviour of the Gram matrices in the infinite width limit.
4.2 STANDARD INFINITE WIDTH LIMITS OF DGPs LACK REPRESENTATION LEARNING
We in a position take viewpoint the DGP analogue of standard NNGP results
(Lee et al., 2017 Matthews et al., 2018 Hron et al.) 2020 Pleiss & Cunningham, 2021). We can then evaluate the log- pestotttrrorr for model written only terms of Gram matrices,
log P (G1,..., G₂ X, log (YlGL) El log P (Gz|GG-1) const.
(8)
Then we take the limit of infinite width.
= N
for
EE (1, (1,...,, L}
with
N - oc.
This limit modifies log P (Ge (Eq. 7. but does not modify GI...., G2 in Eq. (8) as we get choose the values of GI...., GL which to evaluate the log-posterior. Specifically, the log- prior, log P (Gr Gr-1 (Eq. 7). scales with N₂ and hence with N. To get finite limit, we therefore need to divide by N.
lim N log (Go|Gr-1) V1 (log K (G-I) - - (G-.())))))))) + const N
- DEL (N (O, G&)l (G,K()))))) + const. (10)
And remarkably this limit be written as the KL-divergence between multivariate Gaussians. contrast, the log likelihood, log P (Y]G,), is constant wrt N (Eq. 5c). so lim 1 log (Y (Y|Gg) 0. The limiting log-posterior thus,
lim + log Y) - (W (o, G.)lW (G, + const. N
(11)
This form highlights that the log-posterior scales with N, in the limit N co, the posterior converges to a point distribution the global maximum, denoted Gi...., G (see Appendix for formal discussion of weak convergence),
lim P Y = N
O Gū)-
(12)
Moreover, evident from the KL-divergence form for the log-posterior (Eq. that the unique global maximum be computed recursively G K K(G?-1), with e.g. G Po X Thus, the limiting posterior over Gram matrices does not depend on the training targets, there is possibility of representation learning (Aitchison, 2020). This is deeply problematic as the successes of modern deep learning arise from flexibly learning good top-layer representations.
4

---

4.3 THE BAYESIAN REPRESENTATION LEARNING LIMIT
In the previous section, we saw that standard infinite width limits eliminate representation learning because as N
the log-prior terms, log P (Gr|Gr-1), in Eq. (8) dominated the log-likelihood, P P(Y|GL) and the likelihood is the only term that depends on the labels. We therefore introduce the "Bayesian representation learning limit" which retains representation learning. The Bayesian representation learning limit sends the number of output features, NL+1, to infinity as the layer- widths go to infinity,
Ne N M/
for
LE (I, L +1
with
N + co.
(13)
Importantly, the Bayesian representation learning limit gives a valid probabilistic model with a well- defined posterior, arising from the prior, (Eq. ( and likelihood which assumes each output channel is IID,
P PY = II Ni+i N (, + +0'1).
(14)
where Y E BPxNN++1 is infinite width (Eq. 13 whereas the usual DGP data, Y E Px is finite width. Of course, infinite-width data is unusual if not unheard-of. In practice, real data, Y E RPXEL11, almost always has a finite number of features, VL+11 How do apply the DKM to such data? The answer is to define Y as N copies of the underlying data, Y. i.c. Y= (Y As each channel is assumed be IID (Eq. 5c and 14) the likelihood is N times larger,
Y).
log P (G)))) = N log P (Y|GE),
(15)
The pyg-posteroor in the Bayesian representation learning limit very similar to the log- pe-ptttteror in the standard limit (Eq. [T.. The only difference is that the likelihood, log P (Y)Gr.) now scales with N, so it does not disappear as we take the limit, allowing us to retain representation learning.
C(G GL) lim Y log P (G1,..., GL|X, Y) + const, No
(16)
= log P (Y|G,) - El DKL. (N (O. GE)INN (O,K(G))))))
Here, we denote the limiting postterir as L(G...... GL), and this forms the DKM objective. Again, long the global maximum of the DKM objective is unique, the posterior is again point distribution around that maximum (Eq. 12]. Of course, the inclusion of the likelihood term means that the global optimum Gi..... G cannot be computed recursively, but instead need to optimize, e.g. using gradient descent (see Sec. 4.7). Unlike in the standard limit (Eq. [T]. it is no longer possible to guarantee uniqueness of the global maximum. We can nonetheless say that the posterior converges to a point distribution as long as the global maximum of L unique, (i.e. we can have any number of local maxima, as long as they all lie below the unique global maximum). We do expect the global maximum to be unique in most practical settings: we know the maximum is unique when the prior dominates (Eq. [T], Appendix [ we prove uniqueness for linear models, and in Appendix K we give a number of experiments in nonlinear models in which optimizing from very different initializations found the same global maximum, indicating uniqueness in practical settings.
Gr) is
4.4 THE EXACT DGP POSTERIOR OVER FEATURES MULTIVARIATE GAUSSIAN
Above, we noted that the DGP posterior over Gram matrices in the Bayesian representation learning limit is a point distribution, as long as the DKM objective has unique global maximum. Re- markably, in this setting, the corresponding posterior over features is multivariate Gaussian (see Appendix D for the full derivation).
P ( / O,)
(17)
While such a simple result might initially seem remarkable, it should not surprise us too much. In particular, the prior is Gaussian (Eq. []. In addition, in Fig. [ (middle), we saw that the next layer features depend on the current layer features only through the Gram matrices, which are just the raw second moment of the features, Eq. (2). Thus, in effect the likelihood only constrains the raw second moments of the features. Critically, that constraints connected to Gaussian distributions: under the MaxEnt framework, a Gaussian distribution arises by maximizing the entropy under constraints on the raw second moment of the features (laynes) 2003).
the raw second moment are tightly
5

---

Thus it is entirely plausible that Gaussian prior combined with likelihood that "constrains" the raw second moment would give rise to Gaussian posteriors (though of course this is not a proof; see Appendix D for the full derivation).
Finally, note that we appear to use G, or G2 in two separate senses: as 1 F in Eq. 2 and as N:
the posterior covariance the Bayesian representation learning limit (Eq. 17). In the infinite limit, these two uses consistent. In particular, consider the value of G. defined by Eq. (2) under the posterior,
(18) The second equality arises by noticing that we are computing the average of infinitely many terms, fx ((X)), which are IID under the true posterior (Eq. 17). so we can apply the law of large numbers, and the final expectation arises by computing moments under Eq. (7
G
=
lim N. E = ((( [{())) =
N-oo
4.5 THE DKM ORJECTIVE GIVES INTUITION FOR REPRESENTATION LEARNING
The form for the DKM objective in Eq. 16 gives a strong intuition for how representation learning occurs in deep networks. In particular, the likelihood, log (Y L,,,, encourages the model find representation giving good performance on the training data. At the same time, the KL--divergence terms keep the posterior over features, N (0, Gt), (Eq. 17p close to the prior A (0. K (G,-1)) (Eq. Tap. This encourages the optimized representations, G2 to lie close to their value under the standard infinite-width limit, K K(G,11). We could use any form for the likelihood including classi- fication and regression, but to understand how the likelihood interacts with the other KL-divergence terms, is easiest to consider regression (Eq. S5). as this oig-liielihood can also be written as a KL-divergence,
log logP -PL+1 DKL (N (O, GL+1)||N (O=K + const (19)
Thus, the likelihood encourages K (Gt) I be close to the covariance of the data, GL+1 I YY",, while the DGP prior terms encourage all G, to lie close to K(G,-1). In combi- 0z-11 nation, we would expect the optimal Gram matrices to "interpolate" between the input kernel, Go = = _ and the output kernel, GL+1.
To make the notion of interpolation explicit. we consider o² = 0 with linear kernel, K(Gx-1)
G 1. so named because it corresponds to linear neural network layer. With this kernel and with all
V/ v, there is an analytic solution for the (unique) optimum of the DKM objective ( (Appendix J.1
G = Go (G (G) e/(L+1) (G++1)(((+)))
(20)
which explicitly geometrically interpolates between Go and G1+11 Of course, this discussion was primarily for DGPs, but the exact same intuitions hold for BNNs, in that maximizing the DKM
objective finds sequence of Gram matrices, Gi...,, Gj, that interpolate between the input kernel, Go and the output kernel, G GL+11 The only difference is in details of P (Ge(G. and specifically as slight differences in the KL-divergence terms (see below).
4.6 THE DKM OBJECTIVE MIRRORS REPRESENTATION LEARNING IN FINITE NETWORKS
Here, we confirm that the optimizing DKM objective for an infinite network matches doing inference in wide but finite-width networks using Langevin sampling (see Appendix F for details).
We began by looking DGPs, and confirming that the posterior marginals are Gaussian (Eq. 17 Fig. Bab). Then, we confirmed that the representations match closely for infinite-wine wiitt DKMs (Fig. top and bottom rows) and finite-width DGPs (Fig. [ middle two rows), both initialization (Fig. top two rows) and after training to convergence (Fig. 2 bottom two rows). Note that the first column, Kn is a squared exponential kernel applied to the input data, and G3 = yy" is the output Gram matrix (in this case, there is only one output feature).
To confirm that the match improves as the DGP gets wider, we considered the RMSE between elements of the Gram matrices for networks of different widths (x-axis) for different UCI datasets (columns) and different numbers of layers (top row is one-layer, bottom row is two-layers; Fig. 3b). In most cases, we found a good match long as the width was at least 128, which is around the width of typical fully connected neural network, but is a little larger than typical DGP widths (e.g. Damianou & Lawrence 2013 Salimbeni & Deisenroth 2017

---

G1
K(G:)
G2
K(G;)
Init DKM
K(Go)
Init
Index
trained DKM trained DGP
50
1
50
1
50
index
index
index
index
index
index
Figure 2: A two hidden layer DGP with 1024 units per hidden layer and DKM with squared ex- ponential kernels match closely. The data was the first 50 datapoints of the yacht dataset. The first column, Kg is fixed squared exponential kernel applied the inputs, and the last column, T is the fixed output Gram matrix. The first row is the DKM initialization at the prior
G3 = yy Gram matrices and kernels, and the second row is the DGP, which is initialized by sampling from the prior. As expected, the finite width DGP prior closely matches the infinite-width DKM initial- ization, which corresponds the standard infinite width limit. The third row the Gram matrices and kernels for the trained DGP, which has changed dramatically relative to its initialization (second row) in order to better fit the data. The fourth row is the Gram matrices and kernels for the optimized DKM, which closely matches those for the trained DGP.
4.7 THE SPARSE DEEP KERNEL MACHINE AS A DEEP GENERALISATION OF KERNEL METHOOSS
DGPs in the representation learning limit constitute deep generalisation of kernel methods, with very flexible learned kernel, which we call the deep kernel machine (DKM; which was introduced earlier just in the context of the objective). Here, we design sparse DKM, inspired by sparse methods for DGPs (Damianou & Lawrence] 2013 Salimbeni & Deisenroth 2017) (Appendix D. The sparse DKM scales linearly in the number of datapoints, P. opposed to cubic scaling of the plain DKM (similar to the cubic scaling in most naive kernel methods).
We compared DKMs (Eq, I6 and MAP over features (Sec. Ep for DGPs. In addition, we considered a baseline, which was a standard, shallow kernel method mirroring the structure of the deep kernel machine but where the only flexibility comes from the hyperparameters. Formally, this model can be obtained by setting, G K(G and is denoted "Kernel Hyper" in Table We applied these methods to UCI datasets (Gal Ghahramani] (0166) using two hidden layer architecture, with kernel inspired by DGP sxip-connections, K(G,) w/ G + w Keeupp G.).. Here, wf, wh and are hyperparameters, and Kegesp(Gr is a cwwuurdd-xxpoenntii kernel.
We used 300 inducing points fixed to random subset of the training data and not optimised during training. We used the Adam optimizer with a learning rate of 0.001, full-batch gradients and 5000 iterations for smaller datasets and 1000 iterations for larger datasets (kin8nm, naval and protein).
We found that the deep kernel machine objective gave better performance than MAP, or the hyper- parameter optimization baseline (Tab. [.. Note that these numbers are not directly comparable to those in the deep GP literature (Salimbeni & Deisenroth] (017). as deep GPs have a full posterior so offer excellent protection against overfitting, while DKMs give only point estimate.

---

0.6
02
densky
Imput
0.3
1
0.0
0.0
-5
feature
feature
Input
inpur
bos ton
ya
concrete
energy
RMREE
1.0
0.5
width
wieth
width
width
Figure 3: Wide DGP posteriors converge to the DKM. Here, trained DGPs with Langevin sampling (see Appendix F). and compared to trained DKM. a Marginal distribution over features for input datapoint for two-layer DGP trained on subset of yacht. We used width of Ni.... = 1024 and VI... = 5 in all plots to ensure that the data had a strong effect on the learned representations. The marginals (blue histogram) are very close to Gaussian (the red line shows the closest fitted Gaussian). Remember that the true posterior over features is IID (Eq, 31 so each column aggregates the distribution over features (and over 10 parallel chains with 100 samples from each chain) for single input datapoint. The 2D marginal distributions for the same DGP for two input points (horizontal and vertical axes). e Flement-wise RMSE (normalized Frobenius distance) between Gram matrices from trained DKM compared to trained DGPs of increasing width. The DGP Gram matrices converge to the DKM solution as width becomes larger.
Table 1: RMSE for inducing point methods. (Equal) best methods displayed in bold. Error bars give two stderrs for a paired tests, which uses differences in performance between that method and best method, (so there meaningful error bars the best performing method itself). The MAP objective was numerically unstable and thus did not run to completion on the boston dataset.
dataset
P
Kernel Hyper
MAP
C
boston
506
4.41 + 0.31
4.35 + 0.51
concrete
1,030
5.38 = 0.098
5.60 1:00=00 0.15
5.10
energy
768
0.83 = 0.076
0.73 + 0.049
0.47
kin8nm
8.192
(7.3 1.06)-10-2
0.0)
6.6-10 2
naval
11,934
(6.4 0.6)
0.)
4.6-10-4
power 9,568
3.81 0.091
3.73 0.14
3.58
protein
45,730
4.21 4.21000 0.029
0.033 4.30 + 0.033
4.10
wine
1,599
0.68 = 0.0084
0.66 + 0.0067
0.64
yacht 308
0.94 + 0.058
1.14 0.077
0.58
5 CONCLUSION
We introduced the Bayesian representation learning limit, a new infinite-wide limit for BNNs and DGPs that retains representation learning. Representation learning in this limit is described by the intuitive DKM objective, which is composed of a log-likelihood describing performance on the task (e.g. classification or regression) and sum of KL-divergences keeping representations at every layer close to those under the infinite-width prior. For DGPs, the exact posteriors are IID across features and multivariate Gaussian, with covariances given by optimizing the DKM objective. Empirically, we found that the distribution over features and representations matched those in wide by finite DGPs. We argued that DGPs in the Bayesian representation learning limit form
8

---

class of practical deep kemne method: DKMs. We introduce sparse DKMs, which scale linearly in the number of datapoints. Finally, we give the extension for BNNs where the exact posteriors are intractable must be approximated.
REFERENCES
Laurence Aitchison. Why bigger is not always better: finite and infinite neural networks. In International Conference on Machine Learning, pp. 156-164. PMLR, 2020.
Laurence Aitchison, Adam Yang, and Sebastian W Ober. Deep kernel processes. arXiv preprint arXiv:2010.01590, 2020.
Joseph M Antognini. Finite size corrections for neural network gaussian processes. arXiv preprint arXiv:/908. 10030, 2019.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions pattern analysis and machine intelligence, 35(8):1798-1828, 2013.
David Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: review for statisti- cians. Journal of the American statistical Association, 112(518):859 (11188888987777 877, 2017.
T Tony Cai and Linjun Zhang. High-dimensional gaussian copula regression: Adaptive estimation and statistical inference. Statistica Sinica, pp. 96399693999 993-993,220 2018.
Youngmin Cho and Lawrence K. Saul. Kernel methods for deep learning. NIPS, pp. 342-350.
Curran Associates, Inc., 2009.
Andreas Damianou and Neil D Lawrence. Deep gaussian processes. In Artificiaa intelligence and statistics, pp. 207-215. PMLR, 2013.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Ad- vances in neural information processing systems, 32, 2019.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv:1909, 1304, 2019.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model
uncertainty in deep learning- In Proceedings of the 33rd International Conference Machine Learning (ICMI.-16), 2016.
Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional net- works shallow gaussian processes. arXiv preprint
005880 2018.
Arjun K Gupta and Daya K Nagar. Matrix variate distributions. Chapman and Hall/CRC, 2018.
James Halverson, Anindita Maiti, and Keegan Stoner. Neural networks and quantum field theory. Machine Learning: Science and Technology. 2(3):035002, 2021.
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint 19XX::19990099909 2019.
Thomas Hofmann, Bernhard Schölkopf, and Alexander J Smola. Kernel methods in machine learn- ing. The annals of statistics, pp. 1171-1220, 2008
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
Roger A Horn, Roger A Horn, and Charles R Johnson. Topics in matrix analysis. Cambridge university press, 1994.
Jiri Hron, Yasaman Bahri, Roman Novak, Jeffrey Pennington, and Jascha Sohl-Dickstein. Exact posterior distributions of wide bayesian neural networks. arXiv preprint arXiv:2006. 10541, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemne:: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.

---

Edwin T Jaynes. Probability theory: The logic of science. Cambridge university press, 2003.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114, 2013.
arXiv preprint
Alex Krizhevsky, lla Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.0011 2017.
Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The back- propagating renormalization group. arXiv preprint arXiv:2012.040300 2020.
Alexander G de Matthews, Mark Rowland, Jiri Hron, Richard Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv:/804 1804.1 1/271, 2018.
Song Mei, Andrea Montanari, and Than-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 1566))) 2018.
Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects in finite cnns. arXiv preprint
04000,0000,. 2021.
Gadi Naveh, Oded Ben-David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite networks trained with noisy gradients. arXiv preprint arXiv:2004.01 1990 2020.
Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29-53. Springer, 1996.
Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multi- layer neural networks. arXiv preprint arXiv:2001.11443, 2020.
Quynh Nguyen. On connected sublevel sets in deep learning. In International Conference on Machine Learning, pp. 4790-4799. PMLR, 2019.
Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. arXiv preprint arXiv: 18/0.05 2018.
Geoff Pleiss and John P Cunningham. The limitations of large width neural networks: A deep gaussian process perspective. Advances in Neural Information Processing Systems, 34, 2021.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint wwXiv:2106.10166 2021.
Hugh Salimbeni and Mare Deisenroth. Doubly stochastic variational inference for deep gaussian processes. arXiv preprint 1705vv7705.08000 2017.
10

---

Inbar Seroussi and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. arXiv preprint arXiv:2112.15383, 2021.
John Shawe-Taye Taylor and Nello Cristianini. Kernel methods for pattern analysis. Cambridge univer- sity press, 2004.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: central limit theorem. Stochastic Processes and their Applications, 130(3):1800-1822 2020a.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):25-752, 2020b.
Alex J Smola and Bernhard Schölkopf. Learning with kernels. MIT Press, 1998.
Christopher Williams. Computing with infinite networks. Advances in neural information process- ing systems, 9, 1996.
Chieh Wu, Aria Masoomi, Arthur Gretton, and Jennifer Dy. Deep layer-wise networks have closed- form weights, arXiv preprint arXiv:2202.01220, 2022.
Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pp. 165-192. PMLR, 2020.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
Greg Yang and Edward J Hu. Feature learning in inftnite-width neural networks. arXiv preprint arXiv:2011. 14522, 2020.
Jacob Zavatone- Veth and Cengiz Pehlevan. Exact marginal prior distributions of finite bayesian neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Jacob A Zavatone Veth, Abdulkadir Canatar, and Cengiz Pehlevan. Asymptotics of representation learning in finite bayesian neural networks. arXiv preprint arXiv:2106. 00651, 2021.

---

A
BAYESIAN NEURAL NETWORK EXTENSION
Consider a neural network of the form,
F XWo
(21a)
F2 W for 12
(21b)
W N 0, (0:
Wº. N (0, m
(21c)
where Wo R We E NNN and LL+1 N2 xM2 are weight matrices with inde- pendent Gaussian priors and ø is the usually pointwise nonlinearity.
In principle, we could integrate out the distribution over We to find (F1)
P ((F(Fe)1) = dW.P (W.) ((F-1)
(22)
where ō is the Dirac delta. In practice, it is much easier to note that conditioned on Ff-1- the random variables interest, Fz are linear combination Gaussian distributed random variables, We. Thus, F, are themselves Gaussian, and this Gaussian is completely characterised by its mean and variance. We begin by writing the feature vectors, f terms of weight vectors, w
T (F
(23)
As the prior over weight vectors is IID, the prior over features, conditioned on Fc-1), is also IID,
Nt P P(W) II
N2
(w))
= II N (w1:0, N A=1
(24)
A=1
Nr
P F = II P (f))
(25)
A=1
The mean of f5 conditioned on Fe-1 is 0,
E F- = E [d(FF-1)(((BP. = (f B( (F(F=11) [w = 0. (26)
The covariance of f conditioned on Fe-1 is,
F (ft)
F
E (6 (F(FF1) (6(Fe-1) [F-2
= = 1) w/ (w ) o² (Fe-1)
1 Ni_ o F (F
(27)
This and variance imply that Eq. 1 captures the BNN prior, as long choose KBNN and GBNN(.) such that,
No KBNN KENN(GENN (F (Fz-1)) = Ne.1 Ax=1 o(F (tT
(28)
Specifically, we choose the kernel function, K-NN(-) to be the identity function, and GBNN(-) to be the same outer product as in the main text for DGPs (Eq. nonlinearity,
(2), except where we have applied the NN
KNNN (aaNN(GG-1--) G
=G;11
(29)
GuNn(Fa-1) EN ((( ( (P
(30)
This form retains the average-outer-producc form for GBNN(.), which important for our deriva- tions.
Now, Eq. (16) only gave the DKM objective for DGPs. To get more general form, we need to consider the implied posteriors over features. This posterior is IID over features (Appendix D.I. and for DGPs, multivariate Gaussian (Appendix D.2.
P (F G+) = -I P G Go) for DCPPS
IN O.G.)
(31)
12

---

G.
Inte
BNN
G3
VOKM (flow)
50
index 50
VOKM (MvC)
index
index
index
index
0.5
DNN
0.0
feature
feature
feature
feature
Figure 4: The variational DKM closely matches the BNN true posterior obtained with Langevin sampling. Comparison of Gram matrices. The first two rows show Gram matrices for BNN, with the first row being a random initialization, and the second row being the posterior. The last two rows show the Gram matrices from variational DKMs with flow approximate posterior (third row) and multivariate Gaussian approximate posterior (fourth row). In optimizing the variational DKM, we used Eq. (34) with 26 Monte-Carlo samples. The Gram matrices for the flow posterior (third row) closely match those from the BNN posterior (second row), while those for multivariate Gaussian approximate posterior (fourth row) do not match. Marginal distributions over features at each layer for one input datapoint estimated using kernel density estimation. The note that the BNN (blue line) marginals are Gaussinin,,, but the variational DKM with flow posterior (red line) is capable of capturing this non-Gaussianity.
Now,
that Eq. (16) is specific example of general expression. In particular, note that the distribution on the left of the KL-divergence in Eq. (16) is the DGP posterior over features (Eq. 31.. Thus, the DKM objective alternatively be written,
C(G GL) log (Y( El DEL (P (f G.)l (O, ((G)))) (32)
and this form holds for both BNNs and DGPs (Appendix D.3 As in DGPs, the log-posterior is N times L(G),.... Gr,) (Eq. 16). so as taken to infinity, the posterior for all models becomes
point distribution (Eq. 12 G Gz) has unique global maximum.
In practice, the true posteriors required to evaluate Eq. (3) are intractable for BNNs, raising the question of how to develop accurate approximations for BNNs. We develop a variational DKM (vDKM) by taking inspiration from variational inference (Jordan et al. 1999; Blei et al. 2017) (Appendix D.4). Of course, variational inference is usually impossible because is impossible to work with infinitely large latent variables. Our key insight is that as the true posterior factorises across features (Appendix D.I.) work with the approximate posterior
infinite width models,
13

---

only single feature vector, Qg, (f)), where dg are the parameters and f E R is finite. This approach allows us to define a vDKM objective, which bounds the true DKM objective,
((G)( ())0))) 2
(33)
Lv 0r.) log P PYY)G((0))) - Eé=re DKI. (Qo, (O)
with equality when the approximate posteriors, Qg, (EX). equal the true posteriors, Gf G.). The only subtlety here is that it is practically difficult to design flexible ap- proximate posteriors Qo, (f)) where we explicitly specify and optimize the Gram matrices. Instead we optimize general approximate posterior parameters, 0, and compute the implied Gram matrices,
(ot)
1. lim EN EX (FX) = Fo., (13) (O(P))
(34)
Nr No
where f are sampled from Qo, (f)), and the second equality arises from the law of large numbers. We can compute the Gram matrix analytically in simple cases (such as a multivariate Gaussian), but in general we can always estimate the Gram matrix using Monte-Carlo estimate of Eq. (34).
Finally, we checked that the vDKM objective closely matched the posterior under neural networks This bit more involved, as the marginal distributions over features no longer Gaussian (Fig. 4p). To capture these non-Gaussian marginals, we used a simple normalizing flow. In par- ticular, we first sampled ž N (Hz Σe) from a multivariate Gaussian with a learned mean, pa, and covariance, Σe then we obtained features, f f(z)). by passing z through f. a learned pointwise function parameterised as in neural spline flow (Durkan et al., (201)). The resulting distribution is a -ige-diinnsional Gaussian copula (e.g. Cai & Zhang 2018). As shown in Fig. 4 vDKM with multivariate Gaussian (MvG) approximate posterior cannot match the Gram matrices learned by BNN (Fig. D. while vDKM with flow able to capture the non-Gaussian marginals (Fig. 4b) and thus match the learned Gram matrices with BNN.
B GENERAL LIKELIHOODS THAT DEPEND ONLY ON GRAM MATRICES
We consider likelihoods which depend only on the top-layer Gram matrix, G.,
P PY(G)) =
(FL+1 Gt)
(35)
where,
N1
P P(FL-i)LL) = II N "
(36)
This family of likelihoods captures regression,
P (yxlr(+)) N
(37)
(which is equivalent to the model used in the main text Eq. lb) and e.g. classification,
P Fy|F)= Categorical (y: softmax (F) (F)+1)),
(38)
amoung many others.
C WEAK CONVERGENCE
Here, we give formal argument for weak convergence of the DGP posterior over Gram matrices to a point distribution in the limit as + oo,
PN (G (,I,
Gz
$ II 8(G G
(39)
E=1
14

---

where have included N in the subscript of the probability distribution reminder that this distribution depends on the width. By the Portmanteau theorem, weak convergence is established all expectations of bounded continuous functions, converge
lim Epx(GG....G..II... G./G..........). [f(G...... - f ...)
(40)
N
To show this in a reasonably general setting (which the DGP posterior is special case of), consider an unnormalized probability density of the form h(g)e N((o), and compute the moment
E [[ Ig dg (())
(41)
Sg dg N(g))
where g (GI,..., GL) is all L positive semi-definite matrices, GG. Thus, g E 9, where g is convex set.
We consider the superlevel set A(A) = {g|[(g) 2 L(g") A}. where g" is the unique global optimum. We select out a small region, A(A), surrounding the global maximum, and compute the integral as,
Jaaa dg f(g)h(g)) Sorvla) dg ((g) JALA) dg h(g)eNZLo) + Sorvray dg h(g)eN2l()
E [f(
(42)
And divide the numerator and denominator by SA(A) dg h(g))
Ja(ay f(a)M(a)) Jaiay dg NCG)
Jorwa ((a)
Jacay dg Mg)))
E [f()))
(43)
JarMay dg (CG) Jacay dg NCg))
Now, we deal with each term separately. The ratio in the denominator can be lower- bounded by zero, and upper bounded by considering smaller superlevel set, A(A/2), in the denominator,
SerALa) dg ((g)
Sorvia) dg ((g))
5 SA(D) dg ((g)))
S
JALA/D) dg ((g))-N
N(C(n SQVA(A) dg h(g)
-(C()))-112)) SATA/2) dg h(g) -NA/2
SorALa) dg h((g)
(44)
JAJA/2) dg h((g)
The upper bound converges to zero (as h(g) independent of N), and therefore by the sandwich theorem the ratio of interest also tends to zero.
The second ratio in the numerator can be rewritten as,
SorA(s) dg f( JorA(a) dg f(ggh(g)e (orA(s) dg NC())
SorAis) dg h(g)e Sa(A) dg NC(o)
(45)
JA(A) dg NCG)
The first term here is expectation of bounded function, (()), so is bounded, while second term converges to zero in the limit (by the previous result).
Finally, we consider the first ratio in the numerator,
SMLS dg N()
(46)
JACAY dg h(g)))V
which can be understood as an expectation over f(g) in the region A(A). As f is continuous, for any >> O, we can find a > such that for all with Ig* - gi 5, we have
- f(g) <((g)) +c.
(47)
15

---

Further, because the continuous function, L(g), has unique global optimum, 9", for every 8 >> 0 we are always able to find a A > 0 such that all points E A(A) are within ó gf i.. Ig" g < 6. always able to find 8 such that Eq. 47 holds for all ! with lg* gl < 6, and given a 6 we are always able to find a A such that all g E A(A) have Ig" gl < 8. Hence for every > 0 we can find A > 0 such that Eq. 47 holds for all E A(A). Choosing the appropriate e-dependent A and substituting Eq. 47 into Eq. 46 also bounds the error the expectation,
Thus combining the previous two facts, given an
f(g') cc SALA) dg flg)h(g))
JA(A) dg h(g)e f(gg)
(48)
Now, we use the results in Eq. (44). Eq. 45 and Eq. 48) to take the limit of Eq. 43 (we can compose these limits by the algebraic limit theorem as all the individual limits exist and are finite),
f(g") -<<
lim B[(()) f(g') +
(49)
N No000
And as this holds for any e, we have,
f(g") lim [[f(g)..
(50)
N
This result is applicable to the DGP posterior Gram matrices, that posterior can be written
as,
N (GI, Gr Y) x V(l))
(51)
where ((g) is the usual DKM objective,
L(y) = [(G Gr.)
(52)
and h(g) is the remaining terms in the log- peeposttrror which do not depend on N,
h(g)) = exp
P+1 Σ log Ge
(53)
(this requires S so that G. is full-rank).
D GENERAL MODELS IN THE BAYESIAN REPRESENTATION LEARNING LIMIT
Overall, our goal is to compute the integral in Eq. (6) in the limit as N - xo. While the integral use variational inference to reason about its properties. In particular, we can bound the integral using the ELBO,
is intractable for general models such as BNNs,
log P (Ge|Ge-1) 2 ELBO; EQ(F,) [log P (G, + log P (Fr/Ge-1) log Q (Fr))
(54)
Note that Q (F,) here is different from Qa, (f)) in the main text, both because the approximate posterior here, Q (Fe) is over all features jointly, Fe, whereas the approximate posterior in the main text is only over single feature, f3. and because in the main text, chose specific family of distribution with parameters 0z, while here we leave the approximate posterior, Q (Fz) completely unconstrained, so that it has the flexibility to capture the true posterior. Indeed, if the optimal
approximate posterior is equal to the true posterior, Q" (F₂) = P (F, G Gy), then the bound tight, so we get log P (Gr|Ge-1) inference to characterise the optimal approximate which is equal to the true posterior Q" (Fz) [Fe/Ge-1 Ge) and use the corresponding ELBO to obtain log P GG-1)
ELBO; Our overall strategy thus use variational
D.1 CHARACTERISING EXACT BNN POSTERIORS
Remember that if the approximate posterior family, Q (Fr) is flexible enough capture the true posterior P (F G Ge). then the Q" (Fz) that optimizes the ELBO is indeed the true poste- rior. the bound is tight, so the ELBO is equal to log P (Ge/Gr-) (Jrdan al. 1999 Blei et al.) 2017]. Thus, we are careful to ensure that our approximate posterior family captures the true pos- terior, by ensuring that we only impose constraints on (Fr) that must hold for the true posterior,
16

---

P ( Gē). In particular, note that P (G in Eq. (b)) constrains the true posterior to give non-zero mass only to F& that satisfy G. = N; o (F F (Fe). However, this constraint is difficult to handle. We therefore consider an alternative, weaker constraint on expectations, which holds for the true posterior (the first equality below) because Eq. (5b) constrains G. Ne (F)) (Fe), and impose the same constraint the approximate posterior,
G. = (..
Nr (FFe) (F+) [F(F))
I (F)))
(F+)
(55)
Now, we can solve for the optimal Q (F₂) with this constraint on the expectation. In particular, the Lagrangian is obtained by taking the ELBO (Eq, 54), dropping the log P (G& F2) term representing
the equality constraint (that G. Nr ( e) (Fe)) and including Lagrange multipliers for the
expectation constraint, A, (Eq. 55) and the constraint that the distribution must normalize 1, A,
L L= dF (Fz) Plog (Fe/Ge))) - log Q (Fz)
+ + Tr A Gi - (GG-(
10.//FF. )) +A +A(1-1/ dF /FFFFF/FFFE Q(F))
(56)
Differentiating wrt Q (Fr), and solving for the optimal approximate posterior, Q" (Fz),
aL
0 = aQ (F (Fx)))
(57)
0 = (log P (Fr(Gr-1) - log Q" (Ft))
(FF))))"((.)) -
(58)
Solving for log Q" (Fz),
log Q" (Fe) = log P Gr _ Tr (Ad(Fs)o"(Fe)) + const.
(59)
Using the cyclic property of the trace,
log Q" (F&) = log (Fz|Gr-1) - Tr (6" (F.)Ao(F.) + const.
(60)
Thus, log (Fe) can be written
features,
log Q" (Fr) [log P (E - (C))A((()) + const ΣX log Q (f))
(61)
so, the optimal approximate posterior IID over features,
(F.) = Ih Q" (r))
(62)
Remember that this approximate posterior only constrained in expectation, and that this straint held for the true posterior (5). Thus, we might think that this optimal approximate pos- terior would be equal to the true posterior. However, remember that the true posterior had tighter equality constraint, that G No (FF) (Fe), while so far we have only imposed a weaker con- straint in expectation (Eq. (5).. We thus need to check that our optimal approximate posterior does indeed satisfy the equality constraint in the limit as N oo. This be shown using the law of large numbers, as r5 are IID under the optimal approximate posterior, and by using Eq. (55) for the final equality,
N o
lim No N. OFFF)O
(F:)
lim No
(F)) Bo(E)) [o(E)) (C))) G
(63)
Thus, the optimal approximate posterior does meet the constraint in the limit as N + oc, so that limit, the true posterior, like the optimal approximate posterior is IID across features.
P (F Gz) Q* (Fx)
- Q" (r)) Ih P Ge Gz).
(64)
17

---

D.2 EXACTLY MULTIVARIATE GAUSSIAN DGP POSTERIORS
For DGPs, we have ((x) = [ so the optimal approximate posterior is Gaussian,
log QDGP (fx) log Dor (f((G(-1) (f Af const
(65)
= - (F)) (A K (G,-1) f + const
(66)
log (FX: 0, ( K (G--)) 7))
(67)
As the approximate posterior and true posterior are IID, the constraint in Eq. 55
becomes,
G =
[EX(E)) = on (Et) [B(()))] = (A (G-)-
(68)
As the Lagrange multipliers are unconstrained, we can always set them such that this constraint holds. that case both the optimal approximate posterior and the true posterior become.
PDGP (f)GG-1Gr) -DOP (rx) = =N 0, Ge),
(69)
as required.
D.3 GENERAL FORM FOR THE CONDITIONAL DISTRIBUTION OVER GRAM MATRICES
Now that we have shown that the true posterior, P (Fa - G:) factorises, we can obtain simple form for log P (Ge(Ge-1). In particular, log P (Ge|Gr-) is equal to the ELBO if we use the true posterior in place of the approximate posterior,
P (F Ge-1) ((Fe Ge-1 G
lim 1 log P (Ge|Ge-1) No-
lim # GF.F.GG..,.G.)) log P (G./Fe) log Ni
(70)
Under the posterior, the constraint represented by log (G is satisfied, so in the limit we can include that term in a constant,
= lim 1 No N BP(F./G. 1.G.) log P GL-11 Ge) ]
(F const. (71)
lim 1 No- N log P
Now, the fact that the prior, P (Fr/Gr-1) and posterior, P Ge Gt), are IID across features,
P (
lim N log (G. = ve (PE No
log P G G+)
const
(72)
and this expectation is KL-divergence,
lim N log P (G,
DKL (P Ge) Gr/): + const,
(73)
No-000
which gives Eq. 32 when we combine with Eq.
D.4 PARAMETRIC APPROXIMATE POSTERIORS
Eq. 64 represents considerable simplification, as we now need to consider only single feature, K. rather than the joint distribution over all features, Fr. However, in the general case, it still not possible to compute Eq. (64) because the true posterior over a single feature is still not tractable. Fol- lowing the true posteriors derived the previous section, we could chose a parametric approximate posterior that factorises across features,
Qo (F...... F2) - Qo, (f))
(74)
Remember that we optimize the approximate posterior parameters, 6, directly, and set the Gram matrices function of (Fq. 34). As before,
bound, log P (G,=G,(0)) Ge-) using the
18

---

ELBO, and the bound is tight when the approximate posterior equals the true posterior,
log P (Gr G(
(75)
(
-B (,(66)
log P (G.=G,(0e)]F))
F2) + log
(76)
P (Fa Gr-1, Gz= Gg()))
2 (Q,(F)) loge (G)=G()))))) Fe) + log
(
(77)
Q0 (F
Now, we can cancel the log P (G, (B:)) terms, as they represent a constraint that holds both under the true posterior, and under the approximate posterior,
P (FF/Gr,-)
P P(Fr(Ge-1)
G.
log G., G.=G,(0e))
RO,, (Fe)
log
Q0 (Fe)
(78)
Using the fact that the prior, posterior and approximate posterior are IID over features, we can write this inequality in terms of distributions over single feature, f3 and divide by No.
P GI-) log P Gei Gg
P Gf
E Br((,iG....GG....0)))
BQ,, (Ex)
log
Qa (fx)
(79)
Noting that both sides of this inequality are negative KL-divergences, we obtain,
DKL (P (fA Ge-1, G,=),(()) [|P GG-)) 2 DxL (Qo, (Ex) |P (f G(-1)),
(80)
which gives Eq. 33 in the main text.
E THEORETICAL SIMILARITIES IN REPRESENTATION LEARNING IN FINITE AND INFINITE NETWORKS
In the main text, considered probability densities of the Gram matrices, G1...., Gz. However, we can also consider probability densities of the features, Fi...., FL, for a DGP,
log P (F (F.(F.1))
N₂ log K (GDGP (Fr-1))) tr (F.K-! (Guue (Fe-1))FF Fz) + const.
(81)
We can rewrite the density such that it is still the density of features, F,, but it is expressed in terms of the DGP Gram matrix,
log P (Fr(Fr-1)
Nr 2 IK-'(GG-I)GA) N/ log [((G-
const.
(82)
Here, we have used the cyclic property of the trace to combine the Fz and F to form G2, and we have used the fact that our kernels can be written function of the Gram matrix. Overall, we can therefore write the posterior over features, (F1..... F, Y), in terms of only Gram matrices,
JG GL) N log P (F,. X, Y) log P (Y))))
Σ t=1
log P (Fr
N
(83)
substituting Eq. 82
JGG G1) log (Y|Gr) - - (log I(G + (K (G-I)
const. (84)
Thus, JGG...., G1) does not depend on N. and thus the Gram matrices that maximize I (G G1) are the same for any choice of N. The only restriction that we need N₂ 2 P. ensure that the Gram matrices are full-rank.
To confirm these results, we used Adam with a learning rate of 10 to optimize full-rank Gram matrices with Eq. (84) and to directly do MAP inference over features using Eq. (8I). As expected, as the number of features increased, the Gram matrix from MAP inference over features converged rapidly that expected using Eq. (84) (Fig.
19

---

boston
yacht
concrete
energy
0.20
0.15
one- -yee-ayer
RMSE
0.10
0.05
0.00
G:
0.15
G2
DW0- poo-layer
RMSE
0.05
2
21
2
29
2
24
2
2
2
2
2
2
2
2
2
2
2
2
2
2
width
width
width
width
Figure 5: RMSE of trained Gram matrices between one-hidden-layer (first row) and two-hidden- layer (second row) DGPs of various width trained by gradient descent and the corresponding MAP limit. Columns correspond to different datasets (trained on subset of 50 datapoints).
ADDITIONAL EXPERIMENTAL DETAILS
To optimize the analytic DKM objective for DGPs and the variational DKM objective for DGPs (Figs. [ we parameterised the Gram matrices (or covariances for the variational approximate square matrix, R. E PPP, with itself transposed, G, R, R and we used Adam with learning rate -3 to learn Rg, To Bayesian inference finite BNNs and DGPs, used Langevin sampling with 10 parallel chains, and a step size of 10-3, Note that certain senarios, Langevin sampling can be very slow, as the features have a Gaussian prior with covariance ((G+1) which has some very small and some larger eigenvalues, which makes sampling difficult. Instead, we reparameterised the model in terms of the standard Gaussian random variables, V. PxNN, We then wrote F, in terms of Vt.
posterior) as the product
F. Le-
(85)
Here, Le-1 the Cholesky of K(G,-1), so -L, This gives an equivalent distri- bution (Fyl Fr-1). Importantly, the prior on is IID standard Gaussian, sampling V, much faster. To ensure that the computational cost of these expensive simulations remained reasonable, we used subset of 50 datapoints from each dataset.
For the DKM objective for BNNs, we used Monte-Carlo to approximate the Gram matrices,
Go((r)
O (C))
(86)
with f drawn from the appropriate approximate posterior, and K 216 We can use the repa- rameterisation trick (Kingma & Welling 2013) Rezende al. 2014) to differentiate through these Monte-Carlo estimates.
G ADDITIONAL. COMPARISONS WITH FINITE- DGPs
Here, we give additional results supporting those Sec. 4.6 Fig. Fig. particular, we give the DGP representations learned by two-layer networks all UCI datasets (boston, concrete, energy, yacht), except those already given in the main text Fig.
20

---

G1
K(G:)
G2
K(G,)
init DKM
Init DGP
K(Go)
Index
trained DGP
50
50
50
index
index
trained DKM
50
50
50
Index
Index
Index
Index
Figure 6: One hidden layer DGP and DKM with squared exponential kernel trained subset of energy. First and second row: initializations of DGP and DKM. Third and fourth row: trained DGP (by Langevin sampling) and DKM Gram matrices and kernels.
G1
K(G
G2
K(Gg)
Init DKM
Init DGP
K(Go)
Index
trained DGP
50
50
1
50
Index
Index
trained DKM
50
50
50
Index
Index
Index
Index
Figure 7: One hidden layer DGP and DKM with squared exponential kernel trained subset of boston. First and second row: initializations of DGP and DKM. Third and fourth row: trained DGP (by Langevin sampling) and DKM Gram matrices and kernels.
21

---

G
K(G.)
G2
G2
Init
Init DGP
K(Go)
Index
trained DGP
50
50
50
Index
Index
trained DKM
50
50
index
Index
index
index
Figure 8: One hidden layer DGP and DKM with squared exponential kernel trained subset of concrete. First and second row: initializations of DGP and DKM. Third and fourth row: trained DGP (by Langevin sampling) and DKM Gram matrices and kernels.
22

---

H THE FLOW POSTERIOR IN A 2-LAYER BNN
Here, we give the 2-layer version (Fig. 9 of Fig. in the main text, which again shows a close match between the variational DKM with flow posterior, and the BNN true posterior.
Go
G3
0.75
Index
0.00
VoRK (fow))
50
-0.75
Index
VOKM (MVG)
50.
50
Index
Index
b
0.5
VDKM (low)
0.0
feature
feature
Figure 9: Two-layer ReLU BNN and variational DKM with flow. a Initialized (first row) and learned Gram matrices of a width 1024 BNN (second row), vDKM with flow (third row) and vDKM with multivariate Gaussian (fourth row) using 2 Monte-Carlo samples. The Gram matrices between BNN and vDKM (flow) match closely after training. (MvG). b Marginal PDF over features at each layer for one input datapoint using kernel density estimation. The marginal PDFs of BNN are non- Gaussian (blue curves), vDKM with flow is able to capture the non- non-Gaussianity and match closely with BNNs marginals (red curves).
23

---

I MULTIVARIATE GAUSSIAN APPROXIMATE POSTERIORS IN DEEPER NETWORKS
There is a body of theoretical work (e.g. (Scroussi & Ringel 2021 ) on BNNs that approximates BNN posteriors over features as Gaussian. While we have shown that this is a bad idea in general we can nonetheless ask whether there are circumstances where the idea might work well. In particular, we hypothesised that depth is an important factor. In particular, in shallow networks, in order to get G close to the required representation, we may need the posterior over F, to be quite different from the prior. In contrast, in deeper networks, might expect the posterior over F, to be closer to its (Gaussian) prior, and therefore we might Gaussian approximate posteriors to work better.
(Fig. 4 and D.
However, we cannot just make the network deeper, because as we do so, we apply the nonlinearity more times and dramatically alter the network's inductive biases. To resolve this issue, derive a leaky relu nonlinearity that allows (approximately) independent control over the inductive biases (or effective depth) and the actual depth (Appendix [.I.. Using these nonlinearities, we indeed find that very deep networks are reasonably well approximated by multivariate Gaussian approximate posteriors (Appendix 1.2).
1.1 LEAKY RELU NONLINEARITIES
Our goal is to find pointwise nonlinearity, o, such that (under the prior),
E Preer (1x (..)) (o(())) (C)) = (EgIG.)....) [relu(())eelu" (r0) + (1 (1--)
(87)
We will set = a/L, L, where a is the "effective" depth of the network and L is the real depth. These networks are designed such that their inductive biases in the infinite width limit similar to a standard relu network with depth a. Indeed, we would take this approach if we wanted well-defined infinite-depth DKM limit.
Without loss of generality, we consider a 2D case, where x and are zero-mean bivariate Gaussian,
. ( N (())
Exyy
0.
(88)
Σxyy
Eny
where T2x v) is the probability density for the joint distribution. Note that use scaled relu,
/2 x
for << x otherwise
relu(z) = 0
(89)
such that E [relu((x)) = Exx. Mirroring Eq.
87
we want the nonlinearity, o, to satisfy,
E [d(+))]= = EE relu 2 (x) (1 E) Exx
(90a)
E [(())))] = P relu 2 (v) + (1 - p)Eny = Eyy
(90b)
E [(())((v) = p [relu ( reeu((() + (1 - p) E)
(90c)
We hypothesise that this nonlinearity has the form,
(xx) a relu(x)) + ba.
(91)
We will write the relu sum of and |as/
relu(z) /2 (x |æ:),
(92)
because E [f y) = 0 for f(x, v) = alyl or f(r, v) = Joly. It turns out that we get zero expec- tation for all functions where f(-z. -y) = -f y), which holds for the two choices above. To show such functions have a zero expectation, we write out the integral explicitly,
E (xx 1 da L dy ((2,)))(z,
(93)
24

---

We split the domain of integration for y zero.
E [[] [ dr [ dy Tτ. v))((.)))
[ dx S dy ((z,))((z,)))
(94)
We substitute y' y and x' =- in the first integral.
[1((,,)) = { de' 1-- --''"" So dx f dy r(., (x,
0
pc
(95)
As the variables we integrate over are arbitrary we can relabel s/ as and x' as a, and we can then merge the integrals as their limits are the same,
E [f = [- dx L dy
00
-v) + x(x,
(96)
Under a zero-mean Gaussian, π(-I, -v) = x(I,))).
E [[r = |- dx (o dy ((, v) ((-x,-1) + [
(97)
Thus, if f(-x,-1) -2) --(( y). then the expectation of that function under bivariate zero-mean Gaussian distribution is zero.
Remember that overall goal to design nonlinearity, (Eq. 91 which satisfied Eq. We therefore compute the expectation,
20
E [o(x))(v)) EE [(a relu(x)) bzr) (a relu(y) + by)]
(98)
E
+ |z]) bz)
( (y (u/) + by
(99)
Using the fact that E [<lyy]
]
EE | under multivariate Gaussian,
EE
(x + |=|) 1 (v + (w|) + (/2bb + b²) xy
(100)
a
E frelu relu(y)] (2ab 8² E [ryy].
(101)
Thus, we can find the value of a by comparing with Eq. (90c).
p p= a
a Vp.
(102)
For b, things are a bit more involved,
1 /2a + b² = /2 b + b²
(103)
where we substitute for the value of a. This can be rearranged to form a quadratic equation b,
0 62 + / 2p b + (p 1),,
(104)
which can be solved,
= 7
1
--Vpp + 2p 4(p - 1)
- 2p
(105)
b b= 7
22 + /4- - 2p
(106)
b v
+ v - 2
(107)
Only the positive root is of interest.
b V 1 2 v
(108)
Thus, the nonlinearity is.
5(z)= vVruu((()) + (V1-
1/1-3 /B)=
(109)
25

---

BNN
density
4 layers
==1
1=2
4
32 layers
density
==
t=
= 24
== 32
-5
5
-5
5
-5
-5
feature
feature
feature
feature
Comparison of posterior feature marginal distributions between BNN of width 1024 (trained by Langevin sampling over features) and variational DKM with 2 Monte-Carlo samples, in 4-layer (row 1) and 32-layer (row 2) network. We give the BNN posterior features from Langevin sampling (blue histogarm) and the best fitting Gaussian (blue line), and compare against the variational DKM approximate posterior Gaussian distribution (red line).
Figure 10:
where we set p = a/L,, and remember we used the scaled relu in Eq. 89 Finally, we established these choices by considering only the cross term, E (
(oy)) We also need to check that the E (x)) and E (v)] terms are required (Eq. 90a and Eq. 90b). In particular,
E [x){ E relu(x) bar)
IE
(
(ar (x|) + bx) )
(110)
using E [x|>]]] tion,
x | is an odd function of æ, and the zero-mean Gaussian is an even distribu-
(x) = a² (x)] + (Vaab + 6²) 2ix
(111)
using Eq. (102) identify a² and Eq. 103 to identify 2ab + 12 E [6² = relu (x) (1 E)x:
(112)
required.
1.2 MULTIVARIATE GAUSSIAN DEEPER NETWORKS
In the main text, we show that more complex approximate posterior can match the distributions in these networks. Here, consider an alternative approach. In particular, we hypothesise that these distributions are strongly non- Gausssan because the networks are shallow, meaning that the posterior needs to far from prior in order to get top-layer kernel close G We could therefore make the posteriors closer to Gaussian by using leaky-relu nonlinearities (Appendix [1] with fixed effective depth (a = 2), but increasing real depth, L. particular, we use multivariate Gaussian approximate posteriors with learned means,
Q₂, (f)) N (f:; Pg, Σ.)
0: 0=:04 Ez)-
(113)
As expected, for depth 32 network, we have much more similar marginals (Fig. 10 top) and learned representations (Fig. top).
26

---

G.
layers
BMN
1000
Inder Index
G.
G16
G24
G
50
50
Index
Index
32 leyers
DNM
50
50
50
index
index
index
index
Figure 11: Comparison of Gram matrices between BNN of width 1024 (trained by Langevin sam- pling over features) and variational DKM, in 4-layer (row 1-3) and 32-layer networks (row 4-6). Initializations are shown in row and 4, trained BNN Gram matrices are shown in row and 5, and trained variational DKM Gram matrices are shown in row 3 and 6. As in Figure I0 the variational DKM is poor match to Langevin sampling in a BNN for a 4-layer network, but is very similar in 32 layer network.
27

---

J UNIMODALITY IN LINEAR DEEP KERNEL MACHINES
J.1 THEORY: UNIMODALITY WITH A LINEAR KERNEL AND SAME WIDTHS
Here, we show that the deep kernel machine objective unimodal for linear kernel. linear kernel simply returns the input Gram matrix,
K (G) G.
(114)
called a linear kernel, because arises in the neural network setting (Eq. 21 by choosing the nonlinearity, & to be the identity, in which case, F2- F (-1We-1- For linear kernel the objective becomes,
C(, Gr)= ue (log G Ge Tr GG)
(115)
where we have assumed there no output noise, o² Taking all My be equal. Pt (see Appendix J.2 for the general case),
C GL) = log Go GL++| X Tr (G
(116)
Note that Go and G fixed by the inputs and outputs respectively. Thus, find the mode, set the gradient wrt G....., G to zero,
0 dc aG
= ( G G
(117)
Thus, at the mode, the recursive relationship must hold,
T= G Ge G. GG
(118)
Thus, optimal Gram matrices are given by,
G =
Gz=GoT² GoT
(119)
and we can solve for T by noting,
GT T
(120)
Importantly, T is the product of two positive definite matrices, T G G.. so must have positive, real eigenvalues (but T does not have to be symmetric (Horn & Johnson 2012)). There is only one solution to Eq. (12)) with positive real eigenvalues (Horn et al. 1994). Intuitively, this can be seen using the eigendecomposition, Go
GLL+1 V-'DV, where D diagonal,
Tv 'VV' 1/(L+1) = V- ---DDDDDIAAA DD/(AIVV)
(121) Thus, finding T reduces to finding the (I. + 1)th root of each positive real number on the diagonal of D. While there are (I. 1) complex roots, there is only one positive real root, and so T and hence G....., Gr are uniquely specified. This contrasts with a deep linear neural network, which has infinitely many optimal settings for the weights.
Note that for the objective to be well-defined, we need K(G) to be full-rank. With standard kernels (such the squared exponential) this is always the case, even if the input Gram matrix is singular. However, a linear kernel will have a singular output if given a singular input, and with enough data X is always singular. To fix this, could e.g. define Go K Xx?) to be given by applying a positive definite kernel (such as a squared exponential) to 1 X This results in positive definite Go. as long as the input points are distinct.
points, Go vg
J.2 THEORY: UNIMODALITY WITH A LINEAR KERNEL AND ARBITRARY WIDTHS
In the main text we showed that the deep kernel machine is unimodal when all ig are equal. Here, we show that unimodality in linear DKMs also holds for all choices of vy. Recall the linear DKM objective in Eq. (15).
LGG Gz) = Σ
L+1 t=1
(log
Tr (G, G Gel
(122)
E f=l
L+l (log [Grl log GG--l - Tr Th G.))
(123)
28

---

To find the mode, again set the gradient wrt G. to zero,
"211 "
G Mt -
2212 G G G
(124)
2
for 8 = 1,,., L. Right multiplying by 2G, and rearranging,
Pe-1 G. G v G. (v++1 m)) I, for l = 1,,,, L.
(125)
Evaluating this expression for { = 1 and f = 2 gives,
". = 1/ G G₁ (V2 11)
(126)
G3 G3 = u² G2 + (V3 - v2) = vi G. G + (v3 - v1)
(127)
Recursing, get,
my G G. vi MI + (ve - v1
(128)
Critically, this form highlights constraints on G In particular, the right hand side, G G2 is the product of two positive definite matrices, so has positive cigenvalues (but may be non-symmetric (Horn & Johnson 2012)). Thus, all cigenvalues of 1/1 Go G. must be larger than 1/1 Me. and this holds true at all layers. This will become important later, as it rules out inadmissible solutions.
Given Gu and G1, compute any G using.
G.
II (G (GrL,Ge Ge) IT " #=1
II (ux G Grr Ge) ==1
(129)
t
II ve G Ge II Gv G1 + (ve I)
(130)
where the matrix products are ordered as I A = A ... Ar. Now, seek to solve for G using knowledge of GL+I. Computing Go GL+1,
[++1 II Pt t=1
G
G
II -
(131)
We write the eigendecomposition of 1 G G1 as,
14 G G VDV-!
(132)
Thus,
L+1 II pt
L+1 II t=1
G GL+1
(VDV-1 + (v (1)- VAV-1
(133)
F=1
where a diagonal matrix,
L+1 A
D (ve- 1)
(134)
Thus, identify V and A by performing eigendecomposition of the known matrix, (ITZtt) G Then, we can solve for D (and hence G₂) in terms of A and V. The diagonal elements of satisfy,
GL+1.
L+1 0 Au II
(Da + (ve (ve- v1)).
(135)
This is polynomial, and remembering the constraints from Eq. 128). we interested in solutions which satisfy,
1/3 Vnin S D
(136)
29

---

where,
P/mio = min (V......PL++I)
(137)
To reason about the number such solutions, we use Descartes' rule of signs, which states that the number of positive real roots is equal to or multiple of two less than the number of sign changes in the coefficients of the polynomial. Thus, if there is one sign change, there must be one positive real root. For instance, in the following polynomial,
= x + -1
(138)
the signs go as (+), so there is only one sign change, and there is one real root. use Descartes' rule of signs, work in terms of D'., which is constrained to be positive,
0 S D/ = Dú (v1 - Vmmin)
Du = D's + (v1 1/min)
(139)
Thus, the polynomial of interest (Eq. 135 becomes,
L+1
L+1
0= A + II (D,, - (m1 - Mmiin (v1 ve)) = AA + II (Di + (ve - v/min)) E=1
(140)
where 0 < V/- Vini as mmm defined to the smallest va (Eq. 137). Thus, the constant term, - Aa is negative, while all other terms, ......, (D,,.,,()))) (D())++1 in the polynomial have positive coefficients. Thus, there is only one sign change, which proves the existence of only one valid real root, as required.
K UNIMODALITY EXPERIMENTS WITH NONLINEAR KERNELS
For the posterior over Gram matrices converge to point distribution, we need the DKM objective C(G,...., G2) ) to have one unique global optimum. As noted above, this is guaranteed when the prior dominates (Eq. []]. and for linear models (Appendix J.. While we believe that it might be possible to construct counter examples, in practice we expect single global optimum in most practical settings. To confirm this expectation, we did number of experiments, starting with many different random initializations of a deep kernel machine and optimizing using gradient descent (Appendix K). In all cases tested, the optimizers converged the same maximum.
We parameterise Gram matrices G, = v.VV with V, E PxPx being trainable parameters. To make initializations with different seeds sufficiently separated while ensuring stability initialize Ge from broad distribution that depends on k(G:-)). Specifically, we first take the Cholesky decomposition K(G:-1) L then set V. E where each entry of E RPx is independently sampled from standard Gaussian, and D, is diagonal scaling matrix with each entry sampled i.i.d. from inverse- Gaamma distribution. The variance of the inverse-Gamma distribution is fixed to 100, and the mean is drawn from a uniform distribution [[0.5,3] for each seed. Since for any random variable x ~ mn--Gama (o:, 3), E(x) = - and V(x)) = (a- 2)
once we fix the mean and variance we can compute and [ as
E(x)² + 2, V(x)
(141)
E(z)(a 1).
(142)
We set P4- 5, and use the Adam optimizer (Kingma & Ba 2014) with learning rate 0.001 to optimize parameters Ve described above. We fixed all model hyperparameters to ensure that any multimodality could emerge only from the underlying deep kernel machine. As we did not use inducing points, we were forced to consider only the smaller UCI datasets (yacht, boston, energy and concrete). For the deep kernel machine objective, all Gram matrices converge rapidly to the same solution, as measured by RMSE (Fig. 12). Critically, we did find multiple modes for the MAP objective (Fig 13), indicating that experiments are indeed powerful enough to find multiple modes (though of course they cannot be guaranteed to find them). Finally, note that the Gram matrices took a surprisingly long time to converge: this was largely due to the high degree of diversity in the initializations; convergence was much faster if we initialised deterministically from the prior.
30

---

yacht
boston
energy
concrete
objective
-490
-850
-110
1130
500
-860
1110
1740
1.0
RMSE
0.5
50000
100000
50000 Iters
100000
50000
100000
0
50000
100000
iters
Iters
iters
RMSE
0.5
0.0
-500
490
-400
860
850
840
-1110
-1100
-1090
-1740
1730
-1720
objective
objective
objective
objective
Figure 12: One-layer DKMs with squared exponential kernel trained on full UCI datasets (through columns) converges the solution, despite very different initializations by applying stochas- tic diagonal scalings described in Appendix F to the standard initialization with different seeds. Standard initialization is shown in dashed line, while scaled initializations are the color lines each denoting different seed. The first row shows the objective during training for all seeds that all converge to the same value. The second row shows the element-wise RMSE between the Gram ma- trix of each seed and the optimized Gram matrix obtained from the standard initialization. RMSE converges to as all initializations converge on the same maximum. The last row plots RMSE ver- objective value, again showing single optimal objective value where all Gram matrices are the same.
This might contradict usual intuitions about huge multimodality in the weights/features of BNNs and DGPs. This can be reconciled by noting that each mode, written in terms of Gram matrices, corresponds to (perhaps infinitely) many modal features.
particular, in Sec. we show that the log-probability for features, P (F& FFz-1) ) (Eq. 82 depends only the Gram matrices, and note that there are many settings of features which give the same Gram matrix. particular, the Gram matrix is the same for any unitary transformation of the features, F' F, U, satisfying UUT I, as N: F G,. For DGPs we can use any unitary matrix, so there
N: F.F N, F? are infinitely many sets of features consistent with particular Gram matrix, while for BNNs we can only use permutation matrices, which are subset of unitary matrices. Thus, the objective landscape must be far more complex in the feature domain than with Gram matrices, single optimal Gram matrix corresponds to a large family of optimal features.
31

---

yacht
boston
energy
concrete
4700
6510
12450
15800
objective
6460
15750
1.0
RMSE
0.0
50000
100000
50000 iters
100000
100000 iters
200000
50000 iters
100000
iters
1.0
RMSE
0.0
4650
4700
6460
6510
12400
12450
15750
15800
objective
objective
objactive
objective
Figure 13:
One-layer DGP with MAP inference features described in Appendix Eq. 184
Rows and columns are the same as in Figure [2] Using the same randomly scaled initializations de- scribed above, are able find multiple modes in energy and conerete showing our initializations are diverse enough, albeit there is still only single global optimum.
32

---

L INDUCING POINT DKMS
To do large-scale experiments on UCI datasets, we introduce inducing point DKMs by extending Gaussian process inducing point methods (Damianou & Lawrence 2013 Salimbeni & Deisenroth, 2017) to the DKM setting. This approach uses the variational interpretation of the deep kernel machine objective described in Appendix D
To do inducing-point variational inference, we need to explicitly introduce top-layer features mir- roring L++1 IR
PX in Appendix B but replicated N times, F FL+1 E IR PR-NNN-1 Formally,
each feature, f[++1
FL-1 NN11 IID, conditioned on Fz.
P (F (FL+1 FL) I (f (,
(143a)
P (Y|F = - F
(143b)
where give the likelihood for regression, but other likelihoods (e.g. for classification) are possible (Appendix [B.
Further, take the total number of points, P, be made up of P. inducing points and P. test/train points, so that P = P; + P.. Thus, we can separate all features, F, E RPx,,, into the inducing features, F. E IR RxN, and the test/train features, F RAXN, Likewise, we separate the inputs, X. and outputs, Y, into (potentially trained) inducing inputs, X,, and trained inducing outputs, Y,, and the real lest/training inputs, X, and outputs, Y,.
F F
F
X, X.
Y Y
F--
FL+1
X
Y
Y
(144)
We follow the usual doubly stochastic inducing point approach for DGPs. In particular, we treat all the features at intermediate layers, F1,..., FL, and the top-layer train/test features, F/++1 as latent ......,F,,
variables. However, deviate from the usual setup in treating the top-layer inducing outputs, F as learned parameters and maximize over them to ensure that the ultimate method does not require sampling, and at the same time allows minibatched training. The prior and approximate posterior over F,...., Fz are given by,
Q Fx 12 (F.(Fe-), P (FF,...,FF.
(145a)
F2 = -1 P FF.(F..)
(145b)
and remember Fo = X, so Go
No X The prior and approximate posterior at each layer factorises into a distribution over the inducing points and a distribution over the test/train points,
(Fr = (F F.-1) 1) Q (F(),
(146a)
P = P
(146b)
Critically, the approximate posterior samples for the test/train points is the conditional prior P (F which is going to lead to cancellation when we compute the ELBO. Likewise, the approximate posterior over F+++1 is the conditional prior,
Q (F Fz) P FF ,F).
(147)
Concretely, the prior approximate posterior over inducing points are given by,
Q (F = II O,)),,
(148a)
P (F -I ((G(((
(148b)
The approximate posterior is directly analogous to Eq. (69) and the prior is directly analogous to Eq. (Ta), but where have specified that this is only over inducing points. Now compute the ELBO
ELBO G.....,
.... GH))
P (F 1 1
F
F2) P(F.....FF//
= EQ log ( + log
(149)
Q ( F2) ((F.....F,.x
33

---

Note that the P Ff, terms are going to cancel in the ELBO, consider them below when come describing sampling), substituting Eq. [45 147 and cancelling P (Ff Fc-1) and
PF T+ FL).
P (F Q (F
EF ........ ) Eq log (Y
log
(150)
So far, have treated the Gram matrices, G parameters of the approximate posterior. However, in the infinite limit + oc, these are consistent with the features generated by the approximate average over infinitely
posterior. In particular the matrix product N; Ff (Ff)* can be written many IID vectors, fix (first equality. and by the law of large numhers, this equal to the expectation of one term (second equality), which is G (by the approximate posterior Eq. (188a).
lim 1 F (Ff)*= N
lim N --000
N. EE (GEN) (((S))) [ (EW)" G
(151)
By this argument, the Gram matrix from the previous layer, G deterministic. Further, DGP, Ff only depends on F through G; (Eq, 5), and the prior and approximate posterior factorise. Thus, the infinite limit, individual terms in the ELBO can be written,
P (F log
C G Q
lim
Eq
= 1 Eq
log
(152)
N -
Q (F.)
ve DKL (N (0, K(G,
IIW (o. G))
(153)
where the final equality arises when notice that the expectation be written KL- divergence. The inducing DKM objective, Cind, is the ELBO, divided by N to ensure that remains finite in the infinite limit,
Lind (F
lim N EF N
4+ G. G
(154)
= EE [log [(+]
Σ vs DEL (N (W (o. G E=1
Note that this has almost exactly the same form the standard DKM objective for DGPs in the main text (Eq. 16). In particular, the second term is a chain of KL-divergences, with the only difference that these KL-divergences apply only the inducing points. The first term is "performance" term that here depends on the quality of the predictions given the inducing points. As the copies are IID, have,
EQ log P (Ý,15(+1)] = N Eq [log (Y]]
(155)
Now that we have simple form for the ELBO, need compute the expected likelihood, EQ [log P [Y,] This requires to compute the full Gram matrices, including test/train points, conditioned on the optimized inducing Gram matrices. We start by defining the full Gram matrix,
G
(156)
G
G
for both inducing points (labelled "I") and test/training points (labelled "t") from just G For clarity, we have G, E PxP G E R RxP. G E RiP. G E I RR R, where P, is the number of inducing points, P. is the number of train/test points and P P. P. is the total number of inducing and train/test points.
The conditional distribution over F. given Ff is,
P FF Ge-1) = Ne K, Kui)
(157)
34

---

Algorithm DKM prediction
Parameters: {pe)é=)
Optimized Gram matrices (
Inducing and train/test inputs: X,, X
Inducing outputs: Ff+t
Initialize full Gram matrix
G
O:T
X X
X,X! X,X²)
GO
G
Propagate full Gram matrix
for e in (1, (...... I.) do
K
K
G
c--1
(G
= K
Kii
K
G Cl-1-
G
Kn-i K
ii
K
G K
G K Kui
G K K
end for
Final prediction using standard Gaussian process expressions
Ki
Ki
G
(GAT
K
Ki
Kn
G
G
Y, N (K F Kn K n
KI' + +02)
where E. is the activation of the Ath feature for all train/test inputs, t feature for all train/test inputs, and f. and
the activation of the Ath
K K
K K
K
(TLL.FFFLL N F K (G.-1)
(158)
K K₁ K K,
(159)
In the infinite limit, the Gram matrix becomes deterministic via the law of large numbers (as in Eq. 151), and such Gi and G. become deterministic and equal to their expected values. Using
Eq. 157 we can write,
F = K. K 1/23.
(160)
where 3 is matrix with IID standard Gaussian elements. Thus,
G 1 E (T
(161)
K E [F(E))
(162)
K. GG
(163)
and,
G i E [
(164)
1 K E [F() K, 'KKK + / E [S="]
K
(165)
K, L KT + Kni
(166)
For the full prediction algorithm, Alg.
35


Reviews:
[{"review_id": "_JETf7REDgU", "reviewer": "Reviewer_JjRZ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Large-width networks converge to GPs with fixed/ unlearnable kernels, making the networks unable to perform representation learning. In this paper, the authors proposed a simple recipe to allow for representation learning: making the logit layer as wide as the hidden layers and replicating the labels accordingly. The authors argue that this approach allows the kernel to evolve and possible learning a representation. The authors also support their claims using synthetical dataset experiments and UCI datasets. \n", "summary_of_the_review": "Update: \nI thank the authors for the clarification/updates of several notations/terminologies. I like the idea from the paper, a Bayesian prospective of feature learning. I have increased my score accordingly. However, there are still a lot of room for improvement and several open questions. To name some,  \n\n1. Even though this is a theory paper, authors should not shy away from more empirical work. Even doing a small scale experiments on Cifar10 (baseline against kernel) could really helpful to both practitioners and theorists. \n\n2. Authors should have more detailed discussion/comparison between the meanfield/feature learning limit of neural networks and the current proposed Bayesian framework, both theoretically and empirically. E.g., baselining performance (bayesian feature learning vs NN feature learning), are they learning similar features, etc. \n\n\n---------------------------------\n\nOverall, I think the idea from the paper is interesting and worth further exploration. However, I also find the paper hard to parse due to clarity issues from notations and terminology, and the theoretical and empirical results from the paper do not support the strong claims from the paper.    \n\n", "strength_and_weaknesses": "# Strength.\n1. The paper proposed simple and interesting ideas that allow feature learning in the infinite-width network, which is worth further exploration. \n2. Some experiments from the paper show non-Gaussian behavior of the learned representation that differs from the NNGP limit.  \n\n# Weakness\n\n1. The notations and terminologies are confusing, which makes it hard to parse the paper. \n2. Missing key empirical comparison (vs. NNGP/NTK) for representation learning. I would like to see a comparison using CIFAR10. \n3. Several strong claims that require further justifications. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "PWgQUn4MUG", "reviewer": "Reviewer_SEBy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This article considers Bayesian inference with fully connected networks of depth L and width n in the infinite width limit. In order to have non-trivial feature learning at infinite width, the authors propose scaling the MSE negative log-likelihood by a factor proportional to the width, causing it to be appear at the same order in n relative to the prior (written  in terms of the feature-feature covariances $G^\\ell$ of layers ($\\ell = 1,…,L$). The authors explicitly find the MAP solution for deep linear networks and provide some numerical investigation.\n\nWhile the paper is clearly written, the authors don’t seem to realize that the likelihood re-scaling they propose is precisely what is called the mean-field scaling of neural networks in which the variance of final layer weights is rescaled by an extra factor or 1 / n (see below for references). In the language of this article, this precisely corresponds to rescaling the prior $P( Y | G^L )$ by a factor of n since it multiplies the variance of Y given $G^L$ by an additional factor of 1 / n. Given the significant overlap with existing work, I believe it is not ready to be published.", "summary_of_the_review": "This article proposes in a Bayesian framework a well-known modification of neural networks, the so-called mean field regime, that makes them capable of learning features at infinite width and any fixed depth. The article also provides an explicit form of the MAP feature-feature covariances for deep linear networks. While a more detailed understanding of Bayesian interference in the mean-field regime is certainly of interest, the article is not sufficiently novel to be published in its current form. ", "strength_and_weaknesses": "Strengths:\n- The paper is well-written\n- Understanding feature learning in Bayesian inference in mean-field networks is certainly an interesting topic\n- The article points out a nice closed form solution for the MAP value of the feature-feature covariances for deep linear networks in the mean-field regime. I believe this point is novel.\n\nWeaknesses:\n- The main contribution of this article has been considered in many prior articles on mean-field neural networks, though it is phrased somewhat differently (in terms of rescaling the likelihood instead of the weight variance in the final layer). Below are a few representative references. The article would be significantly strengthened by connecting to this literature. \n\nYang, Greg, and Edward J. Hu. \"Tensor programs iv: Feature learning in infinite-width neural networks.\" International Conference on Machine Learning. PMLR, 2021.\n\nSirignano, Justin, and Konstantinos Spiliopoulos. \"Mean field analysis of neural networks: A central limit theorem.\" Stochastic Processes and their Applications 130.3 (2020): 1820-1852.\n\nSirignano, Justin, and Konstantinos Spiliopoulos. \"Mean field analysis of neural networks: A law of large numbers.\" SIAM Journal on Applied Mathematics 80.2 (2020): 725-752.\n\nMei, Song, Andrea Montanari, and Phan-Minh Nguyen. \"A mean field view of the landscape of two-layer neural networks.\" Proceedings of the National Academy of Sciences 115.33 (2018): E7665-E7671.\n\nNguyen, Phan-Minh, and Huy Tuan Pham. \"A rigorous framework for the mean field limit of multilayer neural networks.\" arXiv preprint arXiv:2001.11443 (2020).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OvCprBc_e5A", "reviewer": "Reviewer_gVJs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a Bayesian learning algorithm that preserves the representation-learning ability of neural networks in the fixed-depth-infinite-width limit. This result parallels that of the maximal-update parametrization [Yang & Hu (2020)], which attains the similar limit in the gradient-descent setup (in contrast to the Bayesian setup discussed herein). In order to construct such a limit, the paper proposes to amplify the effect of target labels on the posterior representation by replicating labels N times, where N is a typical width of hidden layers. The proposed algorithm is applied to several tasks. ", "summary_of_the_review": "Overall, the recommendation for \"marginal acceptance\" is based on the novelty and possible significance of the proposed algorithm on deep learning theory. The weakness should be addressed.", "strength_and_weaknesses": "As stated in the \"Summary Of The Paper\" section, this paper contributes a nontrivial idea (strength) -- the construction of the Bayesian representation learning limit through label replications -- albeit with the overstated claim on the theoretical tractability (weakness).\n\nSpecifically, the paper claims that the proposed Bayesian algorithm is \"extremely theoretically tractable.\" While it is easy to state the optimization objective, its theoretical analysis doesn't seem to be tractable for general deep neural networks (which is not surprising as the maximal-update parametrization is also hard to theoretically analyze after some gradient-descent iterations). In fact, as explained in the paper itself, to solve for the maximum of the objective in Eq.(17)/Eq.(21), in general one must resort to numerical/variational methods.\n\nSimilarly (but in an opposite direction), the large body of work listed in the second paragraph of Section 2 -- which utilizes perturbation theory to analyze large-but-finite-width neural networks -- is dismissed as \"highly complex.\" However, in these perturbative approaches, one can systematically understand the behavior of networks order by order in 1/width, and get layer-by-layer recursive equations -- much like that for the NNGP in the infinite-width limit -- that determine the depth dependences of these perturbative corrections. This is theoretically analyzable -- hence such a large body of work exists -- and dismissing them as \"highly complex\" seems misleading (especially because the paper uses such a claim to further illustrate the simplicity of the proposed representation learning limit).\n\nAll those said, the proposed Bayesian algorithm remains nontrivial and interesting. But these claims on the theoretical (in)tractability should be clarified/softened/dropped, unless there is a substantive counter argument.\n\n\nMinor comments:\n\nA. It is preferable _not_ to use subjective adjectives/adverbs. For instance:\n\n(A-i) \"elegant objective\" in the abstract: it is not clear what makes Eqs.(17)/(21) \"elegant.\"\n\n(A-ii) \"strong regulariser\" in 3.4: it is not clear what makes it \"strong\" (in fact, the relative width ratio, \\nu_{\\ell}=N_{\\ell}/N, in front of each KL-divergence term controls its strength, so the relative strength of the regularizer can be adjusted).\n\n(A-iii) \"extremely theoretically tractable\" in the abstract: \"extremely\" is subjective (even if we temporarily forget the previous major objection to this claim).\n\nB. A missing period for \"which we call the deep kernel machine (DKM) Here\" in 3.7. (Also, the terminology/abbreviation is already introduced long before and used throughout the paper, so this phrasing doesn't make sense here.)\n\nC. In the first paragraph, for \"(NNGP Matthews et al., 2018; Novak et al., 2018)\" it is probably more standard to cite [Lee et al., 2017] as well (or in place of [Novak et al., 2018]). (I am not an author of these papers.)\n\nD. It may be appropriate to also cite [Dyer and Gur-Ari, arXiv:1909.11304] and [Hanin and Nica, arXiv:1909.05989] in the second paragraph of Section 2. (I am not an author of these papers.)\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject