cs.LGOct 19, 2026

A "feature ODE" describing the learning behavior of shallow MLPs on simple functions

Authors: Joseph Turnbull, Berkan Ottlik, Jamie Simon

Organizations: UC Berkeley · University of Pennsylvania · Imbue

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We study the gradient flow training dynamics of two-layer multi-layer perceptrons (MLPs) on anisotropic Gaussian data. Even in this simple setting, solving the dynamics in full is notoriously difficult, so we try something different. We propose a "feature ODE" that might be intuitively expected to describe the learning dynamics of a shallow MLP. We find that MLPs trained to learn simple target functions (low-order Hermite polynomials and staircase functions) actually roughly follow the loss trajectories from this feature ODE. This agreement is a surprise because of the size of the simplification and the fact that it does not come from a rigorous chain of approximations. We discuss this surprise and speculate on its reasons and uses.

Figures & tables

Figure 1: We suggest a simple ODE of scalar variables whose solution predicts loss trajectories for shallow MLPs learning simple functions of anisotropic data. These plots treat shallow MLPs with shifted SiLU activation functions ( σ(x)=SiLU(x+0.5)\sigma(x)=\text{SiLU}(x+0.5) ) learning various Hermite polynomial target functions data sampled x∼N(0,Γ){\bm{x}}\sim\mathcal{N}(0,{\bm{\Gamma}}) with Γ=diag(γi){\bm{\Gamma}}=\text{diag}(\gamma_{i}) , γi∝(i+6)−1.8\gamma_{i}\propto(i+6)^{-1.8} . We write zi≡xi/γiz_{i}\equiv{\bm{x}}_{i}/\sqrt{\gamma_{i}} for normalized coordinates. Left. MLP loss trajectories closely match the ODE for example linear and quadratic target functions. Right. Effective optimization times η⋅niter\eta\cdot n_{\text{iter}} required for learning to complete are well-predicted by this ODE across many such target functions. For quadratic terms of the form He2(zi)\text{He}_{2}(z_{i}) , the leading term is indicated in the plots ( zi2z_{i}^{2} ). See Appendix A for experimental details.
Figure 2: Comparison between a shallow MLP and deep polynomial regression (DPR). Left. A shallow MLP has nonlinearities applied elementwise, and each nonlinearity contains many polynomial orders. Different elements tend to different components of the preactivation. Right. In DPR, the nonlinear operation is a preactivation expansion which separately generates monomial interactions of the input vector.
Figure 3: The feature ODE predicts loss trajectories for shallow MLPs learning low-degree “staircase functions.” Predicted loss trajectories are staircases with a number of steps equal to the order of the target function. True loss curves for shallow MLPs roughly follow these staircases, with steps at roughly the predicted times. See Appendix A for experimental details.
Figure 4: The feature ODE can predict MLP loss dynamics despite neuron-level specialization in MLPs. We train a two-layer MLP on a two-term staircase f∗(z)=z1+z1z2f^{*}({\bm{z}})=z_{1}+z_{1}z_{2} and overlay the feature ODE prediction. Top. We plot the online loss of the MLP and of the feature ODE. Bottom. Each row of the heatmap tracks one neuron’s preactivation variance over training, and neurons are sorted by their variance at the end of training. The horizontal striations show neuron-level specialization: individual neurons carry distinctly different preactivation variances. Nevertheless, the feature ODE captures the MLP’s staircase loss.
Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: SFN tracks the feature ODE across all four targets. Online MSE for an SFN (colored) and the feature ODE (dashed black) over effective optimization time. The ODE recovers both the plateau lengths and the loss drops for each target, including the staircase structure of the two- and three-term staircase targets.
Figure 6: DPR tracks the feature ODE across all four targets. Same as Figure 5 but with DPR replacing the SFN. The agreement is tighter than for the SFN, consistent with DPR’s explicit polynomial features removing the random-feature concentration error.
Figure 7: Linear target f∗(z)=z1f^{*}({\bm{z}})=z_{1} . Top: MLP. Bottom: SFN. Neurons are sorted by preactivation variance at the end of training in the heatmaps. Loss, NFMii\sqrt{\text{NFM}_{ii}} (MLP) or AiiA_{ii} (SFN), and per-neuron preactivation variance. Both models track the feature ODE, and every neuron gains variance on the same timescale.
Figure 8: Pure-Hermite target f∗(z)=He2(z1)+He2(z2)f^{*}({\bm{z}})=\mathrm{He}_{2}(z_{1})+\mathrm{He}_{2}(z_{2}) . Top: MLP. Bottom: SFN. Neurons are sorted by preactivation variance at the end of training in the heatmaps. A specialized subpopulation of roughly 500 MLP neurons (top, right column) gains preactivation variance orders of magnitude above the bulk, drives the loss down earlier than the feature ODE predicts, and pins NFMii\sqrt{\text{NFM}_{ii}} below the ODE-predicted AiiA_{ii} . The SFN carries no such subpopulation and tracks the ODE.
Figure 9: Two-term staircase f∗(z)=z1+z1z2f^{*}({\bm{z}})=z_{1}+z_{1}z_{2} . Top: MLP. Bottom: SFN. Neurons are sorted by preactivation variance at the end of training in the heatmaps. A faint specialization band in the MLP heatmap accompanies a small late-time drift of NFM2,2\sqrt{\text{NFM}_{2,2}} above the ODE prediction. The staircase loss is otherwise well captured. The SFN matches the ODE throughout.
More in this topic

Neural Network Training Dynamics

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 499 papers

Explore similar work

CardsList
Apr 2, 2026

Plateaus, Optima, and Overfitting in Multi-Layer Perceptrons: A Saddle-Saddle-Attractor Scenario

Vanishing gradients and overfitting are central problems in machine learning, yet are typically analyzed in asymptotic regimes that obscure their dynamical origins. Here we provide a dynamical description of learning in multi-layer perceptrons (MLPs) via a minimal model inspired by Fukumizu and Amari. We show that training dynamics traverse plateau and near-optimal regions, both organized by saddle structures, before converging to an overfitting regime. Under suitable conditions on the data, this regime collapses to a single attractor modulo symmetry. Furthermore, for finite noisy datasets, convergence to the theoretical optimum is impossible, and the dynamics necessarily settle into an overfitting solution.

Alex Alì Maleknia, Yuzuru Sato

Department of Mathematics, Hokkaido University, Kita 12 Nishi 7, Kita-ku, Sapporo, Hokkaido 060-0812, Japan · Present address: IMAG, IROKO, Univ Montpellier, Inria, CNRS · RIES, Hokkaido University, Kita 12 Nishi 7, Kita-ku, Sapporo, Hokkaido 060-0812, Japan

Neural Network Training DynamicsMultilayer PerceptronsNeural Network GeneralizationGradient Descent Dynamics
AskPDF
Sep 30, 2026

The Conflict Between Logic and Memory: Training Conditions for Optimizer-Dependent Rule Acquisition

Optimizers can fit the same task while acquiring different generalizing relations. We study the training conditions governing these differences in single-hidden-layer ReLU networks, combining composite evidence tasks, parameter-level interventions, and a three-seed strict-parity scan. Our central finding is that nuisance-connected trainability reshapes both shared failures and relative optimizer advantages. In a nuisance-heavy task, all twenty tested optimizer configurations remain near chance on the hardest stage. Retaining every input but fixing nuisance-connected first-layer weights at initialization raises that stage's accuracy from approximately 50% to 70.56%, 68.47%, and 69.00% for momentum SGD, Adam, and Muon. Masking the same inputs only after full training does not recover this performance. On a separate pairwise-mode task, background freezing reduces Muon's rare-mode advantage over momentum SGD by 12.48 percentage points, while the target and mode frequencies remain fixed. Each intervention is evaluated under a common validation-selection protocol with condition-specific learning rates and checkpoints. A strict-parity sweep over orders 1--20 provides a complementary reference without spurious cues or extra nuisance coordinates: the optimizers separate at orders 9--11, then approach chance despite substantial remaining Bayes predictability. A mixed task establishes a recovery boundary, and CIFAR-10 supplies an external architecture comparison. Together, these findings connect optimizer comparison to the acquisition and use of specified relations, identifying permitted adaptation as a concrete training variable that changes what a fixed architecture learns.

Gongyue Zhang, Honghai Liu

Shallow Neural NetworksMultilayer PerceptronsNeural Network Training Dynamics
AskPDF
Aug 25, 2026

Revenge of Monosemanticity: Neuron Specialization as a New Form of Feature Learning in MLPs

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 representation. 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. We show that this ability to specialize gives MLPs a provable data-efficiency advantage over feature-learning methods based on a global low-dimensional representation.

Amirhesam Abedsoltan, Enric Boix-Adsera, Fivos Kalogiannis +1

Department of Computer Science and Engineering, UC San Diego · Department of Statistics and Data Science, The Wharton School, University of Pennsylvania · Halıcıoğlu Data Science Institute, UC San Diego

Multilayer PerceptronsRepresentation Learning
AskPDF