Revenge of Monosemanticity: Specialized Neurons Improve Data Efficiency in MLPs

1 Revisiting Feature Learning

• Unlike classical feature learning, which recovers one global low-dimensional predictive subspace, this mechanism handles tasks that are low-dimensional locally but not globally.Feature learning can reduce dependence from ambient dimension d to intrinsic dimension r when a global predictive subspace exists, whereas clustered tasks may require preserving multiple local associations.Source: §1 · pp. 2–3
• MLPs develop monosemantic neurons that align with individual cluster-specific predictive directions, implicitly learning both local features and where those features apply.This resembles mixture-of-experts routing but arises in a standard MLP without an explicit router or expert decomposition.Source: §1 · p. 3

1.3 Paper Structure

• Gated activations such as ReGLU and SwiGLU substantially improve sample efficiency over standard ReLU and GeLU, while the paper develops empirical and theoretical analyses of specialization.The paper’s empirical results compare MLPs with RFM as cluster count grows, and its theory analyzes specialization in simplified settings.Source: §1 · p. 3, §1.3 · p. 6

1.1 Discussion

• MLPs preserve cluster-specific predictive directions as cluster count grows, while RFM’s global feature-learning advantage diminishes when those directions span the ambient space.Each cluster remains one-dimensional, but collectively the predictive directions may fill the ambient space, eliminating a useful global low-dimensional subspace.Source: §1.1 · pp. 4–5

1.2 Prior Work

• Prior work studies global feature recovery, neuron specialization, monosemanticity, and explicit mixture-of-experts routing; this paper isolates a sample-complexity advantage for locally predictive neuron specialization in standard MLPs.The contrast is with adaptive global-representation methods such as the Recursive Feature Machine and architectures using explicit routers and separate experts.Source: §1.2 · pp. 5–6

• Data are drawn from a Gaussian mixture model with K clusters, whose index is uniformly sampled and whose clusters have means µc and covariances Σc.Source: §2 · p. 7
• Distinct clusters are chosen to be well separated, and responses follow a cluster-specific single-index model with additive Gaussian noise.Source: §2 · p. 7
• Each cluster has a unit-norm predictive direction vc and a potentially different nonlinear link function gc.Source: §2 · p. 7
• Regression performance is evaluated using mean squared error (MSE) for the learned predictor f.Source: §2 · p. 7

3.1 MLPs Develop Specialized First-layer Neurons

• Many active first-layer neurons become monosemantic, aligning with individual cluster-specific predictive directions and concentrating their weights on single predictive coordinates.This specialization is observed across standard and gated activations, with gate and value weights evaluated separately for gated architectures.Source: §3.1 · p. 7, Figure 1 · p. 8

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

• As cluster count increases, MLPs outperform other methods and approach oracle baselines without observing cluster identities; ReGLU further improves over ReLU through explicit gating.GELU behaves similarly to ReLU, while SwiGLU shows the same growing-cluster advantage as ReGLU across five seeds.Source: §3.2 · pp. 9–11, Figure 2 · p. 10

3.3 Trained MLPs Encode Cluster Structure

• Clustering trained MLP first-layer representations substantially improves local Laplace and RFM predictors, with local RFM consistently outperforming global RFM and approaching the cluster-aware oracle.The learned representation separates clusters and identifies cluster-specific predictive features that can be reused by a separate predictor.Source: Figure 4 · p. 12, §3.3 · p. 12

4.1 MLP Neurons Specialize When Learning on Gaussian Mixture Model Data

• The specialization theorem applies to two-layer ReLU population gradient flow on a symmetric K-cluster Gaussian mixture in d = 2K dimensions, with cluster means separated by R and cluster-specific predictive directions.The first K coordinates encode cluster separation, while the last K coordinates encode predictive directions; the clusters use isotropic noise and a cubic Hermite link.Source: §4.1 · pp. 13–14
• Under small initialization and well-separated Gaussian mixture data, randomly initialized MLP neurons specialize almost surely to one cluster’s predictive direction, up to orientation and an O(R^-1) routing component.The selected cluster labels are independent and uniformly distributed across clusters.Source: §4.1 · p. 14
• The specialization mechanism follows signed correlation dynamics: cubic Hermite structure makes one-cluster alignments positive local maxima, while mixed-cluster critical points have unstable directions.ReLU homogeneity reduces directional dynamics to gradient ascent, and strict-saddle avoidance yields almost-sure specialization.Source: §4.1 · p. 15

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

• The sample-complexity separation uses noiseless routing, Frobenius-constrained MLP ERM, and rotationally invariant kernels or RFM supplied with the ground-truth AGOP.The MLP construction assigns two ReLU neurons per cluster, achieving zero approximation error within the Frobenius budget B = 3K; kernel lower bounds follow from rotational invariance and high-dimensional harmonic subspaces.Source: §4.2 · pp. 15–18
• MLPs learn cluster-structured data with polynomial sample complexity, while standard kernel methods and RFM fail under every polynomial sample-size scaling because they cannot compute distinct specialized directions for each cluster.The comparison gives RFM the ground-truth population AGOP, isolating the limitation of using one global feature metric.Source: §4.2 · p. 17

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

• The experiments use 10 clusters in 40 dimensions, with cluster positions, cluster-specific predictive directions, and nuisance coordinates assigned to separate coordinate ranges.Coordinates 1:10 encode cluster position, 11:20 contain predictive directions, and 21:40 are nuisance coordinates.Source: §A.1 · p. 26
• Training, validation, and test sets contain 100,000, 4,096, and 2,048 independently generated examples, respectively.The same normalized second-Hermite link is used across clusters with Gaussian observation noise.Source: §A.1 · p. 26
• The study trains width-2,048 one-hidden-layer ReLU, GELU, ReGLU, and SwiGLU networks with trainable biases and balanced feature-learning initialization.Optimization uses Adam, batch size 2,048, cosine learning-rate decay, and zero weight decay.Source: §A.1 · p. 26
• Models are evaluated at the checkpoint with the smallest validation MSE, with GELU, ReGLU, and SwiGLU trained for at most 20,000 steps and ReLU candidates for at most 40,000.The selected ReLU schedule is 0.004 → 0.0004, with its validation-selected checkpoint at step 37,000.Source: §A.1 · pp. 26–27
• Specialization analysis retains the smallest leading set of neurons accounting for at least 99.9% of total importance, excluding neurons with negligible output influence.For gated architectures, the same active-neuron set is analyzed separately for gate and value weights; the threshold is descriptive rather than statistical.Source: §A.1 · p. 27
• Cosine alignment A_j measures the largest absolute alignment with any cluster-specific predictive direction, so larger values indicate weight norm concentrated along one predictive direction.Absolute inner products treat alignment with v_c and −v_c as equivalent.Source: §A.1 · p. 27
• Dominance D_j complements cosine alignment by identifying whether a predictive coordinate is the largest weight coordinate and whether it exceeds every other coordinate by more than twofold.Appendix Fig. 5 shows cases where moderate cosine alignment coexists with a clearly dominant predictive coordinate because of many small coefficients.Source: §A.1 · pp. 27–28
• Fig. 1 reports active-neuron fractions above A_j thresholds 0.71 and 0.90 and dominance thresholds D_j > 1 and D_j > 2, separately for gated weights.Its lower panel displays highest-importance active GELU neurons satisfying A_j ≥ 0.90 using coordinate-wise absolute incoming weights.Source: §A.1 · p. 28

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

• The data-generating setup separates cluster-identifying coordinates from disjoint cluster-specific predictive signals, with both the predictive direction and nonlinear response function allowed to change across clusters.Experiments consider K ∈{1, 2, 10, 50} equally likely clusters and report means over five runs.Source: §A.2 · pp. 29–31
• With sufficient samples, gated architectures—particularly SwiGLU—achieve the lowest MLP errors for K = 10 and K = 50, although their advantage is not uniform in the smallest-data regime.The multiplicative-gating advantage is most pronounced when many cluster-dependent prediction rules must be learned.Source: §A.2 · pp. 33–34

A.3 Additional Details for Trained MLPs Encode Cluster Structures

• A trained ReLU MLP is converted into learned gates by clustering active-neuron contribution profiles, then fitting independent local Laplace or RFM predictors in gate-specific representations.The method supplies the number of clusters K but never observes true cluster assignments; test inputs are assigned to fixed training centroids before local prediction.Source: §A.3 · pp. 36–37
• The MLP assigns a test input using only its contribution profile, without using the response or true cluster identity.This contrasts with oracle predictors, which are diagnostic rather than fair deployable predictors because they receive true cluster identities.Source: §A.3 · p. 37
• The local RFM starts from a local Laplace predictor and performs three metric updates, recomputing bandwidths after each update; systems under 11,000 examples are solved directly and larger ones use EigenPro.The local Laplace bandwidth is initialized from the median pairwise distance within each gate, with ridge parameter 10−6; the same RFM procedure is used for global and oracle baselines.Source: §A.3 · p. 37
• Figure 4 compares the source ReLU MLP with global Laplace and RFM predictors fitted to all training examples, and MLP-gated local Laplace and RFM using extracted gates and projections.Source: §A.3 · p. 37
• The oracle comparison includes per-cluster Laplace and RFM predictors, while MLP-gated methods use only training input–response pairs and the known number K of gates.This establishes the information available to deployable MLP-gated methods versus diagnostic oracle references.Source: §A.3 · p. 37
• Every Figure 4 curve is the arithmetic mean over five independent runs using newly sampled datasets and independently initialized MLPs.The runs share the same generative model.Source: §A.3 · p. 38

B.1 Proof of Theorem 1

• Almost surely, each teacher-only trajectory selects a cluster label J_j and orientation τ_j, as established by Lemmas 7 and 5.The proof associates every trajectory with a cluster and an orientation before controlling its convergence.Source: §B.1 · p. 39
• Because finitely many teacher-only directions converge, for every δ > 0 a finite T_δ gives simultaneous control of all m trajectories.The resulting bound holds simultaneously for every j ∈ [m].Source: §B.1 · p. 39
• Lemma 8 establishes that the labels J_1, . . . , J_m are independent and uniform, together with the stated cluster-coverage probability.These properties complete the probabilistic part of the theorem’s proof.Source: §B.1 · p. 40
• Figure 7 organizes the proof of Theorem 1 around dependencies among Lemmas 3–9, with Lemma 3 also supplying smoothness for the tracking argument.The main proof invokes Lemmas 3–9, whose statements and proofs appear after it.Source: §B.1 · p. 39, Figure 7 · p. 41

B.1.1 The Population-correlation Landscape

• The analysis uses a signed population-correlation objective expressed through correlations with cluster means and predictive directions, exploiting ReLU homogeneity to apply it to small weights.Source: §B.1.1 · pp. 42, 44
• Every constrained local maximum with positive objective value has exactly one nonzero predictive coordinate, and routing coordinates for inactive predictive coordinates vanish.Source: §B.1.1 · pp. 45–46
• A feasible perturbation shows that any critical point with two nonzero predictive coordinates has positive curvature and therefore cannot be a constrained local maximum.Source: §B.1.1 · p. 46
• For a fixed active cluster and orientation, the one-cluster feasible branch is a connected open interval with a unique critical point that is a strict local maximum for sufficiently large R.Source: §B.1.1 · p. 47

The corresponding direction

• The corresponding one-dimensional objective is positive on its interval and has exactly one critical point, which is therefore its unique maximizer.Source: The corresponding direction · p. 49
• Because every global maximizer must lie on a one-cluster branch and all branches share the same maximal value, every corresponding ω⋆ is a global maximizer of Ψζ.Source: The corresponding direction · p. 49
• There are finitely many corresponding maximizers, one for each cluster c ∈[K] and orientation τ ∈ {±1}, making each isolated and a strict local maximum.Source: The corresponding direction · p. 49

Random initialization

• The teacher-only joint dynamics neglect neuron interactions, justified because the network is small at initialization.These dynamics arise when the second-layer weight is initially zero.Source: §B.1.2 · p. 50

Self-selected signs

• Almost surely, self-selected signs drive neurons to specialized maxima under sufficiently large R, without requiring signs to be fixed initially.The sign is learned at early training times, and the resulting dynamics follow spherical gradient ascent on Ψζ = ζΦ.Source: §B.1.2 · pp. 50, 52, 53

Cluster coverage

• Independent uniformly initialized neurons select clusters independently and uniformly from [K].This follows from permutation symmetry of the data, target, and self-selected vector field.Source: §B.1.2 · p. 53
• A fixed cluster is missed with probability (1 −1/K)^m, and a union bound gives the coverage estimate.Source: §B.1.2 · p. 53

B.1.3 Tracking by the Full Small-initialization Dynamics

• For every fixed T, sufficiently small initialization keeps full neuron trajectories within LT mε2 of teacher-only trajectories uniformly over 0≤t≤T.The comparison follows from bounded teacher-only dynamics, O(mε2) interaction terms, Lipschitz vector fields, and Grönwall’s inequality.Source: §B.1.3 · pp. 54–55
• The tracking argument relies on ReLU homogeneity to preserve an invariant and reduce the full dynamics to a teacher-only vector field plus small interaction errors.Smoothness on bounded intervals and the sphere ensures uniform control of the reduced trajectories for sufficiently small ε.Source: §B.1.3 · pp. 54–55
• The specialization mechanism is expected to persist under perturbations whose population-correlation objective remains small near specialized maxima and the strict-saddle regions of the cubic Hermite target.The pure Hermite target makes the correlation landscape exact, while the perturbation claim concerns persistence of maxima and their attraction basins.Source: §B.1.3 · p. 55

B.2 Proof of Theorem 2

• Theorem 2 is proved in the parameterized setting d = 2K and R ≥ 1, with the main-text theorem recovered as the special case R = 1.The construction uses cluster means µc, predictive directions vc, and covariances Σc specified for this setting.Source: §B.2 · pp. 55–56

B.2.1 Exact MLP Representation and Generalization

• An exactly specialized two-layer ReLU MLP represents the clustered regression function on the data support using paired neurons aligned with each cluster’s predictive direction.For nonmatching clusters, paired ReLUs cancel; the matching pair realizes the required cluster-specific function.Source: §B.2.1 · p. 56
• The exact construction uses two hidden weights per cluster, vc and vc − R−2µc, with output weights +1 and −1, yielding a path-norm bound.The two hidden-weight norms are 1 and 1 + R−2, respectively, and the bound sums contributions across clusters.Source: §B.2.1 · p. 57
• For any B ≥ BR, an empirical squared-loss minimizer over clipped FB satisfies the stated path-norm Rademacher generalization bound.The result follows by containing FB in a scaled absolutely convex hull of unit-norm ReLU atoms, then applying clipping, contraction, and ERM symmetrization.Source: §B.2.1 · pp. 57–58

The AGOP kernel loses the cluster identity

• The ground-truth AGOP kernel loses the cluster identity and contains predictive coordinates but no routing coordinates.Under the data distribution, the predictive coordinates are independent standard Gaussians and independent of the uniform cluster index, while the target depends on the coordinate selected by that index.Source: §B.2.2 · p. 58
• Every measurable function of the AGOP representation has risk at least the conditional-expectation lower bound induced by the missing cluster identity.This follows from the conditional-expectation characterization of squared-loss regression.Source: §B.2.2 · pp. 58–59

The AGOP construction

• The result follows by averaging the regression-gradient outer product over the input distribution and the uniform cluster index.On cluster c, the regression function is f⋆(x) = g(⟨x, v_c⟩), and the Gaussian derivative factor satisfies α := E[g′(G)^2] = P(0 < G < 1) > 0.Source: §B.2.2 · p. 58

Rotational invariance

• The rotational-invariance obstruction is closely related to prior lower bounds for high-dimensional kernel methods [Ghorbani et al., 2020a].Both the standard metric and regularized RFM metric commute with rotations of the predictive subspace.Source: §B.2.3 · p. 59

Hermite coefficients

• The clipped ramp has nonzero Hermite coefficients at arbitrarily high orders, so every finite Hermite tail retains strictly positive squared mass.This persistent high-order target energy supplies the positive limiting contribution in the kernel/RFM lower bound.Source: §B.2.3 · pp. 59–60

Harmonic spaces and invariant subspaces

• For fixed harmonic order r, the traceless Gaussian-chaos space has dimension Θ_r(K^r), while directional energy satisfies α_r,K → 1 as K → ∞.Invariant random-subspace bounds limit captured energy by the subspace dimension relative to the irreducible representation dimension.Source: §B.2.3 · pp. 60–61

All-orders lower bound

• The lower bound applies to every rotationally invariant kernel, every sample size, every ρ ≥ 0, and every empirical ridge parameter λ_n ≥ 0.Both standard-kernel and regularized-RFM estimators are constrained to n-dimensional, rotation-invariant spans in each harmonic order.Source: §B.2.3 · pp. 59, 61, 62
• Polynomial sample budgets leave kernel and regularized RFM estimators with strictly positive error, uniformly over regularization and empirical ridge parameters.For n_K = O(K^A) with finite A, the lower bound remains nonzero for both estimators and any sequence ρ_K ≥ 0.Source: §B.2.3 · pp. 62–63