Title: Should Under-parameterized Student Networks Copy or Average Teacher Weights?
Venue: NeurIPS 2023
Accepted: yes
Keywords: shallow neural networks

non-convex optimization

approximation error

loss landscape
Abstract: 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.

Should Under-parameterized Student Networks Copy or Average Teacher Weights?

Abstract

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.

1 Introduction

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.

∗

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 "unitorthonormal" 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 < γ 1 k with a fixed γ 1 near 0 . 46 (Figure 4).

Surprisingly, we find empirically three regimes of training via gradient flow (GF) 2 for underparameterized networks (Figure 4): (i) for n < γ 1 k , GF converges to an optimal-CA point for all seeds, (ii) for n > γ 2 k with a fixed γ 2 near 0 . 6 , GF converges to a point that we call perturbedn -copy for all seeds, (iii) for γ 1 k < n < γ 2 k , GF converges to either an optimalCA point or a perturbedn -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.

1.1 Related Work

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

A1

Gradient flow (Weight space)

Erf (top)

ReLU (bottom) ReLU (bottom) ReLU (bottom) ReLU (bottom) ReLU (bottom)

A2

Order param's (n=2, k=3)

B

Order param's (n=4, k=8)

A1

Gradient flow (Weight space)

Erf (top)

A2

Order param's (n=2, k=3)

B

Order param's (n=4, k=8)

A1

Gradient flow (Weight space)

Erf (top)

A2

Order param's (n=2, k=3)

B

Order param's (n=4, k=8)

A1

Gradient flow (Weight space)

Erf (top)

A2

Order param's (n=2, k=3)

B

Order param's (n=4, k=8)

A1

Gradient flow (Weight space)

Erf (top)

A2

Order param's (n=2, k=3)

B

Order param's (n=4, k=8)

1.00

1.00

0.75

3

ri

0.75

ri

0.50

aj

0.50-

4

aj

0.25

0.25

3

0.00

1

0.00

2

0.25

-0.25

1

-0.50-

0

-0.50

0

-0.75

0.75

-1.00

ul

u2

-1

-1.00

ul u2 u3 u4

n1n2n3n4

1.00

1.00

5

0.75

ri

0.75

ri

0.50

aj

0.50

4

!e

0.25

0.25

3

0.00

0.00

2

0.25

0.25

1

0.50

-0.50

0.75

-0.75

0

-1.00

ul

u2

-1

-1.00

-1

neuronlneuron2

ul u2 u3 u4

n1n2n3n4

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.

2 Setup

Neural network: Consider a two-layer (student) network function f : R d → R with n neurons

where w i ∈ R d is the incoming vector, a i ∈ R is the outgoing weight of neuron i , and the activation function σ is twice differentiable unless it is specified to be ReLU, i.e. σ 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

where ⊕ denotes the concatenation of two vectors into one vector. We use the notation ⊕ , since the network function can be seen as a sum of its hidden neurons. Sometimes θ is written explicitly in the network function f ( x | θ )= f ( x ) .

Loss function: We assume that the input distribution is a standard d -dimensional Gaussian D = N (0 , I d ) . The target function is denoted by f ∗ : R d → R . Using the square cost, the loss function L : R P → R (also known as the risk or the generalization error) is defined as

Orthogonal teacher network: We assume that the target function is a neural network (also known as the teacher network or a multi-index model)

̸

where its outgoing weights are non-zero and its incoming vectors v 1 , . . . , v k ∈ R d are orthogonal to each other, that is, v i · v j = 0 for i = j . This implies that the input dimension satisfies d ≥ 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 · v j = δ ij .

Optimal loss: We study the optimal solution(s) of the following non-convex optimization problem

for under-parameterized (student) networks, i.e. n < k , and orthogonal teachers. For n ≥ 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 ) := 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

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 · x and V 2 · x are centered Gaussian random variables, the above expectation can be expressed in terms of the covariance of the two-dimensional Gaussian

where r i := ∥ V i ∥ for i = 1 , 2 is the ℓ 2 -norm and, assuming r i > 0 , u := V 1 · V 2 / ( r 1 r 2 ) is the correlation. The covariance entries Q ii = r 2 i , 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.

due to the Cauchy-Schwarz inequality. We introduce the interaction function g σ : R 2 ≥ 0 × [ -1 , 1] → R

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

for all u ∈ [ -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 ∈ ( -1 , 1) , we assume that the interaction function g σ satisfies either the first or both of the following properties

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 σ with respect to the correlation

Hence, if σ is monotonic (increasing or decreasing) 3 , the integrand on the right-hand side is positive; satisfying Assumption 3.1 (i). The ReLU activation function σ relu ( x ) = max(0 , x ) also satisfies it because of the known analytical expression of the interaction [4, 27]

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 · v j / ( r i ∥ v j ∥ ) ,

correlations between the incoming vectors of the student ρ ii ′ = w i · w i ′ / ( r i r i ′ ) ;

where we assumed r i > 0 for all i ∈ [ 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

̸

̸

A

Weight-space

B1

Primary domain of the order parameters

B2

Secondary domain

IRd

.W1

↑Rd

Rd

IB

k

Bk

B

k

R2。

P23

[R(2)

r3

P12

.u1

.W3

.U3

.W2

r2

P13

neuron1

neuron2

neuron1

neuron2

normspace

neuroni-neuroni'

̸

where u i = ( u i 1 , ..., u ik ) and C = E x ∼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 ≥ 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 ≥ 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 ∗ = 0 and a ∗ = 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 θ ∗ = ( w ∗ , a ∗ ) , i.e. ∇ L 1 ,k ( θ ∗ ) = 0 , ∥ w ∗ ∥ ̸ = 0 , a ∗ = 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 ∂ u g σ ( r, ∥ v j ∥ , u 1 j ) = 0 for all j ∈ [ 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.

Finding the optimal solution for the multi-neuron network is challenging. Natural candidates are concatenation of student neurons where each one of them is a critical point of the loss function L 1 ,ℓ i where ℓ i is the number of the subgroup of teacher neurons. More precisely, let us pick a partition ℓ 1 + ... + ℓ n ≤ k such that ℓ i ≥ 1 , and define s m = ∑ m i =1 ℓ i for m ≤ n and s 0 = 0 . We denote a one-neuron critical point by θ ∗ i = ( w ∗ i , a ∗ i ) , when learning from a part of the teacher network

Since f ∗ i is a unit-orthonormal teacher, w ∗ i is in the span of v s i -1 +1 , ..., v s i due to Proposition 4.1.

We use the term copy-average (CA) point or configuration to refer to the concatenation of such one-neuron critical points in the student network with n neurons: if ℓ i = 1 , the student neuron copies one of the teacher neurons ( v j , 1) ; if ℓ i > 1 , it averages a group of teacher neurons in the sense of approximating their sum with one neuron. For odd activation functions, the one-neuron network problems decouple from each other, as the cross-terms E [ σ ( w 1 · x ) σ ( w 2 · x )] vanish for w 1 ⊥ w 2 . For the specific case of erf, we prove that all the copy-average configurations are critical points.

x

where ( w ∗ i , a ∗ i ) is a non-trivial critical point when learning from a unit-orthonormal teacher f ∗ i with the incoming vectors v s i -1 +1 , ..., v s i shown in Eq. 11. Then θ ∗ is a critical point of the loss function L n,k where the target function is f ∗ ( x ) = ∑ k σ ( v j · x ) .

In particular, all neurons are equivalent to each other in a unit-orthonormal teacher network. Therefore, the copy-average configurations where n -1 student neurons each copy a distinct teacher neuron and the n -th student neuron takes an average are also equivalent and called ( n -1) -copy1 -average, or ( n -1) -C1 -A in short. Another interesting configuration is where n student neurons each copy a distinct teacher neuron, which is called n -copy, or n -C in short.

For general activation functions the copy-average parameter vectors are not critical points (see Eq. 56). Nevertheless, we numerically find that the gradient flow converges to similar configurations for the ReLU activation function (see Figure 1, see Appendix C.3 for more experiments).

5 Approximation Error of Underparameterized Networks

The target function is assumed to be a unit-orthonormal teacher network in this section. In Subsection 5.1, we show for the one-neuron network that there is a unique non-trivial critical point up to symmetries, which is necessarily the global minimum (Theorem 5.1). Furthermore, we give the analytic expression of the optimal solution and its loss for erf (Corollary 5.2) and ReLU (Corollary G.5) activation functions. In Subsection 5.2, we provide for the under-parameterized student with n > 1 neurons the exact loss of copy-average critical points for the erf activation function and show that the ( n -1) -copy1 -average configurations reach the lowest loss among CA-critical points (see also Appendix E.1 for the combinatorial number of the equivalent copy-average configurations related to the landscape complexity calculations [39]).

5.1 One-Neuron Network

Using the constrained optimization formulation in 10, we first prove that at any non-trivial critical point of the one-neuron network, the incoming vector aligns equally with all teacher's incoming vectors for unit-orthonormal teachers for activation functions satisfying Assumption 3.1 (see Theorem 5.1). This is related to the symmetric solution visited during the learning plateaus studied in Saad and Solla [2] for erf activation and in Tian [5] for ReLU activation (see Appendix B for a detailed comparison). Our proof works for a broad class of activation functions and does not use the analytic expression of the interaction function.

Theorem 5.1. Assume that the activation function satisfies Assumptions 3.1 (i) and (ii). At any non-trivial critical point ( w ∗ , a ∗ ) of the loss L 1 ,k for the unit-orthonormal teacher network, the incoming vector satisfies

where u ∗ is either 1 √ k or -1 √ k .

Proof Sketch. There is no student-student interaction term since we have a single neuron; therefore we write u j instead of u 1 j and the constrained optimization problem in 10 simplifies to

A

erf

B

softplus

1

2

3

4

5

6

Teacher width (

k

)

0.00

0.02

0.04

0.06

0.08

0.10

0.12

0.14

0.16

L

*

(

k

)

erf

soft

tanh

sig

C

one-neuron approx. error

truefunction

networkfunction

truefunction

networkfunction

6

2

5

1

4

0

3

-1

2

2

0

From Proposition 4.1, we have that ∥ u ∥ = 1 for any non-trivial critical point. Therefore, the constraint of 14 on the correlations u = ( u 1 , ..., u k ) is satisfied. The mapping of any non-trivial critical point to the order-parameter space is a critical point of the Lagrangian loss (see Appendix Lemma G.2). Hence every u j satisfies

for fixed ( r, a ) . Assumption 3.1-(ii) implies that 1 u ∂ u g σ ( r, 1 , u ) is one-to-one hence all u j are equal. End of Proof Sketch.

We show in Lemma F.2 that the interactions of the common activation functions such as erf, tanh, sigmoid, and softplus (respectively)

with β ∈ (0 , 2] satisfy Assumption 3.1-(ii). The interaction of the ReLU activation function, i.e. σ relu ( x ) = max(0 , x ) also satisfies Assumptions 3.1 (with a slight modification in the domain for (ii); see the proof of Corollary G.5).

Thanks to Theorem 5.1, the loss in Eq. 14 can be reduced to a two-dimensional loss in a and r , which can be solved explicitly for ReLU (Corollary G.5) and erf.

Corollary 5.2. Assume that the activation function is σ erf. The optimal solution ( w ∗ , a ∗ ) is given by

or, equivalently, by ( -w ∗ , -a ∗ ) . The optimal loss is then given by

The proof of Corollary 5.2 is presented in Section G.3. For general activation functions, the twodimensional loss does not admit an analytical expression. For this case, from the partial derivatives, we obtain a fixed-point equation in r which we solve numerically for the activation functions listed above (see Appendix Section G.2.2). For softplus, specifically, we prove in addition the following

Theorem 5.3. Assume that the activation function is σ β soft ( x ) with β ≤ 2 . The optimal solution ( w ∗ , a ∗ ) satisfies ∥ w ∗ ∥ ≤ 1 / √ k and a ∗ ≥ k .

We use the FKG inequality to prove Theorem 5.3 (see Appendix G.5). Although the proof requires specific properties of softplus, we show that the above bounds hold also for tanh and sigmoid by numerically solving the fixed point equation (Figure 10, also Figure 5). Note that the incoming vector norm is bounded above by 1 / √ k , hence w ∗ is a damped average of the teacher incoming vectors.

n=1

convergence

8'0

n=2

n=4

=u

0.1339

sim

theory

0.110 -

0.0580

0.0578

0.0230

0.0225

at

0.6 -

0.4

n=8

n=12

n=16

SS07

0.1338

0.1337

0.108

0.106

0.0576

0.0572

0.0570

0.0574

0.0220

0.0215

0.0210

LoSS

0.2

0.1

0.1336

0.104

0.0568

0.0205

0.0200

0

10

0

10

0

10

0

10

0.0

0.0

5

10

n=8,k=16

n=8,k=14

n=8,k=11

n=8,k=9

0

10

20

30

40

50

60

B1In all seeds, GF

B2Insomeseeds,GF

B3Inallseeds,GF

B4In all seeds, GF

ASolutionsofGradientflow(GF)

finds CA-optimal

finds CA-optimal

findsP-n-C solution

findsP-n-C solution

5.2 Multi-Neuron Network

In this subsection, we assume that the activation function is erf such that CA-configurations are critical points (see Theorem 4.2). For a student network with n = 2 and a partition ( ℓ 1 , ℓ 2 ) we can decompose the loss of a CA critical point as

where L ∗ erf (0 , ℓ 0 ) := E x ∼D [ f ∗ ℓ 0 ( x ) 2 ] is the error made by a student with vanishing output when representing a reduced unit-orthonormal teacher network with ℓ 0 neurons. This decomposition is possible because the cross-terms between orthogonal vectors are zero for odd activation functions. Furthermore, for the erf activation function, we show that

(see the proof of Lemma E.2). Therefore, we should search for the minimum loss configuration among the partitions with ℓ 1 + ℓ 2 = k . Among such partitions, Lemma E.2 shows that the optimum CA-point has the partition (1 , k -1) . In words, the optimum is a 1 -copy1 -average point.

For general n , using Lemma E.2 and a small trick, we prove the following.

Theorem 5.5. Consider a unit-orthonormal teacher network f ∗ ( x ) = ∑ k j =1 σ ( v j · x ) and the erf activation function. For an under-parameterized student network with n neurons, the minimum-loss copy-average configuration up to permutations (of the student and teacher neurons) is

where ϵ i ∈ {± 1 } and ( w ∗ n , a ∗ n ) is given by Corollary 5.2 after substituting k with k -n +1 .

See Appendix E.3 for the proof. Because copy-average critical points are not necessarily the only critical points of the loss function for students with n > 1 , we investigate in simulations, if they are found by gradient flow where the weights are initialized as Gaussian with a fixed standard deviation (see Fig. 4).

Interestingly, gradient flow converges to the CA-optimal solution for all random seeds in a broad regime of under-parameterization (green in Fig. 4). Only when n > γ 1 k for n = 8 , 12 , 16 and γ 1 ∼ 0 . 46 , gradient flow converged in some seeds to points close to the n -copy critical point. We call these newly found points 'perturbed n -copy' (Pn -C) points. In gray-blue regimes, the Pn -C points have higher loss than the optimal CA critical point (Fig. 4).

However, this is not always the case: when the student width is close to the teacher width (low compression regime), the Pn -C point has a slightly lower loss the lowest amongst the CA critical points (red in Fig. 4). When k -n is fixed, we found that the ( n -1) -C1 -A solution turns from a minimum for small n to a saddle for large n (see App. Fig. 6); which explains why the gradient flow escapes it in this regime and converges to another minimum at a lower loss.

Finally, based on our theory and experiments, we conjecture that there exists a γ 0 ∈ (0 , γ 3 ) such that when n < γ 0 k and when the activation function is erf, the global optimum of the non-convex loss in Eq. 5 is a ( n -1) -C1 -A configuration. Therefore, if our conjecture holds, the exact approximation error, i.e. the optimal loss, is identical to that of a one-neuron network approximating a teacher with k -n +1 neurons and is given by

6 Conclusion & Future Directions

We studied the learning of under-parameterized student networks from orthogonal teacher networks for standard Gaussian input data and vanishing thresholds. For erf activation function, we introduced a new family of critical points that arise from the decoupling of the problem into one-neuron networks that can be solved separately. Moreover, the exact parameters of copy-average (CA) critical points are given which can be used to study escape behavior near saddles and to determine convergence of first and second-order optimization algorithms [16, 41].

Furthermore, we showed that the optimal CA point is that n -1 neurons copy teacher neurons and the n -th neuron averages the remaining k -n +1 neurons. In simulations, gradient flow converges to a CA-optimal solution for n < γ 1 k where γ 1 is near 0 . 46 . However, for n > γ 2 k where γ 2 is near 0 . 6 , we observe another phase where the gradient flow finds a perturbed copy solution. For the ReLU activation function and the onset of under-parameterization, gradient flow converges to qualitatively similar solutions; however, at the crossing point from under-parameterization to overparameterization (i.e. n = k ), gradient flow is known to get stuck in spurious local minima [27]. On another note, determining the CA-optimal solution of the two-neuron network plays a critical role in our analysis. Still, there is only little literature on two-neuron networks [5, 42] compared to the well-studied one-neuron case [5-7, 22, 23]. The two-neuron network is possibly the simplest model with interactions between neurons, hence it is important to understand the global minimum and gradient flow dynamics of this challenging problem.

On the practical side, our analysis of under-parameterized networks gives a recipe for how to warmstart smaller neural networks for distilling unit-orthonormal teacher networks. If one desires low compression ( n > γ 3 k ), then we recommend initializing the student network in a configuration where each neuron copies a different teacher neuron, to be close to a Pn -C point. However, for higher compression, we recommend initializing the student network in a configuration where n -1 neurons are each copied and the n -th neuron is initialized as an average neuron to be close to a ( n -1) -copy1 -average point. It remains an open question whether this recipe applies to non-idealized scenarios such as non-isotropic input distribution, teacher networks with non-orthogonal incoming vectors, or non-unit outgoing weights. More generally, it is natural to expect that the optimal distillation strategy changes from low compression levels to high compression levels. How exactly and where this change happens is a very intriguing question of theory and practice.

References

Ken-Ichi Funahashi. On the approximate realization of continuous mappings by neural networks. Neural networks , 2(3):183-192, 1989.

David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E , 52(4):4225, 1995.

Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems , 32, 2019.

Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. Advances in neural information processing systems , 22, 2009.

Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International conference on machine learning , pages 3404-3413. PMLR, 2017.

Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics , 46(6A):2747-2774, 2018.

Gilad Yehudai and Shamir Ohad. Learning a single neuron with gradient methods. In Conference on Learning Theory , pages 3756-3786. PMLR, 2020.

Charles M Stein. Estimation of the mean of a multivariate normal distribution. The annals of Statistics , pages 1135-1151, 1981.

Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory , pages 1887-1930. PMLR, 2018.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research , 22(1):4788-4838, 2021.

Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems , 35:97689783, 2022.

Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055 , 2023.

Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at a time. arXiv preprint arXiv:2305.18270 , 2023.

Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. arXiv preprint arXiv:2305.10633 , 2023.

Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. Advances in neural information processing systems , 29, 2016.

Johanni Brea, Flavio Martinelli, Berfin ¸ Sim¸ sek, and Wulfram Gerstner. Mlpgradientflow: going with the flow of multilayer perceptrons (and finding minima fast and accurately). arXiv preprint arXiv:2301.10638 , 2023.

Michael Biehl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general , 28(3):643, 1995.

David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters , 74(21):4337, 1995.

Peter Riegler and Michael Biehl. On-line backpropagation in two-layered neural networks. Journal of Physics A: Mathematical and General , 28(20):L507, 1995.

Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. arXiv preprint arXiv:2202.00293 , 2022.

Weihang Xu and Simon Du. Over-parameterization exponentially slows down gradient descent for learning a single neuron. In The Thirty Sixth Annual Conference on Learning Theory , pages 1155-1198. PMLR, 2023.

Gal Vardi, Gilad Yehudai, and Ohad Shamir. Learning a single neuron with bias using gradient descent. Advances in Neural Information Processing Systems , 34:28690-28700, 2021.

Lei Wu. Learning a single neuron for non-monotonic activation functions. In International Conference on Artificial Intelligence and Statistics , pages 4178-4197. PMLR, 2022.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. arXiv preprint arXiv:2206.04030 , 2022.

Flavio Martinelli, Berfin ¸ Sim¸ sek, Johanni Brea, and Wulfram Gerstner. Expand-and-cluster: Exact parameter recovery of neural networks. arXiv preprint arXiv:2304.12794 , 2023.

Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

Itay Safran and Ohad Shamir. Spurious local minima are common in two-layer relu neural networks. In International Conference on Machine Learning , pages 4433-4441. PMLR, 2018.

Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory , 65(2):742-769, 2018.

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

Yossi Arjevani and Michael Field. Analytic study of families of spurious minima in two-layer relu neural networks: a tale of symmetry ii. Advances in Neural Information Processing Systems , 34:15162-15174, 2021.

Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposition. arXiv preprint arXiv:2209.10652 , 2022.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2015.

Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision , 129:1789-1819, 2021.

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks , 2(5):359-366, 1989.

Andrew R Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information theory , 39(3):930-945, 1993.

Itay M Safran, Gilad Yehudai, and Ohad Shamir. The effects of mild over-parameterization on the optimization landscape of shallow relu neural networks. In Conference on Learning Theory , pages 3889-3934. PMLR, 2021.

Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209.14863 , 2022.

Berfin ¸ Sim¸ sek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances. In International Conference on Machine Learning , pages 9722-9732. PMLR, 2021.

Joan Bruna, Loucas Pillaud-Vivien, and Aaron Zweig. On single index models beyond gaussian data. arXiv preprint arXiv:2307.15804 , 2023.

Charles G Frye, James Simon, Neha S Wadia, Andrew Ligeralde, Michael R DeWeese, and Kristofer E Bouchard. Critical point-finding methods reveal gradient-flat regions of deep network losses. Neural computation , 33(6):1469-1497, 2021.

Chenwei Wu, Jiajun Luo, and Jason D Lee. No spurious local minima in a two hidden unit relu network. Workshop track - ICLR , 2018.

Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 249-256. JMLR Workshop and Conference Proceedings, 2010.

Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

A Summary of Results

In the table above, UO means unit-orthonormal, n.i.g. stands for 'not in general' and yes* follows as a special case from the results with yes on the same row.

B Further Comparison to Literature

In this section, we compare the symmetric solutions found in erf [2] and ReLU networks [5] to our one-neuron solution ( n = 1 ). The main difference is that both earlier studies constrain the search space to the symmetric subspace whereas we first prove that the non-trivial critical points are contained in this subspace in Theorem 5.1 for a broad class of activation functions, including erf and ReLU. Solving the low-dimensional loss, we recover the same solution for ReLU and erf as in [2, 5] for unit-orthonormal teachers.

Symmetric Solution of Saad and Solla [2] for erf activation. The authors focus on the 'symmetric subspace' parameterized as

In this case, the loss is parameterized by three values, that is Q,C,R , hence can be expressed analytically in terms of these values. Solving the fixed point equations, they find the following critical/fixed point (their Eq.22)

which implies r i = 1 / √ 2 k -1 and ρ ij = 1 in our parameterization. This selection of parameters forces all student vectors to be equal therefore reducing the system to a one-neuron network. There are two main improvements in our analysis

We prove that student-teacher correlations u ij are equal to each other at a non-trivial critical point, and give necessary conditions on the activation function (Assumption 3.1) to satisfy this property. We show in Lemma F.2 that not only erf but a large class of common activation functions satisfy Assumption 3.1.

Our student network has a flexible outgoing weight (shallow neural network) as opposed to a fixed outgoing weight +1 (soft-committee machine) in Saad and Solla [2]. It is instructive to compare the generalization errors of the one-neuron network

which are identical since arcsin(0 . 5) = π/ 6 (Saad and Solla [2] uses ϵ g ( k ) = 1 2 L ∗ erf; soft ( k ) that's why there is a factor 0 . 5 difference with respect to their Eq. (23)). However, if we set teacher outgoing weights to say a t , the shallow network adapts and reaches the generalization error a 2 t L ∗ erf ( k ) but the error of the soft-committee machine is

which has a worse coefficient 1 3 compared to 1 3 -1 π as expected.

Symmetric Solution of Tian [5] for ReLU activation. The authors focus on a particular twodimensional subspace ( x, y ) that allows the specialization of student neurons, namely

̸

In particular, they consider the 'symmetric subspace' x = y which is the case when all student neurons collapse to one neuron, and show that the dynamics converge to the following fixed point

Summing over k neurons then produces the following one-neuron due to the positive homogeneity

Our formula (Corollary G.5) gives the identical result due to

In this case, there is no difference between the optimal solution of the soft-committee machine and the shallow network since ReLU is positive-homogeneous as expected.

A

B

C

5.0

theory (ReLU)

0.7

1/√k

k

4.5

10

e

4.0

ReLU

0.6

weigh

3.5

Softplus

0.5

8

E

iul

9

3.0

2.5

2.0

Erf

Sigmoid

lou

0.4

0.3

buobno

6

4

1.5

1.0

0.2

2

2

3

4

7

8

6

10

2

3

4

5

6

8

6

10

2

3

4

5

6

7

8

6

10

#of teacher neurons (k)

# of teacher neurons (k)

#of teacher neurons (k)

C Further Experiments

All experiments in this paper are implemented using the gradient flow package implemented by Brea et al. [16] which is particularly suited to studying gradient flow on the population loss. For activation functions for which there is an analytic formula, it is already implemented in the package; for the others, we used the approximator option for a speed-up compared to the numerical integration option. This method uses a neural network in the background fitted to approximate Gaussian integrals. We trained for 10 5 ode iterations for erf and relu experiments; 10 3 ode iterations for softplus, tanh, and sigmoid. For erf experiments, all seeds converged to configurations with gradient norm below 5 · 10 -8 . For ReLU experiments, a fraction of seeds failed to converge (large gradient norm at the end of training). In Appendix C.3, we report among the seeds that succeeded in converging. Weights initialized as Gaussians with zero mean and standard deviation 0 . 1 or with Glorot initialization [43]. This is in contrast with Saad and Solla [2], Tian [5] where (order) parameters are initialized with positive values (as opposed to the rotationally symmetric initializations done in practice). For each ( n, k ) pair, we implemented 10 or 20 seeds of random initializations.

C.1 One-Neuron Network

Empirically, gradient flow converges to the point where all student-teacher correlations are 1 k 4 .

We know from Theorem 5.1 that at the non-trivial critical point all correlations are equal at correlation 1 / √ k and there is possibly another critical point at correlation -1 / √ k . Depending on the activation function, the point where correlations are -1 / √ k might be

either an equivalent of the optimum solution (for odd activation functions),

or a saddle (for ReLU),

or does not exist (for softplus).

The details can be found in the proofs for individual cases.

C.2 Erf Experiments

In this section, we first numerically investigate whether the CA-optimal critical point is a saddle or minimum in Figure 6. Surprisingly, we find that the point turns from a saddle point to a minimum

5

10

15

20

25

30

n

0.10

0.08

0.06

0.04

0.02

0.00

0.02

min. eig.

k

=

n

k

=

n

+1

k

=

n

+2

k

=

n

+3

5

10

15

20

25

30

n

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

0.0175

min. eig.

k

=

n

k

=2

n

k

=3

n

0

20

40

60

itr

1.0

0.5

0.0

0.5

1.0

corrs

0

20

40

60

itr

0

20

40

60

outgoing weigths

0

20

40

60

itr

0.25

0.50

0.75

1.00

1.25

1.50

norms

0

20

40

60

itr

0

5

10

15

20

losses

20

40

60

0.96

0.98

1.00

0

20

40

60

itr

1.0

0.5

0.0

0.5

1.0

corrs

0

20

40

60

itr

10

0

10

20

30

40

50

60

outgoing weigths

0

20

40

60

itr

0.2

0.4

0.6

0.8

1.0

1.2

1.4

norms

0

20

40

60

itr

0

5

10

15

20

losses

25

50

1.0

1.1

0

20

40

60

itr

1.0

0.5

0.0

0.5

1.0

corrs

0

20

40

60

itr

0

10

20

30

40

50

60

outgoing weigths

0

20

40

60

itr

0.2

0.4

0.6

0.8

1.0

1.2

norms

0

20

40

60

itr

0

5

10

15

20

losses

25

50

1.0

1.2

1.4

point in some regimes of ( n, k ) , despite that the configuration has the same structure of copying n -1 teacher neurons and taking an average of the remaining teacher neurons. When k -n is fixed, and for large n , Figure 6 explains why gradient flow does not converge to the CA-optimal point.

0

20

40

60

itr

1.0

0.5

0.0

0.5

1.0

corrs

0

20

40

60

itr

1.0

0.5

0.0

0.5

1.0

outgoing weigths

0

20

40

60

itr

0.2

0.4

0.6

0.8

1.0

norms

0

20

40

60

itr

0.0

0.5

1.0

1.5

2.0

2.5

3.0

losses

25

50

0.05

0.10

A1

A2

A3

n

= 2

, k

= 5

n

= 2

, k

= 10

n

= 2

, k

= 20

B1

B2

B3

n

= 4

, k

= 5

n

= 4

, k

= 8

n

= 4

, k

= 16

C1

C2

n

= 12

, k

= 16

n

= 12

, k

= 24

n

= 12

, k

= 32

1.00

4

0.75

0.50

3

0.25

2

0.00

0.25

1

0.50

0

0.75

1.00

-1

ulu2

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

ul

u2

1.00

4

0.75

0.50

3

0.25

2

0.00

0.25

1

0.50

0

0.75

1.00

-1

ulu2

1,00

3

0.75

0.50

2

0.25

0.00

1

0,25

0.50

0

0.75

1.00

ulu2u3u4

1.00

3

0.75

0.50

2

0.25

0.00

=

0.25

0.50

0

0.75

-1,00

一

1

ulu2u3u4

n

1.00

0.75

3

0.50

2

0.25

0.00

1

0.25

0.50

0

0.75

1.00

u1u2u3u4

.00

0.75

3

0.50

2

0.25

0.00

1

0.25

0.50

0

0.75

.00

0

5

10

1.00

3

0.75

0.50

2

0.25

0.00

1

0.25

0.50

0

0.75

1.00

0

10

1.00

3

0.75

0.50

2

0.25

0.00

0.25

0.50

0

0.75

1.00

0

10

<latexi sh

1_b

64="3/k

L9Pq

ZM

YR8

j

Up+

H

>A

B7

c

V

NSw

EJ

r

IvgQ

u

y

D

2T

m

z

5

W

O

o

fK

G

X

0

C

d

F

n

<latexi sh

1_b

64="

jWAmZYH

DfTX/cdJ

I

p7

E

>

B8

V

NSwM

2

+

CR

g5

F

UP

L9

z

y

r

v

q

K

0

o

G

k

u3O

n

Q

<latexi sh

1_b

64="g

QJF7

MpEukW

9TwP

Y5

ZNS

>A

B

3

c

V

X

O

8

m

L

rG

/o

K

2

dn vR

f

j

D

z

U

H

q

I

0

y

C

+

<latexi sh

1_b

64="Pd+

ME0X

J

O

fI

>

B73

c

V

NSw

2

T

CR

g5

r

Q

G

9q

Z

pk

y

K

/

H

j

8n

UD

Y

u

F

v

<latexi sh

1_b

64="u

/KnAI5rWdNH29VZ

OD+7

wE

>

B8X

c

S

J3U

q

o

L

g

QkF

y

P

Y

z

T

G

R

j

M

p

C

0

m

f

v

<latexi sh

1_b

64="

8

Vz

ZNn

/

w

>A

B

H

S

M

EJ2

X7

+

T

CR

g5

d

r

Q

9q

p

y

KW

j

UD

u

fF

I

v

O

<latexi sh

1_b

64="T+XkSp

K7

W

E

V

/qH

Rv

0

>A

B8

c

N

wM

J3

9

L

C

Z

j

F

uU

Q5P

m

Gg

d

f

2

r

o

O

D

I

Y

z

y

n

<latexi sh

1_b

64="o

Tvr

K7

y

d

u

2

YPk

R

I

>A

B8H

c

V

NSwM

EJ3

9

q

L

/Z

j

W

F+

U

p

Q5

m

X

Gg

f

n

O

0

D

z

C

<latexi sh

1_b

64="G/5nY

FSc

0p8T

v

N

I

>A

B

X

V

EJ

Ur

q

o

9L

g

Qk

y

P

w

z

u

H

R

7

2OWj

+

C

d

m

Z

f

We show some representative trajectories of gradient flow, in the regime when the CA-optimal critical point is a minimum in Figure 7 and in the regime when it is a saddle point in Figure 8.

C.3 ReLU Experiments

In this subsection, we will present the structure of the minimum loss configuration found by gradient flow. In the regime n ≪ k , the minimum loss configuration is qualitatively similar to the optimal-CA solution but without perfect decoupling. For k that is slightly bigger than n , the gradient flow finds an "all-copy" configuration for n = 4 , 12 . The overall trend of the minimum loss configuration is

similar to the case of erf activation; however, as we do not have the analytic formula of correlations, a theoretical prediction for the optimal solution is left for future work.

D Constrained Optimization Formulation

Expressing the loss function in terms of the order parameters yields the 'projected' loss function

̸

where the constant term C = E x ∼D [ f ∗ ( x ) 2 ] . L proj has n ( k +2) + ( n 2 ) parameters, that is n output weights and n ( k +1) + ( n 2 ) order parameters, instead of n ( d +1) parameters of the original loss function. For d ≫ k +1+ n -1 2 , L proj has significantly less number of parameters.

In the special case d = k , the incoming vectors can be expressed as a linear combination of the teacher's incoming vectors, hence the correlations between them are not free (see Appendix D.1).

Each normalized incoming vector can be expressed as a sum of its projection on the span of the teacher's incoming vectors and an orthogonal component

and the inequality constraint pops up since ∥ u i ∥ 2 = 1 -∥ v ⊥ i ∥ 2 where u i = ( u i 1 , . . . , u ik ) . For the correlations between the incoming vectors, we get

which yields the second set of constraints on the optimization problem

̸

since we have | v ⊥ i · v ⊥ i ′ | ≤ ∥ v ⊥ i ∥∥ v ⊥ i ′ ∥ due to the Cauchy-Schwarz inequality.

We note if all incoming vectors are in the span of the teacher's incoming vectors, we have that ∥ u i ∥ = 1 . As a result, the second set of inequalities in Eq. 28 collapse onto equalities, hence the secondary constraints are in fact equality constraints (see Appendix Eq. 34). We show that this is indeed the case for the non-trivial critical points of the one-neuron network in Section G.

In general, the constrained optimization formulation is possible for non-orthogonal teacher networks.

Remark D.1. (Non-orthogonal teacher network) We can relax the assumption of orthogonality between v 1 , . . . , v k to linear independence. Let us collect the incoming vectors into a matrix V = [ v 1 , . . . , v k ] ∈ R d × k . The expansion in Eq. 26 can be rewritten as

where Γ i = ( γ i 1 , . . . , γ ik ) ∈ R k and v ⊥ i · v j = 0 for all j ∈ [ k ] . The normalized vector has a unit norm, hence we have

The correlation vector can be written as u i = V T V Γ i ∈ R k which yields the following constraint

D.1 Equality Constraints in the Case d = k

In this case, the teacher incoming vectors v 1 , . . . , v k span the input domain R d . Each incoming vector (of the student) can be expressed as a linear combination of the teacher's incoming vectors

and the equality constraint pops up since the normalized vector has a unit ℓ 2 norm. The correlations between the incoming vectors are then expressed in terms of the student-teacher correlations

Therefore the optimization problem is equivalent to

̸

Since ( n 2 ) student-student correlations terms are not free, the problem has only n ( k + 2) free parameters and k equality constraints, yielding n ( k +1) effective parameters, which is the same number as the number of parameters of the original problem in the weight-space.

D.2 Binary-Equality Constraints in the Case d = k +1

In this case, there is only one direction orthogonal to the span of the teacher's incoming vectors (i.e. v ⊥ i ∥ v ⊥ i ′ ). Therefore the general inequality constraint on ρ ii ′ reduces to

D.3 Three-Neuron Network

We present a case for the three-neuron network where the optimal solution of the constrained optimization problem may not be projected back to the weight space. In particular, let us consider a positive and monotonic activation function, i.e. σ ( x ) > 0 . This implies that g σ > 0 and that g σ is increasing in correlation. It is natural to expect that all outgoing weights are positive since this would bring the network function closer to the target function. If the network is overparameterized, some outgoing weights may be zero or even negative (balanced by a positive outgoing weight corresponding to the same incoming vector). Let us pick three positive outgoing weights a 1 , a 2 , a 3 > 0 and three corresponding student-student interaction terms

Note that each ρ ii ′ is decoupled from each other. Since g σ is increasing in correlation, the minimum of each term above is achieved when

This implies that the inequality is tight and therefore v ⊥ i ∥ v ⊥ i ′ moreover,

up to a sign flip. However, it is not possible that the three vectors all have pairwise flipped directions to each other as we would have ( -) · ( -) = (+) . If the optimal solution of the constrained optimization problem verifies ∥ u i ∥ 2 < 1 , then we conclude that it cannot be mapped back to the weight space; in this case, the optimal L proj would only give a lower bound on the optimal loss of the weight-space.

E Copy-Average Critical Points

Theorem E.1. Assume that σ is the erf activation function. We pick a copy-average parameter vector

where ( w ∗ i , a ∗ i ) is a non-trivial critical point when learning from a unit-orthonormal teacher f ∗ i with the incoming vectors v s i -1 +1 , ..., v s i shown in Eq. 11. Then θ ∗ is a critical point of the loss function L n,k where the target function is f ∗ ( x ) = ∑ k j =1 σ ( v j · x ) .

Proof. Let us write down the partial derivatives with respect to the outgoing weights and incoming vectors

We will show that they are equivalent to the following

which implies that the partial derivatives are zero, since ( w ∗ i , a ∗ i ) is a critical point of the loss

Since ( w ∗ i , a ∗ i ) is the optimal solution of the teacher network generated by v s i -1 +1 , ..., v s i , from Theorem 5.1, we have that w ∗ i is in the span of v s i -1 +1 , ..., v s i . We have that

since the two incoming vectors are in the span of two orthogonal subspaces respectively and w ∗ i is orthogonal to all other teacher incoming vectors that are outside of the span of v s i -1 +1 , ..., v s i . For two orthogonal vectors say w ∗ i and v , we have that

which is zero if at least one of σ 1 and σ 2 is odd. This implies the first equation in 41 for the partial derivatives with respect to the outgoing weights. In order to show the second equation for the partial derivatives with respect to the incoming vectors, let us define S j = [ s j -1 +1 , s j ] ,

̸

and note that it suffices to show that W i,j = 0 for all i = j . This is true if and only if W i,j · ¯ v ℓ = 0 where { ¯ v 1 , ..., ¯ v d } form an orthogonal basis of R d . Let us choose ¯ v 1 = v 1 , ..., ¯ v k = v k and that ¯ v k +1 , ..., ¯ v d completes the basis if d > k . One can observe that W i,j · ¯ v ℓ = 0 for k +1 ≤ ℓ ≤ d since x · ¯ v ℓ is an independent Gaussian from the others. Hence, the expectation, i.e. W i,j · ¯ v ℓ , factorizes with a factor of E [ x · ¯ v ℓ ] which is zero.

It remains to check W i,j · v l = 0 for all v l 's. We split the analysis into two cases. If v l / ∈ S j , then x · v l is independent from x · w ∗ j and from x · v k for k ∈ S j . Hence W i,j splits into

̸

where the second term in the product is zero because w ∗ j · x and v k · x are centered Gaussian and σ is odd. For the second case, where v l ∈ S j , using the fact that x · v l is independent from w ∗ i · x and from x · v k for l = k , we have

Applying Stein's Lemma to both terms on the right we have

where Z is standard Gaussian and r ∗ j = √ 1 2 k -1 , u ∗ = √ 1 k , a ∗ j = k (parameters of erf). Hence we want to show that

To show this, we use the following relation [2, 3]

Differentiating with respect to the correlation u we have

In particular, at correlation zero, we get

which implies 51 and the proof is complete.

For general activation functions, using the substitution in Eq.41, the first partial derivatives in Eq. 40 reduce to

̸

which is in general non-zero if σ is not odd.

Lemma E.2. Assume ℓ 2 > ℓ 1 ≥ 1 . We have that

Proof. We first show that the function x 2 arcsin( 1 2 x ) is increasing for x ≥ 1 and convex for x > 0 . Using the Taylor expansion of arcsin , we have that

where the higher-order terms all have positive coefficients. The first derivative is

Therefore, we have

which is positive for x ≥ 1 since we have

The second derivative is

which is positive for positive x .

First, let us show that Eq. 18 holds. Plugging in the analytic expressions for L ∗ erf (0 , · ) and L ∗ erf ( · ) , it is equivalent to

Since f ( x ) is increasing for x ≥ 1 , the second term inside the parenthesis is positive, hence the inequality holds.

We will now prove the statement of the Lemma. It suffices to show the following for all ℓ = ℓ 2 ≥ 2

since then we can continue to decrease ℓ by one, i.e. ℓ -1 , ℓ -2 , ... , until we reach ℓ 1 . Eq. 63 is equivalent to

that is the second-order finite difference, similar to the second-derivative of a continuous function. The proof is completed by observing that L ∗ erf ( · ) is a discrete-concave function since its continuous interpolation

is concave for x > 0 since it can be written as L ∗ erf ( x ) = αx -f ( x ) where f is convex for x > 0 .

Lemma E.2 tells us that if we add one neuron to the teacher, then it is better to approximate it by the student neuron that already approximates many teacher neurons. Applying Lemma E.2 iteratively, we get

Theorem E.3. Consider a unit-orthonormal teacher network f ∗ ( x ) = ∑ k j =1 σ ( v j · x ) and the erf activation function. For an under-parameterized student network with n neurons, the minimum-loss copy-average configuration up to permutations (of the student and teacher neurons) is

where ϵ i ∈ {± 1 } and ( w ∗ n , a ∗ n ) is given by Corollary 5.2 after substituting k with k -n +1 .

Proof. We will conclude with a simple argument that the minimum-loss CA configuration for a multi-neuron network with n neurons is ( n -1) -C1 -A. In particular, if there are two averaging neurons inside the student network, we can redistribute the teacher neurons shared between these two to a lower-loss CA configuration by ensuring that one student neuron copies and the other student neuron averages (see Lemma E.2). The minimum-loss CA point is then achieved among CA configurations where at least n -1 neurons each copy a single teacher neuron (of the k possible ones). The remaining student neuron can be treated as a single-neuron network learning from a teacher with k -n +1 neurons - for which we know the optimal solution is to average (Theorem 5.1).

E.1 Number of CA Critical Points

There is a combinatorial number of ( ℓ 1 , ..., ℓ n ) -CAC critical points, that is

where c n := c ( ℓ 1 , ..., ℓ n ) counts distinguishable permutations between the neurons of the student network, and the binomial coefficients stand for grouping teacher neurons into n non-empty buckets.

If ℓ 1 = ... = ℓ n , permutation between the student neurons is already counted when distributing the teacher neurons, hence c n = 1 . If all ℓ 1 , ..., ℓ n are distinct from each other, we have that c n = n ! since we swap all pairs of student neurons after assigning groups of teacher neurons. In general, let c i denote the number of i 's among ℓ 1 , ..., ℓ n for all i = 1 , .., k ; the formula for the permutation-factor is given by

F General Properties of the Interaction Function

In this Section, we introduce some general properties of the interactions. We use these only for the one-neuron network in this paper (see Section G), however, these properties are likely to play a role in studying the networks with two or more neurons.

We first present the partial derivative of a general interaction function, i.e. two activation functions may be different, for example, if the student activation function does not match the teacher, with respect to the correlation in a simple expression in Lemma F.1. In the second part, we present a property of the activation function sufficient for Assumption 3.1 (ii), and show that the differentiable activation functions studied in this paper satisfy this property in Lemma F.2.

Lemma F.1. Assume that functions σ 1 and σ 2 are differentiable. The partial derivative of the following Gaussian integral term E [ σ 1 ( r 1 x ) σ 2 ( r 2 y )] with respect to the correlation E [ xy ] = u is

We apply the Lemma for σ 1 = σ 2 = σ in the main text in Eq. 9.

Proof. We compute the derivative of E [ σ 1 ( r 1 x ) σ 2 ( r 2 y )] by making the correlation u explicit. Denote u ′ = √ 1 -u 2 and y = ux + u ′ z . After the computation, we use Stein's lemma to reach the desired formula.

where x and z are independent standard Gaussians. Here is a reminder for Stein's Lemma for a standard Gaussian z

To remove x in the first term, we apply Stein's formula for v ( x ) = σ 1 ( r 1 x ) σ ′ 2 ( r 2 ( ux + u ′ z )) yielding

To remove z in the second term, we apply Stein's formula for v ( z ) = σ ′ 2 ( r 2 ( ux + u ′ z )) by considering fixed x which yields

Summing up the two terms completes the proof.

For softplus that is increasing and convex, using Lemma F.1 for σ 1 = σ 2 = σ twice, we infer that the interaction g is also increasing and convex in u . Hence, for u < 0 , Assumption 3.1 (ii) holds for softplus. However, for the other activation functions, using second-order derivatives does not suffice to show the assumption. We will propose a new property of the activation function that implies that the interaction satisfies Assumption 3.1 (ii) and prove that softplus with β ≤ 2 , sigmoid, tanh, and erf satisfy this property.

Lemma F.2. If the activation function σ is thrice-differentiable and it satisfies

then its interaction satisfies Assumption 3.1 (ii) for all u ∈ ( -1 , 1) . Softplus with β ∈ (0 , 2] , sigmoid, tanh, and erf activation functions satisfy the above inequality.

Proof. Let us first write out Assumption 3.1 (ii) explicitly using Lemma F.1

where ¯ σ ( x ) = σ ′ ( x ) . Using Stein's Lemma for v ( x ) = ¯ σ ( r 1 x )¯ σ ′ ( y ) , we get

The desired inequality is equivalent to

Let us introduce f ( x ) = ¯ σ ( x ) -x ¯ σ ′ ( x ) + ¯ σ ′′ ( x ) . For y = ux + u ′ z where u ′ = √ 1 -u 2 , we have the conditional average of y fixing x (we drop conditioning on the right-hand terms for convenience)

where second last equality comes from Stein's Lemma for v ( z ) = ¯ σ ′ ( ux + u ′ z ) . Hence the desired inequality is equivalent to

By straightforward calculus, we will show that f ( x ) > 0 , or that f ( x ) ≥ 0 and f ( x ) = 0 if and only if x = 0 . In the latter case, the expectation in Eq. 79 is positive since f ( y ) > 0 for some y values of the integrand. First, for the sigmoid and tanh activation functions, for which we have

Hence, we can explicitly write f as

Therefore showing f ( x ) > 0 is equivalent to showing that the factor on the right, that is,

is positive. For x < 0 , we have e x < 1 which implies -x (1 -e 2 x ) > 0 and (1 -e x ) e x ≤ 1 / 4 due to the inequality of arithmetic and geometric means hence the first term is upper bounded by -1 / 2 and since we have +2 , the whole term is positive. For x ≥ 0 , we have e x ≥ 1 , hence we can rewrite the inequality as a sum of non-negative terms

Let us now handle the case of erf. Its first three derivatives are given by

Hence, we can explicitly write f as

that is non-negative for all x and zero iff x = 0 .

Finally, for the softplus activation function with β ∈ (0 , 2] , we have the following derivatives

Plugging in the function f , we get

Therefore showing f ( x ) > 0 is equivalent to showing that the factor on the right, that is,

is positive. For x ≤ 0 , we have that -xβ > 0 and 2 -β 2 ≥ -2 since β ≤ 2 , hence it is sufficient to show that the following is positive

which is a sum of squares. For x > 0 , in the rest of the proof we will show that

for β ∈ (0 , 2] . Using e βx ≥ ( βx ) 2 / 2 + βx +1 , it suffices to show that

If ( βx ) 2 / 2 + 3 -β 2 ≥ 1 , then the first term is bigger than βx +1 hence the above term is positive. The remaining possibility is that we have

β ≤ 2 implies x < 1 and x 2 > 0 implies β > √ 2 . Hence we have -xβ + β 2 > 0 since β > x . Therefore, if we have ( βx ) 2 / 2 + 3 -β 2 ≥ 0 , Eq. 92 is positive. Assuming the opposite, we get

β ≤ 2 implies x< 1 / √ 2 and x 2 > 0 implies β> √ 3 .

Going back to Eq. 92, what remains to show is that it is positive in the domain x< 1 / √ 2 , β ∈ ( √ 3 , 2] . It suffices to show that e βx +2 -xβ -β 2 > 0 . Assuming the contrary implies e βx < xβ +2 since β ≤ 2 . We can then deduce that xβ < c = 1 . 2 since otherwise we would have

which implies a contradiction. c can be chosen smaller but this will be enough for our purposes.

Assuming e βx +2 -xβ -β 2 ≤ 0 , let us expand Eq. 92

where in the last inequality we used xβ < 1 . 2 . We note that this inequality holds for slightly larger β using the same technique, however, for significantly larger β , the property breaks down.

G The One-Neuron Network

We study the critical points of the following loss function

in particular, the optimal solution. For n = 1 , all configurations of order parameters are realizable in the weight space, therefore, the optimal solution of the following loss (repeating Eq. 14)

is equivalent to the optimal solution in the weight space. Let us denote the unit ball by B = { ( u 1 , ..., u k ) | u 2 1 + ... + u 2 k ≤ 1 } . Its interior is denoted by int B and its boundary is denoted by ∂B .

We will present the results for the one-neuron network in five parts

In Subsection G.1, we give a proof of Proposition 4.1: any non-trivial critical point ( w,a ) of L 1 ,k satisfies that w is in the span of the teacher's incoming vectors if the activation function satisfies Assumption 3.1 (i). Moreover, we show in Lemma G.2 that the corresponding order parameters should satisfy a Lagrangian condition (Eq. 105).

In Subsection G.2, we characterize the topology of the loss landscape in terms of its critical points for the activation functions studied in this paper and for the unit-orthonormal teacher. Our results for the one-neuron network are strong in the sense that it gives all possible critical points of the loss landscape.

(a) In Subsection G.2.1, we give a proof of Theorem 5.1: for general activation functions satisfying Assumption 3.1, any non-trivial critical point of the one-neuron network attains equal correlations that are either 1 / √ k or -1 / √ k .

(b) In Subsection G.2.2, we study the two-dimensional loss obtained after applying Theorem 5.1. From its derivative constraints, we get a fixed point equation (Eq. 120) that needs to be satisfied by the incoming vector norm r at any non-trivial critical point with equal correlations u . Numerically, we show that there is a unique solution of the fixed point equation for u > 0 for differentiable activation functions studied in this paper (Fig. 10). Finally, we give some sufficient conditions in Eq. 121 to prove uniqueness based on log-concavity; numerically, these are shown to be satisfied by softplus and sigmoid activation functions.

In Subsection G.3, we give a proof of Corollary 5.2 by solving the two-dimensional loss for the erf activation function. Moreover, from the proof, we conclude that there are exactly two non-trivial critical points identical up to the mirror symmetry of erf (since it is odd); and these are the optimal solutions for the loss landscape.

In Subsection G.4, we present and prove Corollary G.5 by solving the two-dimensional loss for the ReLU activation function. We find that there are two non-trivial critical points of the loss function: a saddle 'point' at correlation -1 / √ k and the optimal 'solution' at correlation 1 / √ k . Due to the positive homogeneity of ReLU, these are not two points but two equal-loss hyperbolas in the loss landscape.

In Subsection G.5, we study the two-dimensional loss for the softplus activation function and give a proof of Theorem 5.3. Absence of analytical expression for the Gaussian integral terms make the problem challenging; we use several non-trivial steps in the proof. The proof shows that there is no critical point at correlation -1 / √ k ; and a non-trivial critical point ( w ∗ , a ∗ ) at correlation 1 / √ k satisfies the following bounds: a ∗ ≥ k and ∥ w ∗ ∥ ≤ 1 / √ k . Numerically, we find that these bounds hold for other activation functions studied in this paper (tanh and sigmoid; Fig. 5).

G.1 Any Non-Trivial Critical Point Satisfies the Lagrangian Condition

̸

We add a reminder here for the definition of the non-trivial critical point θ = ( w,a ) : it is a critical point of the loss function that satisfies a = 0 and ∥ w ∥ ̸ = 0 .

Proposition G.1. Assume that f ∗ is an orthogonal teacher network of width k . If the activation function satisfies Assumption 3.1 (i), any non-trivial critical point θ ∗ = ( w ∗ , a ∗ ) , i.e. ∇ L 1 ,k ( θ ∗ ) = 0 , satisfies that w ∗ is in the span of the teacher's incoming vectors.

̸

From Assumption 3.1 (i), we have that ∂ u g σ ( r 1 , r 2 , u ) > 0 for u ∈ ( -1 , 1) which yields a contradiction. Thus, each critical point of the projected loss is on the boundary, i.e. u ∈ ∂B , which implies that the incoming vector is in the span of the teacher's incoming vectors.

We will next show that any non-trivial critical point satisfies a Lagrangian condition since it is on the boundary of a constrained optimization problem.

Lemma G.2. Let θ = ( w,a ) be a non-trivial critical point of L 1 ,k . Then the corresponding order parameters p = ( r, u, a ) satisfy the following Lagrangian condition

Proof. We will first show that for any differentiable path ( r, γ ( t ) , a ) on the boundary such that γ ( t ) ∈ ∂B for t ∈ ( -ϵ, ϵ ) for some ϵ > 0 and γ (0) = u , the following holds

Let us assume the contrary. We construct the corresponding following path in the weight space

Thanks to the equivalence of the losses along the path, we have that

since θ (0) = θ is a critical point in the weight space. Therefore, Eq. 106 holds for any differentiable path on the boundary and implies that ∇ u L ( p ) is orthogonal to all γ ′ (0) . The vector that is orthogonal to all γ ′ (0) is the gradient of the surface, that is 2( u 1 , ..., u k ) . Hence we get ∇ u L ( p ) ∥ u which is written explicitly as the Lagrangian condition in Eq. 105. This is equivalent to setting the partial derivatives of the Lagrangian loss with respect to u j to zero where the Lagrangian loss is given by

We set λ = λ ′ / a in Eq. 105.

G.2 General Activation Functions

Before we present our results, let us take a detour to check the applicability of the convex optimization framework. For a convex and twice-differentiable activation function such as softplus, applying Lemma F.1 twice implies that the interaction g σ ( r 1 , r 2 , · ) is a convex function of the correlation u ∈ ( -1 , 1) for r 1 , r 2 > 0 . Let us consider a fixed a< 0 and r> 0 and consider the loss parameterized by u j 's. It is convex since its Hessian is a diagonal matrix with entries

Proof. We will prove by contradiction. Let us assume that w is outside of the span of the teacher's incoming vectors. We will show that ( w,a ) is not a critical point for any a = 0 . Mapping ( w,a ) to the order parameter space, we get that ( r, u, a ) where u = ( u 1 , ..., u k ) ∈ int B which implies u j ∈ ( -1 , 1) . Since u ∈ int B and r > 0 , we have that ( r, u, a ) is a critical point of L 1 ,k proj since the boundaries are not seen near the neighborhood of this point. Therefore the partial derivatives of L 1 ,k proj are all zero including

Since the constraint on the correlations (Eq. 14) is also convex, we get a convex optimization problem that has a unique global minimum (see Boyd et al. [44], Section 4.2). Swapping a pair of u j does not change the loss, thus it is permutation symmetric. If any two u j were distinct from each other at the minimum, then its permutation would also be a minimum which would violate the unicity. We conclude that at the unique minimum point, the correlations are equal to each other. However, for the case a > 0 , and for other activation functions, the objective is not convex.

We instead use Lagrange multipliers for proving Theorem 5.1.

G.2.1 Proof of Theorem 5.1

Theorem G.3. Assume that the activation function satisfies Assumption 3.1. At any non-trivial critical point ( w ∗ , a ∗ ) of the loss L 1 ,k for the unit-orthonormal teacher network, the incoming vector satisfies

where u is either 1 / √ k or -1 / √ k .

Proof. From Proposition 4.1 and Lemma G.2, we get that any non-trivial critical point should satisfy the Lagrangian condition in Eq. 105. In particular for unit-orthonormal teacher, setting ∥ v j ∥ = 1 and b j = 1 , we get the following Lagrangian condition

If u j = 0 , we get ∂ u g σ ( r, 1 , 0) = 0 which is not possible since g σ ( r, 1 , u ) is increasing due to Assumption 3.1 (i). Hence we have

Let us observe that ∂ u g σ ( r, 1 , u ) /u is decreasing for u ∈ ( -1 , 1) \ { 0 } if and only if

which is equivalent to Assumption 3.1 (ii) for u ∈ ( -1 , 1) \{ 0 } (we included u = 0 in Assumption 3.1 (ii) for a simpler statement which is already implied from Assumption 3.1 (i) at u = 0 ).

̸

Taken together, we conclude that ∂ u g σ ( r, 1 , u ) /u is one-to-one in u ∈ ( -1 , 1) \ { 0 } . We need to consider the remaining case u i ∈{-1 , 1 } . For k ≥ 2 , necessarily, we have u j = 0 for j = i , which is not possible as we have shown. For k = 1 , u i ∈ {-1 , 1 } is the only option that satisfies the boundary condition. For k ≥ 2 , Eq. 113 implies that all correlations are equal. Combining it with the boundary condition, we get u 1 = ... = u k = u with ku 2 = 1 , which completes the proof.

G.2.2 Two-Dimensional Loss, The Derivative Constraints, Uniqueness

At any non-trivial critical point, we proved in Theorem 5.1 that all correlations are equal and denoted by u that is either 1 / √ k or -1 / √ k . The projected loss at a critical point reduces to

Moreover, at a critical point, the partial derivatives with respect to the outgoing weight and norm should also be zero which gives the following two constraints

1.0

erf

1.0

softplus

1.0

sigmoid

1.0

tanh

1.0

gelu

0.8

0.8

0.8

0.8

0.8

0.2

0.6

0.6

0.6

0.6

0.6

=

=

- 0.0

0.4

0.4

0.4

0.4 -

0.4

0.2

0.2

0.2

0.2

0.2 -

- -0.2

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0

5

10

r

r

r

r

r

̸

which can be rearranged into the following (assuming g σ ( r, r, 1) = 0 and ∂ r g σ ( r, r, 1) = 0 )

̸

The second equality between the two ratios of Gaussian integral terms gives a fixed point equation on the norm r . Writing the interactions in Eq. 117 explicitly and rearranging the ratios, we get

where x and y are standard Gaussians with correlation E [ xy ] = u . Let us define the following helper functions

which yields

̸

Let us consider the case u > 0 . We want to show that for any given u ∈ (0 , 1] there is a unique r ∈ (0 , 1] such that f ( r, u ) = 0 . Under the assumption σ (0) = 0 , if the following three conditions are satisfied for all u ∈ (0 , 1] ,

then we have a unique r solving Eq. 120 as we explain next. Note that the first two conditions are equivalent to f (0 , u ) < 0 and f (1 , u ) > 0 , respectively. The tricky part is the third condition which is equivalent to showing that

is log-concave in r . We note that marginalization properties of log-concave functions may be helpful here. In this paper, we were not able to prove the sufficient conditions listed above for general activation functions that do not admit an analytic formula of the interaction, even for softplus which we studied in detail (see Subsection G.5). Instead, we present the numerical integration results, which show that for any given u ∈ (0 , 1] , there is a unique r ∈ (0 , 1] such that f = 0 (see Fig. 10). Once r is shown to be unique, then the matching outgoing weight a follows from Eq. 117.

G.3 Closed-Form Solution for Erf Activation

Corollary G.4. Assume that the activation function is σ erf. The optimal solution ( w ∗ , a ∗ ) is given by

or, equivalently, by ( -w ∗ , -a ∗ ) . The optimal loss is given by

Proof. Since erf is an odd activation function, it suffices to find parameters of the non-trivial critical points satisfying u ≥ 0 . For any such critical point ( w ∗ , a ∗ ) , its mirror symmetry ( -w ∗ , -a ∗ ) is an equivalent critical point due to Eq. 41.

For u = 0 , we have that g erf ( r, 1 , u ) = 0 which implies g erf ( r, r, 1) = E [ σ ( rx ) 2 ] = 0 due to the first derivative constraint in Eq. 116 which holds if and only if r = 0 . This gives a possible trivial critical point yielding the zero predictor function.

For a given u = 1 √ k > 0 , from Fig. 10, we observe that there is a unique r ∈ (0 , 1] satisfying the fixed point equation in Eq. 120. For uniqueness, we rely on numerical integration. We will find one solution to the derivative constraints given below

for a given k , and equivalently u = 1 √ k > 0 ; and due to uniqueness, conclude that it is the only non-trivial critical point up to symmetries.

In particular, we will use the analytic formula for the interaction function [2, 3]

Let us find r where we have

that is satisfied if we have that the arguments of arcsin match, which happens at

Interestingly, at this value of r , we also have

which can be seen by inserting the guessed values in the following equation

Setting a = k in Eq. 124 completes the order parameters of the non-trivial critical point. Finally, let us compute the loss at r = 1 / √ 2 k -1 , u = 1 / √ k , a = k ;

G.4 Closed-Form Solution for ReLU Activation

Corollary G.5. Assume that the activation function is σ relu. Any optimal solution ( w ∗ , a ∗ ) satisfies

forming an equal-loss hyperbola. The optimal loss is given by

We will first show that the interaction of ReLU satisfies

where u 0 = 1 / √ 2 . Note that property (i) is equivalent to Assumption 3.1 (i), and property (ii) is almost equivalent to Assumption 3.1 (ii) except that it holds in the interval ( -1 , u 0 ] ; property (iii) covers up for the missing piece of the interval in the property (ii).

ReLU interaction satisfies Properties 130; Proof. Let us write the first two derivatives of h :

Property (i) easily comes from noting that the derivative of h is positive for u ∈ ( -1 , 1) . Property (ii) holds for u ∈ ( -1 , 0] since both the first and second derivatives are positive. Let us show that Property (ii) holds for u ∈ (0 , u 0 ] , that is equivalent to

Let us note that the left-hand side is smaller than 1 since

Note that arcsin( u ) > 0 for u > 0 ; and π/ 2 > 1 . This completes the proof of Property (ii).

For Property (iii), we first show that h ′ ( u ) /u is convex in u ∈ (0 , 1) . The first two derivatives are

Thus, it is equivalent to showing

Using the Taylor series of arcsin and u > 0 , we have that arcsin( u ) > u . Hence, it suffices to show

where we dropped the positive term π u which holds due to the inequality of arithmetic and geometric means

Let us assume the contrary of Property (iii), that there exists u ∈ ( u 0 , 1) such that

Note that h ′ ( u 0 ) /u 0 > h ′ (1) because π (1 -u 0 ) -arccos( u 0 ) > 0 holds at u 0 = 1 / √ 2 . Since h ′ ( u ) /u is left-continuous at u = 1 , there exists ϵ > 0 such that

Finally, there exists α ∈ (0 , 1) such that u = α (1 -ϵ ) + (1 -α ) u 0 which gives due to the convexity of h ′ ( u ) /u the following

This yields a contradiction since the left-hand side is strictly smaller than h ′ ( u 0 ) /u 0 hence the proof of Property (iii) is complete. ReLU interaction satisfies Properties 130; End of Proof.

Proof. First, we replicate the proof steps of Theorem 5.1 to show that any non-trivial critical point must be on the boundary and attain equal correlations. From Property 130 (i), we get that there is no non-trivial critical point in int B . For k = 1 , this implies that u 1 = -1 or u 1 = 1 .

For general k , let us recall that we get the Lagrangian condition for non-trivial critical points

which is equivalent to Eq. 112 for ReLU activation function. u j = 0 is not possible since we have h ′ (0) = 0 . Hence, we get

Property 130 (ii) implies that f ( u ) = h ′ ( u ) /u is decreasing for u ∈ ( -1 , u 0 ) \{ 0 } . Moreover, f is negative for u < 0 and positive for u > 0 .

If λ/r < 0 , we get that all u j are equal and negative, hence they are equal to -1 / √ k due to the boundary condition.

If λ/r = 0 , we get u j = -1 for all j which implies that k = 1 which is already covered above.

If λ/r > 0 , Property 130 (iii) gives that f ( u 0 ) > f ( u ) for u ∈ ( u 0 , 1) . Since f is decreasing we have also f ( u ) > f ( u 0 ) for u ∈ (0 , u 0 ) ; hence f ( u j ) are equal only when all u j < u 0 or u j > u 0 ; however, the latter case is not possible for k ≥ 2 since it breaks the ball constraint, i.e. u 2 1 + u 2 2 > 1 .

Hence, we get that u j ∈ (0 , u 0 ] and are equal since f is decreasing in this interval. This completes the proof of replica of Theorem 5.1 for the ReLU activation function.

For the ReLU activation function, there is at least one non-differentiable critical point at a = 0 or r = 0 . The careful analysis of this critical point is beyond the scope of this work. For any such 'trivial' point, the error of zero-function is equivalent to

We will next show that ( -1 / √ k ) k j =1 and (1 / √ k ) k j =1 are the global minimum and the global maximum of the following loss function

̸

Due to the Lagrange condition, there is no other critical point, hence these are the only two critical points of the constrained objective in Eq. 140. The objective then reduces to kh ( u ) which is minimized at u = -1 / √ k and maximized at u = 1 / √ k .

Next, we will give the closed-form solution of the remaining order parameters. Plugging in the correlation in the loss and using the factorization of the interaction in Eq. 14, we get

Let us set ˜ a = ar . The loss is a second-order polynomial in ˜ a

where we made the constant explicit. Since the coefficient of the leading term is positive, there is a minimizer and it is the only critical point. Taking the derivative, the minimum is attained at

Finally, plugging in ˜ a ∗ , we get

For u = 1 / √ k and u = -1 / √ k , h ( u ) is non-zero; hence l ( u ) is smaller than the loss of the zero function (trivial critical points). The smallest loss is attained at u = 1 / √ k which is, therefore, the optimal solution. We conclude that the critical point at u = -1 / √ k is a saddle point since it is a maximum in u and a minimum in ˜ a .

G.5 Bounds on Incoming Vector Norm and Outgoing Weight for Softplus

Unlike ReLU and erf, the interaction function does not have a known analytic expression for softplus, hence the proof involves some techniques to compare ratios of Gaussian integral terms.

FKG Inequality. We will use a special case of the FKG inequality repeatedly, that is,

if both f, g are increasing (or decreasing) implying that f and g are positively correlated. The inequality changes direction if f is increasing and g is decreasing (or vice versa) implying that f and g are negatively correlated.

We will rely on some specific properties of the softplus family that are developed in Section G.5.4. Unfortunately, some of these properties do not apply to other activation functions. As a first example of managing interactions that do not have an analytic formula, the proof may inspire generalizations to other activation functions. Below we present the proof sketch for Theorem 5.3. In the following Subsections G.5.1, G.5.2, G.5.3, and G.5.4, the components of the proof are presented in detail.

Proof Sketch. We want to characterize the zero(s) of f introduced in Section G.2.2 that is

For r ∈ [0 , 1] , there is a unique correlation u ∈ [0 , 1] such that f ( r, u ) = 0 . Denoting this correlation by h ( r ) , we have a map h : [0 , 1] → [0 , 1] with boundary conditions h (0) = 0 and h (1) = 1 . For r > 1 , there is no solution of f . As a consequence, no r ≥ 0 solves f ( r, u ) = 0 for negative u , hence there is no non-trivial critical point at u = -1 / √ k .

In Section G.5.2, we prove the inequality h ( r ) ≥ r , which gives us the upper bound on the norm since we have that the correlation at a non-trivial critical point is h ( r ) = 1 / √ k . Using this inequality and Stein's Lemma, we give a lower bound on the outgoing weight, that is a ≥ k , in Section G.5.3. In summary, any non-trivial critical point of the loss has equal correlations that are u = 1 / √ k , the norm satisfies r ≤ u , and the lower bound on the outgoing weight follows. End of Proof Sketch.

G.5.1 Constraining the Zeros of f

In this subsection, we will describe all zero-crossings of f : [0 , ∞ ) × [ -1 , 1] → R . We need to check four cases (i) r = 0 , (ii) r = 1 , (iii) r > 1 , and (iv) r ∈ (0 , 1) .

(i) r = 0 : Note that G (0) = ˜ G (0 , 0) = 0 . Since ˜ G is increasing in correlation for u ∈ [0 , 1] and ˜ G ( u, 0) < ˜ G (0 , 0) for u < 0 (Lemma G.6), the only solution is u = 0 .

(ii) r = 1 : Note that G (1) = ˜ G (1 , 1) since y = x due to correlation one in Eq. 119. Since ˜ G is increasing in correlation for u ∈ [0 , 1] and ˜ G ( u, 1) < ˜ G (0 , 1) for u < 0 (Lemma G.6), the only solution is u = 1 .

(iii) r > 1 : We will show that there is no zero in this case. Let us first show that G ( r ) > ˜ G (1 , r ) for r > 1 , which is equivalent to

Changing the measure of x from the standard Gaussian p ( x ) to ˜ p ( x ) = p ( x ) σ ( rx ) 2 / E [ σ ( rx ) 2 ] , we get the following equivalent inequality

From the property (iv) of Lemma G.8, we have that σ ′ ( rx ) x/σ ( rx ) is increasing after a substitution x ← rx . We need to show σ ( x ) /σ ( rx ) is decreasing in x for r > 1 . We take the derivative

Since σ ′ ( x ) x/σ ( x ) is increasing ∀ x ∈ R , we have

which yields σ ′ ( x ) σ ( rx ) < σ ( x ) σ ′ ( rx ) r , hence we conclude that σ ( x ) /σ ( rx ) is decreasing. Thanks to the FKG inequality, σ ′ ( rx ) x/σ ( rx ) and σ ( x ) /σ ( rx ) are negatively correlated which completes the argument. Since from Lemma G.6, ˜ G is increasing in correlation and ˜ G (0 , r ) > ˜ G ( u, r ) for u < 0 , we have ˜ G (1 , r ) > ˜ G ( u, r ) for all u ∈ [ -1 , 1) , therefore there is no solution of f .

(iv) r ∈ (0 , 1) : We want to show that ∀ r ∈ (0 , 1) , there is a unique u ∈ (0 , 1) such that f ( r, u ) = 0 . It suffices to show

since ˜ G is continuous and increasing in correlation for u ∈ [0 , 1] (Lemma G.6), it then crosses G ( r ) at a unique u ∈ (0 , 1) .

First inequality; ˜ G (0 , r ) < G ( r ) . In this case, x and y are Gaussians with zero correlation, hence independent. We can expand ˜ G (0 , r ) by factorizing the integrals

We want to show

which is equivalent to the following inequality after changing the measure from standard Gaussian p ( x ) to ˜ p ( x ) = p ( x ) σ ( rx ) / E [ σ ( rx )]

This follows from the FKG inequality since we have that both σ ′ ( x ) x/σ ( x ) and σ ( x ) are increasing from the properties (iv) and (i) of softplus (Lemma G.8).

Second inequality; G ( r ) < ˜ G (1 , r ) . This is equivalent to the Ineq. 145, but the direction is reversed since in this case r < 1 . We showed that σ ( x ) /σ ( r ′ x ) is decreasing in x for all r ′ > 1 , therefore its reciprocal σ ( r ′ x ) /σ ( x ) is increasing in x . Substituting x ← rx where r = 1 /r ′ < 1 , we get that σ ( x ) /σ ( rx ) is increasing in x for r < 1 . This yields a positive correlation between σ ′ ( rx ) x/σ ( rx ) and σ ( x ) /σ ( rx ) from the FKG inequality and completes the argument.

Overall, we showed that there are no zeros of f for r > 1 . For r ∈ [0 , 1] , there is a unique correlation u , that we will denote by h ( r ) , such that f ( r, h ( r )) = 0 . Furthermore, h : [0 , 1] → [0 , 1] satisfies the following

h (0) = 0 and h (1) = 1 ,

for r ∈ (0 , 1) , we have h ( r ) ∈ (0 , 1) .

G.5.2 Bound on the Norm

In this subsection, we will show that h ( r ) ≥ r for all r ∈ (0 , 1) . Let us assume the contrary, which implies

due to Lemma G.6. It suffices to show that for all r ∈ (0 , 1) , we have

which yields a contradiction since G ( r ) = ˜ G ( h ( r ) , r ) . Showing this is equivalent to

where r ′ = √ 1 -r 2 . After a change of measure from standard Gaussian p ( x ) to

this is equivalent to the following inequality

What remains to show is that

is non-increasing in x since then we can conclude by the FKG inequality. Since r > 0 we can drop it up to a change in the standard deviation of x . We want to show that its derivative is non-positive:

From the property (iii) of softplus (Lemma G.8), we have that R ( x ) = σ ( x ) /σ ′ ( x ) is convex. Applying Jensen, we get

What remains to show is that

Note that E [ σ ′ ( x + r ′ z ) | x ] is increasing in x since σ ′ is increasing. Moreover, the function

is increasing in x since its integrand R is increasing from the property (ii) of softplus (Lemma G.8). Then we conclude by the FKG inequality that Eq. 154 holds. Therefore, for a solution ( r, u ) of the fixed point Eq. 117, we have r ≤ u = 1 k .

G.5.3 Bounding the Outgoing Weight

To get a bound on a , let us analyze the ratio of interactions in Eq. 117

Using the convexity of softplus (property (i) of Lemma G.8), we get

We can transform the numerator using Stein's lemma with v ( x ) = σ ( rx ) σ ′ ( rx )

which is positive since softplus is positive, increasing, and convex. Combining it with u ≥ r , we get that the ratio is bounded below by 1 which yields a ≥ k .

G.5.4 Helper Lemmas

In this subsection, we provide helper lemmas used in the proof of Theorem 5.3. We present Lemma G.6 which shows that ˜ G is increasing in correlation and Lemma G.7 used in the proof of the former. Finally, we present several properties of the softplus family in Lemma G.8 that are used throughout the proof.

Lemma G.6. The following function is increasing in u ∈ [0 , 1]

for any r ≥ 0 , where x and y are standard Gaussians with correlation E [ xy ] = u . Moreover, ˜ G ( u, r ) < ˜ G (0 , r ) for u < 0 .

Proof. Let us assume 0 ≤ u 1 < u 2 ≤ 1 . For the first part of the statement, we want to show

where E [ xy 1 ] = u 1 and E [ xy 2 ] = u 2 . Changing the measure from the standard Gaussian p ( x ) to

we get the following equivalent inequality

Thanks to the property (iv) of softplus (Lemma G.8), we have that σ ′ ( rx ) x/σ ( rx ) is increasing in x after a substitution x ← rx for r > 0 . For r = 0 , the function reduces to γx with some γ > 0 , hence increasing. We will next show that (all integrations are w.r.t z hereafter, hence we drop the conditioning on x )

is decreasing in x . Computing the derivative w.r.t x , we want to show that it is negative

Note that this is equivalent to showing

for all u ∈ [0 , 1) and x ∈ R . Changing the order of derivatives, it is sufficient to show

is shown to be increasing in x in Lemma G.7 where we need to substitute x → xu 1 for u 1 > 0 , and for u 1 = 0 , we have s 1 ( x ) = γx for some γ > 0 hence it is increasing. The remaining part is to show that the second function

is decreasing. We will consider z ← u ′ z and x ← ux in what follows. We have

due to first applying Stein's Lemma to the numerator and then inverting the ratio. Using the chain rule, it is sufficient to show that

are increasing, since both functions are positive.

Interestingly, f 1 is increasing in x if σ is a log-concave function. Because its derivative is positive

if E [ σ ( x + z )] is log-concave. This is the case since a centered Gaussian distribution p ( z ) is logconcave, therefore σ ( x + z ) p ( z ) is jointly log-concave, and marginalization preserves log-concavity.

Similarly, f 2 is increasing since σ ′ is also log-concave due to property (v) of softplus (Lemma G.8). Hence we showed that

is increasing for u ∈ [0 , 1) . The derivative of r explodes at 1 , however, we can conclude by contradiction that r (1) > r ( u ) for u < 1 : if r ( u ) ≥ r (1) for some 0 ≤ u < 1 , then there exists u 0 ∈ ( u, 1) where the function is decreasing. Therefore, r is increasing for u ∈ [0 , 1] . We can conclude the first part of the proof by the FKG inequality σ ′ ( rx ) x / σ ( rx ) and E [ σ ( y 1 )] / E [ σ ( y 2 )] are negatively correlated.

For the second part of the statement, we need to show

for u < 0 . Changing the measure from standard Gaussian p ( x ) to

The first function the above inequality is equivalent to

This holds since σ ′ ( rx ) x/σ ( rx ) is increasing in x , however, σ ( ux + u ′ z ) is decreasing in x since u is negative which implies a negative correlation due to the FKG inequality.

Lemma G.7. The following function

thus it suffices to show

We will now show the following

is increasing in x where the integrations are w.r.t a centered Gaussian z .

Proof. Since all integrals are w.r.t z , we drop the conditioning with respect to x in the proof. Taking the derivative w.r.t x , and arranging the terms, it suffices to show

which is equivalent to the following due to the property σ ′′ ( z ) = βσ ′ ( z )(1 -σ ′ ( z ))

In the case x ≥ 0 , the LHS is bigger than E [ σ ( x + z )] since σ ′ ( · ) is upper bounded by 1 . Moreover, since σ ( x ) > x and from the convexity of softplus, we get E [ σ ( x + z )] > x . This yields the above inequality by again noting that E [ σ ′ ( x + z )] is upper bounded by 1 .

In the case x < 0 , we need another strategy. We have thanks to Cauchy-Schwartz

for which it suffices to show that v ( z ) := σ ′ ( x + z ) + σ ′ ( x -z ) ≥ 2 σ ′ ( x ) for all z since the centered Gaussian measure p ( z ) is even and the integration can be done over the integrand v ( z ) . We have

that is zero iff either x + z = x -z or x + z = -x + z where the latter is not possible since x < 0 . Hence we get that a critical point of v ( z ) at z = 0 which is a minimizer since v ′′ (0) = 2 σ ′′′ ( x ) > 0 for x < 0 . Hence v ( z ) ≥ v (0) = 2 σ ′ ( x ) for all z which completes the argument.

Finally, it remains to show

From the proof of Lemma G.8, we have that βσ ( x ) > σ ′ ( x ) , which in combination with the following trivial observation for all x < 0 (note that +1 is not needed for the following to hold)

shows that Eq. 177 holds, hence the proof is complete.

Lemma G.8. The softplus family has the following properties i. σ ( x ) is increasing and convex,

ii. σ ( x ) is log-concave (equivalently, σ ( x ) /σ ′ ( x ) is increasing), iii. σ ( x ) /σ ′ ( x ) is convex,

iv. σ ′ ( x ) x/σ ( x ) is increasing, v. σ ′ ( x ) is log-concave.

Proof. For the property (i), see the formulas of σ ′ and σ ′′ in the proof of Lemma F.2. We next prove each one of the properties one after the other. Let us start with property (ii). First note that σ ( x ) is log-concave if and only if σ ( x ) /σ ′ ( x ) is increasing since

where the second inequality is a characterization of log-concavity. We will prove that R ( x ) := σ ( x ) /σ ′ ( x ) is increasing.

Let us write out the ratio explicitly

The first derivative of R is given by

Since log is concave, expanding it around 1 we get log( y + 1) < y for all y > 0 . Substituting y = e βx , we get that the numerator of R ′ is positive, thus R is increasing. This completes the proof of property (ii). Computing the second derivative of R , we get

What remains to show is that βσ ( x ) > σ ′ ( x ) . Using the fundamental theorem of calculus, we get

since 1 / ( y +1) is a lower bound of the integrand which completes the proof of the property (iii). Let us prove the property (iv) by taking the derivative of the function of interest

Using σ ′′ ( x ) = βσ ′ ( x )(1 -σ ′ ( x )) and dropping the positive term σ ′ ( x ) , the numerator of the derivative is

For the case x ≥ 0 , we have σ ( x ) > x and ( βx +1) /e βx > 0 , hence the derivative is positive. For the case x < 0 , we want to show

If e βx + βx +1 > 0 , it is done since the LHS is positive. If e βx + βx +1 ≤ 0 , we have

since log( e βx +1) /e βx is positive. We will next show that log( e βx +1) /e βx is a decreasing function therefore its supremum is achieved at x →-∞ . From the integral expression in Eq. 183, we deduce that log( y +1) /y is a decreasing function since adding smaller terms in the average decreases it. Thus the following limit gives us the supremum using L'Hôpital's rule

Combining it with the Eq. 188 after the substitution y = e βx , we get the desired Ineq. 187 which implies that the derivative is positive in this case too. This completes the proof of property (iv).

For the property (v), we first give a formula for the third derivative of softplus

σ ′ is log-concave if and only if we have

This is equivalent to (1 -y ) log( y +1) < log( y +1) < y where y = e βx > 0 ; the second inequality holds due to y +1 < e y .

which is equivalent to

Peer reviews:
comments: SUMMARY:
This paper is a theoretical study of whether an under-parameterized student network should copy or average teacher weights, there teacher is a one-hidden-layer network with k hidden units and student is the same but with n<k hidden units. The paper derived the optimal “copy-average” configuration under a set of conditions, and also provided a closed-form solution for the case of n=1.

STRENGTHS:
derived some theoretical results in terms of whether smaller two-layer student network should copy or average the weight of a larger two-layer teacher network under certain conditions

WEAKNESSES:
The conditions for the results to hold seems to be very strict, far from practical. I struggle to understand the practical significance of this theory

intuition is not well explained 

mathematical derivation is presented in a straightforward way

QUESTIONS:
Why are the theoretical results presented important when there are so many unrealistic assumptions? 

Why would one “expect that the incoming and outgoing weights of a student neuron are either identical to those of a 29 teacher neuron or that they are closely aligned with the weights of a group of teacher neurons”? What’s the rationale of this expectation?

Could you elaborate how “the loss function in Eq. 5 can be expanded as a weighted sum of 109 the following Gaussian integral terms”, and further expressing “the integral in terms of the covariance of the two-dimensional Gaussian”

LIMITATIONS:
practical significance of the results
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies a shallow teacher-student setting in which a student network being trained attempts to recover the orthogonal weights of a teacher network, assuming normally distributed data. Specifically, the setting being studied is where the number of student neurons $n$ is far smaller than that of the teacher, $k$. An analysis of the optimization landscape is given for the case of a single student neuron, and assuming an odd activation function, it is shown that among the family of critical that are combinations of averages of the teacher's neurons (CA points), the optimal such point is the one where $n-1$ student neurons copy $n-1$ teacher neurons, and the last student neuron averages the remaining teacher neurons.\n\nWhile I'm overall impressed by the technical results in the paper and how well it is written, I find the results a bit limited. On the other hand, this paper studies a technical and challenging problem, and with this in mind, I do believe its merits outweigh its drawbacks.\n\n\nPost-rebuttal:\n\nI have decided to increase my score following the rebuttal and my discussion with the authors. My overall opinion of the paper remains positive and I'm leaning more towards acceptance.\n\nSTRENGTHS:\n- The paper is well structured and written.\n\n- The problem studied is very technical and making interesting observations on it is challenging.\n\n- The technical idea behind the paper (reformulation as a constrained optimization problem) seems interesting and innovative.\n\nWEAKNESSES:\n- The results provided in the paper only partially solve the problem it addresses. For example, the optimality of the copy $n-1$ and average $k-n+1$ teacher neurons point is only established within the family of CA points, and it's possible that there are other non CA points which achieve an even better loss. While a more thorough understanding of the single neuron case is provided, it is of course still limited since it's just a single neuron.\n\n- The analysis only applies to the under parameterized regime where potentially the class of student networks is not even capable of expressing a remotely useful solution. It is far more interesting to study the regime where there are local minima that can achieve small (but not necessarily zero) loss, since this regime might have stronger practical implications. Moreover, the orthogonality assumption which implies $k\\leq d$ further restricts this setting to $n\\ll k\\leq d$, which needless to say does not capture a setting similar to how neural networks are trained in practice in terms of the magnitudes of these quantities.\n\nQUESTIONS:\n- Throughout the paper, in the experimental parts, it is mentioned that gradient flow converges to certain CA or close to CA points. How can you simulate gradient flow? Did you happen to use a very small step size gradient descent? If so please clarify this, since this doesn't imply that gradient flow will necessarily exhibit the same behavior.\n\n\n\nMinor comment:\n\nLine 266: \"span of teacher's\" -> \"span of the teacher's\"\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nhis paper proposes a constrained optimization approach to learning\nstudent networks $g(x) = \\sum_{j=1}^n a_j \\sigma (v_j \\cdot x)$ that\nhave fewer neurons than a two-layer teacher network $f^* (x) =\n\\sum_{i=k} b_i \\sigma( w_i \\cdot x )$. (with $ n \\le k $ as the student network has\nfewer neurons than the teacher )\n\nStarting with the squared-loss between the student and teacher, the authors\npropose a relaxed constrained optimization problem in terms of the student parameter norms\nand correlations to themselves and the teacher parameters.\nThe authors show that in the case\nof general $b$ and odd-activations, for any $n$-partition of the $k$\nteacher neurons, if $(a_j^* , v_j^* )$ are a critical point for a\none-neuron student network, then $v_j^*$ lies in the span of the\nteacher weights $w_j$ corresponding to the partition. They further\nshow that for a multi-neuron student network, concatenating the\ncritical points from one-neuron student networks will produce a\ncritical point for the overall multi-neuron network problem.\n\nThen the authors jump to unit-norm teacher networks, where $b_i = 1$,\nwhere the critical points are unique, and closed-form solutions exist\nfor the erf and ReLU activations.\n\nSTRENGTHS:\n- The stated problem is interesting and I find the results compelling.\n\n- There are several non-trivial aspects to the paper. The authors\n  first convert the student-teacher minimization problem to a relaxed\n  version where they try to optimize variables corresponding to: (a)\n  student parameter norms, (b) correlations between student vectors,\n  and (c) correlations between student and teacher vectors. As the\n  input distribution is Gaussian, the pairwise correlations can be\n  analyzed for a wide class of activation functions.  While some of\n  the techniques are standard (like odd activation functions to\n  decouple neurons in a multi-neuron network), they are compiled well.\n\n- The analysis is fairly slick and the assumptions are reasonable.\n\nWEAKNESSES:\n- The constrained optimization is a relaxation of the original\n  problem of learning a student network, and the critical points of\n  the relaxed problem are analyzed. There is no discussion about the\n  gap between the original and relaxed problem. There is some\n  discussion about small-scale and medium scale problems in the\n  conclusion -- are these related to the gap in the relaxation?\n\n- The teacher network has orthogonal weights, and hence is in a regime\n  where the number of neurons is smaller than the input dimension. I'm\n  not familiar with the literature on learning student networks -- is\n  this standard?\n\n- The authors say that \"for larger n that is close to k, we observe a\n  transient phase where the gradient flow finds a perturbed copy\n  solution that has a lower loss than the optimal CA point.\" How is\n  this possible if your program is a relaxation of the gradient flow\n  objective, and why does it not contradict your claim that \"For\n  shallow neural networks with odd activation function, we prove that\n  “copy-average” configurations are critical points, and the optimum\n  is reached when $n-1$ student neurons with erf activation each copy\n  one teacher neuron and the n-th student neuron averages the\n  remaining $k -n +1$ teacher neurons\"?  I've complained about this elsewhere, but \n  having a Theorem statement for section 5.2 would help alleviate this concern.\n  Also, please maintain consistency between erf and odd-activations, as the sentence I \n  quoted switches between them.\n\n- The writing can be slightly confusing from time to time because of a\n  lot of context switching between one-neuron student networks and\n  multi-neuron odd activation function networks, and then to unit-norm\n  teacher networks, where the outgoing weights are $b_i = 1$.\n\n- Figure 1 is very hard to parse. How should a reader understand that\n  the results in Figure A2 implies convergence to a copy-average of\n  the true solution? And why are the blocks for the\n  pairwise-correlation $u_1$ and $u_2$ weird rectangular blocks?\n  Wouldn't the pairwise correlations just be $(1, \\rho, \\rho, 1)$?\n  Overall, A2, A3, A4 are supposed to explain different settings for\n  $n, k$, but in their current form, they are all equally\n  incomprehensible. The caption is also half a page. There's too much\n  going on, it needs to be simplified.\n\n- Why is there no theorem statement for section 5.2? It seems like a\n  strange choice to not include it given that it's your main result.\n  The lemma currently stated is for a two-neuron network.\n\n- Minor:\n\n   - The first line of the abstract and introduction says that any\n     continuous function can be approximated arbitrarily well given\n     enough neurons in the teacher network, which is true. However,\n     all the results in this paper can only hold when the number of\n     neurons ($k$) is atmost the dimension of the input ($d$), which\n     strongly contradicts the first statement.\n\nQUESTIONS:\nSee points 1, 2, 3 in the weaknesses section above. Overall, I like this paper, but I feel that the presentation makes it hard to parse the results and some statements seem contradictory.\n\nLIMITATIONS:\nCould be discussed in more detail. It's merged with the conclusions section as a monolithic paragraph. The gap between the relaxation and original problem are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe goal of the paper is to understand the following approximation problem: given a student network and teacher network, how should one construct the student network such that the approximation error is minimized. The paper considers the underparameterized regime where the student network has fewer parameters than the teacher network. The first half of the paper considers a constrained optimization problem which is equivalent to the original approximation problem if the student network consists of one neuron (otherwise, the constrained optimization problem is a relaxation of the original approximation problem). The second half provides a characterization of a subset of critical points for the approximation problem. Given additional assumptions on the activation function used in the networks, they also provide an analytical formula for non-trivial critical points of the approximation problem and some empirical verification of the paper's claims.\n\nSTRENGTHS:\nThe paper is mostly well written. In addition, to the best of my knowledge, the proof (in restricted settings) and the empirical observations (in more general settings) of the copy average of phenomenon seems to be novel. Furthermore, the constrained optimization problem that serves as a relaxation of the original approximation problem is interesting. In addition, the paper provides a nice family of critical points for the approximation problem that can be constructed by concatenating critical points of smaller approximation problems.\n\nWEAKNESSES:\nMany of the results of the paper requires that the student network consist of one neuron. While section 5.2 gives a characterization of critical points for the multi-neuron case, it requires that the activation function is erf, an uncommon activation function. In addition, the input data is assumed to be gaussian.\n\nQUESTIONS:\nIn the context of section 5.2 multi-neuron network, is there a nice way to characterize which of the n - 1 teacher neurons will be copied other than trying every (n - 1) copy configurations and computing their approximation error?\n\nFor figure 1.A2, it is not clear which pair of student and teacher neurons are being used to compute the correlation (perhaps label the y-axis with the teacher neurons v1, ..., vk). Furthermore, I don't really see what insight the plots of r_i and a_i give (perhaps add corresponding teacher r'_i and a'_i ? )\n\nDoes lemma 5.4 hold for other activation functions? This seems to be the major workhorse for characterizing critical points for student networks for multiple neurons and thus it would be nice to address if a similar lemma holds for common activation functions.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-24"}