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

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

    Apr 2, 2026Alex Alì Maleknia, Yuzuru SatoNeural Network Training DynamicsMultilayer Perceptrons

    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

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

    Sep 30, 2026Gongyue Zhang, Honghai LiuShallow Neural NetworksMultilayer Perceptrons

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

    Aug 25, 2026Amirhesam Abedsoltan, Enric Boix-Adsera, Fivos Kalogiannis +1Multilayer PerceptronsRepresentation Learning

    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