Free Probability Theory (FPT) provides rich knowledge for handling mathematical difficulties caused by random matrices appearing in research related to deep neural networks (DNNs). However, the critical assumption of asymptotic freeness of the layerwise Jacobian has not been proven mathematically so far. In this work, we prove asymptotic freeness of layerwise Jacobians of multilayer perceptron (MLP) in this case. A key to the proof is an invariance of the MLP.

Asymptotic Freeness of Layerwise Jacobians Caused by Invariance of Multilayer Perceptron

Comm. In Math. Phys. 2022, https://link.springer.com/article/10.1007/s00220-022-04441-7.

Joint work with Benoit Collins. Supported by JST ACT-X and JSPS Sakura Program.

Tomohiro Hayase.
Affiliation: Cluster, Inc., Metaverse Lab. Talk at 2022/Dec./16—18, Japan-China International Conference on Matrix Theory with Applications .

---

Table of Contents

1. Overview
   Deep neural network and Gaussian process
2. Jacobian
   Stability of DNN and random matrices
3. NTK
   Training dynamics and random matrices
4. Asymptotic Freeness Main theorem: asymptotic freeness of Jacobians
5. Summary

---

Overview

---

Multilayer Perceptron

[Figure: https://www.javatpoint.com/multi-layer-perceptron-in-tensorflow] Let $n_0, n_1, \dots, n_L \in \N.$ Parameters:

$$\theta = (W_\ell, b_\ell)_{\ell=1, \dots, L}, W_\ell \in \R^{n_\ell \times n_{\ell-1}},b_\ell \in \R^{n_\ell}.$$

Forward propagation: for $x \in \R^{n_0},$ set $x_0 = x$ and inductively

$$h_\ell = W_\ell x_{\ell-1} + b_\ell, x_\ell = \varphi(h_\ell):= \varphi(h_{\ell,i})_{i\in[n_\ell]}.$$

Finally, deinfe the output by $f_\theta(x) = h_L.$

$\varphi$: Activation Function

Deep Learning

Generally, a standard formulation of supervised deep learning is as follows:

1. We are given a finite set of pairs of input/output data $\mathcal{D}$.
2. We are given a deep neural network (DNN). They have (one of ) the following conditions:
   1. Parameterized family of transformations, which maps a real vector to a real vector.
   2. It is a composition of brief parametrized transformations (e.g. linear transformations and non-linear elementwise function)
3. We are given an object function : e.g. mean squared loss

$$L(x,y, \theta) = \frac{1}{2B}\sum_{j} ( f_\theta(x)_j - y_j)^2$$

---

Optimization

Gradient descent

$$\theta_{t+1} = \theta_t - \eta_t \frac{\partial}{\partial \theta}L (x,y, \theta)$$

---

Initialization of Parameters

e.g. Gaussian (Ginibre) random matrix

$$(W_\ell)_{i,j} \sim \mathcal{N}(0, \sigma_w^2/N), \mathrm{\ i.i.d.}$$

e.g. Haar distributied orthogonal matrix

$$W_\ell= \sigma_w O, O\sim \mathrm{Haar \ Orthogonal \ Prob.}$$

---

The Inifnite-dimensional Limit is Gaussian

[Figure: https://ai.googleblog.com/2020/03/fast-and-easy-infinitely-wide-networks.html]

---

Neural Network Gaussian Process (NNGP)

Consider two inputs $x, x^\prime$ and corresponding hidden units $x_\ell, x^\prime_\ell$ and $h_\ell, h^\prime_\ell$ in MLP. Taking the wide limit, we have [Lee+ICLR2018]

$$(h_\ell, h_\ell^\prime)\sim \mathcal{N}(0, \sigma_w^2 K_\ell(x,x^\prime) + \sigma_b^2)$$

where

$$K_\ell(x, x^\prime) := \lim_{n_\ell \to \infty} \frac{1}{n_\ell} \sum_{j=1}^{n_\ell} x_{\ell,j} x^\prime_{\ell,j}.$$

We have the following Kernel Propagation:

$$K_{\ell+1}(x, x^\prime) = \int \varphi(z_1)\varphi(z_2) p_\mathcal{N}(z) dz$$

where,

$$p_\mathcal{N}= \mathcal{N}( 0, \sigma_w^2\begin{pmatrix} K_\ell(x,x) & K_\ell(x,x^\prime)\\ K_\ell(x,x^\prime)& K_\ell(x^\prime,x^\prime)\end{pmatrix} + \sigma_b^2)$$

• For some activation, we can compute the integral explicitly.

---

NNGP Estimation

Generally, consider $B$ samples. Set $X=(x(a))_{a=1,\dots, B}, Y=(y(a))_{a=1, \dots, B}$ be input/output samples.

$$K(x^*, X) := (K_L(x^*, x(a)))_{n=1, \dots, a} \in \R^B$$ $$K(X,X):= ( K_L(x(a), x(b)) )_{a,b} \in M_B(\R)$$

Posteriror mean/ var: For new input $x^*,$

$$m(y^*) = K(x^*,X)K(X,X)^{-1}Y$$ $$v(y^*) = K(x^*, x^*) - K(x^*,X)K(X,X)^{-1}K(x^*, X)$$

[Lee et.al., Deep Neural Networks as Gaussian Process, ICLR 2018]

Jacobian

---

Vanishing/Exploding Gradients

The optimization of DNN needs its parameter derivations. Since a DNN is a composition of functions, the parameter derivations are computed by  the chian rule. The input-output Jacobian is defined as

$$J = \frac{\partial f_\theta(x)}{\partial x} = \frac{\partial h_L}{\partial x}$$

In the case of MLP, we have

$$J = W_L D_{L-1} \dots W_2 D_1 W_1,$$

where

$$D_\ell = \frac{\partial x_\ell}{\partial h_\ell} = \mathrm{diag}( \varphi^\prime(h_{\ell,1}), \dots, \varphi^\prime(h_{\ell,n_\ell}))$$

---

Dynamical Isometry

A DNN is said to achive dynamical isometry If the eigenvalue distribution of $JJ^\top$ is concentrated aound one. Dynamical Isotmetry prevents the exploding/vanshing gradients.

[Pennington+, AISTATS2018, CH, CIMP2022] If we set the initialization of parameters to be Haar orthgonal and choose appropriate activation function, then we can make the DNN to achieve the dynamical isometry.

Set $\mu_L, \nu$ be limit spectral distributions of $JJ^\top, D^2$ as wide limits respectively.

Under the assumption of asymptotic freeness of Jacobians,

$$\mu_L = [(\sigma^2 \cdot )_* \nu ]^{\boxtimes L}$$

where $\boxtimes$ is the free multiplicative convolution,

Distribution of $D^2$

[Figure: Pennington,  Schoenholz, Ganguli, AISTATS2018]

---

The Limit Spectral Distribution $\mu_L$ of $JJ^\top$

[Figure: Pennington,  Schoenholz, Ganguli, AISTATS2018]

---

Neural Tangent Kernel

---

Neural Tangent Kernel

Under continual vertion of GD, learning dynamics of parameters is given by:

$$\frac{d\theta_t}{dt} = \eta (\nabla_\theta f_{\theta_t})^\top (y - f_{\theta_t})$$

( * The learning rate $\eta$ is fixed.) Then learning dynamics of DNN is given by:

$$\frac{df_{\theta}}{dt} = \eta_t \Theta_t(y-f_{\theta_t})$$

where

$$\Theta_t = \nabla_\theta f_{\theta_t}( \nabla_\theta f_{\theta_t})^\top$$

**Informal[Jacot+NeurIPS2018, Lee+NeruIPS2019]**Under the wide limit $n \to \infty$, the learning dynamics of DNN is approximated by

$$\frac{df_{\theta}}{dt} = \eta \Theta(y-f_{\theta_t})$$

where the neural tangent kernel is defined as

$$\Theta := \lim_{n_2, \dots, n_{L-1} \to \infty} \Theta_0$$

---

Neural Tangent Kernel is A Surrogate Model of DNN+GD

[Figure from Google “Fast and Easy Infinitely Wide Networks with Neural Tangents”]

Appliable to CNN/ResNet

Figure from [Google “Fast and Easy Infinitely Wide Networks with Neural Tangents”]

Moreover, NTK is appliable to Attention: Infinite attention: NNGP and NTK for deep attention networks [https://arxiv.org/abs/2006.10540\]

Eigenvalue Spectrum of NTK

“Spectra of the Conjugate Kernel and Neural Tangent Kernel for linear-width neural networks” https://arxiv.org/abs/2005.11879 They treats the standard formulation: Gaussian Initialization x Multi-samples  x Small output dimension, and they get:

---

One-sample NTK

“The Spectrum of Fisher Information of Deep Networks Achieving Dynamical Isometry” https://arxiv.org/abs/2006.07814 [HK20] When the DNN achieves dynamical isometry,  the spectrum of the (one-sample x high-dim output)”NTK” concentrates around the maximal value, and the maximal values is O(L). Under an assumption on Asymptotic Freeness, we have the following recursive equations:

$$\Theta_{\ell+1} = q_\ell + W_{\ell+1} D_\ell \Theta_\ell D_\ell W_{\ell+1}^\top$$ $$\mu_{\ell + 1} = (q_\ell + \sigma_{\ell+1}^2 \cdot)_* (\nu_\ell \boxtimes \mu_\ell)$$

NTK & Learning Rate

The spectrum (eigenvalues) of the NTK has vital role in tuning the learning dynamics. e.g. $\eta > 1/ \lambda_\mathrm{max} (\Theta) \Longrightarrow$ The learning dynamics does not converge. e.g. The conditional number $c = \lambda_\mathrm{min}/ \lambda_\mathrm{max}$ detemines the converges speed.

Red line  (the boarder line of the exploding gradients) : This line is expected by our theory !

---

Asymptotic Freeness

---

Asymptotic Freeness and Free Probability Theory

Definition(Asymtptotically freeness, C$^*$-version)[Voiculescu’85] Let $(A_j(n), A_j(n)^*)_{j \in J}$ be a family of $n \times n$ random matrices and adjoints. The family is said to be asymptotic free almost surely, if there exists C$^*$-probability spaces $(\mathfrak{A}_j, \tau_j)_{j \in J}$ and elements $(a_j \in \mathfrak{A}_j)_{j \in J}$ so that for any $Q \in \mathbb{C} \langle X_j, X_j^* \mid j \in J \rangle$, the following holds:

$$\lim_{n \to \infty} \mathrm{tr}_n \left[Q(A_j(n), A_j(n)^* \mid j \in J ) \right] \\ = (*_{j \in J} \tau_j) \left[ Q(a_j, a_j^* \mid j \in J) \right]$$

---

Example

For $N \in \N,$ let

• $W(N)$ be Ginibre or Haar orthogonal random matrix,
• $D(N)$ be constant diagonal matrix with a limit distirubution as $N \to \infty.$ Then (W, W^*) and  D are a.s. asymptotically free as N \to \infty.

---

Asymptotic Freeness of Jacobians

Let $W_\ell, D_{\ell} (\ell=1,2,\dots, L)$ be random matrices detemined by MLP. Theorem [CH22] Assuming that $D_1, \dots, D_{L-1}$ have limit joint moments. Then

$$(W_1, W_1^\top), \dots, (W_L, W_L^\top), (D_1, \dots, D_{L-1})$$

are asymptotically free  as $n \to \infty$  almost surely. Difficulty: Entries of $D_\ell \text{\ and }W_\ell$ are not independent. $D_\ell = \mathrm{diag}(\varphi^\prime(h_\ell))$ $h_\ell = W_\ell x_\ell$

(Sketch of Proof) Invariance of MLP + Taking submatrix Construct orthogonal matrix $U_\ell$ fixing $x_{\ell}$ , i.e.

$$U_\ell x_\ell = x_\ell,$$

and

$$U_\ell |_{{\mathbb{C}x_\ell}^\bot} \sim N-1 \times N-1 \text{\ Haar Orthogonal}$$

with

$$(U_0, \dots, U_\ell) \perp\!\!\!\!\perp(W_{\ell+1}, \dots, W_L).$$

for $\ell=0,\dots, L-1.$ Then we only need to show the asymptotic freeness of $N-1 \times N-1$ submatrices of

$$(U_0W_1, \dots, U_{L-1}W_L), (D_1, \dots, D_{L-1}).$$

---

Summary

---

Summary

Considering neural networks with random paramters… $\Longrightarrow$

1. Tuning initializaiton and learning rate
2. Bayesian Estimation with NNGP
3. Understanding Dynamics with NTK

Background: Random Matrix Theory and Free Probability Thoery

---

Future Work

MLP-like NNs were previously only used for toy models, but are now being applied to real-world images and 3D data. e.g. gMLP, NeRF, etc. Theoretically and practically easy to compute positively, just the right next research target!

[Mildenhall, et.al., ***“*NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis”, ECCV 2020 ]