Revenge of Monosemanticity: Specialized Neurons Improve Data Efficiency in MLPs

Amirhesam Abedsoltan$^{1}$
Enric Boix-Adsera$^{2}$
Fivos Kalogiannis$^{1}$
Mikhail Belkin$^{3,1}$

$^{1}$ Department of Computer Science and Engineering, UC San Diego
$^{2}$ Department of Statistics and Data Science, The Wharton School, University of Pennsylvania
$^{3}$ Halıcıoğlu Data Science Institute, UC San Diego

Abstract

Understanding how neural networks learn and organize features is central to understanding their behavior. Much existing theory of feature learning has focused on the emergence of a global low-dimensional predictive geometry. We show that this picture is incomplete. In regression problems with clustered data, we demonstrate that multilayer perceptrons (MLPs) naturally develop monosemantic specialized neurons: individual neurons become strongly aligned with a specific predictive feature relevant to a particular region of the input space. Rather than learning a single global low-dimensional representation, MLPs learn a collection of local low-dimensional representations that can collectively span a high-dimensional space. This specialization provably gives MLPs a data-efficiency advantage over feature-learning methods based on a global low-dimensional representation.

Executive Summary: Neural networks often succeed by learning internal features rather than relying on fixed representations. A common explanation is that they discover a single global low-dimensional subspace of predictive directions in the input. This view is incomplete for data drawn from multiple clusters, where each cluster has its own predictive direction and response function. In such cases, prediction can be simple locally within every cluster while the directions collectively span the full ambient dimension, leaving no useful global low-dimensional structure.

The paper examines how multilayer perceptrons handle exactly this clustered regression setting. It shows that standard MLPs develop a different form of feature learning in which individual neurons become monosemantic: each aligns strongly with the predictive direction of one specific cluster. The network thereby discovers both the latent cluster structure and the local predictive rules without ever receiving cluster labels. Experiments on synthetic Gaussian-mixture data with up to 50 clusters demonstrate that this specialization produces sample complexity close to that of oracle predictors given the true clusters, while global kernel methods and the Recursive Feature Machine degrade as the number of clusters grows. Gated activations such as ReGLU and SwiGLU further improve efficiency over ReLU and GELU.

Theory supports these observations. For two-layer ReLU networks trained from small random initialization on well-separated clusters, population gradient flow drives almost every neuron to align with a single cluster-specific direction. A matching sample-complexity result shows that an appropriately constrained two-layer MLP achieves polynomial risk, whereas kernel ridge regression and RFM with the ground-truth average-gradient-outer-product metric both leave a positive constant error for any polynomial sample size.

These findings matter because they identify a concrete mechanism—neuron-level specialization—that preserves the association between each predictive direction and the cluster in which it applies. The resulting advantage appears precisely when global low-dimensional methods lose their leverage. The work is limited to synthetic, well-separated clusters and focuses its proofs on two-layer networks under specific initialization regimes; whether the same specialization scales reliably to deeper practical models and naturalistic data remains open.

1. Revisiting Feature Learning

Section Summary: Neural networks outperform classical methods largely because they can adapt their internal representations during training to match the structure of a prediction task, rather than relying on a fixed representation chosen in advance. Earlier work has focused on how this feature learning helps uncover a single global low-dimensional subspace that captures the key predictive directions in the input. This section instead describes a distinct mechanism in multilayer perceptrons, in which neurons spontaneously specialize to handle multiple distinct clusters of data, each with its own local predictive directions, allowing effective learning even when no useful global low-dimensional structure exists.

The empirical success of neural networks has motivated substantial work aimed at understanding what enables neural networks to outperform classical learning methods ([1, 2, 3]). One prominent explanation is their ability to perform feature learning: rather than optimizing a predictor over a representation fixed before training, neural networks can use supervision to adapt their internal representations to the structure of the prediction task ([4]).

The distinction between fixed representations and feature learning can be formalized through the neural tangent kernel (NTK). Under a particular initialization scaling and in the infinite-width limit, neural-network training converges to kernel gradient descent with an NTK that remains fixed throughout training ([5, 6]). In this regime, the neural network is effectively equivalent to a classical kernel method operating over a fixed representation. Away from this regime, however, the NTK can evolve substantially during training, and the resulting dynamics can no longer be described by a kernel fixed at initialization ([7, 8, 9, 10]).

This raises a fundamental question:

A major advance in understanding neural networks was the recognition that neural networks can exploit low-dimensional structure in the target, where prediction depends on only a small number of directions in the input space. In such settings, neural networks can adapt their representations to identify these predictive directions, whereas fixed-feature like kernel methods cannot ([11, 12, 13, 14, 2, 15, 16, 17, 18, 19, 20, 21, 22]).

In this setting, the target takes the form

$ f(x)

g!\left(U^{\top}x\right), \qquad U\in\mathbb{R}^{d\times r}, \qquad r\ll d. $

Although $x\in\mathbb{R}^d$, the response depends only on the $r$ -dimensional predictive subspace spanned by the columns of $U$.

A feature-learning method can use the responses of the model to learn the representations corresponding to the predictive subspace, effectively reducing the problem from the ambient dimension $d$ to the intrinsic dimension $r$ . By contrast, a representation fixed independently of the labels would still be dependent on $d$ dimensions.

Taken together, this line of work views feature learning primarily as a mechanism for discovering low-dimensional predictive structure: supervision reshapes the representation so that learning is governed by the intrinsic dimension of the task rather than the ambient dimension of the input.

In our work, we identify a qualitatively different mechanism of feature learning in multilayer perceptrons (MLPs) that is not fundamentally tied to recovering a single global low-dimensional predictive subspace.

In this setting, the data are drawn from multiple clusters, and each cluster can have its own predictive directions and its own link function.

Thus, while the prediction problem could be low-dimensional locally within each cluster, there may be no low-dimensional structure globally.

We show that MLPs can simultaneously discover the cluster structure and learn the predictive functions relevant to each cluster. This uncovers a new form of feature learning that goes beyond the standard paradigm of recovering a single global low-dimensional predictive subspace.

Our main contributions are as follows.

We show that in trained MLPs, a substantial fraction of individual neurons become monosemantic, specializing by aligning predominantly with a single cluster-specific predictive direction. This specialization allows the MLP to learn both the relevant local low-dimensional features and an implicit clustering that determines where each feature is useful. The resulting behavior resembles mixture-of-experts routing, but emerges within a standard MLP without an explicit routing module or expert decomposition.

While this specialization happens for standard ReLU and GeLU ([23]), modern architectures with multiplicative gating, such as ReGLU and SwiGLU ([24]), can substantially improve sample efficiency in this setting.

Furthermore, we show that this specialization leads to increasingly favorable sample-complexity scaling as the number of clusters grows, allowing MLPs to maintain strong performance even when the cluster-specific predictive directions collectively span the ambient space and no useful global low-dimensional structure exists.

1.1 Discussion

The form of feature learning identified in this work differs from the classical feature learning, where the main goal is to recover a single global low-dimensional predictive subspace. In such settings, methods designed for learning low-dimensional features can perform as well or better than MLPs. A prominent example is the Recursive Feature Machine (RFM), a supervised kernel-based method that uses input-response pairs to learn a global low-dimensional predictive subspace ([25, 26]). RFM is therefore a particularly informative comparison: like an MLP, it learns features from the data rather than operating with a fixed representation.

However, as the number of clusters increases, the span of cluster-specific predictive directions grows until it equals the ambient space, so they no longer lie in a single low-dimensional global subspace. This demonstrates a key advantage of MLPs – they can use neuron specialization to preserve the association between each predictive direction and the cluster in which it is relevant, allowing MLPs to remain effective as the number of clusters grows, while RFM feature learning ability diminishes with the number of clusters.

To study this phenomenon systematically, we consider a mixture of single-index models. Suppose the data are drawn from $K$ clusters. For an input $x\in\mathbb{R}^d$ belonging to cluster $c\in[K]$, the target is

$ f(x)

g_c!\left(v_c^\top x\right), \qquad v_c\in\mathbb{R}^d. $

Thus, prediction is one-dimensional within each cluster, but both the relevant direction $v_c$ and the link function $g_c$ may vary across clusters.

The classical single-index model corresponds to $K=1$ . When $K>1$, however, the predictive directions of different clusters may collectively span the entire ambient space:

$ \operatorname{rank} \left(\begin{bmatrix} v_1 & \cdots & v_K \end{bmatrix} \right)

d. $

Hence, prediction can be low-dimensional within every cluster even though there is no useful global low-dimensional predictive subspace to recover.

1.2 Prior Work

Feature learning of global low-dimensional structure.

A large body of work studies feature learning for multi-index targets. These works have shown that neural-network models can adapt to an unknown low-dimensional subspace of the inputs, succeeding where fixed-kernel methods are sample-inefficient ([11, 12, 13, 15]). Recently, the literature has characterized how gradient-based training succeeds in recovering this subspace for isotropic data ([17, 18, 20, 19, 21, 27, 22]), and in the related setting of classifying Gaussian-mixture data with a constant number of clusters ([28, 29]).

However, in our work, we prove that neuron specialization is a different form of feature learning, allowing MLPs to efficiently tackle settings where a single global low-dimensional representation is insufficient.

Neuron specialization.

Emergence of specialized neurons has been observed in several settings. For instance, training neural networks to do modular arithmetic leads to specialized neurons each representing different Fourier components ([30, 31, 32, 33]). When learning XOR-type targets, it has also been shown that neurons specialize to four clusters rather than being distributed evenly across a predictive low-dimensional subspace ([34, 35]). And in teacher-student settings the student's neurons have been shown to often align to the teacher's neurons ([36, 37, 38]), although this depends on initialization ([39]).

These works mainly study specialization to globally relevant features or to cluster directions that themselves determine the target. In our work, neurons specialize to locally predictive features. Thus, an MLP must not only learn the predictive directions, but also preserve their association with the clusters in which they are relevant. Moreover, these works do not isolate a sample-complexity advantage of specialization over an adaptive feature-learning method based on a single global representation, such as the Recursive Feature Machine ([25]). We establish such an advantage: when cluster-specific predictive directions collectively span a high-dimensional space, neuron specialization allows an MLP to preserve these local feature-cluster associations and achieve better sample complexity.

Monosemanticity.

Interpretability works on language models have found that some individual neurons are monosemantic, responding primarily to a single interpretable concept or pattern, while others are polysemantic, responding to multiple distinct concepts or patterns ([40, 41]). More recent work has argued that monosemantic features need not align with individual neurons, and has used sparse dictionary learning to recover more monosemantic feature directions from neural representations ([42, 43]). Our work shows that, at least in simple MLPs, monosemantic features can emerge directly at the level of individual neurons and that this neuron-level specialization can improve sample efficiency compared with methods that are limited to global feature discovery.

Mixture of Experts models.

A complementary line of work studies specialization and latent-cluster recovery in mixture-of-experts (MoE) architectures. [44] analyze a clustered classification problem and show that an MoE router can learn cluster-center features that partition the problem into simpler subproblems handled by different experts. [45] likewise show, theoretically and empirically, that a learned router can route inputs according to latent clusters. Particularly close to our statistical setting, [46] study nonlinear regression with an underlying cluster structure of single-index models and show that an MoE trained by SGD can detect the latent organization and divide the task into cluster-specific subproblems; under their assumptions, a vanilla neural network does not detect this organization within their polynomial complexity regime.

These works build specialization into the architecture through an explicit router and separate experts. In contrast, we study a standard MLP with no explicit routing or expert decomposition, and show that both the implicit clustering and the corresponding local predictive-feature specialization can emerge through the specialization of individual neurons.

1.3 Paper Structure

The remainder of the paper develops the view that MLPs learn cluster-dependent predictive structure through neuron specialization, both empirically and theoretically. In Section 2, we introduce the setting used throughout the paper. In Section 3, we present our empirical results, showing that MLPs learn this cluster structure and achieve better sample-complexity over RFM as the number of clusters grows. We also show that gated activations such as ReGLU and SwiGLU can further improve sample complexity compared with standard activations such as ReLU and GeLU. Finally, in Section 4, we theoretically analyze the emergence of neuron specialization in simplified settings.

2. Preliminary

Section Summary: The paper models its data as coming from a mixture of K Gaussian clusters whose means are placed far apart so that the groups are easy to tell apart. Within each cluster the outcome y is produced by feeding a linear projection of the input vector x through its own nonlinear function and then adding a small amount of random noise. Performance is measured by the usual mean-squared-error loss between the predicted and true values of y.

Throughout our synthetic experiments and theoretical analysis, we consider data drawn from a Gaussian mixture model with $K$ clusters. The cluster index $c$ is drawn uniformly from ${1, \ldots, K}$ . For each cluster $c$, let $\mu_c\in\mathbb{R}^d$ denote its mean and $\Sigma_c\in\mathbb{R}^{d\times d}$ its covariance matrix. Conditioned on cluster $c$, the covariates $x\in\mathbb{R}^d$ are distributed as

$ x \mid c \sim \mathcal{N}(\mu_c, \Sigma_c). $

In all settings considered in this paper, we choose the cluster means so that distinct clusters are well separated. The response follows a cluster-specific single-index model with additive Gaussian noise,

$ y = g_c!\left(\langle x, v_c\rangle\right) + \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, \sigma^2), $

where $\varepsilon$ is independent of $x$ and $c$, $v_c\in\mathbb{R}^d$, with $|v_c|_2=1$, is the cluster-specific predictive direction, and $g_c:\mathbb{R}\to\mathbb{R}$ is a nonlinear "link function" that may vary across clusters. Throughout, we use mean squared error (MSE) as the regression loss,

$ \mathcal{L}(f)

\mathbb{E}!\left[(f(x)-y)^2\right], $

where $f:\mathbb{R}^d\to\mathbb{R}$ denotes the learned predictor.

3. Experiments

Section Summary: The experiments demonstrate that multilayer perceptrons can simultaneously discover hidden cluster structure in data and learn distinct predictive rules within each cluster, enabling a richer form of feature learning than simply identifying one shared low-dimensional subspace. In one set of tests, many neurons in the first layer of trained MLPs become highly specialized, each focusing on a single cluster-specific direction rather than mixing signals broadly. In related regression tasks with varying numbers of clusters, the networks achieve accuracy close to idealized methods that are given the cluster labels in advance, with performance improving further when gated activations are used to help isolate cluster-dependent features.

In this section, we empirically show that MLPs can learn both the cluster structure and the predictive function within each cluster. This gives rise to a form of feature learning that goes beyond recovering a single low-dimensional predictive subspace. We demonstrate these phenomena through three complementary experiments, each presented in a separate subsection.

3.1 MLPs Develop Specialized First-layer Neurons

Main finding.

Figure 1 shows that a significant number of first-layer neurons in trained MLPs become monosemantic, specializing to individual cluster-specific predictive directions. We measure specialization among the active neurons that contribute most strongly to the output, as defined in Appendix A.1. The bottom panel illustrates this directly for selected GELU neurons, whose weights concentrate on a single predictive coordinate.

Data setting.

Following the notation introduced in Section 2, in Figure 1 we use $K=10$ clusters in $d=40$. We choose the cluster centers and predictive directions to lie along separate coordinate axes:

$ [\mu_c] _j = \begin{cases} 20, & j=c, \ 0, & \text{otherwise}, \end{cases} \qquad [v_c] _j = \begin{cases} 1, & j=K+c, \ 0, & \text{otherwise}, \end{cases} \qquad j\in{1, \ldots, d}. $

We further set

$ \Sigma_c = I_d/d, \qquad \varepsilon \sim \mathcal{N}(0, 0.005^2). $

Thus, the cluster centers occupy the first $K$ coordinates, while the predictive directions occupy the next $K$ coordinates.

The cluster centers and predictive directions lie along separate coordinate axes. Since each $v_c$ is a standard basis vector, specialization is directly visible as a large weight on the corresponding coordinate in Figure 1. We use a second-Hermite nonlinearity $g_c$ with small Gaussian response noise.

Specialization measures.

For a neuron's first-layer weight $w_j$, we measure specialization using the maximum absolute cosine similarity with the cluster-specific predictive directions:

$ \max_{1\leq c\leq K} \frac{|\langle w_j, v_c\rangle|} {|w_j|_2, |v_c|_2}. $

Figure 1 reports the fraction of active neurons for which this cosine similarity is at least $\frac{1}{\sqrt{2}}$ or $0.90$ . We use $1/\sqrt{2}\approx 0.71$ as a moderate-alignment threshold, corresponding to at least half of the squared weight norm lying along a single predictive direction, and $0.90$ as a stricter measure of strong alignment. The table also reports two additional measures of specialization, since cosine similarity can underestimate specialization when one predictive coordinate is dominant but many small coefficients contribute to the overall weight norm.

Full experimental details, including the definitions and motivation for the additional specialization measures, are given in Appendix A.1.

3.2 MLPs Jointly Learn Cluster Structure and Cluster-specific Predictive Functions

Following the clustered single-index model introduced in Section 2, we consider regression problems in $d=20$ in which both the predictive direction $v_c$ and nonlinear response function $g_c$ may differ across clusters. We vary the number of clusters over $K\in{1, 2, 10, 50}$ . The MLPs and global kernel baselines, Laplace and RFM, are trained only on input–response pairs and are never given the cluster identities. As a reference, we also consider oracle baselines that are given the true cluster identities and fit a separate Laplace or RFM predictor within each cluster. These oracles therefore measure the performance achievable when the cluster structure is known in advance.

Main findings.

Figure 2 highlights two main conclusions,

1. Sample complexity as the number of clusters grows. MLPs outperform other methods as the number of clusters increases. Remarkably, MLPs achieve performance close to that of "oracle" baselines that are given the true cluster identities and fit a separate RFM or Laplace predictor within each cluster. Thus, MLPs recover much of the benefit of knowing the cluster structure without ever observing the cluster identities.
2. Gated activations enhance cluster-dependent feature learning. ReLU already learns useful cluster-dependent structure and substantially outperforms global RFM when many clusters are present. ReGLU improves further in this regime. Its separate gate and value branches allow the network to select cluster-specific predictive features more directly than an ordinary ReLU hidden layer.

Data setting.

Following Section 2, we consider $K\in{1, 2, 10, 50}$ clusters in $d=20$. We divide the input into ten cluster-identifying coordinates and ten predictive coordinates. The cluster centers are chosen as well-separated unit vectors $s_c\in\mathbb{S}^9:={s\in\mathbb{R}^{10}:|s|_2=1}$, selected sequentially from a large random candidate set so as to maximize separation from the previously selected centers. They are embedded as

$ \mu_c= \begin{bmatrix} s_c\ 0 \end{bmatrix}. $

The cluster-specific predictive directions are sampled independently and uniformly from the unit sphere in the predictive subspace, with

$ \widetilde{v}c\sim\operatorname{Unif}(\mathbb{S}^9), \qquad v_c= \begin{bmatrix} 0\ \widetilde{v}c \end{bmatrix}, \qquad \Sigma_c= \begin{pmatrix} \sigma_K^2 I{10} & 0\ 0 & I{10} \end{pmatrix}, \qquad \varepsilon\sim\mathcal N(0, 0.02^2). $

Here, $\sigma_K$ is adjusted with $K$ so that the clusters remain well separated. Full details of the center construction and the choice of $\sigma_K$ are provided in Appendix A.2.

Each cluster is also assigned a nonlinear response function $g_c$, sampled from the normalized second-order Hermite polynomial, $\sin$, and $\tanh$. Thus, both the predictive direction $v_c$ and nonlinear prediction rule $g_c$ may vary across clusters. The models observe only $(x, y)$, and the cluster identities are not used as predictive inputs or to fit cluster-specific predictors, except for the oracle baselines.

Additional activations.

To test whether the observed behavior extends beyond ReLU and ReGLU, we evaluate GELU and SwiGLU on the same mixed-link data model and the same $K\in{1, 2, 10, 50}$ settings. Across five seeds, GELU behaves similarly to ReLU, and SwiGLU exhibits the same advantage as ReGLU as the number of clusters grows. Complete results are provided in Appendix A.2 and Figure 6.

3.3 Trained MLPs Encode Cluster Structure

We use the mixed-link functions clustered model from Section 3.2 and vary the number of clusters over $K \in {2, 5, 10, 50}$ . Figure 3 summarizes the complete procedure. We train a ReLU MLP on the clustered data and extract its first-layer representations. We then apply K-means to these representations to obtain learned clusters. For each learned cluster, we use the corresponding first-layer weights to construct features for an independent local Laplace or RFM predictor. At test time, each input is assigned to its nearest learned centroid and evaluated using the corresponding local predictor.

Across all numbers of clusters, using the clusters extracted from the MLP first-layer representation substantially improves both Laplace and RFM over their global counterparts (Figure 4). In particular, the resulting local RFM consistently outperforms global RFM and approaches the cluster-aware oracle. These results show that the first-layer representation contains information that both separates the clusters and identifies cluster-specific predictive features, which can be extracted and reused by a separate predictor.

4. Theory

Section Summary: In this theory section, the authors prove that two-layer ReLU networks trained on data from well-separated Gaussian mixture clusters cause their neurons to specialize, with each neuron aligning to the predictive direction of one specific cluster under small random initialization. They further show that this specialization allows MLPs to achieve better sample complexity than kernel methods or recursive feature machines as the number of clusters grows large, since those alternatives cannot develop cluster-specific features. The results are established via analysis of population gradient flow on a cubic target function, demonstrating that neurons converge to one-cluster solutions almost surely.

In this section, we prove two main results describing how MLPs learn data with cluster structure. In Theorem 1, we prove that neurons in the MLP specialize during training. In Theorem 2, we prove that, as the number of clusters tends to infinity, MLPs outperform standard kernel methods and Recursive Feature Machines (RFM) in terms of sample complexity. This comes as a consequence of the fact that the latter two methods cannot compute specialized features for each cluster.

4.1 MLP Neurons Specialize When Learning on Gaussian Mixture Model Data

In this subsection, we show that, under small initialization, neurons trained on well-separated Gaussian mixture data specialize to cluster-specific predictive directions.

We consider $K$-cluster Gaussian mixture data in $d=2K$ dimensions. Let $e_1, \ldots, e_K$ denote the standard basis of $\mathbb{R}^K$. We use the first $K$ coordinates to separate the cluster means and the last $K$ coordinates for the cluster-specific predictive directions.[^1]

[^1]: Note that independent orthogonal changes of basis in the two blocks preserve the isotropic Gaussian noise and all conclusions below.

Symmetric Gaussian mixture

Let $R > 0$ be a cluster separation parameter and, for $c\in[K]$, let

$ \mu_c= \begin{bmatrix} Re_c\ 0 \end{bmatrix}, \qquad v_c= \begin{bmatrix} 0\ e_c \end{bmatrix}. $

Set each cluster to be isotropic with $\Sigma_c=I_d$ and use the common link function $g_c=h_3$, where $h_3(t)=(t^3-3t)/\sqrt6$ is the third Hermite polynomial.

In the notation of Section 2, the data distribution is

$ c\sim\operatorname{Unif}([K]), \qquad x\mid c\sim\mathcal N(\mu_c, \Sigma_c), \qquad y=g_c!\left(\langle x, v_c\rangle\right) =h_3!\left(\langle x, v_c\rangle\right). $

Thus the routing coordinates encode cluster identity through the means $Re_c$, where $R$ should be thought of as a large cluster-separation parameter, so that the clusters are well separated, while within cluster $c$ the response depends only on the cluster-specific predictive direction $v_c$. The cubic Hermite target is chosen because there is an explicit expression for the expected product $\mathbb{E}[y, \phi(\omega^\top x)]$ between the target response and the ReLU neuron's activation, allowing us to characterize the directions to which neurons converge. Next, we consider training a neural network to learn this data distribution.

Two-layer ReLU population gradient flow

We train the two-layer ReLU network

$ f_\theta(x)=\sum_{j=1}^m a_j\phi(w_j^\top x), \qquad \phi(t)=t_+, $

by population gradient flow on the squared loss

$ \mathcal L(\theta)

\frac12\mathbb{E}\left[\bigl(f_\theta(x)-y\bigr)^2 \right]. $

Small random hidden weights and zero output layer

The hidden directions are initialized independently and uniformly at random, scaled by a parameter $\varepsilon > 0$, while the output layer is initialized at zero:

$ \omega_j^0 \stackrel{\mathrm{iid}}{\sim} \operatorname{Unif}(\mathbb{S}^{d-1}), \qquad w_j(0)=\varepsilon\omega_j^0, \qquad a_j(0)=0. $

Under this initialization and data distribution, we are able to prove that the neurons in the MLP specialize to the clusters of the Gaussian mixture. This result is consistent with, and provides theoretical support for, our empirical observations of neuron specialization in Section 3.1.

########## {caption="Theorem 1: Randomly initialized neurons specialize"}

Under Section 4.1, Section 4.1 and Section 4.1 there are universal constants $R_0, C<\infty$ such that the following holds for every $R\geq R_0$. For almost every draw of $\omega_1^0, \ldots, \omega_m^0$, there are cluster labels $J_1, \ldots, J_m\in[K]$ and orientations $\tau_1, \ldots, \tau_m\in{\pm1}$ such that, for every $\delta>0$, there exist a finite time $T_\delta$ and $\varepsilon_0>0$ for which

$ 0<\varepsilon\leq\varepsilon_0 \quad\Longrightarrow\quad \left| \frac{w_j(T_\delta)}{|w_j(T_\delta)|2} -\tau_j v{J_j} \right|_2 \leq \delta+\frac{C}{R} \qquad\text{for every }j\in[m]. $

Thus every neuron specializes, up to orientation and a vanishing $O(R^{-1})$ routing component, to the predictive direction of one cluster.

The selected labels $J_1, \ldots, J_m$ are independent and uniformly distributed on $[K]$. Consequently,

$ \mathbb{P}\left(\text{every cluster is covered by a specialized neuron} \right) \geq 1-Ke^{-m/K}. $

In particular, if

$ m\geq K\log\left(\frac{K}{\eta}\right), $

then all $K$ clusters are covered with probability at least $1-\eta$.

Proof sketch.

We employ a proof strategy of studying feature learning through effectively independent neuron dynamics, which has been used in prior work on early-time and small-initialization regimes ([17, 47, 35]). At initialization, the network output is zero and the hidden weights are stationary. The initial derivative of each output weight is proportional to the population correlation

$ \Phi(\omega) =\mathbb{E}\left[y, \phi(\omega^\top x)\right] $

of its own randomly initialized hidden direction. Except on a measure-zero set, this correlation is nonzero, so the output weight immediately acquires the appropriate sign. ReLU homogeneity then turns the neuron's subsequent directional dynamics into a positive time reparameterization of gradient ascent on the corresponding signed correlation objective.

For the cubic target, the Hermite calculation makes this objective a sum of clusterwise cubic terms. Its positive local maxima are one-cluster solutions: a neuron aligns with $\pm v_c$ and uses only an $O(R^{-1})$ component in the associated routing direction to place its ReLU threshold within cluster $c$. Indeed, we prove that any positive mixed-cluster critical point has an unstable direction. Analytic-gradient-flow convergence and strict-saddle avoidance therefore imply that a random neuron specializes almost surely.

Permutation symmetry makes the selected cluster uniform on $[K]$, and independence of the initial hidden directions makes the selected labels independent. The coverage estimate is then the usual coupon-collector union bound. Finally, on every fixed early-time interval the full network output is $O(m\varepsilon^2)$, so the coupled network dynamics are a vanishing perturbation of the isolated-neuron dynamics. A Grönwall argument transfers the specialization result to population gradient flow. The complete proof is given in Appendix B.1.

4.2 A Sample-complexity Gap Between MLPs, Kernel Methods, and RFM

In this subsection, we establish a sample-complexity separation between two-layer MLPs and both standard kernel methods and Recursive Feature Machines (RFM). We use a data distribution similar to that in the previous subsection, but with full separation between clusters for ease of analysis.

Gaussian mixture with noiseless routing

We set the means $\mu_c=[e_c;0]$ and predictive directions $v_c=[0;e_c]$ in $d=2K$ dimensions, but make the first $K$ coordinates noiseless by setting

$ \Sigma_c= \begin{bmatrix} 0_{K\times K} & 0\ 0 & I_K \end{bmatrix}. $

We use the common clipped-ramp link $g_c=g$, where

$ g(t)=t_+-(t-1)_+. $

Thus, in the notation of Section 2,

$ c\sim\operatorname{Unif}([K]), \qquad x\mid c\sim\mathcal N(\mu_c, \Sigma_c), \qquad y=g_c!\left(\langle x, v_c\rangle\right) =g!\left(\langle x, v_c\rangle\right). $

The routing block now reveals the cluster exactly, while the response within cluster $c$ depends only on its cluster-specific predictive direction.

Instead of studying a network trained through gradient-based updates, we study a two-layer MLP fit by empirical risk minimization subject to a Frobenius-norm constraint on its weights.

Frobenius-constrained MLP ERM

We fit the two-layer ReLU network

$ f_\theta(x)=\sum_{j=1}^m a_j\phi(w_j^\top x), \qquad \phi(t)=t_+, $

using empirical risk minimization (ERM). For a budget $B>0$, consider the set of functions representable by Frobenius-constrained networks,

$ \mathcal F_B := \left{ f_\theta: |\theta|{\mathrm F}^2 =\frac12\sum{j=1}^m\left(a_j^2+|w_j|_2^2\right) \leq B \right}. $

Given $n$ samples, we consider the clipped-ERM estimator

$ \widehat{f}{\mathrm{MLP}} \in \arg\min{f\in\operatorname{clip}\circ\mathcal F_B} \frac1n\sum_{i=1}^n\bigl(f(x_i)-y_i\bigr)^2, \qquad \operatorname{clip}(t)=\min{1, \max{0, t}}. $

We compare this MLP ERM with standard rotationally invariant kernel ridge regression and with RFM, a supervised kernel method ([25, 26]). We give RFM the ground-truth population average gradient outer product (AGOP), thereby removing metric-estimation error and isolating the limitation of using one global feature metric.[^2]

[^2]: Note that an analysis of a limitation for RFM using an empirically estimated AGOP would require bounding empirical fluctuations, so it is not implied by this theorem.

Kernel methods and RFM with ground-truth AGOP

Let $\mathcal K$ be a rotationally invariant kernel. Given $n$ samples, $\widehat{f}_{\mathrm{Kernel}}$ is kernel ridge regression with kernel $\mathcal K(x, x')$. We give RFM the ground-truth population AGOP

$ M :=\mathbb{E}\left[\nabla_x f^\star(x)\nabla_x f^\star(x)^\top \right]. $

For $\rho\geq0$, define the regularized AGOP and corresponding kernel

$ M_\rho:=M+\rho I_d, \qquad \mathcal K_{M_\rho}(x, x') :=\mathcal K\bigl(\sqrt{M_\rho}x, \sqrt{M_\rho}x'\bigr). $

$\widehat{f}{\mathrm{RFM}}$ is kernel ridge regression with kernel $\mathcal K{M_\rho}$. Both estimators may use any kernel ridge parameter $\lambda_n\geq0$.

In the theorem below, we show that a two-layer MLP learns the cluster-structured data with polynomial sample complexity. In contrast, standard kernel methods and RFM provably fail under every polynomial sample-size scaling because they cannot compute a distinct specialized direction for each cluster. In particular, RFM is constrained to learn a single global feature geometry, which cannot capture the cluster-dependent specialization required by the target.

########## {caption="Theorem 2: The MLP succeeds while kernel methods and RFM fail"}

Under Section 4.2 and Section 4.2, the MLP estimator with $B=3K$ satisfies

$ \mathbb{E}\left[|\widehat{f}{\mathrm{MLP}}-f^\star|{L^2(P_{x})}^2 \right] \leq C\frac{K^{3/2}}{\sqrt n} $

for a universal constant $C$.

Under Section 4.2 and Section 4.2, for every fixed $A<\infty$, every sequence $n_K=O(K^A)$, and every sequence $\rho_K\geq0$,

$ \liminf_{K\to\infty} \mathbb{E}\left[|\widehat{f}{\mathrm{Kernel}}-f^\star|{L^2(P_x)}^2 \right]

0, \qquad \liminf_{K\to\infty} \mathbb{E}\left[|\widehat{f}{\mathrm{RFM}}-f^\star|{L^2(P_x)}^2 \right] 0. $

Proof sketch.

We construct an MLP solution that assigns two ReLU neurons to each cluster. On every other cluster, the pair's outputs cancel exactly, while on the selected cluster their difference equals the clipped ramp. This construction has zero approximation error and satisfies the Frobenius budget $B=3K$. Since the MLP class has bounded Frobenius budget, it is contained in a corresponding path-norm ball. Norm-based capacity and Rademacher-complexity bounds for neural networks ([48, 11, 49]), together with standard Rademacher-complexity risk bounds ([50]), then yield the MLP risk estimate.

The kernel lower bound exploits the same rotational-invariance obstruction identified in prior work on the limitations of kernel methods in high dimensions ([3]). In our setting, the ground-truth AGOP is a scaled projector onto the predictive subspace. Hence both the standard kernel metric and the regularized RFM metric $M_\rho$ act by scalars on the routing and predictive subspaces, so their kernels remain rotationally invariant within the predictive coordinates. The representer theorem restricts the predictor to the span of $n$ kernel sections. At Hermite order $r$, the clusterwise target contains a component in an irreducible harmonic subspace of dimension $\Theta_r(K^r)$. The span of the projected kernel sections is an invariant random subspace of dimension at most $n$, so for $n=O(K^A)$ every order $r>A$ leaves asymptotically all of its target energy unrecovered. This argument is uniform over $\rho$ and therefore covers standard kernels and ridge-regularized RFM simultaneously. The complete proof is given in Appendix B.2.

Acknowledgements

Section Summary: The authors thank several U.S. government agencies, including the National Science Foundation, the Office of Naval Research, and the Defense Advanced Research Projects Agency, for providing financial support through specific grants and contracts. They also used advanced computing resources at a national supercomputing center, which were accessed through a broader NSF-supported program. In addition, artificial intelligence tools helped with portions of the experiments and theoretical analysis.

We gratefully acknowledge support from the National Science Foundation (NSF) under grants CCF-2112665 and MFAI 2502258, the Office of Naval Research (ONR N000142412631), and the Defense Advanced Research Projects Agency (DARPA) under Contract No. HR001125CE020. This work used the Delta system at the National Center for Supercomputing Applications through allocation TG-CIS220009 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296.

AI tools were used to assist with aspects of the experiments and theoretical analysis.

Appendix

Section Summary: The appendix provides extra details on experiments showing how simple neural networks can develop specialized neurons in their first layer when trained on synthetic data. It describes generating clustered input data in high dimensions, where specific coordinates encode both cluster identity and predictive signals for a quadratic target function, along with nuisance noise dimensions. The text then covers the training setups for various one-hidden-layer networks, the criteria used to focus analysis on the most influential active neurons, and quantitative measures of how strongly those neurons' weights align with the predictive directions rather than spreading across other coordinates.

A. Additional experimental details

A.1 Additional Details for MLPs Develop Specialized First-layer Neurons Experiments

Data generation.

We use $K=10$ clusters in $d=40$ dimensions, with the cluster index $c$ sampled uniformly from ${1, \ldots, K}$ . We choose the cluster centers and predictive directions to lie along separate coordinate axes:

$ [\mu_c] _j = \begin{cases} 20, & j=c, \ 0, & \text{otherwise}, \end{cases} \qquad [v_c] _j = \begin{cases} 1, & j=K+c, \ 0, & \text{otherwise}, \end{cases} \qquad j\in{1, \ldots, d}. $

For each cluster $c$, we sample

$ x=\mu_c+z, \qquad z\sim\mathcal N(0, I_d/d). $

Thus, coordinates $1{:}10$ encode cluster position, coordinates $11{:}20$ contain the cluster-specific predictive directions ${v_c}_{c=1}^K$, and coordinates $21{:}40$ are nuisance coordinates.

We use the same normalized second-Hermite link in every cluster and add Gaussian observation noise:

$ y= \frac{\operatorname{He}_2!\left(\sqrt d, \langle v_c, x\rangle\right)}{\sqrt{2!}} +\varepsilon, \qquad \operatorname{He}_2(t)=t^2-1, \qquad \varepsilon\sim\mathcal N(0, 0.005^2). $

Training, validation, and test sets are generated independently and contain $100{,}000$, $4{,}096$, and $2{,}048$ examples, respectively.

Models and optimization.

We train one-hidden-layer ReLU, GELU, ReGLU, and SwiGLU networks of width $2{,}048$, with trainable hidden and output biases and balanced feature-learning initialization. For example, the hidden representation of a ReGLU network is

$ \operatorname{ReLU}(W_gx+b_g)\odot(W_vx+b_v). $

All models are trained with Adam under the corresponding width-aware parameterization, using a batch size of $2{,}048$, cosine learning-rate decay, and zero weight decay.

For ReLU, we search initial learning rates

$ {0.004, 0.008, 0.012, 0.016, 0.024, 0.032}, $

with each schedule decaying to one tenth of its initial value. Model selection uses validation MSE only. The selected ReLU schedule is $0.004\to0.0004$, with the validation-selected checkpoint at step $37{,}000$. The schedules used for the remaining architectures are

$ \begin{aligned} \text{GELU:};&0.016\to0.0016, \ \text{ReGLU:};&0.0016\to0.00016, \ \text{SwiGLU:};&0.0256\to0.00256. \end{aligned} $

GELU, ReGLU, and SwiGLU are trained for at most $20{,}000$ steps, while the ReLU candidates are trained for at most $40{,}000$ steps. For every model, we select the checkpoint with the smallest validation MSE before evaluating on the test set.

Active-neuron criterion.

We restrict the specialization analysis to neurons that contribute materially to the network output. For an ordinary ReLU or GELU neuron, we define its importance as

$ I_j=|a_j|, |w_j|_2, $

where $a_j$ is its output weight and $w_j$ is its incoming weight. For a gated ReGLU or SwiGLU neuron, we use

$ I_j=|a_j|, |w_{g, j}|2, |w{v, j}|_2, $

where $w_{g, j}$ and $w_{v, j}$ denote its gate and value weights.

We sort neurons by $I_j$ and retain the smallest leading set accounting for at least $99.9%$ of the total importance. We refer to these retained neurons as active. For gated architectures, the same active-neuron set is used when analyzing the gate and value weights separately.

The $99.9%$ threshold is a descriptive sparsification rule rather than a statistical significance threshold. Its purpose is to prevent the specialization statistics from being dominated by neurons with negligible influence on the network output.

Specialization measures.

For an incoming weight vector $w_j$, we define its maximum absolute cosine similarity with the cluster-specific predictive directions as

$ A_j = \max_{1\leq c\leq K} \frac{|\langle w_j, v_c\rangle|}{|w_j|_2}.\tag{1} $

Because each $v_c$ is a unit vector, $A_j$ is the largest absolute cosine similarity between $w_j$ and any cluster-specific predictive direction. A large value therefore indicates that a substantial fraction of the weight norm is concentrated along a single predictive direction.

Cosine alignment can be conservative when one predictive coordinate is much larger than every individual competing coordinate but the weight vector also contains many small coefficients. To capture this behavior, define

$ \widehat{c}j \in \arg\max{1\leq c\leq K} |\langle w_j, v_c\rangle| $

and the predictive-coordinate dominance score

$ D_j = \frac{|[w_j] _{K+\widehat{c}j}|} {\displaystyle \max{\substack{1\leq \ell\leq d\ \ell\neq K+\widehat{c}_j}} |[w_j] _\ell|}.\tag{2} $

Since the predictive direction $v_c$ is supported on coordinate $K+c$, $D_j>1$ means that the largest-magnitude coordinate of $w_j$ is a predictive coordinate, while $D_j>2$ means that this predictive coordinate is more than twice as large as every other coordinate.

For ReGLU and SwiGLU, both $A_j$ and $D_j$ are computed separately for the gate weights $w_{g, j}$ and value weights $w_{v, j}$. Absolute inner products are used throughout because alignment with $v_c$ and $-v_c$ represents the same predictive direction.

Reported specialization statistics.

The summary table in Figure 1 reports the fractions of active neurons satisfying

$ A_j\geq0.71, \qquad A_j\geq0.90, \qquad D_j>1, \qquad D_j>2. $

These correspond, respectively, to the two cosine-threshold columns, "Max is target-aligned, " and "Target-aligned max

The individual neurons displayed in the bottom panel of Figure 1 are the highest-importance active GELU neurons satisfying $A_j\geq0.90$ . Their coordinate-wise absolute incoming weights are shown so that specialization is visible independently of sign.

Figure 5 in Appendix illustrates why the dominance score complements cosine alignment. In these examples, a neuron has only moderate cosine alignment because its norm contains many small coefficients, yet one cluster-specific predictive coordinate remains clearly dominant.

A.2 Additional Details for MLPs Jointly Learn Cluster Structure and Cluster-specific Predictive Functions

Cluster and predictive geometry.

We choose $d=20$ and divide the input into two ten-dimensional subspaces:

$ \mathbb{R}^{20}

\mathcal S_{\mathrm{cluster}} \oplus \mathcal S_{\mathrm{pred}}. $

The first contains information identifying the cluster, while the second contains the cluster-specific predictive signal. We consider $K\in{1, 2, 10, 50}$ equally likely clusters.

For each $K$, we construct well-separated cluster-center directions in $\mathbb{R}^{10}$ . We first draw

$ L_K=\max{20{,}000, 500K} $

independent Gaussian vectors and normalize them to unit norm. We select the first direction uniformly from this collection and choose each subsequent direction to minimize its largest inner product with the directions already selected. This farthest-point procedure produces unit vectors

$ s_1, \ldots, s_K\in\mathbb{S}^9, \qquad \mathbb{S}^9

{s\in\mathbb{R}^{10}:|s|_2=1}, $

with large pairwise separation. The candidate collection and selected directions are generated separately for each value of $K$ .

For $K>1$, define the largest pairwise similarity

$ \rho_K = \max_{c\neq c'} s_c^\top s_{c'}. $

We set $\rho_1=0$ . The center of cluster $c$ in the full ambient space is

$ \mu_c

\begin{bmatrix} s_c\ 0 \end{bmatrix} \in\mathbb{R}^{20}, \qquad K>1. $

For $K=1$, we set $\mu_1=0$, since no cluster identification is required.

The standard deviation of the noise in the cluster-identifying coordinates is

$ \sigma_K

\frac{1-\rho_K}{2\gamma}, \qquad \gamma=4.\tag{3} $

Thus, the noise level is adjusted using the least-separated pair of cluster centers, maintaining clear cluster separation as $K$ changes.

Each cluster also receives an independently and uniformly sampled unit direction in the predictive subspace,

$ \widetilde{v}_c \sim \operatorname{Unif}(\mathbb{S}^9). $

We embed this direction into the full ambient space by zero-padding the cluster-identifying coordinates:

$ v_c

\begin{bmatrix} 0\ \widetilde{v}_c \end{bmatrix} \in\mathbb{R}^{20}. $

Thus, $v_c$ lies entirely in $\mathcal S_{pred}$, whereas $\mu_c$ lies entirely in $\mathcal S_{cluster}$ . Consequently, the coordinates identifying the cluster are disjoint from those determining its response.

Input distribution.

The cluster identity $C$ is sampled uniformly:

$ C\sim\operatorname{Unif}{1, \ldots, K}. $

Conditional on $C=c$, the input is

$ x

\mu_c + \begin{bmatrix} \sigma_K z_{\mathrm{cluster}}\ z_{\mathrm{pred}} \end{bmatrix}, \qquad z_{\mathrm{cluster}}, z_{\mathrm{pred}} \stackrel{\mathrm{iid}}{\sim} \mathcal N(0, I_{10}).\tag{4} $

Equivalently, in the notation of Section 2,

$ x\mid C=c \sim \mathcal N(\mu_c, \Sigma_c), \qquad \Sigma_c

\begin{pmatrix} \sigma_K^2 I_{10} & 0\ 0 & I_{10} \end{pmatrix}. $

In the finite datasets, the number of examples assigned to each cluster is equal whenever the sample size is divisible by $K$ and otherwise differs by at most one.

Mixed local prediction rules.

Each cluster receives a nonlinear link function independently and uniformly, with replacement, from

$ \mathcal G

\left{ t\mapsto\frac{t^2-1}{\sqrt{2}}, \quad t\mapsto\sin(t), \quad t\mapsto\tanh(t) \right}.\tag{5} $

Sampling with replacement means that two clusters may receive the same link. For a sample from cluster $c$, the response is

$ y

g_c!\left(\langle x, v_c\rangle\right) +\varepsilon, \qquad \varepsilon\sim\mathcal N(0, 0.02^2).\tag{6} $

Since $v_c$ is supported only on the predictive coordinates,

$ \langle x, v_c\rangle = \widetilde{v}c^\top z{\mathrm{pred}}. $

Thus, both the relevant predictive direction $v_c$ and the nonlinear response function $g_c$ may change across clusters.

Training and test sets.

For each $K$ and seed, we generate a training pool containing $200{,}000$ examples. We use the nested training sizes

$ n \in {20{,}000, ;50{,}000, ;100{,}000, ;200{,}000}, $

so that the dataset at a smaller value of $n$ is contained in every larger training set. Each run uses a separate test set of $4{,}096$ examples. All methods evaluated for the same $K$, $n$, and seed use exactly the same training and test examples.

We use five seeds indexed by $r\in{0, \ldots, 4}$ . For a fixed $K$, the cluster centers and predictive directions are shared across the five runs. The link assignments, training and test samples, label noise, and model initialization vary across runs. Specifically, the model and mixed-link seed is $1000+r$, the training-pool seed is $12345+r$, and the test-set seed is $54321+r$ . We report the mean over these five runs, with shaded regions showing one empirical standard deviation.

MLP architectures and optimization.

We evaluate one-hidden-layer ReLU, GELU, ReGLU, and SwiGLU networks, all with width $4{,}096$ . For ReLU and GELU, the hidden representation has the form

$ h(x)=\phi(Wx+b). $

For the gated architectures, it has the form

$ h(x)

\phi(W_gx+b_g) \odot (W_vx+b_v), $

where $\phi$ is ReLU for ReGLU and SiLU for SwiGLU. All hidden weights, readout weights, and biases are trainable.

We optimize the networks using Adam with batch size $4{,}096$, cosine learning-rate decay, and zero weight decay. The activation-specific schedules are

$ \begin{array}{c|c} \text{Activation} & \text{Learning-rate schedule}\ \hline \text{ReLU} & 10^{-2}\to10^{-4}\ \text{GELU} & 3\times 10^{-2}\to3\times10^{-4}\ \text{ReGLU} & 3\times 10^{-3}\to3\times10^{-5}\ \text{SwiGLU} & 10^{-2}\to10^{-4}. \end{array} $

Training continues for at least $20{,}000$ updates and for at most $200{,}000$ updates. After the minimum number of updates, optimization stops when either the training MSE reaches

$ 4\times 10^{-4}, $

which equals the label-noise variance, or the training loss plateaus. The plateau rule stops after 60 consecutive evaluations, performed every 500 updates, without a relative training-MSE improvement of at least 5 x 10^-4 . We restore the checkpoint attaining the lowest training MSE. Stopping and checkpoint selection use training MSE only and never use test performance.

Global Laplace and RFM.

RFM uses the metric-dependent Laplace kernel

$ k_t(x, x')

\exp!\left(-\frac{ |M_t^{1/2}(x-x')|_2 }{h_t} \right),\tag{7} $

where $M_t$ is the metric at RFM iteration $t$ and $h_t$ is the kernel bandwidth. The initial metric is

$ M_0=I_{20}, $

so iteration zero is the isotropic Laplace-kernel baseline.

For numerical stability across different values of $K$, the bandwidth is computed from within-cluster distances in the transformed space. We estimate a median pairwise distance separately within each cluster and take the median of these cluster-level values. The bandwidth is recomputed after every metric update.

This bandwidth calibration uses the cluster identities, but only to select one scalar bandwidth shared by the entire global kernel. The global Laplace and RFM predictors are still fit jointly to all training examples, do not use the cluster identities in their regression objective, and do not use them when making test predictions. Thus, they do not fit separate cluster-specific predictors.

After each kernel fit, RFM estimates the average gradient outer product

$ \widehat{G}_t

\frac{1}{m} \sum_{i=1}^{m} \nabla\widehat{f}_t(x_i) \nabla\widehat{f}_t(x_i)^\top $

using at most $m=20{,}000$ training examples. A normalized version of $\widehat{G}_t$ becomes the metric for the next iteration. We perform three metric updates, producing iterations $0, 1, 2, 3$ .

Exact kernel solves are used for fewer than $11{,}000$ training examples and include a ridge parameter of $10^{-6}$. Larger problems are solved using EigenPro ([51, 52, 53, 54]) without an explicit ridge penalty. EigenPro optimization stops upon reaching training MSE 4 x 10^-4 or its optimization limit. Its effective regularization therefore arises primarily from early stopping.

The displayed global RFM curve reports the smallest test MSE among iterations $0{:}3$ . It is therefore an optimistic diagnostic envelope rather than a deployable iteration-selection procedure.

Per-cluster oracle methods.

The per-cluster Laplace and RFM oracle methods are given the true cluster identity during both training and testing. They partition the training data by cluster, fit one independent predictor within each partition, and evaluate each test example using the predictor associated with its true cluster. The reported MSE is pooled across all test examples.

Iteration zero of each local RFM is the per-cluster Laplace oracle. Each local RFM subsequently performs three AGOP metric updates. For $K=1$, the global and per-cluster methods coincide.

These oracle methods are not fair predictive baselines because they receive the true cluster identities. Their purpose is to measure how easily the local prediction problems can be solved once cluster identification is provided. The ability of an MLP to approach these oracle methods without receiving the cluster identities provides evidence that it learns cluster identification and the corresponding local prediction rules jointly.

Bayes reference.

The dotted reference line reports the test MSE of the Bayes predictor under the known data-generating distribution. It computes the posterior probability of each cluster given the input and averages the corresponding noiseless local predictions. Because the clusters have small but nonzero overlap, its realized test MSE can be slightly larger than the label-noise variance $0.02^2=4\times10^{-4}$ .

Additional activations.

Figure 6 compares all four MLP architectures on the same mixed-link datasets. The training and test examples, cluster geometry, link assignments, and random seeds are shared across activations for every $K$ and $n$ .

The advantage of multiplicative gating is most pronounced when many cluster-dependent prediction rules must be learned. The improvement is not uniform in the smallest-data regime, but with sufficient samples the gated architectures, particularly SwiGLU, obtain the lowest MLP errors for $K=10$ and $K=50$ .

A.3 Additional Details for Trained MLPs Encode Cluster Structures

This appendix describes how we extract cluster gates from a trained ReLU MLP and reuse them to construct independent local Laplace and RFM predictors. The experiment uses the same mixed-link clustered distribution as Section 3.2.

Mixed-link clustered data.

We decompose the ambient space into ten cluster-position coordinates and ten predictive coordinates,

$ \mathbb{R}^{20}

\mathcal S_{\mathrm{gate}} \oplus \mathcal S_{\mathrm{pred}}. $

We consider $K\in{2, 5, 10, 50}$ equally represented clusters. Each cluster $c$ has a unit cluster-position direction $s_c\in\mathbb{S}^9$ and an independently generated unit predictive direction $v_c\in\mathbb{S}^9$ .

The directions $s_c$ are selected by farthest-point sampling from a large set of random unit vectors. Let

$ \rho_K = \max_{c\neq c'} s_c^\top s_{c'} $

be the largest pairwise inner product. We set the cluster-center radius to $R=1$ and choose the standard deviation in the cluster-position coordinates as

$ \sigma_K

\frac{R(1-\rho_K)}{2\gamma}, \qquad \gamma=4. $

Thus, the noise level adapts to the closest pair of cluster centers, keeping the clusters well separated as $K$ changes.

For an example from cluster $c$, we sample

$ z_{\mathrm{gate}}, z_{\mathrm{pred}} \stackrel{\mathrm{iid}}{\sim}\mathcal N(0, I_{10}) $

and construct

$ x

\left(R s_c+\sigma_K z_{\mathrm{gate}}, ;z_{\mathrm{pred}} \right).\tag{8} $

The first ten coordinates therefore identify the cluster, while the last ten coordinates contain the variables used for prediction.

Each cluster independently receives a nonlinear response function $g_c$, drawn uniformly from

$ \mathcal G

\left{ t\mapsto\frac{t^2-1}{\sqrt 2}, \quad t\mapsto\sin(t), \quad t\mapsto\tanh(t) \right}.\tag{9} $

The response is

$ y

g_c!\left(v_c^\top z_{\mathrm{pred}}\right) +\varepsilon, \qquad \varepsilon\sim\mathcal N(0, 0.02^2).\tag{10} $

Consequently, both the predictive direction $v_c$ and the nonlinear response function $g_c$ may differ across clusters. The link functions are sampled independently, so a particular realization need not contain all three functions, especially when $K$ is small.

Datasets and repetitions.

For every $K$, we generate a balanced training pool containing $200{,}000$ examples and use nested prefixes of sizes

$ n\in{20{,}000, 50{,}000, 100{,}000, 200{,}000}. $

Evaluation uses an independently generated test set containing $4{,}096$ examples. Within each run, exactly the same serialized training and test tensors are used by the MLP, global kernel methods, oracle methods, and MLP-gated methods.

Results are averaged over five seeds. For a fixed $K$, the cluster-position and predictive directions are held fixed, while each seed independently resamples the cluster-specific link assignments, training and test examples, label noise, and MLP initialization. No test labels are used to construct the gates.

Source ReLU MLP.

For each $K$, sample size, and seed, we train a one-hidden-layer ReLU network of width $4{,}096$,

$ \widehat{f}(x)

\sum_{j=1}^{4096} a_j\operatorname{ReLU}(w_j^\top x+b_j) +b_{\mathrm{out}}. $

All weights and biases are trainable. Optimization uses Adam with minibatches of size $4{,}096$, zero weight decay, and cosine learning-rate decay from $10^{-2}$ to $10^{-4}$ . Training runs for at least $20{,}000$ steps and at most $200{,}000$ steps. It terminates when the full training MSE reaches 4 x 10^-4, or when the training loss plateaus. We restore the checkpoint with the smallest full training MSE. Neither test MSE nor cluster identity is used for checkpoint selection.

Selecting active neurons.

After training, we freeze the MLP. For neuron $j$, define its average output-weighted activation on the training set as

$ m_j

\frac{1}{n} \sum_{i=1}^{n} |a_j|, \operatorname{ReLU}(w_j^\top x_i+b_j).\tag{11} $

We apply two-means clustering to ${\log_{10}m_j}_{j=1}^{4096}$ and retain the group with the larger average contribution. The resulting active-neuron set is denoted by $\mathcal J$ . If fewer than $K$ neurons are retained, we instead keep the $K$ neurons with the largest $m_j$ .

Contribution profiles and learned gates.

For every training input $x$, we compute its normalized contribution profile over the active neurons:

$ p_j(x)

\frac{ |a_j|\operatorname{ReLU}(w_j^\top x+b_j) }{ \displaystyle \sum_{\ell\in\mathcal J} |a_\ell|\operatorname{ReLU}(w_\ell^\top x+b_\ell) }, \qquad j\in\mathcal J.\tag{12} $

Thus $p(x)$ describes which active first-layer neurons contribute to the MLP prediction on input $x$, independent of the overall magnitude of the prediction.

We apply $K$-means++ to the training contribution profiles and use $M=K$ learned gates. Knowledge of the number of clusters is therefore provided, but the true cluster assignments are never observed. We use eight initializations and at most 30 Lloyd iterations, retaining the solution with the smallest training inertia. Every training example is assigned to exactly one gate through its nearest centroid.

Let $q_c$ denote the centroid of learned gate $c$, with coordinate $q_{cj}$ corresponding to active neuron $j$ . Although examples receive hard gate assignments, the centroid coordinates provide a soft association between neurons and gates: the same neuron may contribute to several gates.

Constructing gate-specific features.

For each learned gate $c$, we form the positive-semidefinite metric

$ G_c

\frac{ \displaystyle \sum_{j\in\mathcal J}q_{cj}w_jw_j^\top }{ \displaystyle \sum_{j\in\mathcal J}q_{cj} }.\tag{13} $

The corresponding gate-specific representation is

$ z_c(x)=G_c^{1/2}x.\tag{14} $

The MLP biases affect the contribution profiles and hence the gate assignments, while $G_c$ itself is constructed from the incoming first-layer weight directions.

For every learned gate, we project only the training examples assigned to that gate and fit an independent local Laplace or RFM predictor in the resulting representation.

Test-time prediction.

For a new test input $x$, we:

1. pass $x$ through the frozen MLP and compute its contribution profile $p(x)$;
2. assign $x$ to the nearest fixed training centroid;
3. transform it using the corresponding representation $z_c(x)=G_c^{1/2}x$ ; and
4. evaluate the local predictor fitted for that learned gate.

The test input is therefore assigned using only its MLP contribution profile. Neither its response nor its true cluster identity is used.

Local Laplace and RFM predictors.

The local Laplace predictor for gate $c$ uses

$ k_c(x, x')

\exp!\left(-\frac{ |G_c^{1/2}(x-x')|_2 }{h_c} \right), $

where $h_c$ is initialized using the median pairwise distance within that gate. We use ridge parameter $10^{-6}$ .

The local RFM begins from this Laplace kernel and performs three metric updates, producing iterations $0, 1, 2, 3$, where iteration $0$ is the local Laplace predictor. After each update, the bandwidth is recomputed using distances under the updated metric. Kernel systems with fewer than $11{,}000$ local training examples are solved directly; larger systems use EigenPro. The same procedure is used for the global and oracle RFM baselines.

Compared methods.

Figure 4 compares:

• the source ReLU MLP;
• one global Laplace predictor fitted to all training examples;
• one global RFM fitted to all training examples;
• MLP-gated local Laplace, using the gates and projections extracted from the ReLU MLP;
• MLP-gated local RFM, using the same extracted gates and projections;
• per-cluster Laplace (oracle), which is given the true cluster identities at training and test time;
• per-cluster RFM (oracle), defined analogously; and
• the irreducible noise level, whose expected MSE is $0.02^2=4\times10^{-4}$ .

The oracle methods are diagnostic references rather than fair deployable predictors because they receive the true cluster identity. By contrast, the MLP-gated methods receive only the training input–response pairs and the known number $K$ of gates.

Aggregation.

Every curve in Figure 4 is the arithmetic mean over five independent runs. Each run uses a newly sampled dataset from the same generative model and an independently initialized MLP.

B. Proofs for Section 4

B.1 Proof of Theorem 1

We first prove Theorem 1 invoking Lemma 3, Lemma 4, Lemma 5, Lemma 6, Lemma 7, Lemma 8, and Lemma 9, whose statements and proofs are provided after the main proof. A roadmap of their dependencies is shown in Figure 7.

Throughout this subsection, the normalized means and predictive directions

$ \frac{\mu_1}{R}, \ldots, \frac{\mu_K}{R}, v_1, \ldots, v_K $

form an orthonormal basis of $\mathbb{R}^d$, where $d=2K$. In the notation of Section 2, we may write

$ c\sim\operatorname{Unif}([K]), \qquad x=\mu_c+z, \qquad z\sim\mathcal N(0, I_d), $

and

$ y=h_3\bigl(v_c^\top z\bigr), \qquad h_3(t)=\frac{t^3-3t}{\sqrt6}. $

For $\omega\in\mathbb{S}^{d-1}$, define the weights' correlations to the cluster means and predictive directions

$ b_c(\omega)

\langle\omega, \mu_c\rangle, \qquad \rho_c(\omega)

\langle\omega, v_c\rangle. $

Write the ReLU activation function and the Gaussian p.d.f. as

$ \phi(t)=t_+, \qquad \varphi(t)

\frac{1}{\sqrt{2\pi}}e^{-t^2/2}, $

and define the unsigned and signed population-correlation objectives, respectively, by

$ \Phi({\omega})

\mathbb{E}\left[{y}, \phi({\omega}^\top{x})\right], \qquad \Psi_\zeta({\omega})

\zeta\Phi({\omega}), \qquad \zeta\in{\pm1}. $

Proof of Theorem 1: For each initial direction ${\omega_j^0}$, let $(\bar{u}_j, {\bar\omega_j})$ be the corresponding self-selected teacher-only trajectory from Lemma 7.

By Lemma 7 and Lemma 5, almost surely there are a cluster label $J_j\in[K]$ and an orientation $\tau_j\in{\pm1}$ such that

$ \limsup_{t\to\infty} \left| {\bar\omega_j(t)-\tau_j v_{J_j}} \right|_2 \leq \frac{C}{R} $

for a universal constant $C$.

Since $m<\infty$ and every teacher-only direction converges, for every $\delta>0$ there is a finite time $T_\delta$ such that, simultaneously for all $j\in[m]$,

$ \left| {\bar\omega_j(T_\delta)-\tau_j v_{J_j}} \right|_2 \leq \frac{\delta}{2}+\frac{C}{R}. $

Apply Lemma 9 on $[0, T_\delta]$. For sufficiently small $\varepsilon$,

$ \left| \frac{w_j(T_\delta)}{|w_j(T_\delta)|_2}

{\bar\omega_j(T_\delta)} \right|2 \leq L{T_\delta}m\varepsilon^2 \leq \frac{\delta}{2} $

for every $j\in[m]$. The triangle inequality therefore gives

$ \left| \frac{w_j(T_\delta)}{|w_j(T_\delta)|_2}

\tau_j {v_{J_j}} \right|_2 \leq \delta+\frac{C}{R}. $

The independence and uniformity of the labels $J_1, \ldots, J_m$, together with the cluster-coverage probability, follow from Lemma 8.

B.1.1 The Population-correlation Landscape

In this subsection, we study the critical points of the optimization problem

$ \max_{\omega \in \mathbb{S}^{d-1}} \Phi_{\zeta}(\omega), .\tag{15} $

The use of homogeneity to characterize parameter directions is part of a broader literature on the implicit bias of gradient methods, beginning with homogeneous linear predictors and extending to homogeneous neural networks ([55, 56, 57, 58]). These results motivate the directional viewpoint used here, but do not directly yield the squared-loss, small-initialization approximation in our setting; we establish the required approximation directly in Lemma 9. (In the subsequent section of the appendix, we will show how the sign $\zeta$ is determined for each neuron at very early times of training.) We show that the critical points of equation 15 are specialized to clusters of the data distribution.

First, we provide an explicit formula for the signed population-correlation objective in terms of a weight's correlations with the cluster means and predictive directions. This lemma is stated for weights $\omega$ on the unit sphere, but is later applied to small weights $\omega$ because of the homogeneity of ReLU networks.

########## {caption="Lemma 3: Cubic correlation identity"}

For every $\omega\in\mathbb{S}^{d-1}$ and $\zeta\in{\pm1}$,

$ \Psi_\zeta(\omega)

-\frac{\zeta}{K\sqrt6} \sum_{c=1}^K b_c(\omega)\varphi\bigl(b_c(\omega)\bigr)\rho_c(\omega)^3. $

Proof: Fix $\omega\in\mathbb{S}^{d-1}$. Conditional on $c$,

$ \omega^\top x =b_c(\omega) +\omega^\top z, \qquad y=h_3\bigl(v_c^\top z\bigr). $

Hence the contribution of cluster $c$ to $\Phi(\omega) = \mathbb{E}\left[y, \phi(\omega^\top x)\right]$ is

$ \mathbb{E}\left[h_3\bigl(v_c^\top z\bigr) \bigl(b_c(\omega) +\omega^\top z\bigr)_+ \right]. $

Set

$ U=v_c^\top z, \qquad V=\omega^\top z, \qquad \rho=\rho_c(\omega) =\langle v_c, \omega\rangle. $

Since $z\sim\mathcal N(0, I_d)$ and $v_c, \omega$ are unit vectors, $U$ and $V$ are standard Gaussian, and

$ \mathbb{E}[UV]

v_c^\top \mathbb{E}[zz^\top]\omega

\langle v_c, \omega\rangle

\rho. $

Thus $(U, V)$ is a jointly standard Gaussian pair with correlation $\rho$.

For $|\rho|<1$, define

$ Z:=\frac{U-\rho V}{\sqrt{1-\rho^2}}. $

Because $(U, V)$ is jointly Gaussian, $(Z, V)$ is also jointly Gaussian. Moreover,

$ \mathbb{E}[Z]=0, \qquad \mathbb{E}[Z^2]=1, $

and

$ \mathbb{E}[ZV]

\frac{\mathbb{E}[UV]-\rho\mathbb{E}[V^2]} {\sqrt{1-\rho^2}} =0. $

Jointly Gaussian random variables with zero covariance are independent, so $Z\sim\mathcal N(0, 1)$ is independent of $V$. Therefore

$ U=\rho V+\sqrt{1-\rho^2}, Z. $

When $|\rho|=1$, the same representation holds with the second term equal to zero.

Conditioning on $V$ and using the independence of $Z$ gives

$ \mathbb{E}[U\mid V]=\rho V $

and

$ \begin{aligned} \mathbb{E}[U^3\mid V] &= \mathbb{E}\left[\left(\rho V+\sqrt{1-\rho^2}, Z\right)^3 , \middle|, V \right]\ &= \rho^3V^3+3\rho(1-\rho^2)V, \end{aligned} $

since $\mathbb{E}[Z]=\mathbb{E}[Z^3]=0$ and $\mathbb{E}[Z^2]=1$. Recalling that

$ h_3(x)=\frac{x^3-3x}{\sqrt6}, $

we obtain

$ \begin{aligned} \mathbb{E}[h_3(U)\mid V] &= \frac{1}{\sqrt6} \left(\mathbb{E}[U^3\mid V] -3\mathbb{E}[U\mid V] \right)\ &= \rho^3 h_3(V). \end{aligned} $

Since $(b_c(\omega)+V)_+$ depends only on $V$, the tower property now gives

$ \begin{aligned} &\mathbb{E}\left[h_3(U)\bigl(b_c(\omega)+V\bigr)+ \right]\ &\qquad= \mathbb{E}\left[\bigl(b_c(\omega)+V\bigr)+ \mathbb{E}[h_3(U)\mid V] \right]\ &\qquad= \rho_c(\omega)^3 \mathbb{E}{G\sim\mathcal N(0, 1)} \left[h_3(G)\bigl(b_c(\omega)+G\bigr)+ \right]. \end{aligned} $

It remains to compute the one-dimensional expectation. For any $b\in\mathbb{R}$,

$ \mathbb{E}\left[h_3(G)(b+G)_+\right]

\frac1{\sqrt6} \int_{-b}^{\infty} (x^3-3x)(b+x)\varphi(x), dx. $

Using

$ \frac{d}{dx} \left[(x^2-1)\varphi(x)\right]

-(x^3-3x)\varphi(x), $

integration by parts yields

$ \int_{-b}^{\infty} (x^3-3x)(b+x)\varphi(x), dx

\int_{-b}^{\infty} (x^2-1)\varphi(x), dx, $

where the boundary term vanishes because $b+x=0$ at $x=-b$ and $\varphi(x)\to0$ as $x\to\infty$. Since

$ \frac{d}{dx}\bigl(x\varphi(x)\bigr)

(1-x^2)\varphi(x), $

we further obtain

$ \int_{-b}^{\infty} (x^2-1)\varphi(x), dx

-b\varphi(b). $

Therefore

$ \mathbb{E}\left[h_3(G)(b+G)_+\right]

-\frac{b\varphi(b)}{\sqrt6}. $

Applying this with $b=b_c(\omega)$, the contribution of cluster $c$ to $\Phi(\omega)$ is

$ -\frac{1}{\sqrt6} b_c(\omega)\varphi\bigl(b_c(\omega)\bigr)\rho_c(\omega)^3. $

Since $c$ is uniform on $[K]$,

$ \Phi(\omega) = -\frac1{K\sqrt6} \sum_{c=1}^K b_c(\omega)\varphi\bigl(b_c(\omega)\bigr)\rho_c(\omega)^3. $

Finally, $\Psi_\zeta(\omega)=\zeta\Phi(\omega)$, and hence

$ \Psi_\zeta(\omega)

-\frac{\zeta}{K\sqrt6} \sum_{c=1}^K b_c(\omega)\varphi\bigl(b_c(\omega)\bigr)\rho_c(\omega)^3, $

as claimed.

Because the $\mu_c/R$ and $v_c$ vectors form an orthonormal basis, the sphere constraint is

$ \sum_{c=1}^K\frac{b_c^2}{R^2} + \sum_{c=1}^K\rho_c^2

1. $

Next, we consider the optimization problem of signed population-correlation maximization between the value of the neuron and the ground truth: $\max_{\omega\in\mathbb{S}^{d-1}}\Psi_\zeta(\omega)$. This is written and analyzed below under the linear change of coordinates

$ b_c=\langle\omega, \mu_c\rangle, \qquad \rho_c=\langle\omega, v_c\rangle, , $

where we show that the neuron maximizes correlation with the target by specializing to one cluster. Here, we use the fact that the target is the cubic Hermite polynomial. (A similar argument would hold for all higher-degree Hermite polynomials as well.)

########## {caption="Lemma 4: Positive local maxima use one cluster"}

Consider the constrained maximization problem

$ \max_{b_1, \ldots, b_K, \rho_1, \ldots, \rho_K} F(b_1, \ldots, b_K, \rho_1, \ldots, \rho_K), $

where

$ F(b_1, \ldots, b_K, \rho_1, \ldots, \rho_K) := -\frac{\zeta}{K\sqrt6} \sum_{c=1}^K b_c\varphi(b_c)\rho_c^3, $

subject to

$ \sum_{c=1}^K\frac{b_c^2}{R^2} + \sum_{c=1}^K\rho_c^2

1. $

Every constrained local maximum of $F$ with positive objective value has exactly one nonzero predictive coordinate $\rho_c$. Moreover, if $\rho_c=0$, then the corresponding routing coordinate $b_c$ also vanishes.

Proof: At a constrained local maximum, first-order stationarity of the Lagrangian

$ F(b_1, \ldots, b_K, \rho_1, \ldots, \rho_K)

\lambda\left(\sum_{c=1}^K\frac{b_c^2}{R^2} + \sum_{c=1}^K\rho_c^2 -1 \right) $

with respect to $b_c$ and $\rho_c$ gives, for every $c\in[K]$,

$ \begin{aligned} -\frac{\zeta}{K\sqrt6} (1-b_c^2)\varphi(b_c)\rho_c^3 &= \frac{2\lambda}{R^2}b_c, \quad\text{(a)}\quad\text{(B.1)}\ -\frac{3\zeta}{K\sqrt6} b_c\varphi(b_c)\rho_c^2 &= 2\lambda\rho_c. \quad\text{(b)}\quad\text{(B.2)} \end{aligned}\tag{16} $

Here we used

$ \frac{d}{db}\bigl(b\varphi(b)\bigr)

(1-b^2)\varphi(b). $

Multiplying equation 16b by $\rho_c$ and summing over $c$ gives

$ 3F(b_1, \ldots, b_K, \rho_1, \ldots, \rho_K) = 2\lambda\sum_{c=1}^K\rho_c^2. $

Since the objective value is positive, at least one predictive coordinate is nonzero. Hence both the left-hand side and $\sum_c\rho_c^2$ are positive, so

$ \lambda>0. $

If $\rho_c=0$, then equation 16a reduces to

$ 0=\frac{2\lambda}{R^2}b_c. $

Since $\lambda>0$, this implies $b_c=0$.

It remains to rule out two nonzero predictive coordinates. Suppose that $\rho_c\neq0$ and $\rho_r\neq0$ for two distinct clusters $c\neq r$. Keep all routing coordinates and all other predictive coordinates fixed, and consider

$ \rho_c(t)

\operatorname{sign}(\rho_c) \sqrt{\rho_c^2+t}, \qquad \rho_r(t)

\operatorname{sign}(\rho_r) \sqrt{\rho_r^2-t}. $

For sufficiently small $|t|$, this perturbation is well defined, preserves the signs of the two coordinates, and satisfies

$ \rho_c(t)^2+\rho_r(t)^2

\rho_c^2+\rho_r^2. $

Thus it preserves the constraint exactly.

Differentiating $F$ along this feasible curve at $t=0$ gives

$ F'(0)

-\frac{3\zeta}{2K\sqrt6} \left(b_c\varphi(b_c)\rho_c

b_r\varphi(b_r)\rho_r \right). $

Since $\rho_c, \rho_r\neq0$, dividing equation 16b by the corresponding predictive coordinate gives

$ -\frac{3\zeta}{K\sqrt6} b_c\varphi(b_c)\rho_c

2\lambda, \qquad -\frac{3\zeta}{K\sqrt6} b_r\varphi(b_r)\rho_r

2\lambda. $

Therefore

$ F'(0)=0. $

The second derivative along the same feasible curve is

$ F''(0)

-\frac{3\zeta}{4K\sqrt6} \left(\frac{b_c\varphi(b_c)}{\rho_c} + \frac{b_r\varphi(b_r)}{\rho_r} \right). $

Using the same KKT identities gives

$ F''(0)

\frac{\lambda}{2} \left(\frac{1}{\rho_c^2} + \frac{1}{\rho_r^2} \right)

0. $

Thus a critical point with two nonzero predictive coordinates has a feasible direction of positive curvature and cannot be a constrained local maximum.

Hence at most one predictive coordinate is nonzero. Since the objective value is positive, at least one is nonzero, so exactly one is active. As shown above, every routing coordinate corresponding to an inactive predictive coordinate also vanishes.

Returning to our objective $\max_{\omega \in \mathbb{S}^{d-1}} \Psi_{\zeta}(\omega)$, let $c$ denote the unique active cluster and let

$ \tau=\operatorname{sign}(\rho_c)\in{\pm1}. $

Then

$ b_j=\rho_j=0 \qquad \text{for every }j\neq c, $

and the constraint becomes

$ \frac{b_c^2}{R^2}+\rho_c^2=1. $

Hence

$ \rho_c

\tau\sqrt{1-\frac{b_c^2}{R^2}}. $

Substituting this into $F$, and writing $b=b_c$, reduces the problem to the one-dimensional objective

$ -\frac{\tau\zeta}{K\sqrt6} b\varphi(b) \left(1-\frac{b^2}{R^2}\right)^{3/2}, \qquad |b|<R. $

Let $c$ be the unique active cluster and let

$ \tau=\operatorname{sign}(\rho_c)\in{\pm1}. $

Then

$ b_j=\rho_j=0 \qquad \text{for every }j\neq c, $

and

$ \rho_c

\tau\sqrt{1-\frac{b_c^2}{R^2}}. $

Writing $b=b_c$, the remaining one-dimensional objective is

$ -\frac{\tau\zeta}{K\sqrt6} b\varphi(b) \left(1-\frac{b^2}{R^2}\right)^{3/2}, \qquad |b|<R. $

########## {caption="Lemma 5: Specialized maxima"}

Fix $c\in[K]$ and $\zeta, \tau\in{\pm1}$. On the one-cluster feasible branch, write

$ f(b) := -\frac{\tau\zeta}{K\sqrt6} b\varphi(b) \left(1-\frac{b^2}{R^2}\right)^{3/2}, \qquad |b|<R, $

and let

$ I:={b\in(-R, R):f(b)>0}. $

Then $I$ is a connected open interval. For all sufficiently large $R$, $f$ has a unique critical point $b^\star\in I$, which is a strict local maximum and satisfies

$ b^\star=-\tau\zeta+O(R^{-2}). $

The corresponding direction

$ \omega^\star

\tau\sqrt{1-\frac{(b^\star)^2}{R^2}}, v_c + \frac{b^\star}{R^2}\mu_c $

is a strict local maximum of $\Psi_\zeta$ and obeys

$ \omega^\star

\tau v_c

\frac{\tau\zeta}{R^2}\mu_c + O(R^{-2}). $

Proof: For every $b\in(-R, R)$,

$ \frac{1}{K\sqrt6}>0, \qquad \varphi(b)>0, \qquad \left(1-\frac{b^2}{R^2}\right)^{3/2}>0. $

Hence the sign of $f(b)$ is determined entirely by $-\tau\zeta b$:

$ f(b)>0 \quad\Longleftrightarrow\quad -\tau\zeta b>0. $

Since $\tau\zeta\in{\pm1}$, it follows that

$ I= \begin{cases} (-R, 0), & \tau\zeta=1, \ (0, R), & \tau\zeta=-1. \end{cases} $

Thus $I$ is connected.

For $b\in I$, set

$ t=-\tau\zeta b. $

The preceding characterization of $I$ shows that this is a bijective linear change of variable from $I$ onto $(0, R)$. Since $\varphi$ is even and $(\tau\zeta)^2=1$, we obtain

$ f(b) = \frac{1}{K\sqrt6} t\varphi(t) \left(1-\frac{t^2}{R^2}\right)^{3/2}. $

Therefore the critical points of $f$ on $I$ correspond exactly to the critical points on $(0, R)$ of

$ t\longmapsto t\varphi(t) \left(1-\frac{t^2}{R^2}\right)^{3/2}. $

This function is strictly positive on $(0, R)$, so its critical points can be found from its logarithmic derivative:

$ \frac{d}{dt} \log\left[t\varphi(t) \left(1-\frac{t^2}{R^2}\right)^{3/2} \right]

\frac1t-t-\frac{3t}{R^2-t^2}. $

Thus a critical point satisfies

$ \frac1t-t-\frac{3t}{R^2-t^2}=0. $

Multiplying by $t(R^2-t^2)$ gives

$ t^4-(R^2+4)t^2+R^2=0. $

Solving this quadratic equation in $t^2$ gives

$ t^2 = \frac{R^2+4\pm\sqrt{R^4+4R^2+16}}{2}. $

The root with the plus sign is larger than $R^2$, whereas the root with the minus sign lies in $(0, R^2)$. Hence there is exactly one critical point in $(0, R)$, and it satisfies

$ t^2

\frac{ R^2+4-\sqrt{R^4+4R^2+16} }{2}

1+O(R^{-2}). $

Since $t>0$,

$ t=1+O(R^{-2}). $

Returning to $b=-\tau\zeta t$ gives

$ b^\star

-\tau\zeta+O(R^{-2}). $

The function

$ t\varphi(t) \left(1-\frac{t^2}{R^2}\right)^{3/2} $

is positive on $(0, R)$ and tends to zero as $t\to0$ or $t\to R$. Since it has exactly one critical point in $(0, R)$, this point is its unique strict maximum. Hence $b^\star$ is the unique maximizer of $f$ on $I$.

It remains to show that the corresponding point is a local maximum of the full constrained objective $F$. Since the feasible set is compact, $F$ attains a global maximum. The point constructed above has positive objective value, so the global maximum is positive. By Lemma 4, every global maximizer must have exactly one nonzero predictive coordinate and therefore lies on one of the one-cluster feasible branches.

For every choice of the active cluster $c$ and orientation $\tau$, the change of variable

$ t=-\tau\zeta b $

reduces the positive part of the corresponding one-cluster objective to

$ \frac{1}{K\sqrt6} t\varphi(t) \left(1-\frac{t^2}{R^2}\right)^{3/2}. $

Thus all one-cluster branches have the same maximal value, attained uniquely at the point identified above. Consequently, every corresponding $\omega^\star$ is a global maximizer of $\Psi_\zeta$.

There are only finitely many such maximizers, one for each $c\in[K]$ and $\tau\in{\pm1}$. Hence each is isolated, and therefore each $\omega^\star$ is a strict local maximum of $\Psi_\zeta$.

B.1.2 Random Initialization, Self-selected Signs, and Cluster Coverage

In the previous section, we considered the signed population-correlation objective $\Phi_{\zeta}$ for each neuron. In this section, we note that the sign $\zeta$ for each neuron is learned at very early times of training, and does not need to be fixed. Consider the teacher-only joint dynamics

$ \dot{u}=\Phi({\omega}), \qquad {\dot\omega} =\tanh(u)\nabla_{\mathbb{S}^{d-1}}\Phi({\omega}).\tag{17} $

These equations arise from the two-layer parameterization when the second-layer weight is initially zero. Crucially, they neglect interactions between neurons, which we will later show is a fine approximation because the network is small at initialization.

########## {caption="Lemma 6: Positive signed-correlation flows specialize"}

Let ${\omega^0}$ have an absolutely continuous distribution on $\mathbb{S}^{d-1}$, fix $\zeta\in{\pm1}$, and assume $\Psi_\zeta({\omega^0})>0$. For all sufficiently large $R$, spherical gradient ascent

$ {\dot\omega} =\nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega}), \qquad {\omega(0)}={\omega^0}, $

converges almost surely to one of the specialized maxima in Lemma 5.

Proof: Along spherical gradient ascent,

$ \frac{d}{dt}\Psi_\zeta({\omega(t)})

\left| \nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega(t)}) \right|_2^2 \geq 0. $

Since the sphere is compact, $\Psi_\zeta$ is bounded above. Hence

$ \int_0^\infty \left| \nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega(t)}) \right|_2^2, dt <\infty. $

Moreover, $\Psi_\zeta$ is smooth on the sphere, so its gradient and Hessian are bounded. It follows that $|\nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega(t)})|_2^2$ has bounded derivative. A nonnegative integrable function with bounded derivative must converge to zero, and therefore

$ \left| \nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega(t)}) \right|_2 \longrightarrow 0. $

We next note that $\Psi_\zeta$ has only finitely many positive critical points. Indeed, at any positive critical point, the stationarity conditions equation 16a–equation 16b imply that $\lambda>0$, and hence $b_c=0$ whenever $\rho_c=0$. For every active coordinate $\rho_c\neq0$, eliminating $\lambda$ from the two stationarity equations gives

$ (1-b_c^2)\rho_c^2 = \frac{3b_c^2}{R^2}. $

Thus $0<|b_c|<1$ and

$ \rho_c^2

\frac{3b_c^2}{R^2(1-b_c^2)}. $

Equation 16b further shows that

$ |b_c|\varphi(b_c)|\rho_c| $

has the same value for every active coordinate. Substituting the preceding expression for $\rho_c^2$, this quantity is proportional to

$ \frac{|b_c|^2\varphi(|b_c|)} {\sqrt{1-|b_c|^2}}, $

which is strictly increasing for $0<|b_c|<1$, since

$ \frac{d}{dt} \log\left(\frac{t^2\varphi(t)}{\sqrt{1-t^2}} \right)

\frac{2}{t}+\frac{t^3}{1-t^2}

0. $

Hence all active $|b_c|$ are equal. If there are $s$ active coordinates, the sphere constraint then gives

$ s, \frac{b_c^2(4-b_c^2)} {R^2(1-b_c^2)} =1, $

which has a unique solution for $b_c^2\in(0, 1)$. Thus, for each choice of the active coordinates and their signs, there is at most one positive critical point. Since there are only finitely many such choices, the set of positive critical points is finite.

Because $\Psi_\zeta({\omega(0)})>0$ and the objective is nondecreasing, every accumulation point of the trajectory has positive objective value. By compactness, accumulation points exist, and the convergence of the gradient to zero implies that every accumulation point is a positive critical point. Since there are only finitely many such points, the trajectory must converge to one of them.

Finally, the argument in Lemma 4 gives an unstable tangent direction at every positive critical point with more than one active predictive coordinate. Since there are only finitely many such critical points, the center-stable manifold theorem implies that the set of initial conditions converging to any of them has measure zero. Lemma 5 shows that the remaining positive critical points are precisely the specialized maxima. Hence an absolutely continuous initialization converges almost surely to a specialized maximum.

For a single neuron, write

$ \omega(t)=\frac{w(t)}{|w(t)|_2}. $

ReLU homogeneity implies that gradient flow preserves

$ |w(t)|_2^2-a(t)^2

|w(0)|_2^2-a(0)^2

\varepsilon^2. $

We therefore define the scalar amplitude coordinate

$ u(t) := \operatorname{arsinh}!\left(\frac{a(t)}{\varepsilon}\right), $

so that, equivalently,

$ a(t)=\varepsilon\sinh u(t), \qquad |w(t)|_2=\varepsilon\cosh u(t). $

Under this parametrization, the target-only part of the directional dynamics is

$ \dot{u}=\Phi(\omega), \qquad \dot\omega

\tanh(u)\nabla_{\mathbb{S}^{d-1}}\Phi(\omega).\tag{18} $

########## {caption="Lemma 7: Self-selected sign and almost-sure specialization"}

Let $\omega^0$ have an absolutely continuous distribution on $\mathbb{S}^{d-1}$, and let $(u(t), \omega(t))$ solve equation 18 with

$ u(0)=0, \qquad \omega(0)=\omega^0. $

For almost every $\omega^0$, define

$ \zeta = \operatorname{sign}(\Phi(\omega^0)) \in{\pm1}. $

Then, for every $t>0$,

$ \zeta u(t)>0, \qquad \zeta\Phi(\omega(t)) \geq |\Phi(\omega^0)|. $

Moreover, $\omega(t)$ is a positive time reparameterization of spherical gradient ascent on $\Psi_\zeta=\zeta\Phi$, and, for all sufficiently large $R$, it converges almost surely to one of the specialized maxima in Lemma 5.

Proof: Since $\Phi$ is real analytic and not identically zero, its zero set has spherical measure zero. Thus, for almost every ${\omega^0}$, $\Phi({\omega^0})\neq0$. Fix such a ${\omega^0}$ and let

$ \zeta=\operatorname{sign}(\Phi({\omega^0})). $

At $t=0$,

$ \frac{d}{dt}\bigl(\zeta u(t)\bigr)\bigg|_{t=0}

\zeta\Phi({\omega^0})

|\Phi({\omega^0})|

0, $

so $\zeta u(t)>0$ for all sufficiently small $t>0$.

As long as $\zeta u(t)>0$, $\tanh(u(t))$ has sign $\zeta$, and hence

$ \zeta\tanh(u(t))=|\tanh(u(t))|. $

Therefore

$ \frac{d}{dt}\bigl(\zeta\Phi({\omega(t)})\bigr)

|\tanh(u(t))| \left| \nabla_{\mathbb{S}^{d-1}}\Phi({\omega(t)}) \right|_2^2 \geq0. $

It follows that

$ \zeta\Phi({\omega(t)}) \geq \zeta\Phi({\omega^0})

|\Phi({\omega^0})|. $

Consequently,

$ \frac{d}{dt}\bigl(\zeta u(t)\bigr)

\zeta\Phi({\omega(t)}) \geq |\Phi({\omega^0})|, $

and therefore

$ \zeta u(t)\geq |\Phi({\omega^0})|t. $

In particular, $\zeta u(t)$ cannot return to zero, so the preceding inequalities hold for every $t>0$. This proves

$ \zeta u(t)>0, \qquad \zeta\Phi({\omega(t)}) \geq |\Phi({\omega^0})|. $

Since $\Psi_\zeta=\zeta\Phi$ and $\tanh(u(t))$ has sign $\zeta$, the directional dynamics satisfy

$ {\dot\omega(t)} = |\tanh(u(t))| \nabla_{\mathbb{S}^{d-1}}\Psi_\zeta({\omega(t)}). $

Thus ${\omega(t)}$ follows the same orbit as spherical gradient ascent on $\Psi_\zeta$, up to a positive reparameterization of time. Moreover, $\zeta u(t)\geq|\Phi({\omega^0})|t$ implies $|u(t)|\to\infty$, and hence $|\tanh(u(t))|\to1$. In particular,

$ \int_0^\infty |\tanh(u(t))|, dt=\infty, $

so this reparameterization covers the entire forward gradient-ascent trajectory.

Finally,

$ \Psi_\zeta({\omega^0})

\zeta\Phi({\omega^0})

|\Phi({\omega^0})|

0. $

Lemma 6 therefore implies that ${\omega(t)}$ converges almost surely to one of the specialized maxima in Lemma 5.

########## {caption="Lemma 8: Uniform selected labels and coverage"}

Let ${\omega_1^0, \ldots, \omega_m^0}$ be independent and uniform on $\mathbb{S}^{d-1}$, and let $J_j$ be the cluster selected by the self-selected flow in Lemma 7. Then $J_1, \ldots, J_m$ are independent and uniform on $[K]$. Consequently,

$ \mathbb{P}\bigl({J_1, \ldots, J_m}\neq[K]\bigr) \leq K\left(1-\frac1K\right)^m \leq Ke^{-m/K}. $

Proof: For a permutation $\pi$ of $[K]$, let $P_\pi$ be the orthogonal map satisfying

$ P_\pi{\frac{\mu_c}{R}}

{\frac{\mu_{\pi(c)}}{R}}, \qquad P_\pi{v_c}

{v_{\pi(c)}}. $

The data distribution, target, and self-selected vector field are equivariant under $P_\pi$. Hence, if ${\omega^0}$ selects cluster $c$, then $P_\pi {\omega^0}$ selects cluster $\pi(c)$. Uniform spherical measure is invariant under $P_\pi$, so all selection probabilities are equal. By Lemma 7, they sum to one, and each is therefore $1/K$. Independence follows because each teacher-only label is a deterministic function of an independent initial direction.

A fixed cluster is missed with probability $(1-1/K)^m$. A union bound yields the coverage estimate.

B.1.3 Tracking by the Full Small-initialization Dynamics

Finally, we put the above ingredients analyzing the trajectories of individual neurons together, and show that they describe the trajectory of the neurons in a neural network (up to rescaling) when the network is initialized small. The proof strategy of studying feature learning through effectively independent neuron dynamics in early-time or small-initialization regimes has been put forward in prior work ([17, 47, 35]). The population loss is

$ \mathcal L(\theta)

\frac12 \mathbb{E}\left[\bigl(f_\theta({x})-{y}\bigr)^2 \right], $

and its gradient-flow equations are

$ \begin{aligned} \dot{a}j &= \mathbb{E}\left[({y}-f\theta({x})) \phi(w_j^\top{x}) \right], \ \dot{w}j &= a_j \mathbb{E}\left[({y}-f\theta({x})) \mathbf 1_{{w_j^\top{x}>0}}{x} \right]. \end{aligned} $

########## {caption="Lemma 9: Early-time tracking from a zero output layer"}

Initialize

$ a_j(0)=0, \qquad w_j(0)=\varepsilon {\omega_j^0}, \qquad {\omega_j^0}\in\mathbb{S}^{d-1}. $

Let $(\bar{u}_j, {\bar\omega_j})$ solve the self-selected teacher-only system

$ \dot{\bar{u}}_j

\Phi({\bar\omega_j}), \qquad {\dot{\bar\omega}_j}

\tanh(\bar{u}j) \nabla{\mathbb{S}^{d-1}}\Phi({\bar\omega_j}), $

with

$ (\bar{u}_j(0), {\bar\omega_j(0)})

(0, {\omega_j^0}). $

For every fixed $T<\infty$, there are constants $L_T<\infty$ and $\varepsilon_T>0$ such that, for $0<\varepsilon\leq\varepsilon_T$,

$ \sup_{0\leq t\leq T} \left(|u_j(t)-\bar{u}_j(t)| + |{\omega_j(t)-\bar\omega_j(t)}|_2 \right) \leq L_Tm\varepsilon^2, $

where

$ a_j(t)=\varepsilon\sinh(u_j(t)), \qquad {\omega_j(t)}=\frac{w_j(t)}{|w_j(t)|_2}. $

Proof: ReLU homogeneity gives the invariant

$ \frac{d}{dt} \left(|w_j|_2^2-a_j^2 \right)

0. $

Under the stated initialization,

$ |w_j(t)|_2^2-a_j(t)^2

\varepsilon^2. $

Since $a_j=\varepsilon\sinh u_j$, it follows that

$ |w_j|_2

\varepsilon\cosh u_j. $

Substituting

$ w_j

\varepsilon\cosh(u_j){\omega_j} $

into the gradient-flow equations and using ReLU homogeneity gives

$ \begin{aligned} \dot{u}_j =& \Phi({\omega_j})

\mathbb{E}\left[f_\theta({x}) \phi({\omega_j}^\top{x}) \right], \ {\dot\omega_j} =& \tanh(u_j) \Bigg(\nabla_{\mathbb{S}^{d-1}}\Phi({\omega_j})\ &

\bigl(I-{\omega_j\omega_j^\top}\bigr) \mathbb{E}\left[f_\theta({x}) \mathbf 1_{{{\omega_j}^\top {x}>0}}{x} \right] \Bigg). \end{aligned} $

Lemma 3 shows that the teacher-only vector field is smooth on bounded $u$-intervals and on the sphere. On every fixed interval $[0, T]$, the quantities $u_j(t)$ remain uniformly bounded for sufficiently small $\varepsilon$. Consequently, there is a constant $C_T<\infty$ such that

$ |a_j(t)|+|w_j(t)|_2 \leq C_T\varepsilon, \qquad 0\leq t\leq T. $

Since ${x}$ has finite moments of every order,

$ |f_\theta(t)|{L^2(P{{x}})} \leq C_T \sum_{j=1}^m |a_j(t)||w_j(t)|_2 \leq C_Tm\varepsilon^2. $

Cauchy–Schwarz therefore bounds both interaction terms in the $(u_j, {\omega_j})$ equations by $C_Tm\varepsilon^2$.

The full and teacher-only systems have the same initial conditions, and their vector fields are Lipschitz on the relevant compact set. Grönwall's inequality yields

$ \sup_{0\leq t\leq T} \left(|u_j(t)-\bar{u}_j(t)| + |{\omega_j(t)-\bar\omega_j(t)}|_2 \right) \leq L_Tm\varepsilon^2. $

Perturbations of the pure cubic target.

The pure Hermite target isolates the specialization mechanism and makes the population-correlation landscape exact. However, we could consider perturbations of this objective and expect similar specialization results. Indeed, specialized maxima and their attraction basins should persist under perturbations whose induced population-correlation objective is sufficiently small near the specialized maxima and the strict-saddle regions induced by the cubic Hermite part of the target.

B.2 Proof of Theorem 2

We prove a slightly stronger, parameterized version of the main-text result. Throughout this subsection, $d=2K$ and $R\geq1$ is fixed. We use the means $\mu_c=[Re_c;0]$, predictive directions $v_c=[0;e_c]$, and covariance $\Sigma_c=\operatorname{diag}(0_{K\times K}, I_K)$. The setting in Theorem 2 is the special case $R=1$. Equivalently,

$ c\sim\operatorname{Unif}([K]), \qquad x=\mu_c+\begin{bmatrix}0\z\end{bmatrix}, \qquad z\sim\mathcal N(0, I_K), $

and the regression function is

$ f^\star(x):=\mathbb{E}[y\mid x] =g\bigl(\langle x, v_c\rangle\bigr), \qquad g(t)=t_+-(t-1)_+. $

B.2.1 Exact MLP Representation and Generalization

For a two-layer ReLU network, define the path norm

$ |f|_{\mathcal P} =\inf\left{ \sum_j|a_j||w_j|_2: f(x)=\sum_j a_j\phi(w_j^\top x) \right}. $

For the parameterization $\theta=((a_j, w_j))_{j=1}^m$, define

$ |\theta|{\mathrm F}^2 :=\frac12\sum{j=1}^m \left(a_j^2+|w_j|_2^2\right), $

and, for $B>0$, let $\mathcal F_B$ be the class of two-layer ReLU networks admitting a parameterization with $|\theta|_{\mathrm F}^2\leq B$. Define $\operatorname{clip}(t)=\min{1, \max{0, t}}$.

########## {caption="Lemma 10: Exact specialized MLP"}

On the support of the data distribution,

$ F(x) =\sum_{c=1}^K\left[\phi\bigl(v_c^\top x\bigr) -\phi\left(\left(v_c-\frac{\mu_c}{R^2}\right)^\top x \right) \right] $

equals $f^\star(x)$ exactly. Moreover,

$ |F|_{\mathcal P} \leq B_R :=K\left(1+\sqrt{1+R^{-2}}\right). $

It also admits a parameterization satisfying $|\theta|_{\mathrm F}^2\leq B_R$.

Proof: Suppose $x$ belongs to cluster $r$. Orthogonality gives

$ v_c^\top x=e_c^\top z, \qquad \frac{\mu_c^\top x}{R^2} =\mathbf 1_{{c=r}}. $

For $c\neq r$, the two ReLUs in the $c$ th pair therefore cancel. The pair with $c=r$ equals

$ \phi(e_r^\top z) -\phi(e_r^\top z-1) =g(e_r^\top z). $

Thus $F=f^\star$.

The two hidden weights for cluster $c$ are $v_c$ and $v_c-R^{-2}\mu_c$, with output weights $+1$ and $-1$. Their norms are $1$ and $\sqrt{1+R^{-2}}$, respectively. Summing over clusters gives the path-norm bound.

Finally, positive homogeneity allows each neuron to be rescaled without changing its realized function so that $|a_j|=|w_j|_2$. For this balanced parameterization,

$ \frac12\sum_j\left(a_j^2+|w_j|_2^2\right) =\sum_j|a_j||w_j|_2 \leq B_R. $

The following bound is the two-layer path-norm specialization of norm-based capacity and Rademacher-complexity results for neural networks ([48, 11, 49]); the conversion from Rademacher complexity to a risk bound follows the standard framework of [50].

########## {caption="Lemma 11: Path-norm Rademacher bound"}

For any deterministic sample $x_1, \ldots, x_n\in\mathbb{R}^d$,

$ \widehat{\mathfrak R}n(\mathcal F_B) \leq \frac{2B}{n}\left(\sum{i=1}^n|x_i|_2^2\right)^{1/2}. $

Consequently, for any $B\geq B_R$, if $\widehat{f}_{\mathrm{MLP}}$ is an empirical squared-loss minimizer over $\operatorname{clip}\circ\mathcal F_B$, then

$ \mathbb{E}\left[|\widehat{f}{\mathrm{MLP}}-f^\star|{L^2(P_x)}^2 \right] \leq16B\sqrt{\frac{K+R^2}{n}}. $

Proof: By the arithmetic–geometric mean inequality, every $f\in\mathcal F_B$ has path norm at most $B$. By positive homogeneity, every $f\in\mathcal F_B$ can therefore be represented as

$ f(x)=\sum_j c_j\phi(\omega_j^\top x), \qquad |\omega_j|_2=1, \qquad \sum_j|c_j|\leq B. $

The class $\mathcal F_B$ is therefore contained in $B$ times the absolutely convex hull of unit-norm ReLU atoms. The contraction inequality and Cauchy–Schwarz give

$ \begin{aligned} \widehat{\mathfrak R}n(\mathcal F_B) &\leq \frac{B}{n}\mathbb{E}\epsilon \sup_{|\omega|2\leq1} \left|\sum{i=1}^n\epsilon_i\phi(\omega^\top x_i)\right|\ &\leq \frac{2B}{n}\mathbb{E}\epsilon \left|\sum{i=1}^n\epsilon_i x_i\right|2\ &\leq \frac{2B}{n}\left(\sum{i=1}^n|x_i|_2^2\right)^{1/2}. \end{aligned} $

Clipping does not increase Rademacher complexity. Since both the target and clipped predictions lie in $[0, 1]$, squared loss is $2$-Lipschitz in the prediction. Because $f^\star\in \operatorname{clip}\circ\mathcal F_B$ whenever $B\geq B_R$, standard ERM symmetrization and contraction applied to the class used by the estimator imply

$ \mathbb{E}\left[|\widehat{f}{\mathrm{MLP}}-f^\star|{L^2(P_x)}^2 \right] \leq8, \mathfrak R_n(\mathcal F_B). $

Finally, $\mathbb{E}|x|_2^2=R^2+K$, and Jensen's inequality yields the stated result.

B.2.2 The Ground-truth AGOP Removes the Routing Block

Define the ground-truth population AGOP across all input coordinates by

$ M :=\mathbb{E}\left[\nabla_x f^\star(x)\nabla_x f^\star(x)^\top \right]. $

########## {caption="Lemma 12: The AGOP kernel loses the cluster identity"}

Let $P_V:=\sum_{c=1}^K v_cv_c^\top$ be the orthogonal projector onto the predictive subspace, let $G\sim\mathcal N(0, 1)$, and set $\alpha:=\mathbb{E}[g'(G)^2]=\mathbb{P}(0<G<1)>0$. Then

$ M=\frac{\alpha}{K}P_V. $

Moreover, every measurable function $h$ of $\sqrt Mx$ satisfies

$ \mathbb{E}\left[\bigl(h(\sqrt Mx)-f^\star(x)\bigr)^2\right] \geq \left(1-\frac1K\right)\operatorname{Var}(g(G)). $

Proof: On the support of cluster $c$, the regression function is $f^\star(x)=g(\langle x, v_c\rangle)$ and hence, almost everywhere,

$ \nabla_x f^\star(x) =g'(\langle x, v_c\rangle)v_c. $

Averaging the corresponding outer product over the input and the uniform cluster index gives

$ M =\frac1K\sum_{c=1}^K \mathbb{E}\left[g'(\langle x, v_c\rangle)^2\mid c\right] v_cv_c^\top =\frac{\alpha}{K}\sum_{c=1}^K v_cv_c^\top =\frac{\alpha}{K}P_V. $

Therefore $\sqrt M=\sqrt{\alpha/K}, P_V$. In particular, $\sqrt Mx$ contains the predictive coordinates $S_j:=\langle x, v_j\rangle$, $j\in[K]$, but none of the routing coordinates. Under the data distribution, $S=(S_1, \ldots, S_K)\sim\mathcal N(0, I_K)$ and is independent of the uniform cluster index $c$, while $f^\star(x)=g(S_c)$. Thus, conditional on $\sqrt Mx$,

$ \mathbb{E}[f^\star(x)\mid\sqrt Mx] =\frac1K\sum_{j=1}^K g(S_j). $

The conditional-expectation characterization of squared-loss regression implies that every $h(\sqrt Mx)$ has risk at least

$ \begin{aligned} \mathbb{E}!\left[\operatorname{Var} \bigl(f^\star(x)\mid\sqrt Mx\bigr)\right] &=\mathbb{E}!\left[\frac1K\sum_{j=1}^K g(S_j)^2 -\left(\frac1K\sum_{j=1}^K g(S_j)\right)^2 \right]\ &=\left(1-\frac1K\right)\operatorname{Var}(g(G)). \end{aligned} $

B.2.3 An All-orders Lower Bound for Kernel Methods and Regularized RFM

We next study the full family in Section 4.2. Let $P_R:=I_d-P_V$ denote the projector onto the routing subspace. By Lemma 12,

$ M_\rho =\rho P_R+\left(\rho+\frac{\alpha}{K}\right)P_V. $

Thus both the standard metric $I_d$ and the regularized RFM metric $M_\rho$ commute with every rotation of the predictive subspace. The lower bound below is uniform over $\rho\geq0$, applies to any rotationally invariant base kernel $\mathcal K$, and is independent of the empirical ridge parameter. This rotational-invariance obstruction is closely related to prior lower bounds for kernel methods in high dimensions ([3]).

Let $\operatorname{He}_r$ denote the probabilists' Hermite polynomial and $h_r=\operatorname{He}_r/\sqrt{r!}$ its normalized version. Write

$ g(t)=\sum_{r=0}^\infty\widehat{g}_r h_r(t), \qquad \widehat{g}_r=\mathbb{E}[g(G)h_r(G)]. $

########## {caption="Lemma 13: Hermite coefficients of the clipped ramp"}

For every $r\geq2$,

$ \widehat{g}r =\frac{ \varphi(0)\operatorname{He}{r-2}(0) -\varphi(1)\operatorname{He}_{r-2}(1) }{\sqrt{r!}}. $

In particular, the sequence $(\widehat{g}r){r\geq2}$ is not eventually zero, and therefore

$ \sum_{r>r_0}\widehat{g}_r^{, 2}>0 \qquad\text{for every finite }r_0. $

Proof: For $a\in\mathbb{R}$ and $r\geq2$, the identity

$ \frac{d}{dx}\bigl(\operatorname{He}_{r-1}(x)\varphi(x)\bigr) =-\operatorname{He}_r(x)\varphi(x) $

gives, by integration by parts,

$ \begin{aligned} \mathbb{E}\left[\operatorname{He}r(G)(G-a)+\right] &=\int_a^\infty\operatorname{He}r(x)(x-a)\varphi(x), dx\ &=\int_a^\infty\operatorname{He}{r-1}(x)\varphi(x), dx\ &=\varphi(a)\operatorname{He}_{r-2}(a). \end{aligned} $

Apply this identity at $a=0$ and $a=1$ and divide by $\sqrt{r!}$.

To show that the coefficients are not eventually zero, define

$ a_m =\varphi(0)\operatorname{He}_m(0) -\varphi(1)\operatorname{He}_m(1). $

The Hermite generating function gives

$ \sum_{m=0}^\infty a_m\frac{t^m}{m!} =e^{-t^2/2}\bigl(\varphi(0)-\varphi(1)e^t\bigr). $

The right-hand side is not a polynomial, so $(a_m)$ is not eventually zero. The same is true of $(\widehat{g}r){r\geq2}$. Parseval's identity then shows that every finite Hermite tail has strictly positive squared mass.

For $r\geq2$, let $\mathcal V_{r, K}$ be the traceless order- $r$ Gaussian-chaos space. Equivalently, under the standard isometry between the $r$ th Gaussian chaos and symmetric order- $r$ tensors, $\mathcal V_{r, K}$ corresponds to the kernel of the tensor trace map. It is the irreducible $O(K)$ representation of harmonic degree- $r$ polynomials.

########## {caption="Lemma 14: Harmonic dimension and directional energy"}

The dimension of $\mathcal V_{r, K}$ is

$ H_{r, K} =\binom{K+r-1}{r}-\binom{K+r-3}{r-2}. $

Let $P_{r, K}$ be orthogonal projection onto $\mathcal V_{r, K}$. For every unit vector $q\in\mathbb{R}^K$,

$ \left|P_{r, K}h_r(q^\top z)\right|{L^2}^2 =\alpha{r, K}, $

where, writing $r=2m$ or $r=2m+1$,

$ \alpha_{2m, K} =\prod_{j=0}^{m-1} \frac{K+2j-1}{K+2m+2j-2}, $

and

$ \alpha_{2m+1, K} =\prod_{j=0}^{m-1} \frac{K+2j-1}{K+2m+2j}. $

For every fixed $r$, $H_{r, K}=\Theta_r(K^r)$ and $\alpha_{r, K}\to1$ as $K\to\infty$.

Proof: The dimension of symmetric order- $r$ tensors is $\binom{K+r-1}{r}$. The trace map onto symmetric order- $(r-2)$ tensors is surjective, and its kernel is the traceless subspace. This gives the stated dimension.

Under the Gaussian-chaos/tensor isometry, $h_r(q^\top z)$ corresponds to $q^{\otimes r}$. By rotational invariance, take $q=e_1$. The orthogonal projection of the homogeneous polynomial $x_1^r$ onto harmonic degree $r$ is

$ \mathcal H_r[x_1^r] =\sum_{j=0}^{\lfloor r/2\rfloor} \frac{(-1)^j\Gamma(r-j+K/2-1)} {4^j j!, \Gamma(r+K/2-1)} |x|^{2j}\Delta^j x_1^r. $

Since

$ \Delta^j x_1^r=\frac{r!}{(r-2j)!}x_1^{r-2j}, $

the coefficient of $x_1^r$ in this harmonic projection is

$ \sum_{j=0}^{\lfloor r/2\rfloor} \frac{(-1)^j r!, \Gamma(r-j+K/2-1)} {4^j j!(r-2j)!, \Gamma(r+K/2-1)}. $

In the symmetric-tensor inner product, this coefficient equals $\langle e_1^{\otimes r}, P_{r, K}e_1^{\otimes r}\rangle$, which equals the squared norm of the projection. Simplifying the finite sum gives the two product formulas displayed above. The asymptotic statements follow immediately from those formulas and the dimension expression.

########## {caption="Lemma 15: Invariant random-subspace bound"}

Let $V$ be a $D$-dimensional irreducible orthogonal representation of a compact group. Let $S\subseteq V$ be a random subspace with invariant law and $\dim S\leq N$ almost surely. Then, for every fixed $v\in V$,

$ \mathbb{E}\operatorname{dist}(v, S)^2 \geq\left(1-\frac ND\right)_+|v|_2^2. $

Proof: Let $\Pi_S$ be orthogonal projection onto $S$. Invariance implies that $\mathbb{E}\Pi_S$ commutes with the group action. Schur's lemma gives $\mathbb{E}\Pi_S=\beta I_V$. Taking traces yields $\beta D=\mathbb{E}\dim S\leq N$. Therefore

$ \mathbb{E}\operatorname{dist}(v, S)^2 =|v|_2^2-\mathbb{E}|\Pi_Sv|_2^2 =(1-\beta)|v|2^2 \geq\left(1-\frac ND\right)+|v|_2^2. $

########## {caption="Lemma 16: All-orders lower bound for regularized AGOP kernels"}

For every rotationally invariant kernel $\mathcal K$, every sample size $n$, every $\rho\geq0$, and every empirical ridge parameter $\lambda_n\geq0$, both estimators from Section 4.2 satisfy

$ \mathbb{E}\left[|\widehat{f}-f^\star|{L^2(P_x)}^2 \right] \geq \sum{r=2}^\infty \widehat{g}r^{, 2}\alpha{r, K} \left(1-\frac{n}{H_{r, K}}\right)+, \qquad \widehat{f}\in \left{\widehat{f}{\mathrm{Kernel}}, \widehat{f}_{\mathrm{RFM}}\right}. $

Proof: Let $A=I_d$ for $\widehat{f}{\mathrm{Kernel}}$ and $A=M\rho$ for $\widehat{f}_{\mathrm{RFM}}$. By the representer theorem, for a fixed training sample the predictor belongs to the span of the $n$ transformed-kernel sections centered at the training inputs. For a test point in cluster $c$, restrict these sections to that cluster's support and write them as functions of its predictive argument $\xi\in\mathbb{R}^K$:

$ \psi_{c, i}^{A}(\xi) :=\mathcal K!\left(\sqrt A x_i, \sqrt A \left(\mu_c+\begin{bmatrix}0\\xi\end{bmatrix}\right) \right), \qquad 1\leq i\leq n. $

At Hermite order $r\geq2$, define

$ S_{c, r} :=\operatorname{span}\left{ P_{r, K}\psi_{c, i}^{A}:1\leq i\leq n \right} \subseteq\mathcal V_{r, K}. $

Then $\dim S_{c, r}\leq n$. The predictive coordinates of the training inputs are standard Gaussian, $A$ commutes with predictive rotations, and $\mathcal K$ is rotationally invariant. Hence the law of $S_{c, r}$ is invariant under the $O(K)$ action on $\mathcal V_{r, K}$ for both estimators and every $\rho\geq0$.

The order- $r$ target component in cluster $c$ is

$ \gamma_{c, r} =\widehat{g}r P{r, K}h_r(e_c^\top\xi), $

whose squared norm is $\widehat{g}r^{, 2}\alpha{r, K}$ by Lemma 14. The projected predictor lies in $S_{c, r}$. Since distinct Gaussian-chaos orders are orthogonal, for every finite $L$,

$ |\widehat{f}c-g(e_c^\top!\cdot)|{L^2}^2 \geq \sum_{r=2}^L\operatorname{dist}(\gamma_{c, r}, S_{c, r})^2. $

Apply Lemma 15 and then let $L\to\infty$ by monotone convergence. This gives, for every cluster $c$,

$ \mathbb{E}\left[|\widehat{f}c-g(e_c^\top!\cdot)|{L^2}^2 \right] \geq \sum_{r=2}^\infty \widehat{g}r^{, 2}\alpha{r, K} \left(1-\frac{n}{H_{r, K}}\right)_+. $

Averaging over the uniform test cluster proves the result.

########## {caption="Lemma 17: Polynomial sample budgets leave a nonzero kernel/RFM error"}

Fix $A<\infty$, let $n_K=O(K^A)$, and let $\rho_K\geq0$ be any sequence. Then, for each $\widehat{f}\in {\widehat{f}{\mathrm{Kernel}}, \widehat{f}{\mathrm{RFM}}}$,

$ \liminf_{K\to\infty} \mathbb{E}\left[|\widehat{f}-f^\star|{L^2(P_x)}^2 \right] \geq \sum{\substack{r\geq2\r>A}}\widehat{g}_r^{, 2}>0. $

Proof: Fix an order $r>A$. By Lemma 14, $H_{r, K}=\Theta_r(K^r)$ and $\alpha_{r, K}\to1$, so

$ \frac{n_K}{H_{r, K}}\longrightarrow0. $

Apply Lemma 16. For any finite collection of orders satisfying $r>A$, each corresponding summand converges to $\widehat{g}_r^{, 2}$. Taking the lower limit and then increasing the finite collection gives, by monotone convergence,

$ \liminf_{K\to\infty} \mathbb{E}\left[|\widehat{f}-f^\star|{L^2(P_x)}^2 \right] \geq \sum{\substack{r\geq2\r>A}}\widehat{g}_r^{, 2}. $

The right-hand side is strictly positive by Lemma 13.

Proof of Theorem 2: Under Section 4.2, the parameterized construction above specializes to $R=1$. Lemma 10 then represents $f^\star$ exactly with width $2K$ and Frobenius budget

$ B_1=K(1+\sqrt2)\leq 3K. $

Thus $f^\star\in\operatorname{clip}\circ\mathcal F_{3K}$, so Lemma 11, applied with the same budget $B=3K$ used by $\widehat{f}_{\mathrm{MLP}}$ in Section 4.2, gives

$ \begin{aligned} \mathbb{E}\left[|\widehat{f}{\mathrm{MLP}}-f^\star|{L^2(P_x)}^2 \right] &\leq 48K\sqrt{\frac{K+1}{n}}\ &\leq 48\sqrt2, \frac{K^{3/2}}{\sqrt n}. \end{aligned} $

This is the first claim of the theorem with the universal constant $C=48\sqrt2$.

For the second claim, fix any $A<\infty$, any sequence $n_K=O(K^A)$, and any sequence $\rho_K\geq0$. Applying Lemma 17 separately to $\widehat{f}{\mathrm{Kernel}}$ and $\widehat{f}{\mathrm{RFM}}$ yields

$ \liminf_{K\to\infty} \mathbb{E}\left[|\widehat{f}-f^\star|{L^2(P_x)}^2 \right] \geq \sum{\substack{r\geq2\r>A}}\widehat{g}r^{, 2}>0, \qquad \widehat{f}\in \left{\widehat{f}{\mathrm{Kernel}}, \widehat{f}_{\mathrm{RFM}}\right}. $

The lower bound is uniform over $\rho_K$ and over the empirical ridge parameters allowed in Section 4.2. This is exactly the pair of kernel and RFM conclusions in the main-text statement.

References

Section Summary: This section compiles a bibliography of research papers on the theoretical analysis of neural networks, with a strong emphasis on their training dynamics, feature learning capabilities, and comparisons to kernel-based methods. The cited works explore concepts like lazy training regimes, the neural tangent kernel, and conditions for effective representation learning in overparameterized models, drawing from major venues such as NeurIPS, ICML, and related journals. Spanning publications from 2013 onward, the references reflect active scholarly interest in when and how gradient-based optimization enables neural networks to surpass simpler mathematical alternatives.

[1] Allen-Zhu, Zeyuan and Li, Yuanzhi (2019). What Can ResNet Learn Efficiently, Going Beyond Kernels?. In Advances in Neural Information Processing Systems.

[2] Ghorbani et al. (2019). Limitations of Lazy Training of Two-layers Neural Network. In Advances in Neural Information Processing Systems.

[3] Behrooz Ghorbani et al. (2020). When Do Neural Networks Outperform Kernel Methods?. In Advances in Neural Information Processing Systems. pp. 14820–14830. https://proceedings.neurips.cc/paper/2020/hash/a9df2255ad642b923d95503b9a7958d8-Abstract.html.

[4] Bengio et al. (2013). Representation Learning: A Review and New Perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence. 35(8). pp. 1798–1828. doi:10.1109/TPAMI.2013.50.

[5] Jacot et al. (2018). Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems.

[6] Lee et al. (2019). Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. In Advances in Neural Information Processing Systems.

[7] Lénaïc Chizat et al. (2019). On Lazy Training in Differentiable Programming. In Advances in Neural Information Processing Systems. pp. 2937–2947. https://proceedings.neurips.cc/paper/2019/hash/ae614c557843b1df326cb29c57225459-Abstract.html.

[8] Woodworth et al. (2020). Kernel and Rich Regimes in Overparametrized Models. In Proceedings of Thirty Third Conference on Learning Theory. pp. 3635–3673.

[9] Geiger et al. (2020). Disentangling Feature and Lazy Training in Deep Neural Networks. Journal of Statistical Mechanics: Theory and Experiment. 2020(11). pp. 113301. doi:10.1088/1742-5468/abc4de.

[10] Yang, Greg and Hu, Edward J. (2021). Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks. In Proceedings of the 38th International Conference on Machine Learning. pp. 11727–11737.

[11] Bach, Francis (2017). Breaking the Curse of Dimensionality with Convex Neural Networks. Journal of Machine Learning Research. 18(19). pp. 1–53.

[12] Soltanolkotabi, Mahdi (2017). Learning ReLUs via Gradient Descent. In Advances in Neural Information Processing Systems.

[13] Yehudai, Gilad and Shamir, Ohad (2019). On the Power and Limitations of Random Features for Understanding Neural Networks. In Advances in Neural Information Processing Systems.

[14] Wei et al. (2019). Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel. In Advances in Neural Information Processing Systems.

[15] Ghorbani et al. (2020). When Do Neural Networks Outperform Kernel Methods?. In Advances in Neural Information Processing Systems. pp. 14820–14830.

[16] Malach et al. (2021). Quantifying the Benefit of Using Differentiable Learning over Tangent Kernels. In Proceedings of the 38th International Conference on Machine Learning. pp. 7379–7389.

[17] Abbe et al. (2022). The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks. In Proceedings of the 35th Conference on Learning Theory. pp. 4782–4887.

[18] Abbe et al. (2023). SGD Learning on Neural Networks: Leap Complexity and Saddle-to-Saddle Dynamics. In Proceedings of the 36th Conference on Learning Theory. pp. 2552–2623.

[19] Ba et al. (2022). High-Dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation. In Advances in Neural Information Processing Systems. pp. 37932–37946.

[20] Damian et al. (2022). Neural networks can learn representations with gradient descent. In Conference on learning theory. pp. 5413–5452.

[21] Dandi et al. (2024). How Two-Layer Neural Networks Learn, One (Giant) Step at a Time. Journal of Machine Learning Research. 25(349). pp. 1–65.

[22] Bruna, Joan and Hsu, Daniel (2025). Survey on algorithms for multi-index models. Statistical Science. 40(3). pp. 378–391.

[23] Hendrycks, Dan and Gimpel, Kevin (2016). Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415.

[24] Shazeer, Noam (2020). GLU Variants Improve Transformer. arXiv preprint arXiv:2002.05202.

[25] Radhakrishnan et al. (2024). Mechanism for Feature Learning in Neural Networks and Backpropagation-Free Machine Learning Models. Science. 383(6690). pp. 1461–1467. doi:10.1126/science.adi5639.

[26] Adityanarayanan Radhakrishnan et al. (2025). Linear Recursive Feature Machines Provably Recover Low-Rank Matrices. Proceedings of the National Academy of Sciences. 122(13). pp. e2411325122. doi:10.1073/pnas.2411325122. https://doi.org/10.1073/pnas.2411325122.

[27] Damian et al. (2025). The generative leap: Sharp sample complexity for efficiently learning gaussian multi-index models. arXiv preprint arXiv:2506.05500.

[28] Refinetti et al. (2021). Classifying High-Dimensional Gaussian Mixtures: Where Kernel Methods Fail and Neural Networks Succeed. In Proceedings of the 38th International Conference on Machine Learning. pp. 8936–8947. https://proceedings.mlr.press/v139/refinetti21b.html.

[29] Ben Arous et al. (2024). High-Dimensional SGD Aligns with Emerging Outlier Eigenspaces. In International Conference on Learning Representations. https://proceedings.iclr.cc/paper_files/paper/2024/hash/d10d6b28d74c4f0fcab588feeb6fe7d6-Abstract-Conference.html.

[30] Nanda et al. (2023). Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217.

[31] Gromov, Andrey (2023). Grokking modular arithmetic. arXiv preprint arXiv:2301.02679.

[32] Morwani et al. (2024). Feature emergence via margin maximization: case studies in algebraic tasks. In International Conference on Learning Representations. pp. 29077–29114.

[33] He et al. (2026). On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking. arXiv preprint arXiv:2602.16849.

[34] Frei et al. (2023). Random Feature Amplification: Feature Learning and Generalization in Neural Networks. Journal of Machine Learning Research. 24(303). pp. 1–49.

[35] Margalit Glasgow (2024). SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem. In The Twelfth International Conference on Learning Representations. https://openreview.net/forum?id=HgOJlxzB16.

[36] Tian, Yuandong (2020). Student Specialization in Deep Rectified Networks With Finite Width and Input Dimension. In Proceedings of the 37th International Conference on Machine Learning. pp. 9470–9480.

[37] Oostwal et al. (2021). Hidden Unit Specialization in Layered Neural Networks: ReLU vs. Sigmoidal Activation. Physica A: Statistical Mechanics and its Applications. 564. pp. 125517. doi:10.1016/j.physa.2020.125517.

[38] Zhu et al. (2025). How Gradient Descent Balances Features: A Dynamical Analysis for Two-Layer Neural Networks. In International Conference on Learning Representations. https://openreview.net/forum?id=25j2ZEgwTj.

[39] Jarvis et al. (2025). A Theory of Initialisation's Impact on Specialisation. In The Thirteenth International Conference on Learning Representations.

[40] Bills et al. (2023). Language Models Can Explain Neurons in Language Models. https://openaipublic.blob.core.windows.net/neuron-explainer/paper/index.html.

[41] Elhage et al. (2022). Softmax Linear Units. Transformer Circuits Thread.

[42] Bricken et al. (2023). Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits Thread. https://transformer-circuits.pub/2023/monosemantic-features/index.html.

[43] Templeton et al. (2024). Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet. Transformer Circuits Thread. https://transformer-circuits.pub/2024/scaling-monosemanticity/index.html.

[44] Chen et al. (2022). Towards Understanding the Mixture-of-Experts Layer in Deep Learning. In Advances in Neural Information Processing Systems. pp. 23049–23062. https://proceedings.neurips.cc/paper_files/paper/2022/hash/91edff07232fb1b55a505a9e9f6c0ff3-Abstract-Conference.html.

[45] Dikkala et al. (2023). On the Benefits of Learning to Route in Mixture-of-Experts Models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. pp. 9376–9396. doi:10.18653/v1/2023.emnlp-main.583. https://aclanthology.org/2023.emnlp-main.583/.

[46] Kawata et al. (2025). Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning. In Proceedings of the 42nd International Conference on Machine Learning. pp. 29390–29448. https://proceedings.mlr.press/v267/kawata25a.html.

[47] Min et al. (2024). Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization. In The Twelfth International Conference on Learning Representations.

[48] Neyshabur et al. (2015). Norm-Based Capacity Control in Neural Networks. In Proceedings of the 28th Conference on Learning Theory. pp. 1376–1401.

[49] Golowich et al. (2018). Size-Independent Sample Complexity of Neural Networks. In Proceedings of the 31st Conference on Learning Theory. pp. 297–299.

[50] Bartlett, Peter L. and Mendelson, Shahar (2002). Rademacher and Gaussian Complexities: Risk Bounds and Structural Results. Journal of Machine Learning Research. 3. pp. 463–482.

[51] Ma et al. (2018). The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning. In Proceedings of the 35th International Conference on Machine Learning. pp. 3325–3334.

[52] Ma, Siyuan and Belkin, Mikhail (2019). Kernel Machines that Adapt to GPUs for Effective Large Batch Training. Proceedings of Machine Learning and Systems. 1. pp. 360–373.

[53] Abedsoltan et al. (2023). Toward Large Kernel Models. In Proceedings of the 40th International Conference on Machine Learning. pp. 61–78.

[54] Abedsoltan et al. (2025). Fast Training of Large Kernel Models with Delayed Projections. In Advances in Neural Information Processing Systems.

[55] Soudry et al. (2018). The implicit bias of gradient descent on separable data. Journal of Machine Learning Research. 19(70). pp. 1–57.

[56] Ji, Ziwei and Telgarsky, Matus (2019). The implicit bias of gradient descent on nonseparable data. In Conference on learning theory. pp. 1772–1798.

[57] Lyu, Kaifeng and Li, Jian (2020). Gradient Descent Maximizes the Margin of Homogeneous Neural Networks. In International Conference on Learning Representations.

[58] Ji, Ziwei and Telgarsky, Matus (2020). Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems. 33. pp. 17176–17186.