Exchangeability and Kernel Invariance in Trained MLPs

1 Introduction

Despite the widespread usage of deep learning in applications, current theoretical understanding of deep networks continues to lag behind the pursued engineering outcomes. Much recent theory concerns networks in their randomized initial state, or contains assumptions about the parameters or data during training.

For example, Cho and Saul [2009], Daniely et al. [2016], Bach [2017] and Tsuchida et al. [2018] analyze the kernels of neural networks with random IID weights. Insightful analysis connecting signal propagation in deep networks to chaos have made similar assumptions [Poole et al., 2016; Raghu et al., 2017]. Random matrix theory has recently been applied to neural networks in an attempt to understand the empirical spectral distribution (ESD) of the Hessian [Pennington and Bahri, 2017] and the Gram matrix [Pennington and Worah, 2017], but these works have made strong assumptions on the weight and data distributions.

The most widely used yet unrealistic assumption is that weights remain IID throughout training. Relaxing such unrealistic assumptions makes obtaining meaningful results more challenging. We take a step in this direction by investigating the probabilistic symmetry known as exchangeability, which is a generalization of the IID assumption. We uncover the striking result that the layer-wise kernel of MLPs with ReLU

Supplemental material available at the arXiv version https://arxiv. org/abs/1810.08351.

activations, trained with many optimizers, remains constant up to a scaling factor during training when the network inputs satisfy certain conditions. Otherwise, we are able to bound the absolute difference between layer-wise kernel and the kernel of the network in its random IID state.

2 Background

2.1 Notation

Random variables, vectors and matrices will be denoted by upper case, bold upper case, and bold upper case with overline characters, respectively. Parenthesized superscripts index the layer of the network to which an object belongs. The first and second post-subscripts index the rows and columns of a matrix, respectively. When the row of a matrix is extracted through an index, it will be assumed to be transposed into a column vector. Pre-subscripts will indicate the iteration of an iterative optimizer. Expectation with respect to the distribution of and random variable R is denoted ER .

Consider an MLP with an input layer and L non-input layers. Denote the number of neurons in layer 0 ≤ l ≤ L by n(l) . Denote an input to the network by x. Denote the random weight matrix connecting layer l-1 to layer l byW(l) . Denote the ℓ2 norm by ‖ · ‖. Denote the activation function by σ. We consider ReLU activations throughout.

2.2 Exchangeability

An exchangeable sequence of random variables (Q1, Q2, ...) has the property that the joint distribution of the sequence is invariant to finite permutations. That is, a sequence (Qi)i≥1 is exchangeable if (Q1,Q2, ...) d= (Qπ(1),Qπ(2), ...) for all finite permutations π. To aid in readability we will omit the index set in the subscript, so that (Qi)i≥1 is the same as (Qi)i.

Infinite exchangeable sequences are characterized as mixtures of IID random variables through de Finetti's theorem.

Theorem 1. [Aldous, 1981] An infinite sequence Q= (Qi)i is exchangeable if and only if there exists a measurable function f such that (Qi)i d= (f(A, Bi))i , where A and B are mutually IID random variables uniform on [0, 1].

Generalizations of Theorem 1 to multi-dimensional arrays exist [Kallenberg, 2006]. A matrix Q is row and column exchangeable (RCE) if its joint distribution is invariant to row and column permutations. That is, Q is RCE if (Qji)ji d= (Qπ1(j)π2(i))ji for all finite permutations π1, π2.

Theorem 2. [Aldous, 1981] An infinite array Q= (Qji)ji is RCE if and only if there exists a measurable function f such that (Qji)ji d= (f (A, Bj , Ci, Dji))ji , where A, B, C, and D are mutually IID uniform on [0, 1].

Intuition concerning the strength of exchangeability in the context of probabilistic symmetries may be aided by the implication graph shown in Figure 1.

2.3 Kernels of RandomMLPs

There is a well-studied connection between the feature maps in MLPs (and other neural network architectures) and the kernel of a reproducing kernel Hilbert space (RKHS) [MacKay, 1992; Neal, 1994; Cho and Saul, 2009; Daniely et al., 2016; Bach, 2017; Bietti and Mairal, 2017]. Consider the angle θ(l) between two random signals σ(W(l)x) and σ(W(l)y) in the lth hidden layer of an MLP for inputs x and y. We have

where W(l) j is the jth row of W(l) . We divide the numerator and denominator by n(l) ‖x‖‖y‖ and use the absolutehomogeneity property of the ReLU σ(|a|z) = |a|σ(z) to consider the scaled numerator

Let ˆx = x ‖x‖ . Suppose that each rowW(l) j of W(l) is IID with all other rows (we relax this requirement later) and is defined on some probability space (Ω, Σ, µ). Asymptotically in the number of neurons n(l) , the strong law of large numbers implies that (2) converges almost surely to

which corresponds to an inner-product in feature space. The kernel is positive semi-definite and uniquely defines an RKHS. When µ is the product measure corresponding to an IID Gaussian with variance E[(W (l) 11 )2 ] and 0 mean, the kernel has a closed-form expression known as the arc-cosine kernel (of degree 1) [Cho and Saul, 2009], given by

where θ(l-1) is the angle between x and y. We will refer to (3) as the layer-wise kernel in layer l, denoted k(l) (x, y). When (3) is normalized in the same fashion as (1), we will call the resulting quantity the layer-wise normalized kernel.

2.4 Layer-wise Kernel in IIDMLPs

Our analysis draws upon and extends results concerning the layer-wise normalized kernels of MLPs with IID weights [Tsuchida et al., 2018], which, for completeness, we briefly review here. Construct a sequence {x(m)}m≥2 such that for all m, x(m) ∈ R∞ and coordinates m + 1, m + 2, ... of x(m) are all 0. Define the sequence {y(m)}m≥2 in the same way, and additionally require that the angle θ(l-1) between x(m) and y(m) is constant in m. Denote the randomly initialized weight matrix by 0W(l) . We would like to evaluate

Sufficient conditions for the central limit theorem (CLT) are given below. Let x(m)i denote the ith coordinate of x(m).

This condition is easily satisfied since for data points with many non-zero entries, ‖x(m)‖ will grow like √m when compared to |x(m)i|. Provided E[0W (l) 11 ] = 0 and E∣∣0W (l) 11 ∣∣3 < ∞, Tsuchida et al. [2018] show that under Hypothesis 3,

The integral of the RHS is 2E[(0W (l) 11 )2 ], so the limit may be brought inside the integral in (5) by Theorem 19 of Royden [2010]. The resulting expectation is (4).

Figure 2 shows the normalized kernels for random weights with PDF ∏mi=1 β 2αΓ(1/β) e-|wi/α|β . This PDF generalizes the isotropic Gaussian PDF (β = 2) and the Uniform PDF (β → ∞). The CLT result says nothing about the kernel of trained networks whose weights are not IID. In §4 we extend the CLT result to trained networks. We do this by first exploring exchangeability in MLPs.

3 Exchangeability inMLPs

Suppose that for every l, the matrix (0W (l) ji )ji is IID and then the weights evolve according to SGD over t iterations. The index j (which corresponds to the jth row of the random weight matrix, or the jth neuron in layer l) is an arbitrary labeling; one may permute these indices along with the corresponding connection in layer l + 1 without changing the output of the network or the joint distribution of the weights.

We show this for L = 3; the generalization to any L ≥ 2 will be straightforward. To start our argument, it is clear that there is full exchangeability of the weights when the network has been randomly initialized with IID weights and has not yet been trained. More restrictively, we have the following.

Observation 4. Let a ∈ Rn(0) and b ∈ Rn(3) be inputs and targets of an MLP. Suppose that the initial weights in each layer 0W(l) are IID, and temporarily drop the pre-subscript. Then for any bijective permutations π1 and π2,

This generalizes to any network with one or more hidden layers (L ≥ 2) because the permutation does not affect the non-exchangeable elements a and/or b. Define g(l) qp to be the function that takes a, b and realizations of (0W(m))m∈[L] and calculates realizations of 1W (l) qp according to an online (batch size of 1) backpropagation update rule. Let g(l) be a matrix-valued function, whose qpth element is g(l) qp . We have

for some cost function E(a, b; (w(r) )r∈[L]) and step-size α. Denote the LHS of (6) by U and the RHS of (6) by Uπ. Then by examining the backpropagation equations,

By the continuous mapping theorem, we may apply g(2) to both sides of (6) if g(2) is almost everywhere (a.e.) continuous. Temporarily dropping the 0 pre-subscripts on the weights,

and the first line is equal to (1W (2) qp )qp . This shows that 1W(2) is RCE. tW(1) is row but not column-exchangeable and tW(L) is column but not row-exchangeable.

When any batch size M is used, the inputs a and b may be replaced by sets {ai}i≤M and {bi}i≤M and (7) still holds. If M is the size of the entire finite dataset, this corresponds to gradient descent. We may use any a.e. continuous g(l) whose evaluation commutes with index permutations in the input (such as SGD, Adam [Kingma and Ba, 2015] or RMSprop). Call such an update rule index commuting. By redefining g(2) to calculate the weights at the tth iteration of SGD, one can show that tW(2) is RCE for all t.

Theorem 5. Let L ≥ 3. Suppose that the initial weights in each layer 0W(l) are IID. Suppose the network is trained using an index commuting update rule. Then for all 2 ≤ l ≤ L - 1 and all optimizer iterations t ≥ 0, the weight matrices tW(l) are RCE. For L ≥ 2, tW(1) is row but not column exchangeable and tW(L) is column but not row exchangeable.

4 Kernels of TrainedMLPs

We now extend the results of §2.4 to trained networks using the results of §3. For the remainder of the paper we will drop the pre-subscript t denoting the training iteration on the weights.

4.1 Layer-wise Kernel in TrainedMLPs

We examine the limit inmof the layer-wise kernel in layer l for a network with infinitely many RCE weights. By Theorem 1, there exists some measurable function f and some mutually independent A and B each uniform on [0, 1] such that

where µA is the uniform probability measure on [0, 1] with kA(x(m), y(m)) given by

where (fA(B)i )i = (f(A, Bi))i and µB is the uniform probability measure. We prove the following in Appendix A.

Then (8) is given by

Note that (9) and (4) are the same up to a scaling factor, which cancels out after normalizing.

4.2 The Ergodic Problem

Unfortunately, (8) is not necessarily the inner product in feature space of an infinitely wide network. By Theorem 2,

for some measurable fAC(B1,D1) = (f(A,Bj,Ci,Dji)i, which converges almost surely to the random variable

depending on A and C by the Birkhoff-Khinchin ergodic theorem (see Appendix E). For the purposes of experimenting, we make the following simplifying assumption.

Hypothesis 7. The following holds:

Hypothesis 7 says that taking averages over j of the products of activations in one network is equivalent to taking averages over one fixed neuron of the products of activations in an ensemble of independent networks. A sufficient condition is that the measure is ergodic with respect to the row-shift transformation. This condition is stronger than necessary. In statistical mechanics, an "approximate ergodicity" applied to sum functions is used to compare time averages with phase averages [Khinchin, 1949; Kurth, 2014]. The Ergodic Problem features heavily in the history of statistical mechanics [Moore, 2015]. It is our hope that by introducing this assumption into the analysis of MLPs, we make further progress towards efforts in connecting neural networks to statistical mechanics [Martin and Mahoney, 2017]. In §5 we demonstrate that Hypothesis 7 is not inconsistent with our empirical observations.

5 Experiments

We illustrate our results with selected figures. Other datasets and optimizers are investigated in the supplemental material.

5.1 Verification of Proposition 6

Architecture. We train an autoencoder with 4 layers and 3072 neurons in each layer on CIFAR10 [Krizhevsky and Hinton, 2009] with pixel values normalized to [0, 1] using an ℓ2 objective. Weights are initialized with a variance of 2 nl [He et al., 2015].

Method. In Figure 3 we plot the empirical layer-wise normalized kernel in each layer. The color of the points moves from blue to red as the training iteration t increases. Each sample is generated using Procedure 1. The numerical steps ensure that the desired angle θ(l-1) is obtained between x and y. The alphabetical steps ensure that ∑ xi = ∑ yi = 0.

Procedure 1 Sample θ(l-1)

Inputs datapoint x, θ(l-1) Output y at angle θ(l-1) to x.

• a Set the last two coordinates of x to 0.
• 1 Sample random vector p orthogonal to x: Set all coordinates of p to zero where x is non-zero and sample remaining coordinates of p from U[0, 1]. Set last two coordinates to 0. Normalize p so that ‖x‖ = ‖p‖.
• b Set the second last coordinate of x to the negative sum of all coordinates of x.
• c Set the last coordinate of p to the negative sum of all coordinates of p.
• 2 Return y = cos θ(l-1)x + sin θ(l-1)p.

5.2 Inputs with Non-Zero Sums

Consider a modification of the method described in §5.1: the alphabetical steps of Procedure 1 are not performed. This means that the sums ∑ ˆxi and ∑ ˆ yi are no longer 0. However, if E[W (l) 11 W (l) 12 ] = 0, Proposition 6 still applies. Note that

Also, by the strong law of large numbers,

be "small". We are interested in finding optimizer hyperparameters that result in (En(l) )2 ̸= 0, which in turn results in deviations from (9). We make the following observations:

(1) When Adam, RMSProp or Nadam [Dozat, 2016] are used, as the hyperparameter ε decreases there is a sharp change in (En(l) )2 and the mean squared error (MSE) of the observed normalized kernel to the normalized arc-cosine kernel of degree one measured at iteration t = 19000. When (En(l) )2 is small the kernel is approximately described by (9). See Figures 4 and 5 and Appendices F and G.

(2) SGD using step sizes α that result in stable training generally have smaller (En(l) )2 than Adam, and thus the normalized kernel agrees more closely with Proposition 6. See Figures 4 and 5 and Appendices F and G.

6 Discussion and Conclusion

We identified that the weights in hidden layers of MLPs are RCE. Using this symmetry, we analyzed the kernels of trained networks. Specifically, we found that the normalized kernel remains invariant when the inputs have sums over their coordinates of 0. When the sums are not 0, a bound which depends on E[W (l) 11 W (l) 12 ] applies to the residual of the normalized kernel to the normalized arc-cosine kernel. We derived a measure (E(n(l) )2 which, when close to 0, indicates whether E[W (l) 11 W (l) 12 ] is close to 0 and thus whether the normalized kernel remains approximately invariant during training.

When empirically comparing optimizers, those which result in small E[W (l) 11 W (l) 12 ] have kernels which follow the normalized arc-cosine kernel during training. The parameter ε present in Adam and other optimizers can increase E[W (l) 11 W (l) 12 ], leading to qualitatively different kernels to the normalized arccosine kernel. Changes in other hyperparameters may change E[W (l) 11 W (l) 12 ], although we had difficulty finding instances where changing α in SGD resulted in a kernel that did not roughly match the normalized arc-cosine kernel without also resulting in unstable training.

In contrast with works that analyze weight distributions through an approximation of SGD by a stochastic differential equation [Seung et al., 1992; Watkin et al., 1993; Martin and Mahoney, 2017; Chaudhari and Soatto, 2018], we incorporate very little knowledge of the learning rule into our theory. The result is that our theory is perhaps more general than required. Interestingly, our results still hold if we perform (stochastic) gradient ascent on the weights. We believe our analysis would benefit from including more knowledge of the update rule.

Jacot et al. [2018] analyse a continuous-time approximation of (not stochastic) gradient descent. In this dynamic, it is shown that MLPs in function space follow a linear differential equation in an infinite width limit. A central object of their study is a positive-definite kernel, the neural tangent kernel, which is shown to stay approximately constant during training. Also utilizing a kernel and working in function space, Du et al. [2019] bound distances between functions trained in continuous and discrete time dynamics in special cases and investigate the surprising fact that certain models can achieve zero training loss. Chizat and Bach [2018] present a unified framework of these and other works under a training regime called lazy training. We remark that it is easy to find networks trained with Adam that do not exhibit approximately constant kernels (for example, see the first row of Figure 4). These networks seem to out-perform those that do have constant kernels during training. This is consistent with the view expressed by Chizat and Bach [2018], that "the most competitive neural networks are not trained in this regime [lazy training]."

In future work, we would like to examine the ESD of weight and Hessian matrices without normality assumptions as in previous works [Pennington and Bahri, 2017; Pennington and Worah, 2017], perhaps using results concerning the ESD of exchangeable random matrices [Chatterjee, 2006; Adamczak et al., 2016].

Acknowledgements

Russell Tsuchida and Fred Roosta gratefully acknowledge the generous support given by the Australian Research Council Centre of Excellence for Mathematical & Statistical Frontiers (ACEMS). Fred Roosta was partially supported by the Australian Research Council through a Discovery Early Career Researcher Award (DE180100923).