Expand-and-Cluster: Parameter Recovery of Neural Networks
Can we identify the weights of a neural network by probing its input-output
mapping? At first glance, this problem seems to have many solutions because of
permutation, overparameterisation and activation function symmetries. Yet, we
show that the incoming weight vector of each neuron is identifiable up to sign
or scaling, depending on the activation function. Our novel method
'Expand-and-Cluster' can identify layer sizes and weights of a target network
for all commonly used activation functions. Expand-and-Cluster consists of two
phases: (i) to relax the non-convex optimisation problem, we train multiple
overparameterised student networks to best imitate the target function; (ii) to
reverse engineer the target network's weights, we employ an ad-hoc clustering
procedure that reveals the learnt weight vectors shared between students --
these correspond to the target weight vectors. We demonstrate successful
weights and size recovery of trained shallow and deep networks with less than
10\% overhead in the layer size and describe an `ease-of-identifiability' axis
by analysing 150 synthetic problems of variable difficulty.
Flavio Martinelli, Berfin Simsek, Wulfram Gerstner, Johanni Brea
# Expand-and-Cluster:

###### Abstract

Can we recover the hidden parameters of an Artificial Neural Network (ANN) by probing its input-output mapping? We propose a systematic method, called 'Expand-and-Cluster' that needs only the number of hidden layers and the activation function of the probed ANN to identify all network parameters. In the expansion phase, we train a series of networks of increasing size using the probed data of the ANN as a teacher. Expansion stops when a minimal loss is consistently reached in networks of a given size. In the clustering phase, weight vectors of the expanded students are clustered, which allows structured pruning of superfluous neurons in a principled way. We find that an overparameterization of a factor four is sufficient to reliably identify the minimal number of neurons and to retrieve the original network parameters in \(80\%\) of tasks across a family of 150 toy problems of variable difficulty. Furthermore, shallow and deep teacher networks trained on MNIST data can be identified with less than \(5\%\) overhead in the neuron number. Thus, while direct training of a student network with a size identical to that of the teacher is practically impossible because of the highly non-convex loss function, training with mild overparameterization followed by clustering and structured pruning correctly identifies the target network.

## 1 Introduction

It is known since the 1980s that finding a solution to the XOR problem with gradient descent is easier with a larger hidden layer, even though a minimal network with two hidden neurons is theoretically sufficient to solve the problem [1]. Indeed, even very small networks have a non-convex loss function [2; 3; 4]. During the last decades, advances in the theory of artificial neural networks indicate that the loss function is rough for networks of a minimal size, but becomes effectively convex in the limit of infinitely large hidden layers [5; 6; 7; 8; 9]. In a teacher-student setup, the complexity of the landscape can roughly be estimated from the ratio between the number of symmetry-induced critical-point manifolds at a positive loss and the number of manifolds at zero loss [10]; at zero loss, teacher and student are functionally equivalent. Importantly, as the degree of overparameterization is increased, the loss landscape undergoes a qualitative change from a ratio larger than one to a ratio smaller than one, suggesting that already for mild overparameterization the landscape is dominated by multi-dimensional zero-loss manifolds [10; 11]. Here we ask whether we can use these theoretical insights to construct a network identification algorithm.

Network parameter identification requires the data to be generated by a neural network; therefore we work in the teacher-student framework. It is arduous to train to zero loss if the student has the same number of hidden neurons as the teacher, however, if the size of the student network is increased by a
factor of two to four, training becomes reliable enough to find a solution close to zero loss. Therefore, we propose to first **expand** the number of neurons in the hidden layer of multiple (\(N\)) students until we can train to very low, ideally zero, loss. Using insights on the structure of the zero-loss manifold [10], we then **cluster** similar neurons between different students and prune back to the minimal network size. Thus, the detour via mildly overparameterized networks enables us to reliably find the solution to the original non-convex parameter identification problem (Figure 1).

We present a procedure to extract a functionally-equivalent model of minimal size; this is more than merely matching the teacher's accuracy on a test set. The desired result is a minimal list of parameters that one can send to a friend, and the friend can consult the teacher to verify that the parameters are correct and the list is minimal. To test the Expand-and-Cluster algorithm, we propose a family of problems with artificial data that enable us to generate hundreds of different regression tasks of variable difficulty, dimensionality, and number of parameters; including generalized XOR-like problems. These tasks are harder than random-teacher models which rarely lead to XOR-like situations. To simplify the procedure, we assume that we know the number of hidden layers and the activation function. For larger-scale applications, we use a regression task to extract parameters of a teacher network optimized on MNIST data. Overall, our examples show that the non-convex network identification problem can be successfully addressed.

Our contributions can be summarized as follows:

* We demonstrate that we can achieve exact, functionally-equivalent _and_ minimal size parameter recovery with learning-based methods on toy-sized problems, despite the extreme non-convexity of the problem.
* Mild overparameterization is enough to solve the non-convex problem due to the combinatorial proliferation of global minima predicted by Simsek et al. [10];
* Our method is orthogonal to other parameter recovery works that are not learning-based, but either rely on special properties of ReLU [12; 13] or perform reconstruction in restricted setups such as committee machines with known layer size [14]. We show successful reconstruction for smooth activation functions on experiments of a similar scale and expand with results on deep fully connected networks.

## 2 Background and Methods

### Theoretical foundations

Overparameterization consists in increasing the number of parameters of a neural network such that its expressivity is larger than necessary for representing a given dataset [15; 16]. For teacher-student setups, we call a student 'overparameterized' if it has more hidden neurons than the teacher in at least one layer. If an overparameterized student network replicates the teacher mapping with zero loss, the space of all possible solutions is fully described by the geometry of the global minima manifold [10]. The global minima manifold contains only two types of hidden units, namely duplicate and zero-type neurons (see theorem 4.2 of Simsek et al. [10] and Fig. 2A); under the following assumptions: one-hidden layer network \(\sum_{i=1}^{m}a_{i}\sigma(w_{i}x)\), infinite input data support, population loss limit, zero bias of all teacher neurons and analytical activation function \(\sigma\) with infinite non-zero even _and_ odd derivatives evaluated at zero. The last assumption guarantees that the activation function has no obvious or hidden symmetries around zero. We call activation functions that satisfy these assumptions'symmetry-free'.

Figure 1: **Expand-Cluster**: we overcome the non-convex problem of recovering the \(q\) parameters of an unknown network by: **(i)** expanding the parameter space by a factor \(\rho\) to relax the optimization problem, \(\Theta\rightarrow\hat{\Theta}\); and **(ii)** map the solution to the original parameter space through theory-informed clustering, \(\hat{\theta}^{*}\rightarrow\theta^{*}\). The dataset \(\mathcal{D}\) is generated by an unknown network.
The intuition for the result of Simsek et al. [10] is that, for zero-loss solutions, each teacher hidden neuron \(a^{*}\sigma(w^{*}x)\) must be replicated in the student by a _duplicate-type_ group of one or more units, contributing \(\sum_{i}a_{i}\sigma(w_{i}x)\). The duplicates' input weight vectors are all aligned with the teacher neuron, \(w_{i}=w^{*}\), while their weighted contribution equals the teacher neuron's output weight \(\sum_{i}a_{i}=a^{*}\). In the same student there can also exist _zero-type_ neuron groups with a null contribution to the student input-output mapping, characterized by \(w_{1}=\cdots=w_{q}\), \(\sum_{i}a_{i}=0\). Neuron types from Simsek et al. [10] but in the presence of biases are summarized in Figure 2A. The assumption of symmetry-free activation function \(\sigma\) is critical because, for example, a student network with tanh units could also contain anti-symmetric input vectors with a switch of output signs.

In this paper, we drop several of the above assumptions to fully reconstruct a teacher network by identifying duplicate neurons while pruning all other superfluous hidden neurons.

### Artificial data in teacher-student networks

To investigate the advantages of overparameterization as a function of layer size, we devised a series of very challenging tasks inspired by the parity-bit problem (or multidimensional XOR), known as a difficult problem for neural networks [1]. To create tasks of variable difficulty, we adapted the problem into a regression format using the following procedure:

All tasks with artificial data have \(d\)-dimensional uniformly distributed input data in the range \(x_{i}\in[-\frac{1}{\sqrt{3}},\frac{1}{\sqrt{3}}]\). A specific task is defined by the parameters of a teacher network. Each hidden neuron \(i\) of the teacher is randomly sampled from a set of input weights \(w_{i}\in\{-1,0,1\}^{d_{in}}\), output weights \(a_{i}\in\{-1,1\}\) and biases \(b_{i}\in\{-\frac{2}{3}\sqrt{3},-\frac{1}{3}\sqrt{3},0,\frac{1}{3}\sqrt{3}, \frac{2}{3}\sqrt{3}\}\). We repeat the sampling if two hidden neurons are identical up to output weights signs to avoid two hidden neurons canceling each other. The resulting input weight vectors \(w\) are first normalized to unity and then both \(w\) and \(b\) are multiplied by a factor of 3. The above procedure yields hyperplanes in direction \(w\) located at a distance \(|b|/||w||\) from the origin, and a steeply rising (or falling) activation on the positive side of the hyperplane. Finally, analogous to batch normalization, the output weights and biases are scaled such that the output has zero-mean and unit variance when averaged over the input distribution: \(a\leftarrow a/\mathrm{std}(y)\) and \(b_{2}=-\langle y\rangle/\mathrm{std}(y)\), where \(y\) is the output vector of the network. We study teachers with input dimensionality \(d_{in}\in\{2,4,8,16,32\}\) and hidden layer size \(r\in\{2,4,8\}\). Figure 4A shows examples of different teachers with input dimension \(d_{in}=2\) and a single hidden layer. We use the symmetry-free activation function \(\sigma=\sigma_{sig}(4x)+\sigma_{soft}(x)\), where \(\sigma_{sig}=\frac{1}{1+e^{-x}}\) and \(\sigma_{soft}=log(1+e^{x})\) for all our simulations unless specified otherwise. The above construction of hidden neuron parameter vectors can be generalized to multi-layer teachers by stacking the procedure (further details in appendix A.5.2).

Our construction yields XOR-like and checkerboard-like structures where hyperplanes are parallel to each other and divide the input space into separate regions. In contrast to our approach, constructing shallow networks with randomly drawn input weight vectors yields easy tasks since all weights tend to be orthogonal [17; 18; 19] and it is known that randomly initialized deep networks tend to behave as constant random functions [20], yielding uninteresting data-generator models. In contrast to committee machines [14], our artificial tasks are more difficult since the output neuron is not merely averaging the contributions of the hidden neurons.

### MNIST regression

To extend the procedure beyond artificial tasks, we also recover parameters of networks trained on the MNIST dataset [21]. We pre-trained teacher networks of either 10, 30, or 60 hidden neurons on MNIST (after removal of the uninformative input pixels, Fig. S5) with standard training procedures. We then used input-output pairs generated by the last layer of the teacher network to define a regression task on which we trained student networks. For our method to work it is crucial that we have access to the classifier's probabilistic output (e.g. the values after the softmax operation or multiple accumulated decisions of a stochastic classifier) and not only the most probable class (e.g. values after argmax operation) which would make parameter recovery impossible.
### Related Work

**1. Pruning:** Commonly used pruning methods are heuristic and based on pruning weights [22, 23, 24], whereas we propose a theory-based [10] approach that removes entire units (structured pruning) potentially relevant for hardware implementations [25]. Our Expand-and-Cluster method reconstructs a minimal network by positively selecting commonly occurring hidden neurons across a set of overparameterized student networks, as opposed to other structured pruning algorithms [26, 27, 28] which remove putatively redundant units; Srinivas and Babu [29] also prune neurons based on weight similarity. In contrast to Ankner et al. [30], who link prunability to data dimensionality, we show a link to data complexity for fixed dimensionalities.

**2. Loss landscape and overparameterization:** Neural network landscapes exhibit a cusp-like transition from under- to over-parameterization [31, 32], related to double-descent phenomena [33]. Overparameterized solutions found by different random initializations are similar to each other in the function space [34] and have been _approximately_ mapped to each other by permutation of hidden neurons [35, 36, 37, 38, 39, 40]. In the teacher-student setting, different zero-loss solutions are _exactly identical_ to each other up to duplicate-type neurons, zero-neuron addition and permutation symmetry in the population loss limit [10]. However, convergence to a zero-loss solution can be hindered either due to the emergence of local minima [41] or the flatness at the bottom of the landscape [42].

**3. Non-convex optimization:** Many non-convex optimization problems are tackled with the following strategy: (i) expand, or _lift_, to a higher dimensional space in order to relax the problem and guarantee convergence to global minima; (ii) map, or _project_, the relaxed solution to the original space by exploiting the problem's intrinsic symmetries and geometry [43]. This approach is used in applied mathematics [44, 45], computer vision [46], nonlinear programming [47, 48, 49], control theory [50], machine learning [43, 51] and many others3. Despite the above-mentioned achievements, for neural networks, the picture is far from complete. Unlike infinitely-wide neural networks [5], the loss functions of finite-width neural networks exhibit essential non-convexity so that the gradient flow converges to many fundamentally different solutions depending on the initialization [52, 41, 53]. Shallow networks with polynomial activation functions can be globally optimized under certain guarantees by lifting the optimization problem to tensor decomposition [51, 54]. Even though the mildly overparameterized regime is non-convex, we show that we can exploit its reduced complexity (compared to non-overparameterized problems [10]) to find zero-loss solutions.

Footnote 3: A curated list of solvable non-convex optimization problems can be found here: [https://sunju.org/research/nonconvex](https://sunju.org/research/nonconvex)

**4. Interpretability:** Explaining in qualitative terms the behavior of single neurons embedded in deep networks is a challenging task [55, 56]. For example, in symbolic regression, small networks with vanishing training loss are desirable for interpretability [57]. We provide precise mathematical explanations of all hidden neurons found in zero-loss overparameterized student networks in relation to a teacher network of minimal size; going beyond the notion of'superimposed' features described in Elhage et al. [58]. We follow the terminology of [10] and use the terms duplicate or zero type neurons which are related to 'Monosemantic' [59] and 'Frivolous' [60] units found in various settings.

**5. Functionally Equivalent Model Extraction:** Our paper focuses on functionally equivalent extractions, that is retrieving a model \(\mathcal{M}\) such that \(\forall x\in X,\mathcal{M}(x)=\mathcal{M}^{*}(x)\), where \(\mathcal{M}^{*}(x)\) is the target model. This type of extraction is the hardest achievable goal in the field of Model Stealing Attacks [61], using only input-output pairs [13]. In addition, out of all the functionally equivalent models that we describe in Section 3.1, we extract the one of _minimal size_. Conditions for neural network identifiability and their symmetries have been studied theoretically for different activation functions [62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72], although overparameterized solutions were not considered. To the best of our knowledge, existing functionally equivalent extractions of _trained networks_ rely on identifying boundaries between linear regions of shallow ReLU networks [13, 73], with partial success for deep ReLU networks [12]. Janzamin et al. [51] show a theoretical reconstruction based on third-order derivatives. Fornasier et al. [14], building upon [74, 75, 76, 77], propose an identification method for wide committee machines (shallow networks with unitary second layer weights) and unit norm teacher weights, where knowledge of the teacher layer size is necessary. We are the first to propose an exact recovery method for arbitrary activation functions on shallow and deep fully connected networks of unknown layer widths, with successful reconstructions of networks trained on MNIST. Our method is learning-based and fundamentally different from the known approaches. However, none of these methods were shown to work in large-scale applications.
## 3 Results

### Neuron types near zero-loss

We extend the categorization of neuron types of Simsek et al. [10] to neurons with _bias_, _finite_ input data support and analytical activation functions that contain one _symmetry_ due to infinite even or odd non-zero derivatives. This generalization includes commonly used activation functions such as GELU, sigmoid, tanh, softplus. The catalog of new neuron types is sketched in Figure 2B.

In the case of a teacher neuron with bias \(a^{*}\sigma(w^{*}x+b^{*})\), a group of _duplicate-type_\(\sum_{i}a_{i}\sigma(w_{i}x+b_{i})\) has input weight vectors and biases aligned to the teacher: \(w_{i}=w^{*},\ b_{i}=b^{*},\ \sum_{i}a_{i}=a^{*}\), while the new _zero-type_ group has aligned, but arbitrary, weights and biases \(w_{1}=\cdots=w_{q},\ b_{1}=\cdots=b_{q},\sum_{i}a_{i}=0\). With biases, _constant-type_ student neurons can also arise: they have vanishing input weights \(w=0\) and contribute a constant amount of \(a\sigma(b)\) to the next layer; to keep an exact mapping of the teacher, this constant contribution must be accounted for in the next-layer biases.

**GELU or softplus** teacher neurons can be written as a combination of a linear and an even function: \(a^{*}\sigma(w^{*}x+b^{*})=a^{*}c_{1}\cdot(w^{*}x+b^{*})+a^{*}\sigma_{\text{ even}}(w^{*}x+b^{*})\), where \(c_{1}\) is the slope of the linear approximation around \(0\). The even symmetry allows student neurons to combine in groups of aligned, \(w=w^{*}\), and opposite, \(w=-w^{*}\), input weight vectors and biases (see appendix A.1.1 for details). If in such a group the output weights sum to zero, we obtain a _linear-type_ group that contributes a linear function to the next layer. To guarantee an exact mapping of the teacher, there must be another _linear-type_ group in the same layer contributing the exact opposite linear term. Alternatively, if the sum of the output weights matches the teacher output weight, then the neuron group is a _linear duplicate-type_, replicating a teacher neuron up to a misaligned linear contribution; the latter can be accounted for by another linear group in the same layer. A numerical example of different neurons found by a softplus student is shown in Figure S6.

**Tanh and Sigmoid** can be treated similarly: student neurons can group with aligned and opposite weight vectors and combine to form constant types or duplicate plus constant. They are described in detail in the appendix along with a categorization of all commonly used activation functions (appendix A.1.2).

Finally, in near zero-loss solutions, students may have _offbound-type_ neurons characterized by being almost constant, linearly increasing or decreasing (depending on the asymptotic behavior of the activation function) in regions of the input space with actual data, which is reminiscent of the dead ReLU phenomenon. The output of groups of _offbound-type_ neurons can synergize so as to contribute constant amounts, which can be compensated by bias adjustments in the next layer; see Figure 2.

Figure 2: **Catalog of neuron types near zero loss:** with \(\sigma\) analytical, near-zero loss overparameterized students can only contain a handful of neuron types. On the left, a teacher neuron is defined along with a sketch of its output, the grey bar indicates the finite input support \(x\) to the neuron; the same color-coded letters indicate equal quantities. **A) Neuron types from Simsek et al. [10] adapted with biases:**_duplicate-type_ neurons combine together to replicate a teacher neuron by copying its weight vector \(w^{*}\) and bias \(b^{*}\), their activations \(a_{i}\) sum up to the teacher activation \(a^{*}\). _Zero-types_ have aligned weight vectors and biases but cancel each other via output weights. **B) Additional neuron types:**_constant-types_ contribute a fixed amount to the next layer by learning a null vector. Only in the presence of even + linear activation functions, such as GELU, _linear-type_ groups combine in a way to contribute a linear function, while _linear duplicate types_ replicate the teacher neuron summed to an extra linear component. _Offbound-types_ appear at non-exact zero loss and are characterized by having the nonlinearity placed away from their input support.
### Expand-and-Cluster algorithm

We recall that the only way an overparameterized student can reach zero loss is by representing all the neurons at least once; see Section 3.1. If a student can be trained to exact zero loss, which is possible in small toy setups, it is almost trivial to identify the different neuron types, including the teacher neurons (see a numerical example in Fig. S6). However, in practice and even in larger tasks with artificial data, overparameterized networks are difficult to train to exact zero loss because of computational and memory budgets. Therefore we need ways to identify the teacher neurons from imperfectly trained students. These students present approximate duplicates of teacher neurons and neurons of other types that point in arbitrary directions. In a group of \(N\) imperfectly trained students, neurons that approximate teacher duplicates form clusters while other neurons do not, since they are not aligned between students. Therefore we propose the following clustering procedure (Fig. 3):

**Step 1: Expansion phase.** Rapidly train a sequence of networks with increasing sizes of hidden layers to a fixed convergence criterion. To do so, use teacher-generated input-output pairs, \(\mathcal{D}=\{\mathbf{X},y\}\), and minimize the mean square error loss with standard gradient descent methods. This expansion phase allows finding a network width \(m\) at which convergence to nearly zero loss is possible (Fig. S3B).

**Step 2: Training phase.** Train \(N\) students of \(L\) layers, width \(m\), on \(\mathcal{D}=\{\mathbf{X},y\}\) to minimize the mean square error to the lowest value achievable by the chosen optimizer (Fig. 3A).

**Step 3: Clustering phase.** Collect the first hidden layer neurons of the \(N\) students, then cluster the input weight vectors with hierarchical clustering on the L2 distance. With a threshold selection criteria that maximizes the number of large clusters (size \(\geq\gamma N,\,0\,<\,\gamma\,\leq 1\)), obtain groups of aligned weight vectors; these clusters should include all duplicate teacher neurons. Proceed to filter out clusters whose elements are not aligned in angle (median alignment \(\geq\beta\)), removing eventual zero or constant type neuron clusters. Then, merge each remaining cluster of duplicate neurons into single

Figure 3: **Parameter identification with Expand-and-Cluster.****A) Training scheme:** once an overparameterization factor yields near-zero training losses, train N overparameterized students on the teacher-generated dataset \(\mathcal{D}(\mathbf{X},y)\); **B) Similarity matrix**: \(\ell_{2}\)-distance between hidden neurons’ input weight vectors of layer \(l\) for all \(N\) students. Large-sized clusters (red and green) are good candidate weight vectors. **C) Dendrogram obtained with hierarchical clustering:** the selected linkage threshold is shown in orange. Clusters are eliminated if too small (blue) or unaligned (red), the remaining clusters are shown in green.
hidden neurons to reconstruct the layer. We noticed that higher layers align with the teacher weights only at prohibitively low losses [78]. Therefore, if the student networks have more than one hidden layer left to reconstruct, we go back to Step 2 and train again \(N\) overparameterized students of \(L\gets L-1\) layers, using as input \(\mathbf{X}\) the output of the last reconstructed layer. Repeat this procedure until the last hidden layer is reconstructed (Fig. 3 Algorithm 1, more details in appendix A.5.2).

**Step 4: Fine-tuning phase.** Adjust the final parameters using the training data.

### Artificial data experiments

We trained overparameterized students on the family of teachers of Section 2.2 with a single hidden layer of \(r\) neurons (Figs. 4 and S7). For an overparameterization factor \(\rho\), the student hidden layer has \(\rho r\) neurons. To stay close to the theoretical setting of near-zero loss and to obtain perfect parameter recovery, we trained all the networks with the package MLPGradientFlow.jl [79]. This allowed us to find global and local minima with machine precision accuracy for networks without overparameterization (Fig. 4B, \(\rho=1\)). However, even with second-order ODE solvers and slightly larger networks, it becomes challenging to converge fully to global minima within a reasonable amount of time (Fig. 4C, \(\rho\in\{2,4,8\}\)). Hence, methods to deal with imperfectly trained students are needed. Since training is full-batch (30k data points), the only source of randomness is in the initialization; see

Figure 4: **Artificial data define tasks of variable difficulty.****A) For fixed \(\mathbf{d_{in}}\), teacher complexity increases with number \(\mathbf{r}\) of hidden neurons:** contourplot of the teacher network output \(y=b^{2}+\sum_{i}^{r}a_{i}\ g(w_{i}{}^{T}x+b_{i}^{1})\) for \(d_{in}=2\) input dimensions and different \(r\), \(b^{2}\) denotes the bias of the output unit. Each hidden neuron generates a hyperplane, \(w_{i}{}^{T}x+b_{i}^{1}=0\) (dashed lines); the direction of the weight vector \(w_{i}\) is indicated by an arrow and the sign of the output weight \(a_{i}\) by its color. Top left: generalization of the XOR or parity-bit problem to a regression setting. From left to right: As the number of hidden neurons increases the level lines become more intricate. **B) Non-convexity prevents training to zero loss:** for each combination of \(d_{in}=2,4,8,16,32\) and \(r=2,4,8\) we generated 10 teachers; for each teacher, we trained 20 or 10 students (for \(r=8\)) with different seeds. Each teacher corresponds to one row of dots while each dot corresponds to one seed (see inset bottom right). Dark blue dots indicate loss below \(10^{-14}\). Student networks of the same size as the teacher (\(\rho=1\)) get often stuck in local minima. The effect is stronger for larger ratios \(r/d_{in}\). **C) Effects of overparameterization on convergence:** student networks with overparameterization \(\rho\geq 2\) are more likely to converge to near-zero loss than those without. We report the following general trends: (i) overparameterization avoids high loss local minima, (ii) the dataset complexity, i.e. number of hidden neurons per input dimension \(r/d_{in}\), determines the amount of overparameterization needed for reliable convergence to near-zero loss. For difficult teachers, i.e. overcomplete (\(r/d_{in}\geq 1\)), training is very slow and convergence is not guaranteed in a reasonable amount of time (see Fig. S4).
appendix A.5 for more details. Figure 4C shows a beneficial trend as overparameterization increases, but also highlights a strong dependence on the dataset (or teacher) complexity \(r/d_{in}\): as the number of hyperplanes per input dimension increases, it becomes harder to train overparameterized students to global minima.

We find that direct training of 20 student networks without overparameterization (data generated from a teacher network with \(r=4\) hidden neurons and input dimensionality \(d_{in}=4\)) does not yield a single case of convergence to zero loss (Fig. 5A). For the same teacher, the application of the Expand-and-Cluster algorithm yields student networks that achieve zero loss and hidden-layer size equal to that of the teacher if an overparameterization of \(\rho=2,4\) or 8 is used in Step 2. This suggests that successful retrieval of all parameters of the teacher network is possible (Fig. 5A).

We tested the quality of parameter identification with Expand-and-Cluster for each teacher network of Figure 4 and illustrate the final loss of the reconstructed networks in Figure 5. For example, of the 30 teacher networks with input dimensionality \(d_{in}=8\), all except 2 networks were correctly identified as indicated by a zero-loss solution (for \(\rho=4,8\) and RMSE \(\leq 10^{-14}\), dark blue in Fig. 5B). Of 150 different teacher networks, 118 (\(\sim 80\%\)) were correctly identified with \(\rho=4\). For all correctly identified networks, the maximum angle between the weight vectors of the teacher network and the corresponding vector of the reconstructed student network was less than \(3\cdot 10^{-8}\) radians. In all but 7 out of 118 successful recoveries, the number of neurons found matches that of the teacher; the other cases have at most up to 4 neurons in excess (these can be easily categorized into zero-type and constant-type neurons, see e.g. Fig. S6).

### MNIST shallow and deep experiments

Now we turn to applications using bigger networks where zero-loss solutions are not attainable because of computational time and memory budget, and noise in the optimizer (SGD). Expand-and-Cluster can discern duplicates from zero or constant type neurons that would not be otherwise distinguishable with a simple hard threshold on the weight values and therefore outperforms'magnitude pruning' [23]. We explore the sensitivity of Expand-and-Cluster to the number of students \(N\) on the MNIST regression task described in Section 2.3 (Figure 6B). We found that with two students Expand-and-Cluster works reasonably well but does not reliably identify the minimal network. However, with \(N\geq 4\) student networks, the Expand-and-Cluster algorithm enabled us to identify the teacher network with at most 5 percent of additional neurons.

Figure 5: **A) Expand-and-Cluster applied to mildly overparameterized students reaches zero loss: a total of 80 student networks with 4, 8, 16 or 32 hidden neurons have been trained using data generated by a teacher with \(r=4\) hidden neurons and \(d_{in}=4\) input dimensions. None of the 20 students with 4 hidden neurons reached zero loss (orange dots, \(\rho=1\)), while all overparameterized student networks have zero loss with 4 hidden neurons after reconstruction (large colored stars). B) Loss after Expand-and-Cluster for all teacher networks and student sizes from Figure 4: the color of each small horizontal bar represents the final loss. Only a small fraction of teacher networks (i.e., those in yellow) were not identified correctly.**
Unlike other methods [12; 13; 14], Expand-and-Cluster identifies deep fully connected networks trained on artificial data (see appendix A.3) and larger networks trained on MNIST, Figure 6B. We show a successful reconstruction of three hidden layers of \(30\) neurons each by applying Expand-and-Cluster to students overparameterized by a factor 3. The reconstruction identifies layer sizes up to \(4\) neurons in excess for the last hidden layer and achieves a loss of \(3\) orders of magnitude lower than similar-sized networks trained from random initialization (Fig. 6B). The reconstructed network has a fidelity (fraction of labels matching with the teacher) of \(100\%\) and \(96.75\%\) on train and test set, respectively.

## 4 Conclusions

Even if the data is generated by a teacher of known size and architecture, recovering the parameters of the teacher network by training a student network of the same size is difficult to achieve with gradient-descent methods [13]. Indeed, if a student network of the same size as the teacher is trained, it is rare to find the global minimum because of the extreme non-convexity of the loss function: across 1024 training runs of students that contain the same number of hidden neurons as the teacher the global minimum was found only in 2 cases (Fig. S3A). However, using the same computation budget to train 128 students containing \(\rho=8\) times more neurons than the teacher, we reached near zero-loss solutions in 72 cases (\(56\%\)). For a successful application of the Expand-and-Cluster algorithm, it is desirable to have at least \(N=3\) students near a global minimum. Thus, with the above probability of convergence with an overparameterization of \(\rho=8\), ten students are sufficient to find at least three low-loss solutions with \(\sim 95\%\) probability - yet, using the same computation time for 80 runs of the student with the minimal number of hidden units is unlikely to yield even one zero-loss solution. In summary, the detour of expansion-convergence-clustering is a computationally efficient way to find the unique (up to permutations) global minima of a non-convex loss function where standard solution methods rarely work.

Four generalizations of our approach should be straightforward. First, Expand-and-Cluster should work on convolutional layers just as well as it does on dense layers. Second, we obtained our results in a regression setting with a mean square error loss, but the method should also work with the cross-entropy loss as is commonly used in classification settings. Third, our simulation results are obtained with a symmetry-free activation function but, as discussed in Section 3.1, our results can be generalized to common activation functions, if we correctly account for symmetries. Fourth, the choice of dataset should not bring any difference in difficulty as we are dealing with a regression task of data generated by a smaller teacher network. At this stage, our results are limited to simple setups, primarily due to the high computational budget required to reach low losses in more complicated setups. For shallow relu networks, appendix A.4, our method does not reach the same exact extraction accuracy as other methods [12; 13], probably due to the slow finetuning required. In contrast, our method works for arbitrary smooth activation functions and is extendable to deeper networks.

Figure 6: **Parameter recovery for larger networks**. **A) MNIST shallow teachers**: fraction of excess neurons with respect to the teacher size) clustered as a function of the \(N\) students used for Expand-and-Cluster(\(N,\gamma=0.5,\beta=\pi/6\)). Combined statistics across three shallow teachers of sizes \(r\in\{10,30,60\}\) pre-trained on MNIST data. **B) MNIST deep teacher**: A fully connected teacher of layer sizes 784-30-30-30-10 trained on MNIST is reconstructed with Expand-and-Cluster(\(N=50,\gamma=0.5,\beta=\pi/5\)) applied to students of factor \(\rho=3\) overparameterization (magenta dots) with 4 excess neurons in the last hidden layer (magenta star). Direct training with the teacher network architecture never reached losses below \(0.1\) (orange dots).
## Acknowledgments and Disclosure of Funding

This work was supported by Sinergia Project CRSII5_198612 and SNF Project 200020_207426.
