PaperlayerView exact revision on arXiv

Source-linked AI summary

Revenge of Monosemanticity: Specialized Neurons Improve Data Efficiency in MLPs

Amirhesam Abedsoltan, Enric Boix-Adsera, Fivos Kalogiannis, Mikhail Belkin

arXiv:2608.24007v1cs.LGstat.ML

TL;DR

The paper asks whether feature learning must recover one global low-dimensional representation in clustered regression. It analyzes MLPs on clustered data and finds that neurons specialize to local predictive features, giving MLPs a data-efficiency advantage when those features collectively span a high-dimensional space.

  • Problem

    Existing feature-learning theory largely focuses on recovering a single global low-dimensional predictive subspace, but clustered problems may lack such global structure despite being locally low-dimensional.

  • Method

    The paper studies MLPs on Gaussian-mixture regression data and analyzes how neurons specialize to cluster-specific predictive directions and functions.

  • Results

    MLPs preserve local feature-cluster associations and remain effective as cluster-specific directions span the ambient space, while global-representation feature learning diminishes with cluster count.

  • Takeaways & Limitations

    Neuron specialization is a distinct form of feature learning that lets standard MLPs efficiently handle settings where a single global low-dimensional representation is insufficient.

  • Takeaways & Limitations

    The analysis uses Gaussian-mixture data with cluster-specific assumptions; derivatives are interpreted tangentially along the affine support components.

Abstract

from arXiv · show

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.

1 Revisiting Feature Learning

The paper broadens feature learning beyond recovering one global low-dimensional predictive subspace: on clustered data, MLPs can learn local predictive features together with their cluster associations. This specialization can preserve performance as cluster-specific directions span a high-dimensional ambient space.

  • Feature learning traditionally adapts representations to discover low-dimensional predictive structure rather than using a fixed representation.In the standard picture, supervision reduces learning from ambient dimension d to intrinsic dimension r.
  • In clustered regression, prediction can be low-dimensional within each cluster even when no useful global low-dimensional predictive subspace exists.Different clusters may have distinct predictive directions and link functions.
  • MLPs can jointly discover cluster structure and learn the predictive functions relevant to each cluster.This constitutes a form of feature learning beyond recovering a single global subspace.
  • A substantial fraction of trained MLP neurons become monosemantic, aligning predominantly with one cluster-specific predictive direction.The resulting implicit routing resembles mixture-of-experts behavior without an explicit router or expert decomposition.
  • Multiplicative-gating architectures such as ReGLU and SwiGLU can substantially improve sample efficiency in this clustered setting, although specialization also occurs with ReLU and GeLU.The specialization itself is observed across these activation choices.
  • As the number of clusters grows, specialization gives MLPs increasingly favorable sample-complexity scaling even when cluster directions collectively span the ambient space.Compared with global-representation feature learning such as RFM, MLPs preserve local feature-cluster associations as global low-dimensional structure disappears.

2 Preliminary

The paper studies regression data generated from a clustered Gaussian-mixture model with cluster-specific predictive directions and nonlinear link functions. The response includes independent additive Gaussian noise and is evaluated using mean squared error.

  • The experiments and theory use a Gaussian mixture model with K clusters and well-separated cluster means.Each cluster has its own mean µc and covariance matrix Σc.
  • Within each cluster, the response follows a single-index model using a unit predictive direction vc and a potentially cluster-specific nonlinear link function gc.The model also includes additive Gaussian noise independent of the covariates and cluster index.
  • The regression loss used throughout the paper is mean squared error, with f denoting the learned predictor.

3 Experiments

Across three experiments, MLPs learn cluster-dependent predictive structure through specialized neurons and first-layer representations, without being given cluster identities. This specialization supports strong performance as clusters multiply and enables local predictors to approach cluster-aware oracles.

  • Experiments: MLPs learn both cluster structure and cluster-specific predictive functions from input–response pairs, extending feature learning beyond a single global low-dimensional subspace.The experiments cover specialized neurons, joint cluster-function learning, and extracting cluster structure for separate local predictors.
  • 3.1 MLPs Develop Specialized First-layer Neurons: A significant fraction of active first-layer neurons become monosemantic, aligning predominantly with one cluster-specific predictive direction.Selected GELU neurons concentrate their weights on one predictive coordinate, with little weight on cluster-position and nuisance coordinates.
  • 3.1 MLPs Develop Specialized First-layer Neurons: In Figure 1, the specialization experiment uses K = 10 clusters in d = 40, with cluster centers and predictive directions placed on separate coordinate axes.This construction makes specialization visible as a large weight on the corresponding predictive coordinate.
  • 3.2 MLPs Jointly Learn Cluster Structure and Cluster-specific Predictive Functions: As the number of clusters grows, MLPs outperform other methods and approach oracle baselines that receive true cluster identities and fit separate within-cluster predictors.ReLU substantially outperforms global RFM when many clusters are present, while ReGLU improves further through separate gate and value branches.
  • 3.2 MLPs Jointly Learn Cluster Structure and Cluster-specific Predictive Functions: Across K ∈{1, 2, 10, 50}, GELU behaves similarly to ReLU, while SwiGLU shows the same growing-cluster advantage as ReGLU.These additional-activation results were obtained across five seeds on the same mixed-link data model.
  • 3.3 Trained MLPs Encode Cluster Structure: Clustering trained ReLU first-layer representations and fitting local Laplace or RFM predictors substantially improves both methods over global counterparts, with local RFM approaching the cluster-aware oracle.At test time, inputs are routed to the predictor associated with their nearest learned centroid.

4 Theory

The theory shows that MLP neurons specialize to cluster-specific predictive directions and that this local structure yields a sample-complexity advantage over globally represented alternatives.

  • Sample-complexity separation: As the number of clusters grows, MLPs outperform standard kernel methods and RFM in sample complexity because those methods cannot compute specialized features for each cluster.The comparison gives RFM the ground-truth population AGOP, isolating the limitation of using one global feature metric.
  • Neuron specialization: MLPs specialize to cluster-specific predictive directions when trained on well-separated Gaussian mixture data.The result is established under small initialization for the two-layer ReLU population-gradient-flow setup.
  • Neuron specialization: Every neuron specializes, up to orientation and a vanishing O(R^-1) routing component, to one cluster’s predictive direction.The selected cluster labels are independent and uniformly distributed across clusters.
  • Theoretical setup: The theory analyzes a two-layer MLP fit by Frobenius-constrained empirical risk minimization on Gaussian mixtures with noiseless routing.The routing block reveals cluster identity exactly, while within-cluster responses depend on cluster-specific predictive directions.
  • Sample-complexity separation: The MLP estimator learns the cluster-structured data with polynomial sample complexity, whereas standard kernels and RFM fail under every polynomial sample-size scaling in the analyzed setting.The obstruction is that RFM remains constrained to a single global feature geometry rather than cluster-dependent specialization.

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

The experiments quantify specialization among active first-layer neurons using alignment and dominance measures, with coordinate-aligned synthetic data making cluster-specific features directly visible.

  • Experimental setup: The experiments use K = 10 clusters in d = 40 dimensions, with separate coordinate axes for cluster position and predictive directions.Coordinates 1:10 encode cluster position, coordinates 11:20 contain predictive directions, and coordinates 21:40 are nuisance coordinates.
  • Experimental setup: Models are one-hidden-layer ReLU, GELU, ReGLU, and SwiGLU networks of width 2,048 trained with Adam and validation-MSE checkpoint selection.Training uses batch size 2,048, cosine learning-rate decay, and zero weight decay.
  • Active-neuron analysis: Active neurons are the smallest leading set accounting for at least 99.9% of total importance.The threshold is a descriptive sparsification rule intended to exclude neurons with negligible influence on the output, not a statistical significance threshold.
  • Specialization measures: Specialization is measured by maximum absolute cosine similarity to cluster-specific predictive directions and by predictive-coordinate dominance.Aj > 1 is not used; instead, Aj thresholds include 0.71 and 0.90, while Dj > 1 and Dj > 2 measure coordinate dominance.
  • Observed specialization: A significant number of active first-layer neurons become monosemantic, with selected GELU neurons concentrating their weights on a single predictive coordinate.The dominance score complements cosine alignment when many small coefficients reduce the cosine measure despite one clearly dominant predictive coordinate.

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

The experiments use clustered regression data where cluster identity and cluster-specific predictive signals occupy disjoint subspaces, while response functions vary independently across clusters. MLPs and global or oracle kernel baselines are compared across cluster counts, sample sizes, architectures, and training procedures.

  • Data construction: The synthetic setting uses d = 20 with separate ten-dimensional cluster-identifying and predictive subspaces.Each cluster has a center in the identifying subspace and a unit predictive direction in the predictive subspace.
  • Data construction: K ∈{1, 2, 10, 50} clusters are sampled uniformly, with cluster-specific predictive directions and nonlinear response functions varying across clusters.Functions are sampled with replacement from normalized quadratic, sine, and hyperbolic-tangent responses.
  • Experimental protocol: The experiments use nested training sizes from a 200,000-example pool, separate 4,096-example test sets, and five seeds reporting means with one empirical standard deviation.Methods sharing a setting use identical training and test examples.
  • Experimental protocol: One-hidden-layer ReLU, GELU, ReGLU, and SwiGLU MLPs all use width 4,096 and trainable hidden weights, readout weights, and biases.Optimization uses Adam with batch size 4,096, cosine learning-rate decay, and zero weight decay.
  • Baselines: Global kernel baselines fit all examples jointly without cluster-specific predictors, whereas oracle methods receive true cluster identities and fit independent predictors within clusters.The global methods use cluster identities only to calibrate one shared bandwidth, while oracle MSE is pooled across test examples.
  • Results: At larger training sizes, ReGLU and especially SwiGLU generally outperform corresponding standard activations for K = 10 and K = 50.The gated-activation advantage is not uniform in the smallest-data regime.

B.1 Proof of Theorem 1

The proof of Theorem 1 is organized around lemmas characterizing specialized teacher-only trajectories, their cluster-selection distribution, finite-time approximation, and coverage, with dependencies summarized in a proof roadmap.

  • Proof structure: Theorem 1 is proved using Lemmas 3–9, whose statements and proofs appear after the main proof.Figure 7 records the principal logical dependencies and an additional use of Lemma 3 for smoothness in the tracking argument.
  • Trajectory characterization: Almost surely, each teacher-only trajectory selects a cluster label Jj and orientation τj ∈{±1} associated with a specialized direction.The selected labels and orientations arise from the self-selected teacher-only dynamics.
  • Coverage: The selected cluster labels are independent and uniform, and finite-time tracking can simultaneously approximate all trajectories within any δ > 0.Cluster coverage follows from the resulting selection distribution.

B.1.1 The Population-correlation Landscape

The population-correlation landscape is expressed through neuron correlations with cluster means and predictive directions. Its positive local maxima are one-cluster solutions, while mixed-cluster critical points are unstable.

  • Objective: The signed population-correlation objective maximizes a neuron’s correlation with the target over unit weight directions.The analysis uses correlations with routing coordinates and cluster-specific predictive coordinates.
  • Specialized maxima: For the cubic target, positive local maxima align with ±vc and use only an O(R−1) routing component for cluster placement.The Hermite calculation reduces the objective to clusterwise cubic terms.
  • Specialized maxima: Every positive constrained local maximum has exactly one nonzero predictive coordinate ρc, and its corresponding routing coordinate is also the only potentially active routing coordinate.If ρc = 0, the corresponding bc vanishes.
  • Landscape geometry: Any critical point with two nonzero predictive coordinates has a feasible direction of positive curvature and therefore cannot be a constrained local maximum.This establishes instability of mixed-cluster local maxima.
  • Landscape geometry: For sufficiently large R, each one-cluster branch has a unique strict local maximum, and all branches attain the same maximal value.The resulting maximizers are isolated and indexed by the active cluster and orientation.

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

Random initialization determines each neuron’s sign and cluster selection through early teacher-only dynamics. The resulting directional flow converges almost surely to specialized maxima, while independent neurons provide probabilistic cluster coverage.

  • Self-selected signs: The sign ζ is learned at very early training times rather than fixed in advance.The teacher-only dynamics begin with zero second-layer weight and determine the subsequent signed objective.
  • Almost-sure specialization: Mixed-cluster positive critical points are avoided almost surely because each has an unstable tangent direction and a measure-zero center-stable basin.The remaining positive critical points are the specialized maxima.
  • Almost-sure specialization: For absolutely continuous random initialization, the directional dynamics are a positive time reparameterization of spherical gradient ascent on Ψζ and converge almost surely to specialized maxima.This holds for sufficiently large R.
  • Cluster coverage: Independent uniformly initialized neurons select independent uniformly distributed clusters.Permutation symmetry makes all cluster-selection probabilities equal to 1/K.
  • Cluster coverage: A fixed cluster is missed with probability (1 −1/K)^m, yielding a union-bound estimate for coverage across clusters.The estimate follows from m independent neuron selections.

B.1.3 Tracking by the Full Small-initialization Dynamics

The analysis combines effectively independent neuron dynamics with small-initialization tracking to connect teacher-only trajectories to the full neural-network dynamics. It also constructs an exact specialized MLP and derives a path-norm-based generalization bound.

  • Tracking dynamics: Small-initialization analysis reduces early neural-network feature learning to effectively independent neuron dynamics up to rescaling.The proof studies individual-neuron trajectories and invokes prior work on early-time or small-initialization regimes.
  • Tracking dynamics: For every fixed T, full and teacher-only neuron trajectories remain close, with error bounded by L_T m ε^2.The comparison uses matching initial conditions, Lipschitz vector fields, and Grönwall’s inequality.
  • Tracking dynamics: ReLU homogeneity supplies an invariant parameterization that converts the network equations into the analyzed (u_j, ω_j) dynamics.The transformation w_j = ε cosh(u_j)ω_j is inserted into the gradient-flow equations.
  • Exact representation: An exact specialized MLP uses two hidden weights per cluster, with opposite output weights causing non-relevant cluster pairs to cancel.For cluster c, the hidden weights are v_c and v_c − R^-2 μ_c, with output weights +1 and −1.
  • Generalization: The exact construction yields a path-norm bound by summing the norms of the cluster-specific neuron pairs.The relevant hidden-weight norms are 1 and sqrt(1 + R^-2), and positive homogeneity permits balanced rescaling.
  • Generalization: For B ≥ B_R, empirical squared-loss minimization over clipped path-norm-bounded networks admits the stated Rademacher-based risk guarantee.The argument uses clipping, contraction, Cauchy–Schwarz, and standard ERM symmetrization.

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

The ground-truth AGOP captures predictive coordinates but removes routing information. Consequently, predictors based only on this representation face an irreducible risk lower bound.

  • AGOP construction: The ground-truth AGOP is defined as the population expectation of regression-function input gradients’ outer products.This constructs the feature metric from the ground-truth regression function rather than learned predictions.
  • Routing information: The AGOP kernel is rotationally structured within the predictive subspace and does not retain cluster-specific routing directions.The result is established through the clusterwise regression function and averaging over the uniform cluster index.
  • Routing information: The AGOP representation contains predictive coordinates S_j but no routing coordinates, while the cluster index is independent of S.On cluster c, the target is g(S_c), so the relevant predictive coordinate depends on unavailable cluster identity.
  • Risk consequence: Every predictor based only on the AGOP representation has risk at least the stated positive lower bound.The bound follows from the conditional-expectation characterization of squared-loss regression.

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

Rotational invariance prevents kernel methods and regularized RFM from efficiently representing the cluster-dependent higher-order target structure. The resulting lower bound remains nonzero when the sample budget grows only polynomially with dimension.

  • Rotational obstruction: Both the standard metric and regularized RFM metric are invariant under rotations of the predictive subspace.The resulting lower bound is uniform over regularization, rotationally invariant base kernels, and empirical ridge parameters.
  • High-order structure: The clipped-ramp target has infinitely many nonzero Hermite coefficients, so every finite Hermite tail carries strictly positive squared mass.This supplies the target-side mass needed for a positive aggregate lower bound.
  • High-order structure: For fixed order r, the harmonic-space dimension scales as Θ_r(K^r), while the directional-energy factor α_r,K approaches 1 as K grows.These asymptotics quantify why finite-dimensional sample spans capture little of high-order rotationally symmetric structure.
  • All-orders lower bound: For every rotationally invariant kernel, kernel and regularized-RFM estimators obey an all-orders lower bound on prediction error.The theorem covers every sample size, regularization level, and empirical ridge parameter.
  • All-orders lower bound: At each Hermite order, the estimator’s cluster-restricted predictor lies in a sample-generated subspace whose dimension is at most n.Rotational invariance makes the law of this subspace invariant under the orthogonal action on the harmonic chaos space.
  • Polynomial sample budgets: When n_K = O(K^A), choosing any order r > A leaves a nonzero kernel/RFM error lower bound for every regularization sequence.The conclusion follows by combining the high-order asymptotics with the all-orders bound.
Loading 2608.24007v1…