$\Lambda(1520)$ resonance production with respect to transverse spherocity
using EPOS3 + UrQMD

Nasir Mehdi Malik$\star$ and Sanjeev Singh Sambyal

University of Jammu

*<EMAIL_ADDRESS>
51st International Symposium on Multiparticle Dynamics (ISMD2022)
Pitlochry, Scottish Highlands, 1-5 August 2022 10.21468/SciPostPhysProc.?

## Abstract

Resonances are sensitive to the properties of the medium created in heavy ion
collision. They also provide insight into the properties of the hadronic
phase. Studying the dependence of the yield of resonances on transverse
spherocity and multiplicity allows us to understand the resonance production
mechanism with event topology and system size, respectively. The results
reported pertain to $\Lambda(1520)$, using predictions from EPOS3+UrQMD event
generator.

###### Contents

  1. 1 Introduction
  2. 2 Distribution plots
  3. 3 pT spectra of $\Lambda(1520)$
  4. 4 Conclusion

##  1 Introduction

At collider experiments, a major challenge is to understand
production/interaction of resonances in a hot dense medium because their
lifetime is comparable to that of hadronic phases [1]. Resonance particle
$\Lambda(1520)$ having a life-time between $K^{*0\pm}(980)$ and $\phi(1020)$
make it a good candidate to understand the hadronic phase[2] [3]. Here we are
analysing production of $\Lambda(1520)$ with respect to transverse spherocity
$S_{0}$ using the EPOS3 event generator[4] [5] [6] with UrQMD as an
afterburner to simulate the hadronic rescatterings after hadronisation.
Transverse spherocity is an event shape variable which help to distinguish
between isotropic and jetty events. Its value ranges from 0 to 1 and it is
given by

$S_{0}={\pi^{2}\over
4}{\mathrm{\mathop{min}_{\vec{n}=(n_{x},n_{y},0)}}}\left(\frac{{}_{\sum_{i}\left|\vec{P_{\mathrm{T_{i}}}}\times\hat{n}\right|}}{\mathrm{N_{particles>5}}}\right)^{2}$

Where,

$S_{0}=\begin{cases}0,&\text{"pencil-like" limit (hard events)},\\\
1,&\text{"isotropic" limit (soft events)},\end{cases}$

and $\hat{n}$ is a two dimensional unit vector in the transverse plane,chosen
such that $S_{0}$ is minimised. By restricting it to the transverse plane,
transverse spherocity becomes infrared and collinear safe. Also the jetty
events are usually hard events while the isotropic ones are the result of soft
processes [7, 8, 9, 10].

##  2 Distribution plots

For the present study, events having more than 5 charge paricles are
considered. Moreover, to study in various multiplicity intervals and
transverse spheorcity classes, we separated events into the 0-10%
quantile(high multiplicity) and 30-60% quantile(low multiplicity). To
disentangle the spherocity classes i.e jetty and isotropic events, we have
applied event-selection cuts on the spherocity distribution. The values, cuts
and definition of various multiplicity percentile and spherocity classes can
be found in Table 1(a) and 1(b). We observed that the spherocity distributions
are shifted toward 1 from the low multiplicity quantile to the high
multiplicty one. Figure 1 displays charged particle multiplicity, transverse
spherocity and spherocity as a function of multiplicity.

No. of charged | Percentile | Percentile
---|---|---
particle | (lower) | (upper)
28 - 110 | 90-100% | 0-10%
18-28 | 70-90% | 10-30%
11 - 18 | 40-70% | 30-60%
0-110 | 0-100% | 0-100%
(a) Percentile of multiplicity distribution

Multiplicity(%) | Jetty | Isotropic
---|---|---
| 0-10% | 0-10%
| (lower) | (upper)
0-100 | 0.00-0.381 | 0.834-1.00
0-10 | 0.00-0.644 | 0.896-1.00
30-60 | 0.00-0.442 | 0.824-1.00
(b) Quantile of spherocity distribution

Table 1:  Figure 1: Shows distribution of :- (a) Charge particle
multiplicities. (b) Transverse spherocity. (c): Transverse spherocity vs
multiplicity.

##  3 pT spectra of $\Lambda(1520)$

Figure 2(2(a),2(b)) shows the pT-spectra of $\Lambda(1520)$ for 0-10% and
30-60% multiplicity classes with top and bottom 10% spherocity values
responsible for selecting jetty and isotropic events respectively. The bottom
panels show the ratio with respect to the spherocity integrated class.

(a)

(b)

Figure 2:  (a) pT spectra of $\Lambda(1520)$ evaluated at 0-10% multiplicity
for Isotropic(blue), Jetty (red) and Integrated(black). (b) pT spectra of
$\Lambda(1520)$ evaluated at 30-60% multiplicity for Isotropic(blue), Jetty
(red) and Integrated(black).

##  4 Conclusion

A study of the transverse spherocity of $\Lambda(1520)$ in p-p collisions at
$\sqrt{s}=13$ TeV using EPOS3 generator with UrQMD as afterburner was
reported. The observables are measured using primary charge particles and
reported as a function of charged particle multiplicity for events with
different scales defined by the transverse spherocity. As can be seen in
Figure(1c),jetty events dominate the low-multiplicity region. However, we see
that production of $\Lambda(1520)$ is dominated in both multiplicity regions
by isotropic event shapes, since EPOS is a parton model, with many parton-
parton binary interaction,so highest multiple interaction(MPI) mean dominace
in isotropic event(soft QCD). An experimental investigation in this direction
would be very beneficial to comprehend the event shape dependence of system
dynamics.

## Acknowledgements

We thank Vikash Sumberia, Randhir Singh and Dr. Ranbir Singh for inspiring us
to do this work and providing also valuable advice. Thanks are also due to
Johannes Jahan Ph.D. student in Subatech, Nantes -France for helping us to
understand technical term in EPOS3. This work was supported by Council of
Scientific and Industrial Research (CSIR), India. The authors acknowledge the
HEP Lab. Department of Physics, University of Jammu for providing
infrastructure for computing etc.

## References

  * [1] M. Bleicher and J. Aichelin,  _Strange resonance production: probing chemical and thermal freeze-out in relativistic heavy ion collisions_ ,  Physics Letters B 530(1), 81 (2002),  https://doi.org/10.1016/S0370-2693(02)01334-5.
  * [2] A. G. Knospe, C. Markert, K. Werner, J. Steinheimer and M. Bleicher,  _Hadronic resonance production and interaction in partonic and hadronic matter in the epos3 model with and without the hadronic afterburner urqmd_ ,  Phys. Rev. C 93, 014911 (2016),  10.1103/PhysRevC.93.014911.
  * [3] M. Tanabashi, K. Hagiwara, K. Hikasa, K. Nakamura, Y. Sumino, F. Takahashi, J. Tanaka, K. Agashe, G. Aielli, C. Amsler, M. Antonelli, D. M. Asner _et al._ ,  _Review of particle physics_ ,  Phys. Rev. D 98, 030001 (2018),  10.1103/PhysRevD.98.030001.
  * [4] H. Drescher, M. Hladik, S. Ostapchenko, T. Pierog and K. Werner,  _Parton-based gribov–regge theory_ ,  Physics Reports 350(2), 93 (2001),  https://doi.org/10.1016/S0370-1573(00)00122-8.
  * [5] K. Werner, B. Guiot, I. Karpenko and T. Pierog,  _Analyzing radial flow features in $p$-pb and $p$-$p$ collisions at several tev by studying identified-particle production with the event generator epos3_,  Phys. Rev. C 89, 064903 (2014),  10.1103/PhysRevC.89.064903.
  * [6] K. Werner,  _Core-corona separation in ultrarelativistic heavy ion collisions_ ,  Phys. Rev. Lett. 98, 152301 (2007),  10.1103/PhysRevLett.98.152301.
  * [7] N. Mallick, R. Sahoo, S. Tripathy and A. Ortiz,  _Study of transverse spherocity and azimuthal anisotropy in pb–pb collisions at sNN=5.02 TeV using a multi-phase transport model_ ,  Journal of Physics G: Nuclear and Particle Physics 48(4), 045104 (2021),  10.1088/1361-6471/abeb59.
  * [8] A. Ortiz, E. Cuautle and G. Paić,  _Mid-rapidity charged hadron transverse spherocity in pp collisions simulated with pythia_ ,  Nuclear Physics A 941, 78 (2015),  https://doi.org/10.1016/j.nuclphysa.2015.05.010.
  * [9] S. Prasad, N. Mallick, D. Behera, R. Sahoo and S. Tripathy,  _Event topology and global observables in heavy-ion collisions at the large hadron collider_ ,  Scientific Reports 12(1), 3917 (2022),  10.1038/s41598-022-07547-z.
  * [10] ALICE Collaboration, S. Acharya, D. Adamová, S. P. Adhya, A. Adler, J. Adolfsson, M. M. Aggarwal, G. A. Rinella, M. Agnello, N. Agrawal, Z. Ahammed, S. Ahmad _et al._ ,  _Charged-particle production as a function of multiplicity and transverse spherocity in pp collisions at $\sqrt{\mathbf{s}}=\mathbf{5.02}$ and 13 TeV_,  The European Physical Journal C 79(10), 857 (2019),  10.1140/epjc/s10052-019-7350-y.


marginparsep has been altered.
topmargin has been altered.
marginparwidth has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Implicit Acceleration and Feature Learning in
Infinitely Wide Neural Networks with Bottlenecks

Etai Littwin 1  Omid Saremi 1  Shuangfei Zhai 1  Vimal Thilak 1  Hanlin Goh 1
Joshua M. Susskind 1  Greg Yang 2

††footnotetext:  1Apple 2MSR. Correspondence to: Etai Littwin.
<EMAIL_ADDRESS>
Presented at the ICML 2021 Workshop on Overparameterization: Pitfalls &
Opportunities. Copyright 2021 by the author(s).

###### Abstract

We analyze the learning dynamics of infinitely wide neural networks with a
finite sized bottleneck. Unlike the neural tangent kernel limit, a bottleneck
in an otherwise infinite width network allows data dependent feature learning
in its bottleneck representation. We empirically show that a single bottleneck
in infinite networks dramatically accelerates training when compared to purely
infinite networks, with an improved overall performance. We discuss the
acceleration phenomena by drawing similarities to infinitely wide deep linear
models, where the acceleration effect of a bottleneck can be understood
theoretically.

##  1 Introduction

The study of infinitely wide neural networks is one of the most actively
researched topics in deep learning theory Jacot et al. (2018); Daniely et al.
(2016); Hazan & Jaakkola (2015); Roux & Bengio (2007); Lee et al. (2018);
Matthews et al. (2018); Hinton & Neal (1995); Novak et al. (2019); Yang (2019;
2020a; 2020b); Yang & Littwin (2021); Littwin et al. (2020a; b). Previous work
identified distinct training regimes that are determined by the networks
hyperparameters. In Yang & Hu (2020), it is shown that by correctly scaling
hyperparameters such as learning rate, weight multipliers and initialization
constants, neural network training dynamics under gradient descent exhibit a
limiting behaviour as the width of the network tends to infinity. Two distinct
training regimes are identified in the limit. 1) The kernel regime, where the
network evolves like a linear model during training. In this regime,
intermediate activations in infinite layers change by a vanishing amount,
stripping the network of its ability to learn features. 2) The feature
learning regime, where the intermediate activations change in a nontrivial way
during training, resulting in data dependent feature learning.
In Yang & Hu (2020), a clear dichotomy exists between these regimes in the
infinite width limit, where a network is either in one regime or another, but
not in both at the same time. A network undergoing training by gradient
descent algorithms can be placed in one regime or the other depending on its
parametrization at initialization. The underlying assumption in both cases
however is that all hidden layers are infinitely wide. Some practical
architectures incorporate narrow layers by design, or use varying width layers
of which some are wide, while others are narrow. For example, bottleneck
layers force a network to learn a low dimensional latent representation, and
are typically used as part of a generally wide network. Perhaps the most
prominent architecture which uses such layers is the autoencoder, where the
encoder learns a typically low dimensional representation of its input, while
the decoder reconstructs the input from its latent representation. In these
types of models, the discrepancy between wide and narrow layers is a carefully
hand designed feature of the architecture. Therefore, the standard infinite
width approximation cannot be applied without ”assuming away” a prominent
architectural feature.
In this work, we consider a different type of limit, where a bottleneck layer
of finite width is inserted in an otherwise infinitely wide network. From a
Bayesian perspective, such models have been investigated in Agrawal et al.
(2020); Aitchison (2020), however their learning dynamics under gradient
descent have not yet been explored to the best of our knowledge. As we show,
gradient descent on such models can be understood and simulated exactly in
function space rather than parameter space, at a relatively cheap
computational cost. This is in contrast to the feature learning limit in Yang
& Hu (2020), where it is practically infeasible to simulate exact learning
dynamics in the limit without approximations. Empirically, we show that adding
a bottleneck layer can significantly boost speed of training convergence. We
further speculate on possible reasons for this dramatic boost in convergence
speed by drawing equivalence to infinitely wide deep linear models with
bottlenecks, where acceleration of the bottlenecks is fully tractable.

##  2 Neural Networks With Bottlenecks

A natural way of thinking about wide neural networks with bottlenecks is
through composition of networks. Let
$f_{n}(x;w):\mathbb{R}^{d}\to\mathbb{R}^{d_{r}}$ denote the (vector) output of
a neural network given input vector $x\in\mathbb{R}^{d}$ and parameters $w$,
where $n$ denotes the width of the hidden layers111We use the term ”width”
losely here. For MLPs, this means the size of the hidden layers, or the
feature dimension in convolutional neural nets.. In the present work, we
consider the case where the input $x$ is itself given by the output of a
neural network $x(\xi)=g_{n}(\xi;\theta):\mathbb{R}^{d_{0}}\to\mathbb{R}^{d}$,
given input vector $\xi\in\mathbb{R}^{d_{0}}$ and parameters $\theta$, with
hidden layers width parametrized by $m$. For simplicity, we assume the width
of all hidden layers in both $f$ and $g$ is $n$. The output of the composed
architecture is then given by $\mathcal{F}_{n}(\xi)=(f_{n}\circ g_{n})(\xi)$.
Note that even if $n$ is large, the output of $g_{n}(\xi)$ is still $d$
dimensional, hence $\mathcal{F}_{n}$ can be viewed as a wide network with a
bottleneck of dimension $d$.

##### Setup

To make things concrete, we consider the case where $f_{n}$ implement an MLP
of depths $2$ and width $n$. Due to technical considerations,222Rigorously
generalizing the claims in this paper to deeper MLPs is not straightforward,
and involves a considerable technical challenge. This difficulty arises due to
the structure of the input-output Jacobian, which cannot be implemented as a
Tensor Program Yang (2019; 2020a; 2020b) in its current form. we restrict our
discussion in the current paper to a 1 hidden layer MLP for $f$, and leave the
rigorous analysis of deeper networks to future work. Given an input vector
$g\in\mathbb{R}^{d}$ and parameters $w=\\{u,v\\}$, the output $f_{n}(g)$ is
given by:

$\displaystyle f_{n}(x)=\frac{v}{\sqrt{n}}\phi(\frac{ux}{\sqrt{d}})$ (1)

where $u\in\mathbb{R}^{n\times d},v\in\mathbb{R}^{d_{r}\times n}$ are the
weight matrices sampled from a normal distribution, and $\phi$ is a coordinate
nonlinearity which we assume, for the sake of analysis, is twice
differentiable with bounded derivatives. We let $g$ implement an arbitrary
neural network function with parameter matrices
$\theta=\\{\theta^{l}\\}_{l=1}^{L}$, with suitable dimensions. We are
interested in theoretically understanding the behaviour of the composition
$\mathcal{F}_{n}=f_{n}\circ g_{n}$ during training in the overparametrized
scenario where the width $n$ of the hidden layers tend to infinity, while the
bottleneck dimension $d$ remains fixed in size. Our setup involves the
training of the composition function $\mathcal{F}_{n}$ on a training dataset
$\\{\xi_{i}\\}_{i=1}^{N}$, using a loss function $\mathcal{L}$ implicitly
containing the labels. We state our results assuming infinitesimal learning
rate (aka gradient flow), however we expect our results to carry over to SGD.
Since $\mathcal{F}_{n}$ contains a bottleneck of finite size, we will have to
reason about the evolution of forward and backward signals as they propagate
through a finite sized layer.

##### Notations

We use $x,\tilde{x}$ to denote generic placeholder vectors to $f$, and
$g(\xi)$ as a specific instantiation of $x$ by the output $g(\xi,\theta)$. We
denote the input-output Jacobian $J(x)=\frac{\partial f(x)}{\partial
x}\in\mathbb{R}^{d_{r}\times d}$, with the notation $\tilde{J}=J(\tilde{x})$.
We use $\mathcal{F},g,J$ for $\mathcal{F}(\xi),g(\xi),J=J\big{(}g(\xi)\big{)}$
where $\xi$ is some generic input. To remove clutter, we use subscripts
$n,t,i$ to denote the value of a vector/matrix parametrized by $n$ at time $t$
given sample $\xi_{i}$. (i.e $g_{n,t,i}=g_{n,t}(\xi_{i}),g_{n,t}=g_{n,t}(\xi)$
and same for $J$). We use superscripts to denote coordinates of a
vector/matrix (i.e $g^{\alpha}$ is the $\alpha$ coordinate of vector $g$). The
absence of a subscript $n$ implies we have taken $n\to\infty$. Finally, we use
$\chi_{i}$ to denote the loss derivative for sample $i$ (i.e
$\chi_{i}=\frac{\partial\mathcal{L}}{\partial\mathcal{F}_{i}}\in\mathbb{R}^{d_{r}}$).

###  2.1 Preliminaries

As width increases, pre-activations of initialized neural networks without
bottlenecks approach a centered gaussian process (GP), independent across
coordinates but possibly correlated across inputs. At the output level, the
following hold at initialization:

$\displaystyle\lim_{n\to\infty}g_{n}(\xi)\overset{d}{=}g(\xi),$
$\displaystyle~{}~{}~{}\lim_{n\to\infty}f_{n}(x)\overset{d}{=}f(x)$ (2)

where $g(\xi),f(x)$ are GPs with respect to their respective inputs. For $f$
as implemented in eq. 1, we can write the defining covariance $\Lambda$ of the
GP given input samples $x,\tilde{x}$:

$\displaystyle\forall_{\alpha},\begin{pmatrix}f^{\alpha}(x)\\\
f^{\alpha}(\tilde{x})\end{pmatrix}$
$\displaystyle\sim\mathcal{N}\big{(}0,\Lambda(x,\tilde{x})\big{)}$ (3)
$\displaystyle\Lambda(x,\tilde{x})$
$\displaystyle=\begin{pmatrix}\Sigma(x,x)&\Sigma(x,\tilde{x})\\\
\Sigma(\tilde{x},x)&\Sigma(\tilde{x},\tilde{x})\end{pmatrix}\in\mathbb{R}^{2\times
2}$ (4)

We assume both $f,g$ have empirical NTKs $\mathcal{K}_{n},\Theta_{n}$ defined
as $\mathcal{K}_{n}(x,\tilde{x})=\frac{\partial f(x)}{\partial
w}\frac{\partial f(\tilde{x})^{\top}}{\partial w}\in\mathbb{R}^{d_{r}\times
d_{r}},\Theta_{n}(\xi,\tilde{\xi})=\frac{\partial
g(\xi)}{\partial\theta}\frac{\partial
g(\tilde{\xi})^{\top}}{\partial\theta}\in\mathbb{R}^{d\times d}$ with
corresponding limits
$\mathcal{K}(x,\tilde{x})I_{d_{r}},\Theta(\xi,\tilde{\xi})I_{d}$, where
$I_{d_{r}},I_{d}$ are identity matrices of size $d_{r},d$, and
$\mathcal{K}(x,\tilde{x}),\Theta(\xi,\tilde{\xi})\in\mathbb{R}$. Consistent
with intuition, it was rigorously shown in Agrawal et al. (2020) that as the
width increases the output of an MLP with bottlenecks converges to a
composition of GPs.

##### Training

We consider training the model described in section 2. Under gradient flow,
the weights evolve according to
$\dot{w}_{t}=-\nabla_{w_{t}}\mathcal{L}_{t},~{}~{}~{}\dot{\theta}_{t}=-\nabla_{\theta_{t}}\mathcal{L}_{t}$.
The evolution of the output of the composition function can be described by
the following set of ODEs:

$\displaystyle\dot{\mathbf{\mathcal{F}}}_{n,t}(\xi)$
$\displaystyle=\frac{\partial\mathcal{F}_{n,t}(\xi)}{\partial
w_{t}}\dot{w}_{t}+\frac{\partial\mathcal{F}_{n,t}(\xi)}{\partial\theta_{t}}\dot{\theta}_{t}$
(5)

Substituting the empirical kernel definitions:

$\displaystyle\dot{\mathcal{F}}_{n,t}(\xi)$
$\displaystyle=-\sum_{i=1}^{N}\mathcal{K}_{n,t}\big{(}g_{n,t},g_{n,t,i}\big{)}\chi_{n,t,i}$
(6)
$\displaystyle-\sum_{i=1}^{N}J_{n,t}\Theta_{n,t}(\xi,\xi_{i})J^{\top}_{n,t,i}\chi_{n,t,i}$

where $g_{n,t}=g_{n,t}(\xi),J_{n,t}=J_{n,t}(\xi)$. The above evolutionary
equations can be interpreted as Kernel equations with random, evolving kernels
that depend on the weights $w,\theta$. In Jacot et al. (2018), it was shown
for infinite width networks (without bottlenecks) the output is fully
deterministic at any time $t$ conditioned on the initial GP output at time
$t=0$. In contrast, since both bottleneck embedding $g$ and input-output
Jacobian $J$ are finite, we expect eq. 6 to remain random even when taking the
limit $n\to\infty$. To get a more complete view of the evolution of the
composition function $\mathcal{F}$ during training at the limit, we must
reason about the dynamics of the Jacobian term $J$.

##  3 Dynamics in Function Space

A key observation in our analysis is that the input-output Jacobian $J(x)$
converges to a multivariate GP, and evolves as a linear function in the
infinite width limit, similarly to outputs of an infinitely wide network.
Here, $J(x)$ will have non trivial correlations across its coordinates, unlike
layers in the NTK limit where the coordinates are independent. Our first
result relating to the initial state of $J$ is stated in the following
proposition:

######  Lemma 3.1 (GP behaviour of the Jacobian).

For $f_{n}(x)$ and its limit $f(x)$ as described in eqs. 1 and 2, the
following holds at initialization for every pair of fixed inputs
$x,\tilde{x}$:

$\displaystyle\lim_{n\to\infty}J_{n}(x)=J(x)$ (7)

where $J(x)$ is a multivariate GP with independent rows, and $J(x),f(x)$ are
jointly Gaussian with:

$\displaystyle\begin{pmatrix}J^{\alpha,\beta}(x)\\\
J^{\alpha,\gamma}(\tilde{x})\end{pmatrix}$
$\displaystyle\sim\mathcal{N}(\bm{0},\begin{pmatrix}\Sigma_{(2)}^{\beta,\beta}(x,x)&\Sigma_{(2)}^{\beta,\gamma}(x,\tilde{x})\\\
\Sigma_{(2)}^{\gamma,\beta}(\tilde{x},x)&\Sigma_{(2)}^{\gamma,\gamma}(\tilde{x},\tilde{x})\end{pmatrix}$
(8) $\displaystyle\begin{pmatrix}f^{\alpha}(x)\\\
J^{\alpha,\beta}(\tilde{x})\end{pmatrix}$
$\displaystyle\sim\mathcal{N}(\bm{0},\begin{pmatrix}\Sigma(x,\tilde{x})&\Sigma_{(1)}^{\beta}(x,\tilde{x})\\\
\Sigma_{(1)}^{\beta}(x,\tilde{x})^{\top}&\Sigma_{(2)}^{\beta,\beta}(\tilde{x},\tilde{x})\end{pmatrix}$
(9)

where:

$\displaystyle\Sigma_{(1)}(x,\tilde{x})$
$\displaystyle=\frac{\partial}{\partial
b}\Sigma(x,b)\Big{|}_{b=\tilde{x}}\in\mathbb{R}^{1\times d}$ (10)
$\displaystyle\Sigma_{(2)}(a,b)$ $\displaystyle=\frac{\partial^{2}}{\partial
a\partial b}\Sigma(a,b)\Big{|}_{a=x,b=\tilde{x}}\in\mathbb{R}^{d\times d}$
(11)

lemma 3.1 already illustrates a novel aspect of infinite width networks. As
the input-output Jacobian is frequently used to derive sensitivity to
perturbations, lemma 3.1 shows that sensitivity to perturbations and outputs
can be jointly modeled as a multivariate GP for sufficiently wide models.
In the next theorem, we show that the $J(x)$ evolves linearly in wide models,
in a similar fashion to the network outputs. Now we are ready to characterize
the full training dynamics of $\mathcal{F}$ in the infinite width limit in the
following theorem:

######  Theorem 3.2 (Evolution of composed function).

For networks $f,g$ as in section 2, the following ODEs describe the dynamics
of $\mathcal{F},g,J$ in the limit of $n\to\infty$:

$\displaystyle\dot{g}_{t}(\xi)$
$\displaystyle=-\sum_{i=1}^{N}\Theta(\xi,\xi_{i})J_{t,i}^{\top}\chi_{t,i}$
(12) $\displaystyle\dot{J}_{t}(x)$
$\displaystyle=-\sum_{i=1}^{N}\chi_{t,i}\Xi(x,g_{t,i})^{\top}$ (13)
$\displaystyle\dot{\mathcal{F}}_{t}(\xi)$
$\displaystyle=-\sum_{i=1}^{N}\Big{[}\mathcal{K}\big{(}g_{t},g_{t,i}\big{)}I_{d_{r}}+\Theta(\xi,\xi_{i})J_{t}J_{t,i}^{\top}\Big{]}\chi_{t,i}$
(14)

where $\Xi(-,-)$ is a deterministic function defined as
$\Xi(x,\tilde{x})\in\mathbb{R}^{d}=\lim_{n\to\infty}\frac{\partial^{\top}}{\partial
x}\mathcal{K}_{n}^{\alpha,\alpha}(x,\tilde{x})$.

For $\phi=\text{ReLU}$, we give an explicit form for $\Xi$ in appendix D. ODEs
in eqs. 12, 13 and 14 depend on deterministic, frozen kernels
$\mathcal{K},\Theta,\Xi$. The evolution of $\mathcal{F},g,J$ is hence
completely deterministic after conditioning on initial states, and can
therefore be expressed in functional space.

##  4 Implicit Acceleration by Bottlenecks in Linear Networks

To intuitively understand the training aspects of introducing bottlenecks in
infinite width networks, we draw inspiration from deep linear networks. Assume
$f_{n},g_{n}$ implement deep linear MLPs of depth $L_{f},L_{g}$ respectively
and width $n$, and let $w_{\text{eff}}\in\mathbb{R}^{d_{r}\times
d}=\frac{1}{\sqrt{n^{(L_{f}-1)}d}}w^{L_{f}}w^{L_{f-1}}...w^{1},~{}\theta_{\text{eff}}\in\mathbb{R}^{d\times
d_{0}}=\frac{1}{\sqrt{n^{(L_{g}-1)}d_{0}}}\theta^{L_{g}}\theta^{L_{g-1}}...\theta^{1}$.
Hence, we have that $g(\xi)=\theta_{\text{eff}}\xi$, $f(g)=w_{\text{eff}}g$
and $\mathcal{F}(\xi)=w_{\text{eff}}\theta_{\text{eff}}\xi$. In finite
networks, recent results have shown that in some cases, the stacking of linear
layers produces an acceleration effect, and a low rank bias when optimized by
gradient descent Arora et al. (2018). Moreover, the acceleration effect is
akin to momentum, and cannot be reproduced by adding some regularizer to the
objective. However, the acceleration effect as outlined in Arora et al. (2018)
disappears when considering the NTK regime. This is because in this regime,
training speed is determined by the NTK itself, which does not change
meaningfully with stacking of additional linear layers. Indeed, for a linear
$g$ we have that
$\Theta(\xi,\tilde{\xi})=L_{g}\frac{\xi^{\top}\tilde{\xi}}{d_{0}}$. However,
by introducing a bottleneck, we regain the lost acceleration effect, as
illustrated in the following lemma.

######  Lemma 4.1.

In the limit of $n\to\infty$, optimizing $\mathcal{F}$ by running gradient
flow on the weights $\\{w^{l}\\}_{l=1}^{L_{f}}$ with a learning rate
$\epsilon_{f}$, and $\\{\theta^{l}\\}_{l=1}^{L_{g}}$ with a learning rate
$\epsilon_{g}$, is equivalent to running gradient flow directly on the
effective weights $w_{\text{eff}},\theta_{\text{eff}}$ with learning rates
$L_{f}\epsilon_{f},L_{g}\epsilon_{g}$.

lemma 4.1 suggests an infinitely wide deep linear network with a bottleneck is
essentially reduced to a two layer, finite linear network with weight matrices
$w_{\text{eff}},\theta_{\text{eff}}$ under gradient descent. Therefore, under
mild initialization conditions,333The acceleration effect formally requires a
small initialization and learning rates to hold. Empirically, these conditions
may sometimes be relaxed. a bottleneck brings about accelerated learning in
linear networks. It is worth noting, additional capacity cannot be attained by
stacking additional layers in linear networks. Moreover, shallow and deep
linear functions represent the same function class. Hence training speed can
be directly attributed to trajectories of gradient descent. This does not hold
for nonlinear finite networks, where changes in depth or width also affect
capacity. However, in the infinite width regime, with bottlenecks or not,
capacity is infinite. Hence, we can isolate the effect of bottlenecks on
training and acceleration without confounding capacity.

##  5 Experiments

We provide empirical support for our theoretical contributions in three parts:

  * •

In part 1, we conduct simulations to numerically verify the theory in lemmas
3.1, 4.1 and 3.2. We present these results in appendix E.

  * •

In part 2, we train infinite neural networks with bottlenecks on MNIST LeCun
et al. (1989) and CIFAR-10 Krizhevsky (2009) datasets by simulating SGD on the
loss in function space, investigating training acceleration effects and test
performance. Figure 1 summarizes the results from these experiments. We
observe that the accelerated training predicted by lemma 4.1 for linear models
is also visible when we train infinite width nonlinear networks with
bottlenecks on the two real world datasets.444Figure 1(b) show loss for first
15K steps on CIFAR-10, as some training runs did not complete in time due to
compute resource scarcity. Extended results showing loss trajectories for
models trained longer are available in figure 8 in appendix D. We present
additional results and analysis in appendix D.

  * •

In part 3, we run experiments with finite width networks trained with standard
SGD and verify that the acceleration effect holds in this setting. These
results are presented in appendix F.

(a) MNIST training loss

(b) CIFAR-10 training

Figure 1: Training loss for infinite width bottleneck models with widths from
smallest to largest (inf indicates infinite width bottleneck). Loss is reduced
in function space by simulating SGD with eqs. 12, 13 and 14

##  6 Conclusion

In this work we investigate the effect of applying a bottleneck in an
otherwise infinite width network. We do this by first deriving the ODEs
corresponding to optimization of such a model under gradient flow. Our
theoretical analysis reveals novel insights regarding the behaviour of input-
output Jacobians, both at initialization and training. Though stated for
shallow, single hidden layer networks post bottleneck, we expect our results
to hold in more general cases. Empirically, we observe that infinite width
networks with bottlenecks train much faster than their fully infinite
counterparts, while typically achieving better overall performance. We hope
our results pave the way for new understanding of learning dynamics beyond
kernel regimes.

## References

  * Abadi et al. (2015) Abadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z., Citro, C., Corrado, G. S., Davis, A., Dean, J., Devin, M., Ghemawat, S., Goodfellow, I., Harp, A., Irving, G., Isard, M., Jia, Y., Jozefowicz, R., Kaiser, L., Kudlur, M., Levenberg, J., Mané, D., Monga, R., Moore, S., Murray, D., Olah, C., Schuster, M., Shlens, J., Steiner, B., Sutskever, I., Talwar, K., Tucker, P., Vanhoucke, V., Vasudevan, V., Viégas, F., Vinyals, O., Warden, P., Wattenberg, M., Wicke, M., Yu, Y., and Zheng, X.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015\.  URL https://www.tensorflow.org/.  Software available from tensorflow.org.
  * Agrawal et al. (2020) Agrawal, D., Papamarkou, T., and Hinkle, J.  Wide neural networks with bottlenecks are deep gaussian processes.  _ArXiv_ , abs/2001.00921, 2020.
  * Aitchison (2020) Aitchison, L.  Why bigger is not always better: on finite and infinite neural networks.  _ArXiv_ , abs/1910.08013, 2020.
  * Arora et al. (2018) Arora, S., Cohen, N., and Hazan, E.  On the optimization of deep networks: Implicit acceleration by overparameterization.  _ArXiv_ , abs/1802.06509, 2018.
  * Arora et al. (2019) Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R., and Wang, R.  On exact computation with an infinitely wide neural net.  _ArXiv_ , abs/1904.11955, 2019.
  * Daniely et al. (2016) Daniely, A., Frostig, R., and Singer, Y.  Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity.  In _NIPS_ , 2016.
  * Hazan & Jaakkola (2015) Hazan, T. and Jaakkola, T.  Steps toward deep kernel methods from infinite neural networks.  _ArXiv_ , abs/1508.05133, 2015.
  * Hinton & Neal (1995) Hinton, G. E. and Neal, R.  Bayesian learning for neural networks.  1995\.
  * Huang et al. (2019) Huang, C., Zhai, S., Talbott, W., Martin, M. B., Sun, S.-Y., Guestrin, C., and Susskind, J.  Addressing the loss-metric mismatch with adaptive loss alignment.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2891–2900. PMLR, 09–15 Jun 2019.  URL http://proceedings.mlr.press/v97/huang19f.html.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  In _NeurIPS_ , 2018.
  * Krizhevsky (2009) Krizhevsky, A.  Learning multiple layers of features from tiny images.  2009\.
  * LeCun et al. (1989) LeCun, Y., Boser, B. E., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W. E., and Jackel, L. D.  Handwritten digit recognition with a back-propagation network.  In _NIPS_ , 1989.
  * Lee et al. (2018) Lee, J., Bahri, Y., Novak, R., Schoenholz, S., Pennington, J., and Sohl-Dickstein, J.  Deep neural networks as gaussian processes.  _ArXiv_ , abs/1711.00165, 2018.
  * Littwin et al. (2020a) Littwin, E., Galanti, T., and Wolf, L.  On random kernels of residual architectures.  _arXiv: Learning_ , 2020a.
  * Littwin et al. (2020b) Littwin, E., Galanti, T., and Wolf, L.  On the optimization dynamics of wide hypernetworks.  _ArXiv_ , abs/2003.12193, 2020b.
  * Matthews et al. (2018) Matthews, A., Rowland, M., Hron, J., Turner, R., and Ghahramani, Z.  Gaussian process behaviour in wide deep neural networks.  _ArXiv_ , abs/1804.11271, 2018.
  * Novak et al. (2019) Novak, R., Xiao, L., Bahri, Y., Lee, J., Yang, G., Hron, J., Abolafia, D., Pennington, J., and Sohl-Dickstein, J.  Bayesian deep convolutional networks with many channels are gaussian processes.  In _ICLR_ , 2019.
  * Roux & Bengio (2007) Roux, N. L. and Bengio, Y.  Continuous neural networks.  In Meila, M. and Shen, X. (eds.), _Proceedings of the Eleventh International Conference on Artificial Intelligence and Statistics_ , volume 2 of _Proceedings of Machine Learning Research_ , pp. 404–411, San Juan, Puerto Rico, 21–24 Mar 2007. PMLR.  URL http://proceedings.mlr.press/v2/leroux07a.html.
  * Yang (2019) Yang, G.  Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are gaussian processes.  _ArXiv_ , abs/1910.12478, 2019.
  * Yang (2020a) Yang, G.  Tensor programs ii: Neural tangent kernel for any architecture.  _ArXiv_ , abs/2006.14548, 2020a.
  * Yang (2020b) Yang, G.  Tensor programs iii: Neural matrix laws.  _ArXiv_ , abs/2009.10685, 2020b.
  * Yang & Hu (2020) Yang, G. and Hu, E. J.  Feature learning in infinite-width neural networks.  _ArXiv_ , abs/2011.14522, 2020.
  * Yang & Littwin (2021) Yang, G. and Littwin, E.  Tensor programs iib: Architectural universality of neural tangent kernel training dynamics.  _ArXiv_ , abs/2105.03703, 2021.

##  Appendix A Proof of lemma 3.1

See 3.1

###### Proof.

Intuitively, the derivative process of any centered GP with a covariance
kernel function $\mathcal{K}(x,\tilde{x})$ is another GP, provided that the
kernel function $\mathcal{K}$ is everywhere differentiable. Since the Jacobian
is by definition the derivative of the output with respect to the input, this
implies that eq. 7 holds, along with the stated covariances. However, the
subtlety here is to show that the empirical Jacobian indeed converges in
distribution to the derivative of the NN-GP in the limit. For a single hidden
layer MLP, the convergence of the Jacobian to a GP can be established by a
straightforward application of the standard CLT theorem, as in the feed
forward case.
The Jacobian takes the form:

$\displaystyle J_{n}(x)=\frac{\partial f(x)}{\partial
x}=\frac{v}{\sqrt{nd}}\text{diag}\big{(}\phi^{\prime}(\frac{ux}{\sqrt{d}})\big{)}u$
(15)

Each element $J_{n}^{\alpha,\beta}(x)$ is hence given by:

$\displaystyle
J_{n}^{\alpha,\beta}(x)=\frac{1}{\sqrt{nd}}\sum_{\gamma=1}^{n}v^{\alpha,\gamma}\phi^{\prime\gamma}(\frac{ux}{\sqrt{d}})u^{\gamma,\beta}$
(16)

Note that the sequence
$\\{v^{\alpha,\gamma}\phi^{\prime\gamma}(ux)u^{\gamma,\beta}\\}_{\gamma=1}^{n}$
is a sequence of zero mean iid variables. From the polynomial boundness of
$\phi$, the coordinates of $\phi^{\prime}$ have bounded second moment. We can
then apply the CLT theorem to establish a convergence to a GP with some kernel
function $\Sigma^{\star}\in\mathbb{R}^{3\times 3}$:

$\displaystyle\begin{pmatrix}f^{\alpha}(x)\\\ J^{\alpha,\beta}(\tilde{x})\\\
J^{\alpha,\delta}(x^{\prime})\end{pmatrix}\overset{p}{=}\lim_{n\to\infty}\frac{1}{\sqrt{n}}\begin{pmatrix}\sum_{\gamma=1}^{n}v^{\alpha,\gamma}\phi^{\gamma}(\frac{ux}{\sqrt{d}})\\\
\frac{1}{\sqrt{d}}\sum_{\gamma=1}^{n}v^{\alpha,\gamma}\phi^{\prime\gamma}(\frac{u\tilde{x}}{\sqrt{d}})u^{\gamma,\beta}\\\
\frac{1}{\sqrt{d}}\sum_{\gamma=1}^{n}v^{\alpha,\gamma}\phi^{\prime\gamma}(\frac{ux^{\prime}}{\sqrt{d}})u^{\gamma,\delta}\end{pmatrix}\overset{p}{=}\mathcal{N}(\bm{0},\Sigma^{\star})$
(17)

From the independence of $u,v$, it trivially holds that
$\forall_{\alpha,\beta}\mathbb{E}[J^{\alpha,\beta}(x)]=0$,
$\forall_{\alpha\neq\gamma,\beta,\delta}\mathbb{E}[J^{\alpha,\beta}(x)J^{\gamma,\delta}(\tilde{x})]=0$
and
$\forall_{\alpha\neq\gamma,\beta}\mathbb{E}[J^{\alpha,\beta}(x)f(\tilde{x})^{\gamma}]=0$.
For any two inputs $x,\tilde{x}$ we have:

$\displaystyle\forall_{\alpha,\beta,\delta}\mathbb{E}[J^{\alpha,\beta}(x)J^{\alpha,\delta}(\tilde{x})]$
$\displaystyle=\frac{1}{d}\mathbb{E}\Big{[}\big{(}v^{\alpha,\gamma}\big{)}^{2}\phi^{\prime}(\frac{ux}{\sqrt{d}})^{\gamma}\phi^{\prime}(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}u^{\gamma,\beta}u^{\gamma,\delta}\Big{]}$
(18)
$\displaystyle=\frac{1}{d}\mathbb{E}\Big{[}\phi^{\prime}(\frac{ux}{\sqrt{d}})^{\gamma}\phi^{\prime}(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}u^{\gamma,\beta}u^{\gamma,\delta}\Big{]}$
(19) $\displaystyle=\mathbb{E}\Big{[}\frac{\partial}{\partial
x^{\beta}}\frac{\partial}{\partial\tilde{x}^{\delta}}\phi(\frac{ux}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}\Big{]}$
(20)

From our assumption that $\phi$ is differentiable with bounded derivatives, it
holds that $\mathbb{E}\Big{[}\big{|}\frac{\partial}{\partial
x^{\beta}}\frac{\partial}{\partial\tilde{x}^{\delta}}\phi(\frac{ux}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}\big{|}\Big{]}<\infty$.
By the dominated convergence theorem:

$\displaystyle\mathbb{E}\Big{[}\frac{\partial}{\partial
x^{\beta}}\frac{\partial}{\partial\tilde{x}^{\delta}}\phi(\frac{ux}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}\Big{]}=\frac{\partial}{\partial
x^{\beta}}\frac{\partial}{\partial\tilde{x}^{\delta}}\mathbb{E}\Big{[}\phi(\frac{ux}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}\Big{]}=\frac{\partial^{2}}{\partial
x^{\beta}\partial\tilde{x}^{\delta}}\Sigma(x,\tilde{x})$ (21)

Similarly:

$\displaystyle\forall_{\alpha,\beta,\delta}\mathbb{E}[J^{\alpha,\beta}(x)f^{\alpha}(\tilde{x})]$
$\displaystyle=\mathbb{E}\Big{[}\big{(}v^{\alpha,\gamma}\big{)}^{2}\phi^{\prime}(\frac{ug}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}u^{\gamma,\beta}\Big{]}$
(22) $\displaystyle=\frac{\partial}{\partial
x^{\beta}}\mathbb{E}\Big{[}\phi(\frac{ux}{\sqrt{d}})^{\gamma}\phi(\frac{u\tilde{x}}{\sqrt{d}})^{\gamma}\Big{]}=\frac{\partial}{\partial
x^{\beta}}\Sigma(x,\tilde{x})$ (23)

which concludes the proof. ∎

##  Appendix B Proof of 3.2

See 3.2

###### Proof.

The empirical dynamical equation for $\mathcal{F},g$ are given in eq. 6:

$\displaystyle\dot{g}_{n,t}(\xi)$
$\displaystyle=-\sum_{i=1}^{N}\Theta_{n,t}(\xi,\xi_{i})J_{n,t,i}^{\top}\chi_{n,t,i}$
(24) $\displaystyle\dot{\mathcal{F}}_{n,t}(\tilde{\xi})$
$\displaystyle=-\sum_{i=1}^{N}\mathcal{K}_{n,t}\big{(}g_{n,t},g_{n,t,i}\big{)}\chi_{n,t,i}-\sum_{i=1}^{N}J_{n,t}\Theta_{n,t}(\tilde{\xi},\xi_{i})J^{\top}_{n,t,i}\chi_{n,t,i}$
(25)

Prior results regarding convergence of NTK functions have been established
pointwise on a fixed dataset Yang (2020a). In our scenario, we have both the
bottleneck embeddings and the Jacobian which evolve continuously during
training. This would not pose a major problem to us since a pointwise
convergence implies local uniform convergence over a closed region. All we
need is insure both the bottleneck embedding terms $g$ and Jacobian terms $J$
do explode during training. Assuming this is the case, it is straightforward
that:

$\displaystyle\lim_{n\to\infty}\sum_{i=1}^{N}\Theta_{n,t}(\xi,\xi_{i})J_{n,t,i}^{\top}\chi_{n,t,i}$
$\displaystyle=\sum_{i=1}^{N}\Theta(\xi,\xi_{i})J_{t,i}^{\top}\chi_{t,i}$ (26)
$\displaystyle\lim_{n\to\infty}\sum_{i=1}^{N}J_{n,t}\Theta_{n,t}(\tilde{\xi},\xi_{i})J^{\top}_{n,t,i}\chi_{n,t,i}$
$\displaystyle=\sum_{i=1}^{N}\Theta(\tilde{\xi},\xi_{i})J_{t}J^{\top}_{t,i}\chi_{t,i}$
(27)

We now have to show how the bottleneck outputs evolve by deriving the
dynamical equation for the Jacobian. This is done by taking the input
derivative of the dynamical equation for $f(x)$. Namely:

$\displaystyle\dot{J}_{n,t}(x)=\frac{\partial}{\partial
x}\dot{f}_{n,t}(x)=-\frac{\partial}{\partial
x}\sum_{i=1}^{N}\mathcal{K}_{n,t}(x,g_{n,t,i})\chi_{n,t,i}$ (28)

We will now prove that $x,\tilde{x}$,
$\lim_{n\to\infty}\frac{\partial\mathcal{K}_{n,t}(x,\tilde{x})}{\partial
x}=\frac{\partial\mathcal{K}(x,\tilde{x})I_{d_{r}}}{\partial x}$.

Let $y(x)=\frac{ux}{\sqrt{d}},z(x)=\phi\big{(}y(x)\big{)}$. For our 1 hidden
layer MLP, the kernel $\mathcal{K}_{n}$ is given by:

$\displaystyle\mathcal{K}_{n}(x,\tilde{x})=\frac{\partial f_{n}(x)}{\partial
w}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
w}=\frac{x^{\top}\tilde{x}}{d}\cdot\frac{1}{n}\frac{\partial f(x)}{\partial
y_{n}(x)}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
y_{n}(\tilde{x})}+\frac{z_{n}(x)^{\top}z_{n}(\tilde{x})}{n}I_{d_{r}}$ (29)

Hence, we have that:

$\displaystyle\frac{\partial}{\partial
x^{\alpha}}\mathcal{K}_{n}(x,\tilde{x})=\frac{\tilde{x}^{\alpha}}{d}\cdot\frac{1}{n}\frac{\partial
f_{n}(x)}{\partial y_{n}(x)}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
y_{n}(\tilde{x})}+\frac{x^{\top}\tilde{x}}{d}\cdot\frac{1}{n}\frac{\partial^{2}f_{n}(x)}{\partial
y_{n}(x)\partial x^{\alpha}}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
y_{n}(\tilde{x})}+\frac{1}{n}\frac{\partial z_{n}(x)^{\top}}{\partial
x^{\alpha}}z_{n}(\tilde{x})I_{r_{d}}$ (30)

We now handle each term separately. For the first term, we have the trivial
limit which holds throughout training:

$\displaystyle\frac{\tilde{x}^{\alpha}}{d}\cdot\frac{1}{n}\frac{\partial
f_{n}(x)}{\partial y_{n}(x)}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
y_{n}(\tilde{x})}\to\frac{\tilde{x}^{\alpha}}{d}\dot{\Sigma}(x,\tilde{x})$
(31)

where:

$\displaystyle\dot{\Sigma}(x,\tilde{x})$
$\displaystyle=\mathbb{E}_{a,b\sim\mathcal{N}(\bm{0},\Lambda_{0}(x,\tilde{x}))}\big{[}\dot{\phi}(a)\dot{\phi}(b)\big{]},~{}~{}~{}\Lambda_{0}(x,\tilde{x})=\begin{pmatrix}\frac{\|x\|^{2}}{d}&\frac{x^{\top}\tilde{x}}{d}\\\
\frac{\tilde{x}^{\top}x}{d}&\frac{\|\tilde{x}\|^{2}}{d}\end{pmatrix}$ (32)

For the second term:

$\displaystyle\frac{\partial f_{n}^{\beta}(x)}{\partial
y_{n}^{\gamma}(g)}=v^{\beta,\gamma}\phi^{\prime\gamma}(\frac{ux}{\sqrt{d}}),~{}~{}~{}\frac{\partial^{2}f_{n}^{\beta}(x)}{\partial
y_{n}^{\gamma}(x)\partial
x^{\alpha}}=\frac{1}{\sqrt{d}}v^{\beta,\gamma}u^{\gamma,\alpha}\phi^{\prime\prime\gamma}(\frac{ux}{\sqrt{d}})$
(33)

hence at the limit:

$\displaystyle\frac{1}{n}\frac{\partial^{2}f_{n}^{\beta}(x)}{\partial
y_{n}(x)\partial x^{\alpha}}\frac{\partial
f_{n}^{\delta}(\tilde{x})^{\top}}{\partial y_{n}(\tilde{x})}$
$\displaystyle=\frac{1}{\sqrt{d}}\frac{\sum_{\gamma=1}^{n}v^{\beta,\gamma}v^{\delta,\gamma}u^{\gamma,\alpha}\phi^{\prime\prime\gamma}(\frac{ux}{\sqrt{d}})\phi^{\prime\gamma}(\frac{u\tilde{x}}{\sqrt{d}})}{n}$
(34)
$\displaystyle\to\frac{1}{\sqrt{d}}\mathbb{E}\Big{[}u^{\gamma,\alpha}\phi^{\prime\prime\gamma}(\frac{ux}{\sqrt{d}})\phi^{\prime\gamma}(\frac{u\tilde{x}}{\sqrt{d}})\Big{]}\mathbb{1}(\beta=\delta)$
(35) $\displaystyle=\mathbb{E}\Big{[}\frac{\partial}{\partial
x^{\alpha}}\phi^{\prime\gamma}(\frac{ux}{\sqrt{d}})\phi^{\prime\gamma}(\frac{u\tilde{x}}{\sqrt{d}})\Big{]}\mathbb{1}(\beta=\delta)$
(36) $\displaystyle=\frac{\partial}{\partial
x^{\alpha}}\mathbb{E}\Big{[}\phi^{\prime\gamma}(\frac{ux}{\sqrt{d}})\phi^{\prime\gamma}(\frac{u\tilde{x}}{\sqrt{d}})\Big{]}\mathbb{1}(\beta=\delta)$
(37) $\displaystyle=\frac{\partial}{\partial
x^{\alpha}}\dot{\Sigma}(x,\tilde{x})\mathbb{1}(\beta=\delta)$ (38)

where we used the independence of the coordinates of
$\phi^{\prime},\phi^{\prime\prime}$ to apply the LLN (law of large numbers)
theorem on eq. 34, and used the dominated convergence theorem to get from eq.
35 to eq. 36. We arrive at the result for the second term:

$\displaystyle\frac{x^{\top}\tilde{x}}{d}\cdot\frac{1}{n}\frac{\partial^{2}f_{n}(x)}{\partial
y_{n}(x)\partial x^{\alpha}}\frac{\partial f_{n}(\tilde{x})^{\top}}{\partial
y_{n}(\tilde{x})}\to\frac{x^{\top}\tilde{x}}{d}\frac{\partial}{\partial
x^{\alpha}}\dot{\Sigma}(x,\tilde{x})I_{d_{r}}$ (39)

For the final term we have similarly:

$\displaystyle\frac{1}{n}\frac{\partial z_{n}(x)^{\top}}{\partial
x^{\alpha}}z_{n}(\tilde{x})I_{r_{d}}=\frac{1}{\sqrt{d}}\frac{\sum_{\gamma=1}^{n}u^{\gamma,\alpha}\phi^{\prime\gamma}(\frac{ux}{\sqrt{d}})\phi^{\gamma}(\frac{u\tilde{x}}{\sqrt{d}})}{n}\to\frac{\partial}{\partial
x^{\alpha}}\Sigma(x,\tilde{x})$ (40)

where we again used the independence of the coordinates of $\phi$ and applied
LLN to the sum. Wrapping up all three terms, we arrive at:

$\displaystyle\lim_{n\to\infty}\frac{\partial}{\partial
x^{\alpha}}\mathcal{K}_{n}(x,\tilde{x})=\frac{\partial}{\partial
x^{\alpha}}\Big{(}\frac{x^{\top}\tilde{x}}{d}\dot{\Sigma}(x,\tilde{x})+\Sigma(x,\tilde{x})\Big{)}=\frac{\partial}{\partial
x^{\alpha}}\mathcal{K}(x,\tilde{x})=\Xi^{\alpha}(x,\tilde{x})$ (41)

concluding the proof. ∎

##  Appendix C Proof of lemma 4.1

See 4.1

###### Proof.

We prove the claim by deriving the update equations using 3.2 for both cases
and show their equivalence. Concretely, we must show that the functions
$\Theta,\mathcal{K},\Xi$ are the same in both cases ups to constants. When
optimizing the effective weights $w_{\text{eff}},\theta_{\text{eff}}$
directly, we trivially have that
$\Theta(\xi,\tilde{\xi})=\frac{\xi^{\top}\tilde{\xi}}{d_{0}},\mathcal{K}(x,\tilde{x})=\frac{x^{\top}\tilde{x}}{d}$
and $\Xi(x,\tilde{x})=\frac{\tilde{x}}{d}$.
When optimizing the weights
$\\{\theta^{l}\\}_{l=1}^{L_{g}},\\{w^{l}\\}_{l=1}^{L_{f}}$ instead, we use the
recursive formulas for the NTK of MLPs in Arora et al. (2019), equipped with
linear activations, and arrive at
$\Theta(\xi,\tilde{\xi})=L_{g}\frac{\xi^{\top}\tilde{\xi}}{d_{0}},\mathcal{K}(x,\tilde{x})=L_{f}\frac{x^{\top}\tilde{x}}{d}$
and $\Xi(x,\tilde{x})=L_{f}\frac{\tilde{x}}{d}$, concluding the proof. ∎

##  Appendix D Empirical support for implicit acceleration - setup and
additional results

##### Stochastic Gradient Descent

To perform empirical experiments, we resort to the discrete time version of
eqs. 12, 13 and 14 which will allow simulating SGD on the loss in function
space. Let $\xi_{i_{s}}$ denote the sample fed into the SGD algorithm at step
$s$. The discrete evolution of $\mathcal{F},g,J$ assuming a learning rate
$\mu$ is given by:

$\displaystyle g_{t+1}(\xi)$
$\displaystyle=g_{0}(\xi)-\mu\sum_{s=1}^{t}\Theta(\xi,\xi_{i_{s}})J_{s,i_{s}}^{\top}\chi_{s,i_{s}}$
(42) $\displaystyle J_{t+1}(x)$
$\displaystyle=J_{0}(x)-\mu\sum_{s=1}^{t}\chi_{s,i_{s}}\Xi(x,g_{s,i_{s}})^{\top}$
(43) $\displaystyle\mathcal{F}_{t+1}(\xi)$
$\displaystyle=\mathcal{F}_{0}(\xi)-\mu\sum_{s=1}^{t}\mathcal{K}\big{(}g_{t+1}(\xi),g_{s,i_{s}}\big{)}I_{d_{r}}\chi_{s,i_{s}}$
(44)

###  D.1 Model description

We use 1 hidden layer MLPs with ReLU555While ReLUs dont meet the criteria of a
twice differentiable activations, our results indeed hold still by
approximating the ReLU function using the softplus function
$\phi(a;m)=\frac{1}{m}\log(1+\text{e}^{ma})$, and taking its limit
$m\to\infty$ after taking $n\to\infty$. activations for both $f,g$ in all
experiments. For any pair of inputs $\xi,\tilde{\xi}$ with
$D=\|\xi\|\|\tilde{\xi}\|,\lambda=\frac{\xi^{\top}\tilde{\xi}}{D}$, the NTK
function $\Theta(\xi,\tilde{\xi})$ is given by
$\Theta(\xi,\tilde{\xi})=\frac{\xi^{\top}\tilde{\xi}}{d_{0}}\dot{\Sigma}(\xi,\tilde{\xi})+\Sigma(\xi,\tilde{\xi})$
where:

$\displaystyle\Sigma(\xi,\tilde{\xi})=\frac{D}{d_{0}}\frac{\lambda(\pi-\text{arccos}(\lambda))+\sqrt{1-\lambda^{2}}}{2\pi},~{}~{}~{}\dot{\Sigma}(\xi,\tilde{\xi})=\frac{\pi-\text{arccos}(\lambda)}{2\pi}$
(45)

$\mathcal{K}(x,\tilde{x})$ computes the same function with the replacement
$d_{0}\to d$, and the inputs
$D=\|x\|\|\tilde{x}\|,\lambda=\frac{x^{\top}\tilde{x}}{D}$.
The function $\Xi(x,\tilde{x})$ is computed by taking the derivative of
$\mathcal{K}(x,\tilde{x})$ with respect to the first input $x$. This results
in the following:

$\displaystyle\Xi(x,\tilde{x})=\frac{\tilde{x}}{d}\dot{\Sigma}+\frac{x^{\top}\tilde{x}}{d}\frac{1}{2\pi\sqrt{1-\lambda^{2}}}\big{(}\frac{\tilde{x}}{D}-\lambda\frac{x}{\|x\|^{2}}\big{)}+\frac{D}{d}\dot{\Sigma}(x,\tilde{x})\big{(}\frac{\tilde{x}}{D}-\lambda\frac{x}{\|x\|^{2}}\big{)}+\frac{x\|\tilde{x}\|}{\|x\|D}\Sigma(x,\tilde{x})$
(46)

We consider a range of bottleneck widths from infinite bottleneck (“no
bottleneck”) (labeled as inf in plots) to a maximum value that varies per
dataset as noted in the result section below. Finally, we use initialized
networks with $n=10000$ to sample $g_{0}(\xi),J_{0}(x),F_{0}(\xi)$ for all
inputs. Experiments are implemented using the Tensorflow Abadi et al. (2015)
package.

###  D.2 Tasks and datasets

We consider a regression problem on synthetic dataset in section D.4.1, multi-
class classification on MNIST dataset LeCun et al. (1989) in fig. 5 and a
binary classification problem (dog vs deer) on CIFAR-10 dataset Krizhevsky
(2009) in section D.4.3. The networks are optimized with mean squared error
(MSE) criterion where the targets are one-hot encoded vectors in
classification experiments.

###  D.3 Compute budget

We train all models for a fixed amount of time per set of hyperparameters that
include learning rate and mini-batch size per dataset. This budget is
necessary in order for us to manage the burden imposed on our compute
resources. Consequently narrower finite width bottleneck models train for a
larger number of steps compared to wider finite width bottleneck models.

###  D.4 Results

We present the results of experiments with infinite width bottleneck nonlinear
models in this section. Our main finding is that the acceleration effect
suggested by lemma 4.1 for linear models during optimization is also seen with
nonlinear models on both synthetic and real datasets. Our presentation
includes two types of plots that allows us to make observations about both
training speed as well as test performance:

##### Loss evolution

We present plots that show the evolution of loss and accuracy for both
training and test datasets. These plots allow us to make observations on the
implicit acceleration effect in infinite width bottleneck models as well as
its impact on test performance.

##### Training vs test loss

Additionally, we plot test loss as a function of training loss for several
bottleneck widths. This plot allows us to compare test performance for a given
training loss across different bottleneck widths.

The procedure used to generate this plot is described below for completeness:

  * •

determine a range for training loss values that covers all bottleneck widths
for a given dataset. We set this range to be the minimum and maximum recorded
training losses across all bottleneck widths for each dataset.

  * •

resample training loss values in the above range from training loss values
recorded in our experiments.

  * •

determine test loss values for the training loss values obtained in the above
step again by resampling with linear interpolation.

  * •

smooth both training and test loss values with a moving average filter of
length 5.

####  D.4.1 Synthetic dataset

|
---|---
(a) | (b)
Figure 2: Synthetic dataset: metrics for finite bottleneck models of various widths. Widths are shown from smallest to largest where inf indicates infinite bottleneck (a) training loss (b) test loss |
---|---
(a) | (b)
Figure 3: Synthetic dataset: heatmap view of (a) training loss evolution and (b) test loss evolution for the first 4000 steps |
---|---
(a) | (b)
Figure 4: Synthetic dataset: training loss vs test loss plot for (a) finite
and infinite width bottleneck models and (b) finite width bottleneck models

We first study the behavior of infinite width bottleneck models by running
eqs. 42, 43 and 44 on a regression problem with synthetic data. To this end,
we generate a dataset that consists of 2000 training and test samples where
the:

  * •

inputs are random vectors sampled from a standard multivariate normal
distribution.

  * •

targets are formed by passing these vectors through a deep Gaussian Process
model.

The widths considered for this dataset range from infinite width (inf) to a
width of 500. The models are trained with SGD using a mini-batch size of 20
and learning rate set to 2000.

##### Loss evolution

Figure 2a and figure 2b show the evolution of training and test loss for
models with different bottleneck widths. Figure 2a clearly shows that
architectures with narrower bottleneck train faster than their wider
bottleneck counterparts. Figure 3 shows a heatmap view of the evolution of
training and test loss for the first 4000 steps to zero in on the early phase
of learning. Figure 3 clearly demonstrates the acceleration effect of training
and test loss for small to intermediate bottleneck sizes, while no bottleneck
and large bottleneck models exhibit slower learning.

##### Training vs test loss

We observe from figure 2a that the infinite width bottleneck model exhibits a
higher training loss compared to its finite width bottleneck counterparts. To
make meaningful comparisons, we resort to using two plots for this dataset —
figure 4a shows performance of all models including infinite width bottleneck
(labeled as inf) and figure 4b that shows performance of finite width
bottlenecks only. We make the following observations from the two plots:

  * •

finite width bottleneck models lead to a lower test loss for the same training
loss over infinite width bottlenecks models.

  * •

the widest bottleneck width model (500) has the best test performance as seen
from figure 4b

The above observations suggest that using a bottleneck in infinite width
models can lead to better test performance across a range of bottleneck widths
over infinite width (“no bottleneck”) model. The narrowest bottleneck width
(10) considered here in this experiment appears to show an increase in test
loss. This, however, is an artifact that arises due the training loss range
displayed in figure 4a. We see from figure 4b the the training vs test loss
curve eventually decreases and matches values observed in models with
bottleneck widths of 20, 30 and 50. This non-monotonic behavior may be caused
by sub-optimal hyperparameter selection that we plan to address as future
work.

####  D.4.2 MNIST dataset

(a) train loss

(b) train accuracy

(c) test loss

(d) test accuracy

Figure 5: MNIST dataset: training metrics for bottleneck models with widths
from smallest to largest (inf indicates infinite width bottleneck)

We study the behavior of infinite width neural networks on MNIST dataset LeCun
et al. (1989) in this section. We train models on a multi-class (10-class)
classification problem using the MSE loss in this experiment. The bottleneck
widths used here range from “no bottleneck” (labeled as inf) to 1000 and are
trained with SGD using a learning rate of 250. As with the previous section,
we discuss implicit acceleration during optimization followed by test
performance below:

##### Loss and accuracy evolution

Figure 5 shows the evolution of training and test loss and accuracy over the
course of training. We observe that training loss decays rapidly for networks
with narrower bottlenecks. Figure 6 (appendix D) further examines training and
test loss evolution for the first 4000 steps, confirming that training
accelerates as bottleneck size narrows. Note that while narrower bottleneck
models exhibit lower test loss compared to their wider bottleneck
counterparts, this does not hold for test accuracy, potentially due to a loss-
metric mismatch as is commonly observed in practice Huang et al. (2019). These
results are inline with the analysis of linear networks with bottlenecks, that
wide networks with bottlenecks accelerate training, and confirm empirically
that this effect holds in practical nonlinear settings.

|
---|---
(a) | (b)
Figure 6: MNIST dataset: heatmap view of (a) training loss evolution and (b)
test loss evolution for the first 4000 steps
Figure 7: MNIST dataset: training loss vs test loss plot

##### Training vs test loss

Figure 7 shows a plot of training vs test loss for MNIST dataset from which we
make the following observations:

  * •

finite width bottleneck models achieve lower test loss for a given training
loss over infinite width bottlenecks.

  * •

narrower bottlenecks attain lower test loss compared to their wider finite
width counterparts for all widths considered for this dataset.

The above observations suggest that using bottlenecks in infinite width
nonlinear models can lead to a lower test loss and hence improved test
performance on real datasets as well. However, we repeat our previous
observation that lower test loss does not always lead to an improvement in
test accuracy.

####  D.4.3 CIFAR-10 dataset

Finally, we study minimizing loss in function space on CIFAR-10 Krizhevsky
(2009) which is another instance of real world data. We study a binary
classification problem using two classes from CIFAR-10 in this set of
experiments. As before, we first present our findings on acceleration effect
during training followed by observations on test performance. The models are
trained with SGD using a learning rate of 1000 in these experiments.

##### Loss and accuracy evolution

Figure 8 shows training loss, training accuracy, test loss, test accuracy for
experiments on CIFAR10 dataset. Figure 9 shows the evolution of training loss
for 4000 steps and clearly illustrates the acceleration effect in finite
bottleneck models. The accelerated learning effect is stronger in narrower
bottleneck architectures. Interestingly, we observe that lower loss value
leads to better test accuracy in this dataset unlike our observation for MNIST
in fig. 5.

##### Training vs test loss

We make the following observations from figure 10 that shows a plot of
training vs test loss for CIFAR-10 dataset:

  * •

Larger finite width bottleneck models attain a lower test loss compared to
infinite width bottleneck model.

  * •

Larger finite width bottleneck models achieve a lower test loss compared to
their narrower width counterparts.

(a) train loss

(b) train accuracy

(c) test loss

(d) test accuracy

Figure 8: CIFAR-10 dataset: training metrics for bottleneck models with widths from smallest to largest (inf indicates infinite width bottleneck) |
---|---
(a) | (b)
Figure 9: CIFAR-10 dataset: heatmap view of (a) training loss evolution and
(b) test loss evolution for the first 4000 steps
Figure 10: CIFAR-10 dataset: training loss vs test loss plot

##  Appendix E Experimental verification of the theory

Here we present results of a number of numerical simulations, serving as
sanity checks on the theoretical results at finite but large $n$. In order to
numerically verify the results stated in lemma 3.1, i.e., Jacobian’s GP
behavior at initialization, we numerically estimate the Jacobian covariance
and its cross- covariance with $f$ for a two-layer ReLU network with a large
number of hidden units $n$ by averaging over initializations (100K for $J-J$,
30K for $J-f$ and $f-f$ each). Deviation of the empirical Jacobian covariances
from their theoretical counterparts, given by eq. 8, eq. 9 and eq. 10, is
subsequently measured for randomly generated input pairs as $\lVert
cov_{empirical}-cov_{theory}\rVert_{F}/\lVert cov_{theory}\rVert_{F}$. Each
plot in Figure fig. 11 shows deviations of a given covariance structure for up
to 50 randomly generated input data pairs of dimension two and for different
values of (large) $n$.666The main objective here is to only verify
deviations/fractional errors are small for large $n$. In particular, we do not
attempt to numerically establish convergence as a function of $n$.

Next, we shift our attention to showing the evolution equations in
LABEL:{thm:main} hold numerically for sample parameter values. The following
direct strategy is used: We train four-layer ReLU networks with finite
bottlenecks and large number of hidden units $n$, with architecture described
in section 2, using full batch gradient descent (GD) and a small learning rate
$lr=10^{-3}$ (to approximate a gradient flow) on a regression task with an
L2-loss. Training data is synthetically generated. Instances are taken to be
standard normal-distributed, while labels are assigned by a fixed random
projection of the training data. In order to numerically track the network in
the functional space during training, we checkpoint values of the network
function $f$, $\chi$ (loss response) and the bottleneck embedding $g$, all
evaluated on the training instances, as well as the Jacobian $J$ (for a fixed
sample bottleneck embedding) at every training step during GD. The idea is to
numerically show these snapshots satisfy the concentrated evolution equations
eq. 12, eq. 13 and eq. 14 at their corresponding time steps, if the NTK
limiting ReLU kernels $\mathcal{K}$ and $\Theta$ are used and time derivatives
in the evolution equations are discretized by forward finite differences (with
time step size set to be the learning rate). Equivalently, we plot the
corresponding residual error per step. Figure (12) summarizes the simulation
results for a choice of parameters. Each point in the Figure (12) plots is a
residual error measured at a time step $t$ as follows
$\lVert(\psi[t+1]-\psi[t])/lr-EQ_{\psi}^{rhs}[t]\rVert/\lVert
EQ_{\psi}^{rhs}[t]\rVert$, where $\psi\in\\{f,g,J\\}$.777Besides finite $n$,
continuous time discretization error (finite learning rate) contributes to the
error. Color-coding corresponds to different training samples on which the
functional evolution equations were evaluated.

To verify the main result in lemma 4.1, we train two linear networks with
different architectures of the type described in section 4 for sample
parameter values. One of the networks is taken to have $f_{n}$ and $g_{n}$
maps as linear MLPs of depth two. The second network has the simple form
$\mathcal{F}(\xi)=w_{\text{eff}}\theta_{\text{eff}}\xi$ (as if the depth-two
components of the first network are collapsed into one effective each).
Similar to the previous case, we record the two networks dynamics in
functional space during GD. We then numerically show the evolution equations
eq. 12, eq. 13 and eq. 14 are satisfied by the empirical trajectories in the
functional space, when the appropriate NTK kernels for each linear network is
used. fig. 13 and fig. 14 show the evolution equations’ residual error for the
four and two layer linear networks described above, where kernels
$\Theta(\xi,\tilde{\xi})=L_{g}\frac{\xi^{\top}\tilde{\xi}}{d_{0}},\mathcal{K}(x,\tilde{x})=L_{f}\frac{x^{\top}\tilde{x}}{d}$
for $L_{f}=L_{g}=2$ and
$\Theta(\xi,\tilde{\xi})=\frac{\xi^{\top}\tilde{\xi}}{d_{0}},\mathcal{K}(x,\tilde{x})=\frac{x^{\top}\tilde{x}}{d}$
were used in the evolution equations respectively. The learning rate was set
to $10^{-4}$.

|
---|---
(a) | (b)
Figure 11: Deviation of the experimental covariances from their theoretical predictions in (8) for pairs of random input samples for (a) $cov_{JJ}$ (b) $cov_{ff}$ and (c) $cov_{Jf}$. See appendix E for details. |
---|---
(a) | (b)
Figure 12: Residual error for $f$, $g$ and $J$ evolution equations for $n=50K$ and a bottleneck size of three. See appendix E for details. |
---|---
(a) | (b)
Figure 13: Residual error for $f$, $g$ and $J$ evolution equations for a four-layer linear network with $n=50K$ and a bottleneck size of three. See appendix E for details.. |
---|---
(a) | (b)
Figure 14: Residual error for $f$, $g$ and $J$ evolution equations for a two-
layer effective linear network described in appendix E. Bottleneck size of
three. Refer to appendix E for details.

##  Appendix F Finite approximation of infinite width networks

We set up four-layer neural networks with a large number of dimensions $n$
(100K) to approximate infinite width layers, as well as a significantly
smaller bottleneck layer between the wide layers. We used the standard batched
stochastic gradient descent (SGD) with an L2-loss to train the model for a
multioutput regression task on the MNIST dataset. The batchsize was reduced to
5 to allow for these wider networks to fit within GPU memory. We trained
variants of the network with different bottleneck sizes and analyzed the
change in losses during training (Figure 15). The final loss and error rate
(1-accuracy) at the end of training for 36 different bottleneck sizes: [1 to
9, 10 to 90, 100 to 900, 1000 to 9000] is shown in Figure 16 . For better
clarity, Figure 15 shows the following selection of bottleneck sizes: [1, 2,
4, 8, 10, 20, 40, 80, 100, 200, 400, 800, 1000].

From the results, the best performing models have a relatively small
bottleneck size. We observe that performance improves when the bottleneck size
goes from wide to narrow, before degrading again when the width is extremely
small as the representation power degrades. At the sweet spot, the model not
only achieves lower training and test losses, it also trains faster.
Interestingly, the bottleneck size with the lowest test loss did not result in
the lowest test error (or highest accuracy). These observations are consistent
with those in the infinite width bottleneck experiments (section 5).

(a)

(b)

Figure 15: Progression of training and test losses over training steps for
$n=100K$ and bottleneck sizes from 1 to 1000.

(a)

(b)

Figure 16: Loss and error rate w.r.t. bottleneck size for $n=100K$ and
bottleneck sizes from 1 to 9000.