ICML 2026

Quantifying and Optimizing Simplicity via Polynomial Representations

A general, measurable, and differentiable function-space notion of neural network simplicity.

Tianren Zhang* Xiangxin Li* Minghao Xiao* Guanyu Chen Feng Chen

Department of Automation, Tsinghua University

* Equal contribution

We probe a neural network along interpolation paths between real data points, fit its outputs with orthogonal polynomials, and quantify simplicity from the distribution of polynomial orders.

Abstract

Simplicity bias is a compelling explanation for neural-network generalization, yet “simplicity” is difficult to define in a way that is architecture-agnostic, computable on modern models, and useful for optimization. We introduce a polynomial representation of neural functions along data-dependent interpolation paths and summarize the resulting coefficient spectrum with Effective Degree (ED).

ED is defined in function space, predicts generalization more reliably than several parameter- and sharpness-based measures, and is fully differentiable. Used as a regularizer, it improves generalization across vision, language, vision-language fine-tuning, and reinforcement learning.

  • General: applicable across architectures, modalities, and tasks.
  • Quantifiable: stable and tractable on real, trained networks.
  • Optimizable: differentiable through the polynomial fitting procedure.

Method

1 Slice the function

Connect two data samples and evaluate the model along their interpolation path.

2 Fit a polynomial

Approximate each one-dimensional function slice with a Chebyshev basis.

3 Measure its order

Aggregate coefficient mass across orders to obtain a stable complexity score.

For samples \(x_1,x_2\sim\mathcal D\), we construct \(x(\alpha)=\alpha x_1+(1-\alpha)x_2\) and approximate the model output along the path with

\[ P(\alpha)=\sum_{k=0}^{K}c_kT_k(2\alpha-1). \]

Here, \(T_k\) denotes the degree-\(k\) Chebyshev basis function, and \(c_k\) measures how much of that order appears in the network output along the path.

Low-order coefficients dominate when the model changes smoothly along data paths; oscillatory behavior requires more high-order components. The normalized Effective Degree is

\[ \operatorname{ED}_{\mathrm{norm}}(P) =\frac{\sum_{k=0}^{K}|c_k|k} {\sum_{k=0}^{K}|c_k|}. \]

Averaging over paths yields a distribution-dependent measure of the neural function rather than its particular parameterization. Our theory shows that random interpolation paths preserve the algebraic degree of multivariate polynomials almost surely, providing a foundation for this one-dimensional reduction.

  • Data-dependent paths: measurements stay near the real data distribution rather than probing arbitrary input regions.
  • Numerically stable fitting: an orthogonal Chebyshev basis together with cosine-spaced sampling reduces the ill-conditioning of higher-order polynomial fits.
  • Scalable outputs: optional path-wise PCA retains the dominant output directions before fitting high-dimensional model predictions.
  • Differentiability: polynomial fitting is a least-squares problem, so its coefficients have analytic gradients with respect to network outputs and ED can be backpropagated into model parameters.

Results

Diagnostic Predict generalization gaps ED consistently tracks the gap between training and test performance across model and hyperparameter sweeps.
Dynamics Explain grokking ED reveals the transition from memorization to rule learning when standard complexity measures do not.
Optimization Improve generalization ED regularization improves vision, language, vision-language fine-tuning, and reinforcement learning.

ED predicts generalization

Across hyperparameter sweeps, higher-order behavior along data paths tracks a larger train-test gap more consistently than the tested baselines.

ResNet18 on CIFAR-10. ED has the strongest linear association with the generalization gap among the evaluated measures.
CLIP ViT-B/32 on ImageNet. ED remains positively correlated with the generalization gap while the other measures are negatively correlated in this setting.

ED explains grokking dynamics

During grokking, ED rises as the model memorizes the training set, peaks near the onset of the abrupt generalization transition, and then falls as the model discovers the underlying rule. Parameter norms grow monotonically, while sharpness does not clearly mark this phase change. ED therefore provides a function-space account of how a model moves from memorization to generalization.

Tracking grokking dynamics. ED rises during memorization, peaks at the generalization transition, and decreases afterward. The parameter norm and sharpness-based measures provide a less distinct transition signal.

ED can directly improve generalization

Because the polynomial fit is differentiable, we regularize the model with \(\mathcal L=\mathcal L_{\mathrm{task}}+\lambda\widehat{\operatorname{ED}}\).

Across the full study, ED regularization improves generalization in image and text classification, large-scale training, vision-language fine-tuning under distribution shift, and reinforcement learning. These experiments span training from scratch and adapting pre-trained models, showing that the same function-space principle transfers across architectures, modalities, and learning paradigms.

CLIP fine-tuning. ED regularization improves the in-distribution/out-of-distribution trade-off across interpolation coefficients.
Reinforcement learning. ED-regularized PPO improves learning across four environments.

Prior Work: Why Can Simplicity Help Representation Learning?

Simplicity is usually invoked to explain why a learned predictor generalizes. Our ICML 2025 paper, When Do Neural Networks Learn World Models?, offered a complementary explanation for representation learning: simplicity can help a model recover the hidden variables that generate its observations.

From a world-model perspective, representation learning is a latent identifiability problem. Many alternative latent variables can explain exactly the same observations, so reconstruction or prediction alone cannot identify the true representation. Across multiple tasks, however, the true latent variables support systematically lower-order task functions. A model with a low-order simplicity bias can therefore select the true world variables from many observationally equivalent representations. We formalized this mechanism with a Fourier-Walsh decomposition and proved identifiability under suitable task distributions. The work received an Outstanding Paper Award at the ICLR 2025 World Models Workshop, where it was presented orally.

As related work, the recent paper When Does LeJEPA Learn a World Model? by Klindt, LeCun, and Balestriero uses a related identifiability perspective and a Hermite-polynomial decomposition. The shared idea is to decompose functions by order and analyze how training objectives favor simpler components.

This project follows the shared direction of these theoretical works by turning low-order simplicity from a theoretical explanation for representation learning into a quantity that can be measured and optimized in real neural networks.

Limitations

Simpler is not always better. If a low-complexity feature is a non-robust shortcut, ED may favor it. ED measures function complexity, not causal correctness, semantic validity, or fairness. It is also a distribution-dependent proxy based on sampled paths, and the additional interpolation points introduce training overhead.

Citation

@inproceedings{zhang2026quantifying,
  title     = {Quantifying and Optimizing Simplicity via Polynomial Representations},
  author    = {Zhang, Tianren and Li, Xiangxin and Xiao, Minghao and Chen, Guanyu and Chen, Feng},
  booktitle = {International Conference on Machine Learning},
  year      = {2026}
}