Luck Matters: Understanding Training Dynamics of Deep ReLU Networks

1 Introduction

Although neural networks have made strong empirical progress in a diverse set of domains (e.g., computer vision

alexnet ; vgg ; resnet , speech recognition hinton2012deep ; amodei2016deep

word2vec ; bert , and games alphago ; alphagozero ; darkforest ; mnih2013playing

), a number of fundamental questions still remain unsolved. How can Stochastic Gradient Descent (SGD) find good solutions to a complicated non-convex optimization problem? Why do neural networks generalize? How can networks trained with SGD fit both random noise and structured data

rethinking ; krueger2017deep ; neyshabur2017exploring , but prioritize structured models, even in the presence of massive noise rolnick2017deep ? Why are flat minima related to good generalization? Why does over-parameterization lead to better generalization neyshabur2018towards ; zhang2019identity ; spigler2018jamming ; neyshabur2014search ; li2018measuring ? Why do lottery tickets exist lottery ; lottery-scale ?

In this paper, we propose a theoretical framework for multilayered ReLU networks. Based on this framework, we try to explain these puzzling empirical phenomena with a unified view. We adopt a teacher-student setting where the label provided to an over-parameterized deep student ReLU network is the output of a fixed teacher ReLU network of the same depth and unknown weights (Fig. 1(a)). Here over-parameterization means that at each layer, the number of nodes in student network is more than the number of nodes in the teacher network. In this perspective, hidden student nodes are randomly initialized with different activation regions (Fig. 2(a)). During optimization, student nodes compete with each other to explain teacher nodes. From this setting, Theorem 4 shows that lucky student nodes which have greater overlap with teacher nodes converge to those teacher nodes at a fast rate, resulting in winner-take-all behavior. Furthermore, Theorem 5 shows that in the 2-layer case, if a subset of student nodes are close to the teachers’, they converge to them and the fan-out weights of other irrelevant nodes of the same layer vanishes.

With this framework, we try to intuitively explain various neural network behaviors as follows:

Fitting both structured and random data. Under gradient descent dynamics, some student nodes, which happen to overlap substantially with teacher nodes, will move into the teacher node and cover them. This is true for both structured data that corresponds to small teacher networks with few intermediate nodes, or noisy/random data that correspond to large teachers with many intermediate nodes. This explains why the same network can fit both structured and random data (Fig. 2(a-b)).

Over-parameterization. In over-parameterization, lots of student nodes are initialized randomly at each layer. Any teacher node is more likely to have a substantial overlap with some student nodes, which leads to fast convergence (Fig. 2(a) and (c), Thm. 4), consistent with lottery ; lottery-scale . This also explains that training models whose capacity just fit the data (or teacher) yields worse performance li2018measuring .

Flat minima

. Deep networks often converge to “flat minima” whose Hessian has a lot of small eigenvalues

sagun2016eigenvalues ; sagun2017empirical ; lipton2016stuck ; baity2018comparing . Furthermore, while controversial dinh2017sharp , flat minima seem to be associated with good generalization, while sharp minima often lead to poor generalization hochreiter1997flat ; keskar2016large ; wu2017towards ; li2018visualizing . In our theory, when fitting with structured data, only a few lucky student nodes converge to the teacher, while for other nodes, their fan-out weights shrink towards zero, making them (and their fan-in weights) irrelevant to the final outcome (Thm. 5), yielding flat minima in which movement along most dimensions (“unlucky nodes”) results in minimal change in output. On the other hand, sharp minima is related to noisy data (Fig. 2(d)), in which more student nodes match with the teacher.

Implicit regularization. On the other hand, the snapping behavior enforces winner-take-all: after optimization, a teacher node is fully covered (explained) by a few student nodes, rather than splitting amongst student nodes due to over-parameterization. This explains why the same network, once trained with structured data, can generalize to the test set.

Lottery Tickets. Lottery Tickets lottery ; lottery-scale ; zhou2019deconstructing is an interesting phenomenon: if we reset “salient weights” (trained weights with large magnitude) back to the values before optimization but after initialization, prune other weights (often $> 90 %$ of total weights) and retrain the model, the test performance is the same or better; if we reinitialize salient weights, the test performance is much worse. In our theory, the salient weights are those lucky regions ( $E_{j_{3}}$ and $E_{j_{4}}$ in Fig. 3) that happen to overlap with some teacher nodes after initialization and converge to them in optimization. Therefore, if we reset their weights and prune others away, they can still converge to the same set of teacher nodes, and potentially achieve better performance due to less interference with other irrelevant nodes. However, if we reinitialize them, they are likely to fall into unfavorable regions which cannot cover teacher nodes, and therefore lead to poor performance (Fig. 3(c)), just like in the case of under-parameterization. Recently, Supermask zhou2019deconstructing shows that a supermask can be found from winning tickets. If it is applied to initialized weights, the network without training gives much better test performance than chance. This is also consistent with the intuitive picture in Fig. 3(b).

Figure 1: (a) Teacher-Student Setting. For each node $j$ , the activation region is $E_{j} = {x : f_{j} (x) > 0}$ . (b) node $j$ initialized to overlap substantially with a teacher node $j_{1}^{\circ}$ converges faster towards $j_{1}^{\circ}$ (Thm. 4). (c) Student nodes initialized to be close to teacher converges to them, while the fan-out weights of other irrelevant student nodes goes to zero. (Thm. 5).

Figure 2: Explanation of implicit regularization. Blue are activation regions of teacher nodes, while orange are students’. (a) When the data labels are structured, the underlying teacher network is small and each layer has few nodes. Over-parameterization (lots of red regions) covers them all. Moreover, those student nodes that heavily overlap with the teacher nodes converge faster (Thm. 4), yield good generalization performance. (b) If a dataset contains random labels, the underlying teacher network that can fit to it has a lot of nodes. Over-parameterization can still handle them and achieves zero training error.

Figure 3: Explanation of lottery ticket phenomenon. (a) A successful training with over-parameterization (2 filters in the teacher network and 4 filters in the student network). Node $j_{3}$ and $j_{4}$ are lucky draws with strong overlap with two teacher node $j_{1}^{\circ}$ and $j_{2}^{\circ}$ , and thus converges with high weight magnitude. (b) Lottery ticket phenomenon: initialize node $j_{3}$ and $j_{4}$ with the same initial weight, clamp the weight of $j_{1}$ and $j_{2}$ to zero, and retrain the model, the test performance becomes better since $j_{3}$ and $j_{4}$ still converge to their teacher node, respectively. (c) If we reinitialize node $j_{3}$ and $j_{4}$ , it is highly likely that they are not overlapping with teacher node $j_{1}^{\circ}$ and $j_{2}^{\circ}$ so the performance is not good.

2 Mathematical Framework

Notation. Consider a student network and its associated teacher network (Fig. 1(a)). Denote the input as $x$ . For each node $j$ , denote $f_{j} (x)$ as the activation, $f_{j}^{'} (x)$ as the ReLU gating (for the top-layer, $f_{c}^{'} (x)$ and $f_{c^{\circ}}^{'} (x)$ are always $1$ ), and $g_{j} (x)$

as the backpropagated gradient, all as functions of

x

. We use the superscript

^{\circ}

to represent a teacher node (e.g.,

j^{\circ}

). Therefore,

g_{j^{\circ}}

never appears as teacher nodes are not updated. We use

w_{j k}

to represent weight between node

j

and

k

in the student network. Similarly,

w_{j^{\circ} k^{\circ}}^{*}

represents the weight between node

j^{\circ}

and

k^{\circ}

in the teacher network.

We focus on multi-layered ReLU networks. We use the following equality extensively: $σ (x) = σ^{'} (x) x$ . For ReLU node $j$ , we use $E_{j} \equiv {x : f_{j} (x) > 0}$ as the activation region of node $j$ .

Teacher network versus Dataset. The reason why we formulate the problem using teacher network rather than a dataset is the following: (1) It leads to a nice and symmetric formulation for multi-layered ReLU networks (Thm. 1). (2) A teacher network corresponds to an infinite size dataset, which separates the finite sample issues from induction bias in the dataset, which corresponds to the structure of teacher network. (3) If student weights can be shown to converge to teacher ones, generalization bound can naturally follow for the student. (4) The label complexity of data generated from a teacher is automatically reduced, which could lead to better generalization bound. On the other hand, a bound for arbitrary function class can be hard.

Objective

. We assume that both the teacher and the student output probabilities over

C

classes. We use the output of teacher as the input of the student. At the top layer, each node

c

in the student corresponds to each node

c^{\circ}

in the teacher. Therefore, the objective is:

min w J (w) = \frac{1}{2} E_{x} [∥ f_{c} (x) - f_{c^{\circ}} (x) ∥^{2}]

(1)

By the backpropagation rule, we know that for each sample $x$ , the (negative) gradient $g_{c} (x) \equiv \partial J / \partial f_{c} = f_{c^{\circ}} (x) - f_{c} (x)$ . The gradient gets backpropagated until the first layer is reached.

Note that here, the gradient $g_{c} (x)$ sent to node $c$ is correlated with the activation of the corresponding teacher node $f_{c^{\circ}} (x)$ and other student nodes at the same layer. Intuitively, this means that the gradient “pushes” the student node $c$ to align with class $c^{\circ}$ of the teacher. If so, then the student learns the corresponding class well. A natural question arises:

Are student nodes at intermediate layers correlated with teacher nodes at the same layers?

One might wonder this is hard since the student’s intermediate layer receives no direct supervision from the corresponding teacher layer, but relies only on backpropagated gradient. Surprisingly, the following theorem shows that it is possible for every intermediate layer:

Theorem 1 (Recursive Gradient Rule).

If all nodes $j$ at layer $l$ satisfies Eqn. 2

g_{j} (x) = f_{j}^{'} (x) ⎡ ⎣ \sum j^{\circ} β_{j j^{\circ}}^{*} (x) f_{j^{\circ}} (x) - \sum j^{'} β_{j j^{'}} (x) f_{j^{'}} (x) ⎤ ⎦,

(2)

then all nodes $k$ at layer $l - 1$ also satisfies Eqn. 2 with $β_{k k^{\circ}}^{*} (x)$ and $β_{k k^{'}} (x)$ defined recursively from top to bottom-layer as follows:

β_{k k^{\circ}}^{*} (x) \equiv \sum j j^{\circ} w_{j k} f_{j}^{'} (x) β_{j j^{\circ}}^{*} (x) f_{j^{\circ}}^{'} (x) w_{j^{\circ} k^{\circ}}^{*}, β_{k k^{'}} (x) \equiv \sum j j^{'} w_{j k} f_{j}^{'} (x) β_{j j^{'}} (x) f_{j^{'}}^{'} (x) w_{j^{'} k^{'}}

(3)

And for the base case where node $c$ and $c^{\circ}$ are at the top-most layer, $β_{c c^{\circ}}^{*} (x) = β_{c c^{\circ}} (x) = δ (c - c^{\circ})$ (i.e., $1$ when $c$ corresponds to $c^{\circ}$ , and $0$ otherwise).

Note that Theorem 1 applies to arbitrarily deep ReLU networks and allows different number of nodes for the teacher and student. The role played by ReLU activation is to make the expression of $β$ concise, otherwise $β_{k k^{\circ}}$ and $β_{k k^{\circ}}^{*}$ can take a very complicated (and asymmetric) form.

In particular, we consider the over-parameterization setting: the number of nodes on the student side is much larger (e.g., 5-10x) than the number of nodes on the teacher side. Using Theorem 1, we discover a novel and concise form of gradient update rule:

Assumption 1 (Separation of Expectations).

	$E_{x} [β_{j j^{\circ}}^{*} (x) f_{j}^{'} (x) f_{j^{\circ}}^{'} (x) f_{k} (x) f_{k^{\circ}} (x)]$	$=$	$E_{x} [β_{j j^{\circ}}^{*} (x)] E_{x} [f_{j}^{'} (x) f_{j^{\circ}}^{'} (x)] E_{x} [f_{k} (x) f_{k^{\circ}} (x)]$		(4)
	$E_{x} [β_{j j^{'}} (x) f_{j}^{'} (x) f_{j^{'}}^{'} (x) f_{k} (x) f_{k^{'}} (x)]$	$=$	$E_{x} [β_{j j^{'}} (x)] E_{x} [f_{j}^{'} (x) f_{j^{'}}^{'} (x)] E_{x} [f_{k} (x) f_{k^{'}} (x)]$		(5)

Theorem 2.

If Assumption 1 holds, the gradient dynamics of deep ReLU networks with objective (Eqn. 1) is:

{˙ W}_{l} = L_{l}^{*} W_{l}^{*} H_{l + 1}^{*} - L_{l} W_{l} H_{l + 1}

(6)

Here we explain the notations. $W_{l}^{*} = [w_{1}^{*}, \dots, w_{m_{l}}^{*}]$ is $m_{l}$ teacher weights, $β_{l + 1}^{*} = E_{x} [β_{j j^{\circ}} (x)]$ , $d_{j j^{\circ}}^{*} = E_{x} [f_{j}^{'} (x) f_{j^{\circ}}^{'} (x)]$ and $D_{l}^{*} = [d_{j j^{\circ}}^{*}]$ , $H_{l + 1}^{*} = [h_{j j^{\circ}}] = β_{l + 1}^{*} \circ D_{l}$ , $l_{j j^{\circ}}^{*} = E_{x} [f_{j} (x) f_{j^{\circ}} (x)]$ and $L_{l}^{*} = [l_{j j^{\circ}}^{*}]$ . We can define similar notations for $W$ (which has $n_{l}$ columns/filters), $β$ , $D$ , $H$ and $L$ (Fig. 4(c)). At the lowest layer $l = 0$ , $L_{0} = L_{0}^{*}$ , at the highest layer $l = l_{max} - 1$ where there is no ReLU, we have $β_{l_{max}} = β_{l_{max}}^{*} = H_{l_{max}} = H_{l_{max}}^{*} = I$ due to Eqn. 1. According to network structure, $β_{l + 1}$ and $β_{l + 1}^{*}$ only depends on weights $W_{l + 1}, \dots, W_{l_{max} - 1}$ , while $L_{l}$ and $L_{l}^{*}$ only depend on $W_{0}, \dots, W_{l - 1}$ .

3 Analysis on the Dynamics

Figure 4: (a) Small overlaps between node activations. Figure drawn in the space spanned by the activations of last layer so all decision boundaries are linear. (b) Lipschitz condition (Assumption 2). (c) Notation in Thm. 2.

In the following, we will use Eqn. 6 to analyze the dynamics of the multi-layer ReLU networks. For convenience, we first define the two functions $ψ_{l}$ and $ψ_{d}$ ( $σ$ is the ReLU function):

ψ_{l} (w, w^{'}) = E_{x} [σ (w^{T} x) σ ({w^{'}}^{T} x)], ψ_{d} (w, w^{'}) = E_{x} [I (w^{T} x) I ({w^{'}}^{T} x)] .

(7)

We assume these two functions have the following property .

Assumption 2 (Lipschitz condition).

There exists $K_{d}$ and $K_{l}$ so that:

∥ ψ_{i} (w, w_{1}) - ψ_{i} (w, w_{2}) ∥ \leq K_{i} ψ_{i} (w, w_{1}) ∥ w_{1} - w_{2} ∥, i \in {d, l}

(8)

Using this, we know that $d_{j j^{'}} = ψ_{d} (w_{j}, w_{j^{'}})$ , $d_{j j^{'}}^{*} = ψ_{d} (w_{j}, w_{j^{'}}^{*})$ , and so on. For brevity, denote $d_{j j^{'}}^{* *} = ψ_{d} (w_{j}^{*}, w_{j^{'}}^{*})$ (when notation $j_{1}^{\circ}$ is heavy) and so on. We impose the following assumption:

Assumption 3 (Small Overlap between teacher nodes).

There exists $ϵ_{l} ≪ 1$ and $ϵ_{d} ≪ 1$ so that:

d_{j_{1} j_{2}}^{* *} \leq ϵ_{d} d_{j_{1} j_{1}}^{* *} (o r ϵ_{d} d_{j_{2} j_{2}}^{* *}), l_{j_{1} j_{2}}^{* *} \leq ϵ_{l} l_{j_{1} j_{1}}^{* *} (o r ϵ_{l} l_{j_{2} j_{2}}^{* *}), f o r j_{1} \neq j_{2}

(9)

Intuitively, this means that the probability of the simultaneous activation of two teacher nodes $j_{1}$ and $j_{2}$ is small. If we have sufficient training data to cover the input space, then a sufficient condition for Assumption 3 to happen is that the teacher has negative bias, which means that they cut corners in the space spanned by the node activations of the lower layer (Fig. 4

a). We have empirically verified that the majority of biases in BatchNorm layers (after the data are whitened) are negative in VGG11/16 trained on ImageNet (Sec.

4.1).

3.1 Effects of BatchNorm

Batch Normalization batchnorm has been extensively used to speed up the training, reduce the tuning efforts and improve the test performance of neural networks. Here we use an interesting property of BatchNorm: the total “energy” of the incoming weights of each node $j$ is conserved over training iterations:

Theorem 3 (Conserved Quantity in Batch Normalization).

For Linear $\to$ ReLU $\to$ BN or Linear $\to$ BN $\to$ ReLU configuration, $∥ w_{j} ∥$ of a filter $j$ before BN remains constant in training. (Fig. 15).

See Appendix for the proof. The similar lemma is also in arora2018theoretical . This may partially explain why BN has stabilization effect: energy will not leak from one layer to nearby ones. Due to this property, in the following, for convenience we assume $∥ w_{j} ∥^{2} = ∥ w_{j}^{*} ∥^{2} = 1$ , and the gradient ${˙ w}_{j}$ is always orthogonal to the current weight $w_{j}$ . Note that on the teacher side we can always push the magnitude component to the upper layer; on the student side, random initialization naturally leads to constant magnitude of weights.

3.2 Same number of student nodes as teacher

We start with a simple case first. Consider that we only analyze layer $l$ without over-parameterization, i.e., $n_{l} = m_{l}$ . We also assume that $L_{l}^{*} = L_{l} = I$ , i.e., the input of that layer $l$ is whitened, and the top-layer signal is uniform, i.e., $β_{l + 1}^{*} = β_{l + 1} = 1 1^{T}$ (all $β$ entries are 1). Then the following theorem shows that weight recovery could follow (we use $j^{'}$ as $j^{\circ}$ ).

Theorem 4.

For dynamics ${˙ w}_{j} = P_{w_{j}}^{⊥} (W^{*} h_{j}^{*} - W h_{j})$ , where $P_{w_{j}}^{⊥} \equiv I - w_{j} w_{j}^{T}$ is a projection matrix into the orthogonal complement of $w_{j}$ . $h_{j}^{*}$ , $h_{j}$ are corresponding $j$ -th column in $H^{*}$ and $H$ . Denote $θ_{j} = ∠ (w_{j}, w_{j}^{*})$ and assume $θ_{j} \leq θ_{0}$ . If $γ = cos θ_{0} - (m - 1) ϵ_{d} M_{d} > 0$ , then $w_{j} \to w_{j}^{*}$ with the rate $1 - η ¯ d γ$ ( $η$ is learning rate). Here $¯ d = [1 + 2 K_{d} sin (θ_{0} / 2)] {min}_{j} d_{j j}^{* 0}$ and $M_{d} = (1 + K_{d}) {[1 + 2 K_{d} sin (θ_{0} / 2)]}_{d}^{2} / cos \frac{θ_{0}}{2}$ .

See Appendix for the proof. Here we list a few remarks:

Faster convergence near $w_{j}^{*}$ . we can see that due to the fact that $h_{j j}^{*}$ in general becomes larger when $w_{j} \to w_{j}^{*}$ (since $cos θ_{0}$ can be close to $1$ ), we expect a super-linear convergence near $w_{j}^{*}$ . This brings about an interesting winner-take-all mechanism: if the initial overlap between a student node $j$ and a particular teacher node is large, then the student node will snap to it (Fig. 1(c)).

Importance of projection operator $P_{w_{j}}^{⊥}$ . Intuitively, the projection is needed to remove any ambiguity related to weight scaling, in which the output remains constant if the top-layer weights are multiplied by a constant $α$ , while the low-layer weights are divided by $α$ . Previous works du2017gradient also uses similar techniques while we justify it with BN. Without $P_{w_{j}}^{⊥}$ , convergence can be harder.

Top-down modulation. Note that here we assume the top-layer signal is uniform, which means that according to $β_{k k^{\circ}}^{*}$ , there is no preference on which student node $k$ corresponds to which teacher node $k^{\circ}$ . If there is a preference (e.g., $β_{k k^{\circ}}^{*} = δ (k - k^{\circ})$ ), then from the proof, the cross-term $h_{k k^{\circ}}^{*}$ will be suppressed due to $β_{k k^{\circ}}^{*}$ , making convergence easier. As we will see next, such a top-down modulation plays an important role for 2-layer and over-parameterization case. We believe that it also plays a similar role for deep networks.

3.3 Over-Parameterization and Top-down Modulation in 2-layer Network

In the over-parameterization case ( $n_{l} > m_{l}$ , e.g., 5-10x), we arrange the variables into two parts: $W = [W_{u}, W_{r}]$ , where $W_{u}$ contains $m_{l}$ columns (same size as $W^{*}$ ), while $W_{r}$ contains $n_{l} - m_{l}$ columns. We use $[u]$ (or $u$ -set) to specify nodes $1 \leq j \leq m$ , and $[r]$ (or $r$ -set) for the remaining part.

In this case, if we want to show “the main component” $W_{u}$ converges to $W^{*}$ , we will meet with one core question: to where will $W_{r}$ converge, or whether $W_{r}$ will even converge at all? We need to consider not only the dynamics of the current layer, but also the dynamics of the upper layer. Using a 1-hidden layer over-parameterized ReLU network as an example, Theorem 5 shows that the upper-layer dynamics $˙ V = L^{*} V^{*} - L V$ automatically apply top-down modulation to suppress the influence of $W_{r}$ , regardless of their convergence. Here $V = [\begin{matrix} V_{u} V_{r} \end{matrix}]$ , where $V_{u}$ are the weight components of $u$ -set. See Fig. 5.

Figure 5: Over-parameterization and top-down modulation. Thm. 5 shows that under certain conditions, the relevant weights $W_{u} \to W^{*}$ and weights connecting to irrelevant student nodes $V_{r} \to 0$ .

Theorem 5 (Over-Parameterization and Top-down Modulation).

Consider $˙ W = W^{*} H^{*} - W H$ with over-parameterization ( $n > m$ ) and its upper-layer dynamics $˙ V = L^{*} V^{*} - L V$ . Assume that initial value $W_{u}^{0}$ is close to $W^{*}$ : $θ_{j} = ∠ (w_{j}, w_{j}^{*}) \leq θ_{0}$ for $j \in [u]$ . If (1) Assumption 3 holds for all pairwise combination of columns of $W^{*}$ and $W_{r}^{0}$ , and (2) there exists $γ = γ (θ_{0}, m) > 0$ and $¯ λ$ so that Eqn. 41 and Eqn. 42 holds, then $W_{u} \to W^{*}$ , $V_{u} \to V^{*}$ and $V_{r} \to 0$ with rate $1 - η ¯ λ γ$ .

See Appendix for the proof (and definition of $¯ λ$ in Eqn. 45). The intuition is: if $W_{u}$ is close to $W^{*}$ and $W_{r}$ are far away from them due to Assumption 3, the off-diagonal elements of $L$ and $L^{*}$ are smaller than diagonal ones. This causes $V_{u}$ to move towards $V^{*}$ and $V_{r}$ to move towards zero. When $V_{r}$ becomes small, so does $β_{j j^{'}}$ for $j \in [r]$ or $j^{'} \in [r]$ . This in turn suppresses the effect of $W_{r}$ and accelerates the convergence of $W_{u}$ . $V_{r} \to 0$ exponentially so that $W_{r}$ stays close to its initial locations, and Assumption 3 holds for all iterations. A few remarks:

Flat minima. Since $V_{r} \to 0$ , $W_{r}$ can be changed arbitrarily without affecting the outputs of the neural network. This could explain why there are many flat directions in trained networks, and why many eigenvalues of the Hessian are close to zero sagun2016eigenvalues .

Understanding of pruning methods. Theorem 5 naturally relates two different unstructured network pruning approaches: pruning small weights in magnitude han2015learning ; lottery and pruning weights suggested by Hessian lecun1990optimal ; hassibi1993optimal . It also suggests a principled structured pruning method: instead of pruning a filter by checking its weight norm, pruning accordingly to its top-down modulation.

Accelerated convergence and learning rate schedule. For simplicity, the theorem uses a uniform (and conservative) $γ$ throughout the iterations. In practice, $γ$ is initially small (due to noise introduced by $r$ -set) but will be large after a few iterations when $V_{r}$ vanishes. Given the same learning rate, this leads to accelerated convergence. At some point, the learning rate $η$ becomes too large, leading to fluctuation. In this case, $η$ needs to be reduced.

Many-to-one mapping. Theorem 5 shows that under strict conditions, there is one-to-one correspondence between teacher and student nodes. In general this is not the case. Two students nodes can be both in the vicinity of a teacher node $w_{j}^{*}$ and converge towards it, until that node is fully explained. We leave it to the future work for rigid mathematical analysis of many-to-one mappings.

Random initialization. One nice thing about Theorem 5 is that it only requires the initial $∥ W_{u} - W^{*} ∥$ to be small. In contrast, there is no requirement for small $∥ V_{r} ∥$ . Therefore, we could expect that with more over-parameterization and random initialization, in each layer $l$ , it is more likely to find the $u$ -set (of fixed size $m_{l}$ ), or the lucky weights, so that $W_{u}$ is quite close to $W^{*}$ . At the same time, we don’t need to worry about $∥ W_{r} ∥$

which grows with more over-parameterization. Moreover, random initialization often gives orthogonal weight vectors, which naturally leads to Assumption

3.4 Extension to Multi-layer ReLU networks

Using a similar approach, we could extend this analysis to multi-layer cases. We conjecture that similar behaviors happen: for each layer, due to over-parameterization, the weights of some lucky student nodes are close to the teacher ones. While these converge to the teacher, the final values of others irrelevant weights are initialization-dependent. If the irrelevant nodes connect to lucky nodes at the upper-layer, then similar to Thm. 5, the corresponding fan-out weights converge to zero. On the other hand, if they connect to nodes that are also irrelevant, then these fan-out weights are not-determined and their final values depends on initialization. However, it doesn’t matter since these upper-layer irrelevant nodes eventually meet with zero weights if going up recursively, since the top-most output layer has no over-parameterization. We leave a formal analysis to future work.

Figure 6: Distribution of BatchNorm bias in pre-trained VGG16 on ImageNet. Orange/blue are positive/negative biases. Conv0 corresponds to the lowest layer (closest to the input). VGG11/13/19 in Fig. 18.

4 Simulations

4.1 Checking Assumption 3

To make Theorem 4 and Theorem 5 work, we make Assumption 3 that the activation field of different teacher nodes should be well-separated. To justify this, we analyze the bias of BatchNorm layers after the convolutional layers in pre-trained VGG11/13/16/19. We check the BatchNorm bias $c_{1}$ as these models use Linear-BatchNorm-ReLU architecture. After BatchNorm first normalizes the input data into zero mean distribution, the BatchNorm bias determines how much data pass the ReLU threshold. If the bias is negative, then a small portion of data pass ReLU gating and Assumption 3 is likely to hold. From Fig. 6, it is quite clear that the majority of BatchNorm bias parameters are negative, in particular for the top layers.

4.2 Numerical Experiments of Thm. 5

We verify Thm. 5 by checking whether $V_{r}$ moves close to $0$ under different initialization. We use a network with one hidden layer. The teacher network is 10-20-30, while the student network has more nodes in the hidden layers. Input data are Gaussian noise. We initialize the student networks so that the first $20$ nodes are close to the teacher. Specifically, we first create matrices $W_{ϵ}$ and $V_{ϵ}$ by first filling with i.i.d Gaussian noise, and then normalizing their columns to $1$ . Then the initial value of student $W$ is $W_{u}^{0} = c o l u m n_n o r m a l i z e (p_{W} W^{*} + W_{ϵ})$ , where $p_{W}$ is a factor controlling how close $W_{u}^{0}$ is to $W^{*}$ . For $W_{r}$ we initialize with noise. Similarly we initialize $V_{u}$ with a factor $p_{V}$ . The larger $p_{W}$ and $p_{V}$ , the close the initialization $W_{u}^{0}$ and $V_{u}^{0}$ to the ground truth values.

Fig. 7 shows the behavior over different iterations. All experiments are repeated 32 times with different random seeds, and (mean $\pm$

std) are reported. We can see that a close initialization leads to faster (and low variance) convergence of

V_{r}

to small values. In particular, it is important to have

W_{u}^{0}

close to

W^{*}

(large

p_{W}

), which leads to a clear separation between row norms of

V_{u}

and

V_{r}

, even if they are close to each other at the beginning of training. Having

V_{u}^{0}

close to

V^{*}

makes the initial gap larger and also helps convergence. On the other hand, if

p_{W}

is small, then even if

p_{V}

is large, the gap between row norms of

V_{u}

and

V_{r}

only shifts but doesn’t expand over iterations.

Figure 7: Numerical verification of Thm. 5. Rows are $2 \times$ , $5 \times$ and $10 \times$ over-parameterization. Columns are different initializations of student networks. All figures show $∥ v_{j} ∥$ , i.e., norm of each row of $V$ .

5 Experiments

5.1 Experiment Setup

We evaluate both the fully connected (FC) and ConvNet setting. For FC, we use a ReLU teacher network of size 50-75-100-125. For ConvNet, we use a teacher with channel size 64-64-64-64. The student networks have the same depth but with $10 \times$ nodes/channels at each layer, such that they are substnatially over-parameterized. When BatchNorm is added, it is added after ReLU.

We use random i.i.d Gaussian inputs with mean 0 and std $10$ (abbreviated as GAUS) and CIFAR-10 as our dataset in the experiments. GAUS generates infinite number of samples while CIFAR-10 is a finite dataset. For GAUS, we use a random teacher network as the label provider (with $100$ classes). To make sure the weights of the teacher are weakly overlapped, we sample each entry of $w_{j}^{*}$ from $[- 0.5, - 0.25, 0, 0.25, 0.5]$ , making sure they are non-zero and mutually different within the same layer, and sample biases from $U [- 0.5, 0.5]$ . In the FC case, the data dimension is 20 while in the ConvNet case it is $16 \times 16$ . For CIFAR-10 we use a pre-trained teacher network with BatchNorm. In the FC case, it has an accuracy of $54.95 %$ ; for ConvNet, the accuracy is $86.4 %$ . We repeat 5 times for all experiments, with different random seed and report min/max values.

Two metrics are used to check our prediction that some lucky student nodes converge to the teacher:

Normalized correlation $¯ ρ$

. We compute normalized correlation (or cosine similarity)

ρ

between teacher and student activations¹¹1For

f_{j} = [f_{j} (x_{1}), \dots, f_{j} (x_{n})]

and

f_{j^{\circ}}

ρ_{j j^{\circ}} = {~ f}_{j}^{T} {~ f}_{j^{\circ}} \in [- 1, 1]

, where

{~ f}_{j} = (f_{j} - m e a n (f_{j})) / s t d (f_{j})

. evaluated on a validation set. At each layer, we average the best correlation over teacher nodes:

¯ ρ = {m e a n}_{j^{\circ}} {max}_{j} ρ_{j j^{\circ}}

, where

ρ_{j j^{\circ}}

is computed for each teacher and student pairs

(j, j^{\circ})

¯ ρ \approx 1

means that most teacher nodes are covered by at least one student.

Mean Rank $¯ r$ . After training, each teacher node $j^{\circ}$ has the most correlated student node $j$ . We check the correlation rank of $j$ , normalized to $[0, 1]$ ( $0$

=rank first), back at initialization and at different epochs, and average them over teacher nodes to yield mean rank

¯ r

. Small

¯ r

means that student nodes that initially correlate well with the teacher keeps the lead toward the end of training.

Figure 8: Correlation $¯ ρ$ and mean rank $¯ r$ over training on GAUS. $¯ ρ$ steadily grows and $¯ r$ quickly improves over time. Layer-0 (the lowest layer that is closest to the input) shows best match with teacher nodes and best mean rank. BatchNorm helps achieve both better correlation and lower $¯ r$ , in particular for the CNN case.

Figure 9: Same experiment setting as in Fig. 8 on CIFAR-10. BatchNorm helps achieve lower $¯ r$ .

Figure 10: Ablation studies on GAUS. (a) $¯ ρ$ converges much faster in small models (10-15-20-25) than in large model (50-75-100-125). BatchNorm hurts in small models. (b) $¯ ρ$ stalls using finite samples.

Figure 11: Visualization of (transpose of) $H^{*}$ and $β^{*}$ matrix before and after optimization (using GAUS). Student node indices are reordered according to teacher-student node correlations. After optimization, student node who has high correlation with the teacher node also has high $β$ entries. Such a behavior is more prominent in $H^{*}$ matrix that combines $β^{*}$ and the activation patterns $D^{*}$ of student nodes (Sec. 2).

5.2 Results

Experiments are summarized in Fig. 8 and Fig. 9. $¯ ρ$ indeed grows during training, in particular for low layers that are closer to the input, where $¯ ρ$ moves towards $1$ . Furthermore, the final winning student nodes also have a good rank at the early stage of training, in particular after the first epoch, which is consistent with late-resetting used in lottery-scale . BatchNorm helps a lot, in particular for the CNN case with GAUS dataset. For CIFAR-10, the final evaluation accuracy (see Appendix) learned by the student is often $\sim 1 %$ higher than the teacher. Using BatchNorm accelerates the growth of accuracy, improves $¯ r$ , but seems not to accelerate the growth of $¯ ρ$ .

The theory also predicts that the top-down modulation $β$ helps the convergence. For this, we plot $β_{j j^{\circ}}^{*}$ at different layers during optimization on GAUS. For better visualization, we align each student node index $j$ with a teacher node $j^{\circ}$ according to highest $ρ$

. Despite the fact that correlations are computed from the low-layer weights, it matches well with the top-layer modulation (identity matrix structure in Fig.

11). Besides, we also perform ablation studies on GAUS.

Size of teacher network. As shown in Fig. 10(a), for small teacher networks (FC 10-15-20-25), the convergence is much faster and training without BatchNorm is faster than training with BatchNorm. For large teacher networks, BatchNorm definitely increases convergence speed and growth of $¯ ρ$ .

Figure 12: Effect of different degrees of over-parameterization. Left panel: Teacher FC (10-15-20-25) without batchnorm, right panel: teacher CNN (10-15-20-25) with batchnorm.

Degree of over-parameterization. Fig. 12 shows the effects of different degree of over-parameterization ( $1 \times$ , $2 \times$ , $5 \times$ , $10 \times$ , $20 \times$ and $50 \times$ ). We initialize 32 different teacher network (10-15-20-25) with different random seed, and plot $\pm$ standard derivation with shaded area. We can clearly see that $¯ ρ$ grows more stably and converges to higher values with over-parameterization. On the other hand, $20 \times$ and $50 \times$ are slower in convergence due to excessive parameters.

Finite versus Infinite Dataset. We also repeat the experiments with a pre-generated finite dataset of GAUS in the CNN case (Fig. 10(b)), and find that the convergence of node similarity stalls after a few iterations. This is because some nodes receive very few data points in their activated regions, which is not a problem for infinite dataset. We suspect that this is probably the reason why CIFAR-10, as a finite dataset, does not show similar behavior as GAUS.

6 Conclusion and Future Work

In this paper we propose a new theoretical framework that uses teacher-student setting to understand the training dynamics of multi-layered ReLU network. With this framework, we are able to conceptually explain many puzzling phenomena in deep networks, such as why over-parameterization helps generalization, why the same network can fit to both random and structured data, why lottery tickets lottery ; lottery-scale exist. We backup these intuitive explanations by Theorem 4 and Theorem 5, which collectively show that student nodes that are initialized to be close to the teacher nodes converge to them with a faster rate, and the fan-out weights of irrelevant nodes converge to zero. As the next steps, we aim to extend Theorem 5 to general multi-layer setting (when both $L$ and $H$ are present), relax Assumption 3 and study more BatchNorm effects than what Theorem 3 suggests.

7 Acknowledgement

The first author thanks Simon Du, Jason Lee, Chiyuan Zhang, Rong Ge, Greg Yang, Jonathan Frankle and many others for the informal discussions.

References

[1] Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, et al. Deep speech 2: End-to-end speech recognition in english and mandarin. In
International conference on machine learning
, pages 173–182, 2016.
[2] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. arXiv preprint arXiv:1812.03981, 2018.
[3] Marco Baity-Jesi, Levent Sagun, Mario Geiger, Stefano Spigler, G Ben Arous, Chiara Cammarota, Yann LeCun, Matthieu Wyart, and Giulio Biroli. Comparing dynamics: Deep neural networks versus glassy systems. ICML, 2018.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[5] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1019–1028. JMLR. org, 2017.
[6] Simon S Du, Jason D Lee, Yuandong Tian, Barnabas Poczos, and Aarti Singh. Gradient descent learns one-hidden-layer cnn: Don’t be afraid of spurious local minima. ICML, 2018.
[7] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Training pruned neural networks. ICLR, 2019.
[8] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. The lottery ticket hypothesis at scale. arXiv preprint arXiv:1903.01611, 2019.
[9] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. In Advances in neural information processing systems, pages 1135–1143, 2015.
[10] Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE international conference on neural networks, pages 293–299. IEEE, 1993.
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In
Proceedings of the IEEE conference on computer vision and pattern recognition
, pages 770–778, 2016.
[12] Geoffrey Hinton, Li Deng, Dong Yu, George Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Brian Kingsbury, et al. Deep neural networks for acoustic modeling in speech recognition. IEEE Signal processing magazine, 29, 2012.
[13] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1–42, 1997.
[14] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. ICML, 2015.
[15] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.
[16] Jonas Kohler, Hadi Daneshmand, Aurelien Lucchi, Ming Zhou, Klaus Neymeyr, and Thomas Hofmann. Towards a theoretical understanding of batch normalization. arXiv preprint arXiv:1805.10694, 2018.
[17] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.
Imagenet classification with deep convolutional neural networks.
In Advances in neural information processing systems, pages 1097–1105, 2012.
[18] David Krueger, Nicolas Ballas, Stanislaw Jastrzebski, Devansh Arpit, Maxinder S Kanwal, Tegan Maharaj, Emmanuel Bengio, Asja Fischer, and Aaron Courville. Deep nets don’t learn via memorization. ICLR Workshop, 2017.
[19] Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in neural information processing systems, pages 598–605, 1990.
[20] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. ICLR, 2018.
[21] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems, pages 6389–6399, 2018.
[22] Zachary C Lipton. Stuck in a what? adventures in weight space. arXiv preprint arXiv:1602.07320, 2016.
[23] Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In Advances in neural information processing systems, pages 3111–3119, 2013.
[24] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[25] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. In Advances in Neural Information Processing Systems, pages 5947–5956, 2017.
[26] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of over-parametrization in generalization of neural networks. arXiv preprint arXiv:1805.12076, 2018.
[27] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. ICLR Workshop, 2015.
[28] David Rolnick, Andreas Veit, Serge Belongie, and Nir Shavit. Deep learning is robust to massive label noise. arXiv preprint arXiv:1705.10694, 2017.
[29] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. ICLR, 2017.
[30] Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. ICLR 2018 Workshop Contribution, arXiv preprint arXiv:1706.04454, 2018.
[31] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484, 2016.
[32] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. Nature, 550(7676):354, 2017.
[33] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014.
[34] Stefano Spigler, Mario Geiger, Stéphane d’Ascoli, Levent Sagun, Giulio Biroli, and Matthieu Wyart. A jamming transition from under-to over-parametrization affects loss landscape and generalization. arXiv preprint arXiv:1810.09665, 2018.
[35] Yuandong Tian. A theoretical framework for deep locally connected relu network. arXiv preprint arXiv:1809.10829, 2018.
[36] Yuandong Tian and Yan Zhu. Better computer go player with neural network and long-term prediction. ICLR, 2016.
[37] Lei Wu, Zhanxing Zhu, et al. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv preprint arXiv:1706.10239, 2017.
[38] Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz. A mean field theory of batch normalization. ICLR, 2019.
[39] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. ICLR, 2017.
[40] Chiyuan Zhang, Samy Bengio, Moritz Hardt, and Yoram Singer. Identity crisis: Memorization and generalization under extreme overparameterization. arXiv preprint arXiv:1902.04698, 2019.
[41] Hattie Zhou, Janice Lan, Rosanne Liu, and Jason Yosinski. Deconstructing lottery tickets: Zeros, signs, and the supermask. arXiv preprint arXiv:1905.01067, 2019.

8 Appendix: Proofs

Figure 13:
Teacher-Student Setting, loss function and notations.

8.1 Theorem 1

Proof.

The first part of gradient backpropagated to node $j$ is:

$g_{j}^{1} (x)$	$=$	$f_{j}^{'} (x) \sum j^{\circ} β_{j j^{\circ}}^{*} (x) f_{j^{\circ}} (x)$	(10)
	$=$	$f_{j}^{'} (x) \sum j^{\circ} β_{j j^{\circ}}^{} (x) f_{j^{\circ}}^{'} (x) \sum k^{\circ} w_{j^{\circ} k^{\circ}}^{} f_{k^{\circ}} (x)$	(11)
	$=$	$\sum k^{\circ} ⎡ ⎣ f_{j}^{'} (x) \sum j^{\circ} β_{j j^{\circ}}^{} (x) f_{j^{\circ}}^{'} (x) w_{j^{\circ} k^{\circ}}^{} ⎤ ⎦ f_{k^{\circ}} (x)$	(12)

Therefore, for the gradient to node $k$ , we have:

	$g_{k}^{1} (x)$	$=$	$f_{k}^{'} (x) \sum j w_{j k} g_{j}^{1} (x)$		(13)
		$=$	$f_{k}^{'} (x) \sum k^{\circ} ⎡ ⎣ \sum j j^{\circ} w_{j k} f_{j}^{'} (x) β_{j j^{\circ}}^{} (x) f_{j^{\circ}}^{'} (x) w_{j^{\circ} k^{\circ}}^{} ⎤ ⎦      β_{k k^{\circ}}^{*} (x) f_{k^{\circ}} (x)$		(14)

And similar for $β_{k k^{'}} (x)$ . Therefore, by mathematical induction, we know that all gradient at nodes in different layer follows the same form. ∎

Figure 15: Different BatchNorm Configuration.

8.2 Theorem 2

Proof.

Using Thm. 1, we can write down weight update for weight $w_{j k}$ that connects node $j$ and node $k$ :

	${˙ w}_{j k}$	$=$	$\sum j^{\circ}, k^{\circ} w_{j^{\circ} k^{\circ}}^{} E_{x} [f_{j}^{'} (x) β_{j j^{\circ}}^{} (x) f_{j^{\circ}}^{'} (x) f_{k} (x) f_{k^{\circ}} (x)]      β_{j j^{\circ} k k^{\circ}}^{*}$		(15)
		$-$	$\sum j^{'}, k^{'} w_{j^{'} k^{'}} E_{x} [f_{j}^{'} (x) β_{j j^{'}} (x) f_{j^{'}}^{'} (x) f_{k} (x) f_{k^{'}} (x)]      β_{j j^{'} k k^{'}}$		(15)

Note that $j^{\circ}$ , $k^{\circ}$ , $j^{'}$ and $k^{'}$ run over all parents and children nodes on the teacher side. This formulation works for over-parameterization (e.g., $j^{\circ}$ and $j^{'}$ can run over different nodes). Applying Assumption 1 and rearrange terms in matrix form yields Eqn. 6. ∎

8.3 Theorem 3

Proof.

Given a batch with size $N$ , denote pre-batchnorm activations as $f = [f_{j} (x_{1}), \dots, f_{j} (x_{N})]^{T}$ and gradients as $g = [g_{j} (x_{1}), \dots, g_{j} (x_{N})]^{T}$ (See Fig. 14(a)). $~ f = (f - μ) / σ$ is its whitened version, and $c_{0} ~ f + c_{1}$ is the final output of BN. Here $μ = \frac{1}{N} \sum_{i} f_{j} (x_{i})$ and $σ^{2} = \frac{1}{N} \sum_{i} (f_{j} (x_{i}) - μ)^{2}$ and $c_{1}$ , $c_{0}$

are learnable parameters. With vector notation, the gradient update in BN has a compact form with clear geometric meaning:

Lemma 1 (Backpropagation of Batch Norm [35]).

For a top-down gradient $g$ , BN layer gives the following gradient update ( $P_{f, 1}^{⊥}$ is the orthogonal complementary projection of subspace ${f, 1}$ ):

g_{f} = J^{B N} (f) g = \frac{c_{0}}{σ} P_{f, 1}^{⊥} g, g_{c} = S (f)^{T} g

(16)

Intuitively, the back-propagated gradient $J^{B N} (f) g$ is zero-mean and perpendicular to the input activation $f$ of BN layer, as illustrated in Fig. 14. Unlike [16, 38] that analyzes BN in an approximate manner, in Thm. 1 we do not impose any assumptions.

Given Lemma 1, we can prove Thm. 3. For Fig. 15(a), using the property that $E_{x} [g_{j}^{l i n} f_{j}^{l i n}] = 0$ (the expectation is taken over batch) and the weight update rule ${˙ w}_{j k} = E_{x} [g_{j}^{l i n} f_{k}]$ (over the same batch), we have:

\frac{1}{2} \frac{d ∥ w_{j} ∥^{2}}{d t} = \sum k \in c h (j) w_{j k} {˙ w}_{j k} = E_{x} ⎡ ⎣ \sum k \in c h (j) w_{j k} f_{k} (x) g_{j}^{l i n} (x) ⎤ ⎦ = E_{x} [f_{j}^{l i n} (x) g_{j}^{l i n} (x)] = 0

(17)

For Fig. 15(b), note that $E_{x} [g_{j}^{l i n} f_{j}^{l i n}] = E_{x} [g_{j}^{r l} f_{j}^{r l^{'}} f_{j}^{l i n}] = E_{x} [g_{j}^{r l} f_{j}^{r l}] = 0$ and conclusion follows. ∎

8.4 Lemmas

For simplicity, in the following, we use $δ w_{j} = w_{j} - w_{j}^{*}$ .

Lemma 2 (Bottom Bounds).

Assume all $∥ w_{j} ∥ = ∥ w_{j^{'}} ∥ = 1$ . Denote

p_{j j^{'}}^{*} \equiv w_{j^{'}}^{*} d_{j j^{'}}^{*}, p_{j j^{'}} \equiv w_{j^{'}} d_{j j^{'}}, Δ p_{j j^{'}} \equiv p_{j j^{'}}^{*} - p_{j j^{'}}

(18)

If Assumption 2 holds, we have:

∥ Δ p_{j j^{'}} ∥ \leq (1 + K_{d}) d_{j j^{'}}^{*} ∥ δ w_{j^{'}} ∥

(19)

If Assumption 3 also holds, then:

d_{j j^{'}}^{*} \leq ϵ_{d} (1 + K_{d} ∥ δ w_{j^{'}} ∥) (1 + K_{d} ∥ δ w_{j} ∥) d_{j j}^{*}

(20)

Proof.

We have for $j \neq j^{'}$ :

$∥ Δ p_{j j^{'}} ∥$	$=$	$∥ w_{j^{'}}^{} d_{j j^{'}}^{} - w_{j^{'}} d_{j j^{'}} ∥$	(21)
	$=$	$∥ w_{j^{'}} (d_{j j^{'}}^{} - d_{j j^{'}}) + (w_{j^{'}}^{} - w_{j^{'}}) d_{j j^{'}}^{*} ∥$	(22)
	$\leq$	$∥ w_{j^{'}} ∥ ∥ d_{j j^{'}}^{} - d_{j j^{'}} ∥ + ∥ w_{j^{'}}^{} - w_{j^{'}} ∥ d_{j j^{'}}^{*}$	(23)
	$\leq$	$d_{j j^{'}}^{} K_{d} ∥ δ w_{j^{'}} ∥ + d_{j j^{'}}^{} ∥ δ w_{j^{'}} ∥$	(24)
	$\leq$	$(1 + K_{d}) d_{j j^{'}}^{*} ∥ δ w_{j^{'}} ∥$	(25)

If Assumption 3 also holds, we have:

$d_{j j^{'}}^{*}$	$\leq$	$d_{j j^{'}}^{* *} (1 + K_{d} ∥ δ w_{j^{'}} ∥)$	(26)
	$\leq$	$ϵ_{d} d_{j j}^{* *} (1 + K_{d} ∥ δ w_{j^{'}} ∥$	(27)
	$\leq$	$ϵ_{d} d_{j j}^{*} (1 + K_{d} ∥ δ w_{j} ∥) (1 + K_{d} ∥ δ w_{j^{'}} ∥)$	(28)

∎

Lemma 3 (Top Bounds).

Denote

q_{j j^{'}}^{*} \equiv v_{j^{'}}^{*} l_{j j^{'}}^{*}, q_{j j^{'}} \equiv v_{j^{'}} l_{j j^{'}}, Δ q_{j j^{'}} \equiv q_{j j^{'}}^{*} - q_{j j^{'}}

(29)

If Assumption 2 holds, we have:

∥ Δ q_{j j^{'}} ∥ \leq (1 + K_{l}) l_{j j^{'}}^{*} ∥ δ w_{j^{'}} ∥

(30)

If Assumption 3 also holds, then:

l_{j j^{'}}^{*} \leq ϵ_{l} (1 + K_{l} ∥ δ w_{j^{'}} ∥) (1 + K_{l} ∥ δ w_{j} ∥) l_{j j}^{*}

(31)

Proof.

The proof is similar to Lemma 2. ∎

Lemma 4 (Quadratic fall-off for diagonal elements of $L$ ).

For node $j$ , we have:

∥ l_{j j}^{*} - l_{j j} ∥ \leq C_{0} l_{j j}^{*} ∥ δ w_{j} ∥^{2}

(32)

Proof.

The intuition here is that both the volume of the affected area and the weight difference are proportional to $∥ δ w_{j} ∥$ . $∥ l_{j j}^{*} - l_{j j} ∥$ is their product and thus proportional to $∥ δ w_{j} ∥^{2}$ . See Fig. 16. ∎

8.5 Theorem 4

Proof.

First of all, note that $∥ δ w_{j} ∥ = 2 sin \frac{θ_{j}}{2} \leq 2 sin \frac{θ_{0}}{2}$ . So given $θ_{0}$ , we also have a bound for $∥ δ w_{j} ∥$ .

When $β = β^{*} = 1 1^{T}$ , the matrix form can be written as the following:

{˙ w}_{j} = P_{w_{j}}^{⊥} w_{j}^{*} h_{j j}^{*} + \sum j^{'} \neq j P_{w_{j}}^{⊥} (w_{j^{'}}^{*} h_{j j^{'}}^{*} - w_{j^{'}} h_{j j^{'}}) = P_{w_{j}}^{⊥} p_{j j}^{*} + \sum j^{'} \neq j P_{w_{j}}^{⊥} Δ p_{j j^{'}}

(33)

by using $P_{w_{j}}^{⊥} w_{j} \equiv 0$ (and thus $h_{j j}$ doesn’t matter). Since $∥ w_{j} ∥$ is conserved, it suffices to check whether the projected weight vector $P_{w_{j}^{*}}^{⊥} w_{j}$ of $w_{j}$ onto the complementary space of the ground truth node $w_{j}^{*}$ , goes to zero:

P_{w_{j}^{*}}^{⊥} {˙ w}_{j} = P_{w_{j}^{*}}^{⊥} P_{w_{j}}^{⊥} p_{j j}^{*} + \sum j^{'} \neq j P

Luck Matters: Understanding Training Dynamics of Deep ReLU Networks

Authors

Over-parameterization as a Catalyst for Better Generalization of Deep ReLU network

A theoretical framework for deep locally connected ReLU network

From complex to simple : hierarchical free-energy landscape renormalized in deep neural networks

FitNets: Hints for Thin Deep Nets

Learning performance in inverse Ising problems with sparse teacher couplings

Ternary Neural Networks for Resource-Efficient AI Applications

Hidden Unit Specialization in Layered Neural Networks: ReLU vs. Sigmoidal Activation

This week in AI

1 Introduction

2 Mathematical Framework

Theorem 1 (Recursive Gradient Rule).

Assumption 1 (Separation of Expectations).

Theorem 2.

3 Analysis on the Dynamics

Assumption 2 (Lipschitz condition).

Assumption 3 (Small Overlap between teacher nodes).

3.1 Effects of BatchNorm

Theorem 3 (Conserved Quantity in Batch Normalization).

3.2 Same number of student nodes as teacher

Theorem 4.

3.3 Over-Parameterization and Top-down Modulation in 2-layer Network

Theorem 5 (Over-Parameterization and Top-down Modulation).

3.4 Extension to Multi-layer ReLU networks

4 Simulations

4.1 Checking Assumption 3

4.2 Numerical Experiments of Thm. 5

5 Experiments

5.1 Experiment Setup

5.2 Results

6 Conclusion and Future Work

7 Acknowledgement

References

8 Appendix: Proofs

8.1 Theorem 1

Proof.

8.2 Theorem 2

Proof.

8.3 Theorem 3

Proof.

Lemma 1 (Backpropagation of Batch Norm [35]).

8.4 Lemmas

Lemma 2 (Bottom Bounds).

Proof.

Lemma 3 (Top Bounds).

Proof.

Lemma 4 (Quadratic fall-off for diagonal elements of L).

Proof.

8.5 Theorem 4

Proof.

Lemma 4 (Quadratic fall-off for diagonal elements of $L$ ).