Expand-and-Cluster: Parameter Recovery of Neural Networks

TL;DR

This work shows that the incoming weight vector of each neuron is identifiable up to sign or scaling, depending on the activation function, and can identify layer sizes and weights of a target network for all commonly used activation functions.

Abstract

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.

Figures (22)

• Figure 1: Expand-and-Cluster: we overcome the non-convex problem of recovering the $q$ parameters of an unknown network by: (i) expanding the dimensionality of the parameter space $\Theta$ by a factor $\rho$ to relax the optimisation problem, $\Theta \rightarrow \hat{\Theta}$; (ii) mapping the loss minimiser in expanded space $\hat{\theta}^*$ to the original parameter space via clustering of $N$ overparameterised solutions, $\hat{\theta}^* \rightarrow \theta^*$.
• Figure 2: Catalogue of student neuron types at zero loss: zero loss overparameterised students can only contain a few neuron types. 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 colour-coded letters indicate equal quantities. A) Neuron types from simsek2021geometry adapted with biases:duplicate-type neurons combine 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) Novel neuron types:constant types contribute a fixed amount to the next layer by learning a null vector. For even + linear activation functions, linear type groups combine to contribute a linear function. Linear duplicate type groups copy the teacher vector and its opposite, replicating the teacher neuron up to a linear mismatch. Offbound types: at non-exact zero loss, their input support is placed in the linear or zero region of $\sigma$.
• Figure 3: Parameter identification with Expand-and-Cluster.A) Training scheme: once an overparameterisation factor yields near-zero training losses, train N overparameterised students on the teacher-generated dataset $\mathcal{D}(\mathbf{X}, y)$; B) Similarity matrix: L2-distance between hidden neurons' input weight vectors of layer $l$ for all $N$ students. Large-sized clusters 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. The code is available at https://github.com/flavio-martinelli/expand-and-cluster.
• Figure 4: Synthetic teachers 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. Each hidden neuron generates a hyperplane, ${w_i}^T x + b_i = 0$ (dashed lines); the direction of the weight vector $w_i$ is indicated by an arrow starting from the hyperplane and the sign of the output weight $a_i$ by its colour. Top left: generalization of the XOR or parity-bit problem to a regression setting. From left to right: As the number of teacher hidden neurons $r$ increases the contour 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 overparameterisation on convergence: student networks with overparameterisation $\rho \geq 2$ are more likely to converge to near-zero loss than those without. We report the following general trends: (i) overparameterisation 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 overparameterisation 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. \ref{['fig:supp2']}).
• Figure 5: A) Expand-and-Cluster applied to mildly overparameterised 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 overparameterised student networks have zero loss with 4 hidden neurons after reconstruction (large coloured stars). B) Loss after Expand-and-Cluster for all teacher networks and student sizes from Figure \ref{['fig:4']}: the colour 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.