SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Berfin Şimsek NYU*
bs3736@nyu.edu
Amire Bendjeddou EPFL
amire.bendjeddou@epfl.ch
Wulfram Gerstner EPFL
wulfram.gerstner@epfl.ch
Johanni Brea EPFL
johanni.brea@epfl.ch
Any continuous function  $f^{*}$  can be approximated arbitrarily well by a neural network with sufficiently many neurons  $k$ . We consider the case when  $f^{*}$  itself is a neural network with one hidden layer and  $k$  neurons. Approximating  $f^{*}$  with a neural network with  $n < k$  neurons can thus be seen as fitting an under-parameterized "student" network with  $n$  neurons to a "teacher" network with  $k$  neurons. As the student has fewer neurons than the teacher, it is unclear, whether each of the  $n$  student neurons should copy one of the teacher neurons or rather average a group of teacher neurons. For shallow neural networks with erf activation function and for the standard Gaussian input distribution, we prove that "copy-average" configurations are critical points if the teacher's incoming vectors are orthonormal and its outgoing weights are unitary. Moreover, the optimum among such configurations is reached when  $n - 1$  student neurons each copy one teacher neuron and the  $n$ -th student neuron averages the remaining  $k - n + 1$  teacher neurons. For the student network with  $n = 1$  neuron, we provide additionally a closed-form solution of the non-trivial critical point(s) for commonly used activation functions through solving an equivalent constrained optimization problem. Empirically, we find for the erf activation function that gradient flow converges either to the optimal copy-average critical point or to another point where each student neuron approximately copies a different teacher neuron. Finally, we find similar results for the ReLU activation function, suggesting that the optimal solution of underparameterized networks has a universal structure.
A shallow neural network with a single hidden layer of a large number  $k$  of neurons can approximate any continuous function  $f^{*}$  arbitrarily well on a compact subset of the input space [1]. We consider a related problem, where the function  $f^{*}$  itself is a neural network with a large number  $k$  of neurons, and its approximation is a smaller network with  $n < k$  neurons. In other words, we fit an underparameterized "student" network with  $n$  neurons to a "teacher" network with  $k$  neurons. As the student has fewer neurons than the teacher, it cannot perfectly match the teacher. In the configuration with the lowest loss, where the approximation error is smallest, one may expect that the incoming and outgoing weights of a student neuron are either identical to those of a teacher neuron or that they are aligned with the weights of a group of teacher neurons, but it is unclear what the optimal configuration is.
*Previous address: EPFL.
To answer the question of whether student neurons should "copy" or "average" teacher neurons, and more generally to shed light on the loss landscape of under-parameterized neural networks, we study the theoretically tractable setup with standard Gaussian input data and teacher networks with orthogonal incoming vectors. First, we re-parameterize the loss in terms of interactions between pairs of neurons, similar to [2, 3], and we re-formulate the original optimization problem as a constrained optimization problem. The interactions between neurons can be written as a function expressed in terms of the standard deviation and correlation of two Gaussian random variables, with explicit formulas for the erf and ReLU activation functions [2-4]. Next, we prove several properties of the most extremely under-parameterized student network with a single neuron  $n = 1$ , extending thus the important work of [5-7]. For many commonly used activation functions, we prove for the network with a single hidden neuron that the optimal solution is the only non-trivial critical point of the loss function up to symmetries and is achieved when the incoming vector of the one-neuron student reaches a configuration that can be interpreted as a damped average of all incoming teacher weights.
The proof relies on identifying the critical points of the constrained optimization problem and showing that the common activation functions satisfy the assumptions. We rely in particular on the derivative rule of the interaction function which comes as a pleasant consequence of Stein's Lemma [8] instead of the Hermite basis expansion which is a commonly used technique [9-14]. For the erf and ReLU activation functions we derive additionally a closed-form solution of the optimization problem for  $n = 1$ . Next, we investigate "copy-average" configurations of students with  $n > 1$  neurons, where some student neurons copy teacher neurons and other student neurons average sub-groups of teacher neurons, in the sense that they are at the optimal one-neuron solution for the given sub-group of teacher neurons. Our particular contributions are:
- We propose a constrained optimization formulation of the standard minimization problem in the weight-space in terms of the interaction function (Section 3). The interaction function is a natural generalization of the dual activation [15].
- Applying the constrained optimization formulation for  $n = 1$ , we prove that the incoming vector of the student lies in the span of the incoming vectors of an orthogonal teacher network (Proposition 4.1). For a broad class of activation functions, we prove that the incoming vector aligns with the average of the teacher's incoming vectors for the "unit-orthonormal" teacher network (Theorem 5.1). Using the derivative rule of the interaction function (Lemma F.1), we show that common activation functions such as erf, softplus, tanh, and ReLU satisfy this property (Lemma F.2 and Corollary G.5).
- Assuming a unit-orthonormal teacher network and erf activation function, we prove that the concatenation of critical points of single neurons (of the student network) each approximating a teacher subnetwork is a copy-average critical point (Theorem 4.2).
- Assuming a unit-orthonormal teacher network and erf activation function, we prove that the optimal copy-average (CA) configuration is such that  $n - 1$  student neurons each copy a teacher neuron and the  $n$ -th student neuron approximates optimally the sum of the remaining teacher neurons (Theorem 5.5; see also Fig. 1, top row). Empirically, we find that the gradient flow converges to an optimal-CA point for all seeds when  $n < \gamma_1 k$  with a fixed  $\gamma_1$  near 0.46 (Figure 4).
- Surprisingly, we find empirically three regimes of training via gradient flow  $(\mathrm{GF})^2$  for underparameterized networks (Figure 4): (i) for  $n < \gamma_1k$ , GF converges to an optimal-CA point for all seeds, (ii) for  $n > \gamma_2k$  with a fixed  $\gamma_2$  near 0.6, GF converges to a point that we call perturbed- $n$ -copy for all seeds, (iii) for  $\gamma_1k < n < \gamma_2k$ , GF converges to either an optimal-CA point or a perturbed- $n$ -copy point. Therefore, as the under-parameterized network grows larger, the solution found with gradient flow where the weights are initialized randomly with a fixed standard deviation changes qualitatively. The code to reproduce these findings is available on GitHub, and we refer to Appendix C for details.
The teacher-student setup has been extensively used in the literature to study the evolution of gradient flow trajectories and of the generalization error [2, 3, 17-20]. This series of work gives insight into
2We use a numerical ODE solver for multi-layer networks [16] to simulate the gradient flow in this paper. All "solutions", which are the points at which gradient flow converges, have a gradient norm of at most  $5 \cdot 10^{-8}$ .
ReLU (bottom)
A1 Gradient flow (Weight space)
A2 Order param's  $(n = 2, k = 3)$
B Order param's  $(n = 4, k = 8)$
Figure 1: The gradient flow converges to the copy-average optimum point for erf activation (top), or nearby for ReLU activation (bottom): the first  $n - 1$  neurons copy one teacher neuron each; the  $n$ -th neuron takes an average of the remaining teacher neurons. The teacher network is unit-orthonormal, i.e.  $f^{*}(x) = \sum_{j=1}^{k} \sigma(v_{j} \cdot x)$  where  $v_{j} \in \mathbb{R}^{d}$ 's are orthonormal, and  $d = k + 1$ . A1 The gradient flow trajectory is shown in the weight space for  $n = 2$ ,  $k = 3$ : the positions of the circles (red and green) represent incoming vector  $w_{i}$  projected down to the span of  $v_{1}, v_{2}, v_{3}$  and the sizes of the circles represent outgoing weights  $a_{i}$ . The blue circle represents the one-neuron solution (the position shows  $w^{*}$ , the size shows  $a^{*}$ ). A2 Same setting, the weight-space parameters at convergence are mapped to the order-parameter space;  $u_{i} = (u_{i1}, \dots, u_{ik})$  where  $u_{ij}$  represents the normalized dot product between  $w_{i}$  and  $v_{j}$  and  $r_{i} = \| w_{i} \|$ . B Order parameters shown at convergence for  $n = 4$ ,  $k = 8$ . For erf (top) the point at convergence is exactly an  $(n - 1)$ -copy-1-average point, whereas for ReLU, it is perturbed away from this configuration. Neurons are reordered for clarity.
the solution found at convergence, however, they rely on numerically integrating the equations of dynamics. Tian [5] gives convergence guarantees for ReLU activation function, however, their method only works for one student and one teacher neuron. Xu and Du [21] recently gave the convergence rates for multiple student neurons for the case of one teacher neuron as a prototypical setup for overparameterization. These convergence guarantees were extended to broad input distributions [7, 22] and finite training data [6, 23]. We give the analytical formula of the optimal solution and its generalization error for one student neuron and unit-orthonormal teacher network with multiple neurons for erf and ReLU activation functions and a partial characterization for a broader class of activation functions without relying on the analytic formula of the loss.
The studies cited above showed positive results for a single-neuron teacher or a unit-orthonormal teacher. However, even for settings where the teacher has only a few neurons, hard teachers can be constructed in the sense that the student fails to find a zero-loss solution for a certain fraction of random initializations [24-26]. Moreover, for medium-scale problems, gradient flow often converges to 'non-zero loss' solutions [27-29]. Arjevani and Field [30] characterized some families of local minima using symmetries, for the ReLU activation function and unit-orthonormal teacher network. In this paper, we similarly characterize, for the case that the student has a smaller size than the teacher, a large family of 'copy-average' critical points, but for the erf activation function. Our approach focuses on the important regime of under-parameterized networks which is relevant for the superposition of features [31] and for the distillation of large networks into smaller ones [32, 33].
There is a large history of approximation theory of neural networks that give universal guarantees on the approximation error, e.g. [1, 34-36]. However, these works focus on rates of convergence and provide neither a formula nor an approximation for the error. In this paper, we make a conjecture for the exact formula of the approximation error of under-parameterized student networks which we support both theoretically and numerically.
Neural network: Consider a two-layer (student) network function  $f:\mathbb{R}^d\to \mathbb{R}$  with  $n$  neurons
$$
f (x) = \sum_ {i = 1} ^ {n} a _ {i} \sigma \left(w _ {i} \cdot x\right) \tag {1}
$$
where  $w_{i}\in \mathbb{R}^{d}$  is the incoming vector,  $a_{i}\in \mathbb{R}$  is the outgoing weight of neuron  $i$ , and the activation function  $\sigma$  is twice differentiable unless it is specified to be ReLU, i.e.  $\sigma_{\mathrm{relu}}(x) = \max (0,x)$ , and the dot marks the scalar product.  $P = n(d + 1)$  is the number of parameters.
Parameter vector: The parameter vector is represented as
$$
\theta = \left(w _ {1}, a _ {1}\right) \oplus \dots \oplus \left(w _ {n}, a _ {n}\right) \in \mathbb {R} ^ {P} \tag {2}
$$
where  $\oplus$  denotes the concatenation of two vectors into one vector. We use the notation  $\oplus$ , since the network function can be seen as a sum of its hidden neurons. Sometimes  $\theta$  is written explicitly in the network function  $f(x|\theta) = f(x)$ .
Loss function: We assume that the input distribution is a standard  $d$ -dimensional Gaussian  $\mathcal{D} = \mathcal{N}(0, I_d)$ . The target function is denoted by  $f^* : \mathbb{R}^d \to \mathbb{R}$ . Using the square cost, the loss function  $L : \mathbb{R}^P \to \mathbb{R}$  (also known as the risk or the generalization error) is defined as
$$
L (\theta) = \mathbb {E} _ {x \sim \mathcal {D}} \left[ (f (x | \theta) - f ^ {*} (x)) ^ {2} \right]. \tag {3}
$$
Orthogonal teacher network: We assume that the target function is a neural network (also known as the teacher network or a multi-index model)
$$
f ^ {*} (x) = \sum_ {j = 1} ^ {k} b _ {j} \sigma \left(v _ {j} \cdot x\right) \tag {4}
$$
where its outgoing weights are non-zero and its incoming vectors  $v_{1}, \ldots, v_{k} \in \mathbb{R}^{d}$  are orthogonal to each other, that is,  $v_{i} \cdot v_{j} = 0$  for  $i \neq j$ . This implies that the input dimension satisfies  $d \geq k$ . Following [27, 30, 37], we particularly focus on the unit-orthonormal teacher network where the outgoing weights are all one, that is  $b_{j} = 1$ , and the incoming vectors have unit norm, i.e.  $v_{i} \cdot v_{j} = \delta_{ij}$ .
Optimal loss: We study the optimal solution(s) of the following non-convex optimization problem
$$
L ^ {n, k} \left(\oplus_ {i = 1} ^ {n} \left(a _ {i}, w _ {i}\right)\right) = \mathbb {E} _ {x \sim \mathcal {D}} \left[ \left(\sum_ {i = 1} ^ {n} a _ {i} \sigma \left(w _ {i} \cdot x\right) - \sum_ {j = 1} ^ {k} b _ {j} \sigma \left(v _ {j} \cdot x\right)\right) ^ {2} \right]. \tag {5}
$$
for under-parameterized (student) networks, i.e.  $n < k$ , and orthogonal teachers. For  $n \geq k$  neurons, the network can copy all teacher neurons and set the outgoing weights of the remaining neurons to zero, therefore the optimal loss is trivially zero. If the teacher is unit-orthonormal, then all of its neurons contribute equally; hence the optimal loss is determined by  $n$  and  $k$  only and denoted by  $L^{*}(n,k)$ . If the student neural network has one neuron we use the notation  $L^{*}(k) \coloneqq L^{*}(1,k)$ .
3 Foundations & Constrained Optimization Formulation
In this section, we introduce a constrained optimization problem that is a reformulation of the minimization problem in Eq. 5. This formulation allows us to show that the incoming vector of any non-trivial critical point of the one-neuron network is in the span of the teacher's  $k$  orthogonal (or potentially even non-orthogonal, see Appendix Remark D.1) incoming vectors (see Proposition 4.1). We give the exact solution in the case of a unit-orthonormal teacher (see Corollary 5.2 and Corollary G.5).
Using the linearity of expectation, the loss function in Eq. 5 can be expanded as a weighted sum of the following Gaussian integral terms
$$
\mathbb {E} _ {x \sim \mathcal {D}} [ \sigma (V _ {1} \cdot x) \sigma (V _ {2} \cdot x) ]
$$
where  $V_{1}$  and  $V_{2}$  represent two arbitrary vectors of student and teacher networks such as  $(w_{i}, w_{j})$  or  $(w_{i}, v_{j})$ . As both  $V_{1} \cdot x$  and  $V_{2} \cdot x$  are centered Gaussian random variables, the above expectation can be expressed in terms of the covariance of the two-dimensional Gaussian
$$
\mathbb {E} _ {x \sim \mathcal {D}} \left[ \begin{array}{c c} (V _ {1} \cdot x) ^ {2} & (V _ {1} \cdot x) (V _ {2} \cdot x) \\ (V _ {1} \cdot x) (V _ {2} \cdot x) & (V _ {2} \cdot x) ^ {2} \end{array} \right] = \left[ \begin{array}{c c} r _ {1} ^ {2} & r _ {1} r _ {2} u \\ r _ {1} r _ {2} u & r _ {2} ^ {2} \end{array} \right]
$$
where  $r_i \coloneqq \| V_i \|$  for  $i = 1, 2$  is the  $\ell_2$ -norm and, assuming  $r_i > 0$ ,  $u \coloneqq V_1 \cdot V_2 / (r_1 r_2)$  is the correlation. The covariance entries  $Q_{ii} = r_i^2$ ,  $Q_{12} = r_1 r_2 u$  have been used to study the gradient flow trajectories [2, 3, 24]. We prefer the parametrization with  $r_i$  and  $u$  as it enables us to make the positive definiteness constraint explicit, i.e.
$$
| u | = \frac {\left| V _ {1} \cdot V _ {2} \right|}{r _ {1} r _ {2}} \leq 1, \tag {6}
$$
due to the Cauchy-Schwarz inequality. We introduce the interaction function  $g_{\sigma}:\mathbb{R}_{\geq 0}^{2}\times [-1,1]\to \mathbb{R}$
$$
g _ {\sigma} \left(r _ {1}, r _ {2}, u\right) = \mathbb {E} _ {\left(x _ {1}, x _ {2}\right) \sim \mathcal {N} (0, \Sigma)} \left[ \sigma \left(r _ {1} x _ {1}\right) \sigma \left(r _ {2} x _ {2}\right) \right] \quad \text {w i t h} \Sigma = \left[ \begin{array}{l l} 1 & u \\ u & 1 \end{array} \right], r _ {1}, r _ {2} > 0, \tag {7}
$$
to express the Gaussian integral terms. Note that  $u$  is not well-defined if one of the norms is zero. Extending the formula above, for the case w.l.o.g.  $r_2 = 0$ , we define
$$
g _ {\sigma} \left(r _ {1}, 0, u\right) := \mathbb {E} _ {x \sim \mathcal {N} (0, 1)} [ \sigma \left(r _ {1} x\right) ] \sigma (0)
$$
for all  $u \in [-1, 1]$ . In this paper, we consider the activation functions satisfying the following.
Assumption 3.1. For all  $r_1, r_2 > 0$  and  $u \in (-1, 1)$ , we assume that the interaction function  $g_{\sigma}$  satisfies either the first or both of the following properties
$$
(i) \frac {d}{d u} g _ {\sigma} \left(r _ {1}, r _ {2}, u\right) > 0, \quad (i i) \frac {d ^ {2}}{d u ^ {2}} g _ {\sigma} \left(r _ {1}, 1, u\right) u <   \frac {d}{d u} g _ {\sigma} \left(r _ {1}, 1, u\right). \tag {8}
$$
To check whether a specific activation function satisfies the above properties, we mainly rely on Lemma F.1 which gives us the rule for the partial derivative of  $g_{\sigma}$  with respect to the correlation
$$
\frac {d}{d u} g _ {\sigma} \left(r _ {1}, r _ {2}, u\right) = r _ {1} r _ {2} \mathbb {E} \left[ \sigma^ {\prime} \left(r _ {1} x\right) \sigma^ {\prime} \left(r _ {2} y\right) \right]. \tag {9}
$$
Hence, if  $\sigma$  is monotonic (increasing or decreasing), the integrand on the right-hand side is positive; satisfying Assumption 3.1 (i). The ReLU activation function  $\sigma_{\mathrm{relu}}(x) = \max (0,x)$  also satisfies it because of the known analytical expression of the interaction [4, 27]
$$
g _ {\text {r e l u}} \left(r _ {1}, r _ {2}, u\right) = r _ {1} r _ {2} h (u) \quad \text {w h e r e} h (u) = \frac {1}{2 \pi} \left(\sqrt {1 - u ^ {2}} + (\pi - \operatorname {a r c c o s} (u)) u\right).
$$
Checking Assumption 3.1 (ii) for a given activation function is delicate. We rely on it in Section 5.
Using the interaction function, the loss function can be expressed in terms of the order parameters:
- norms of the incoming vectors of the student  $r_i = \| w_i \|$ ,
- correlations between the incoming vectors of the student and teacher  $u_{ij} = w_i \cdot v_j / (r_i \| v_j \|)$ ,
- correlations between the incoming vectors of the student  $\rho_{ii'} = w_i \cdot w_{i'} / (r_ir_{i'})$ ;
where we assumed  $r_i > 0$  for all  $i \in [n]$ . The constrained optimization formulation is possible for general non-orthogonal teacher networks (see Remark D.1 in the Appendix). For the sake of simplicity, we formulate here the constrained optimization problem for the case of orthogonal teachers and reformulate the objective in Eq. 5 as
$$
\text {m i n i m i z e} \sum_ {i = 1} ^ {n} a _ {i} ^ {2} g _ {\sigma} \left(r _ {i}, r _ {i}, 1\right) + 2 \sum_ {i \neq i ^ {\prime}} a _ {i} a _ {i ^ {\prime}} g _ {\sigma} \left(r _ {i}, r _ {i ^ {\prime}}, \rho_ {i i ^ {\prime}}\right) - 2 \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {k} a _ {i} b _ {j} g _ {\sigma} \left(r _ {i}, \| v _ {j} \|, u _ {i j}\right) + C
$$
subject to  $\| u_{i}\| \leq 1$ ,  $r_i\geq 0$ , for all  $i\in [n]$
$$
\left| \rho_ {i i ^ {\prime}} - u _ {i} \cdot u _ {i ^ {\prime}} \right| \leq \sqrt {1 - \| u _ {i} \| ^ {2}} \sqrt {1 - \| u _ {i ^ {\prime}} \| ^ {2}}, \quad \text {f o r a l l} i \neq i ^ {\prime} \in [ n ]; \tag {10}
$$
3 Increasing (or decreasing) mean strictly increasing (or decreasing) everywhere in this paper.
A Weight-space
Figure 2: Cartoon representation of the mapping of a student with three neurons from the weight space  $\mathbf{A}\mathbb{R}^{nd}$  to order parameter space B1-B2. The mapping between the outgoing weights is an identity mapping hence not shown. A Each axis shows the direction of weights  $v_{i}$  of one teacher neuron  $(k\geq 3)$ . B1 Each incoming vector  $w_{i}\in \mathbb{R}^{d}$  is first transformed into  $(r_i,w_i / r_i)$  and then  $w_{i} / r_{i}$  is projected onto the span of the teacher's incoming vectors, yielding the student-teacher correlation vector  $u_{i} = (u_{i1},\dots,u_{ik})$ . B2 The student-student correlations  $\rho_{ii^{\prime}}$  are in general free parameters bounded in between  $u_{i}\cdot u_{i^{\prime}}\pm \sqrt{1 - \|u_{i}\|^{2}}\sqrt{1 - \|u_{i^{\prime}}\|^{2}}$  hence the box constraint. An activated constraint, w.l.o.g.  $u_{1}\in \mathbb{S}^{k - 1}$ , gives a vanishing  $\pm$  term for the interval of correlation  $\rho_{1i}$  for all  $i\neq 1$ , hence they are no longer free (shown in red). In the case  $d = k$ , all  $u_{i}$  are on the hypersphere due to the problem geometry, hence the correlations  $\rho_{ii^{\prime}}$  are fixed and not free (see Appendix D.1).
B1 Primary domain of the order parameters
B2 Secondary domain
where  $u_{i} = (u_{i1},\dots,u_{ik})$  and  $C = \mathbb{E}_{x\sim \mathcal{D}}[f^{*}(x)^{2}]$ . The constraints in Eq. 10 give tighter bounds than simply bounding correlations with the help of Eq. 6. See Appendix D for the derivation of the constraints and Fig. 2 for a schematic.
The objective above is exact for  $n = 1,2$ , in the sense that its optimal solution is equivalent to the optimal solution in the weight space, since the mapping from the weight-space to the order space is invertible. However, it is a relaxation for  $n \geq 3$ , since there are order-parameter configurations in the domain (see Figure 2) that do not correspond to any weight-space configuration (see Appendix D.3 for a construction). It seems possible to overcome this gap by considering the geometry of the angles between  $n \geq 3$  incoming vectors to tighten the constraints between student-student correlations.
4 Copy-Average Critical Points
In this section, we identify a new family of critical points by 'combining' critical points of one-neuron networks for the unit-orthonormal teacher and the erf activation function. We first show that in a network with  $n = 1$  student neuron, for any "non-trivial" critical point, that is  $w^{*} \neq 0$  and  $a^{*} \neq 0$ , the incoming vector  $w^{*}$  is in the span of the teacher's incoming vectors (Proposition 4.1). Applying this proposition to the special case of the erf activation function, we show that the concatenation of such critical points is also a critical point for multi-neuron networks (Theorem 4.2).
Proposition 4.1. Assume that  $f^{*}$  is an orthogonal teacher network (Eq. 4) of width  $k$ . If the activation function satisfies Assumption 3.1 (i), any non-trivial critical point  $\theta^{*} = (w^{*}, a^{*})$ , i.e.  $\nabla L^{1,k}(\theta^{*}) = 0$ ,  $\| w^{*} \| \neq 0$ ,  $a^{*} \neq 0$ , satisfies that  $w^{*}$  is in the span of the teacher's incoming vectors.
The proof uses the constrained optimization formulation for  $n = 1$  (see Appendix G.1). In short, a critical point mapped to the order parameter space satisfies either  $\| u_1\| = 1$  or  $\partial_{u}g_{\sigma}(r,\| v_{j}\| ,u_{1j}) = 0$  for all  $j\in [k]$ . Under the Assumption 3.1 (i) these partial derivatives are non-zero, hence  $\| u_{1}\| = 1$ . In a recent work [38], the incoming vectors of the student network are also shown to converge to the span of the vectors of the multi-index model using weight-decay.