Research map / Architectures and efficiency / RNNs and normalization
Simplicity Bias of Two-Layer Networks beyond Linearly Separable Data
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Characterizes simplicity bias of two-layer networks trained by gradient flow from small initialization on non-linearly-separable data, proving feature clustering in a few directions independent of width and extreme simplicity bias on XOR-like data, with a MNIST-CIFAR10 domino experiment.
Problem
Simplicity bias (over-reliance on simple features, hurting OOD generalization) has only been theoretically characterized under strong assumptions like linear separability or one-dimensional data.
Approach
Approximates early gradient-flow dynamics from small initialization by a linearized system with a data-dependent function G whose global extrema determine the directions where prominent neurons cluster, independent of hidden width; for XOR-like data it proves extreme simplicity bias in the infinite-time limit via max-margin theory.
Key results
- On the XOR-like 2D experiment with m=4096 neurons, almost all neurons align with the four cluster directions (0, pi/2, pi, -pi/2) and the first-layer scale grows from 0.3 to 5.5.
- On the MNIST-CIFAR-10 domino setup, OOD accuracy peaks around 45% versus about 90% when training on plain CIFAR-10 with the same recipe.
- For the small-regularization logistic regression, OOD accuracy drops by approximately 8.13% ± 0.80% between epochs 64 and 256 (p-value around 1e-6).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train ResNet-18 on the domino data, where MNIST labels are perfectly correlated with CIFAR-10 labels | ResNet-18 · MNIST-CIFAR-10 domino OOD test | accuracy: peaking in accuracy at about 45% | Hurt |
| Train the same ResNet-18 recipe on plain CIFAR-10 instead of domino data | ResNet-18 · CIFAR-10 | accuracy: around 90% accuracy | Helped |
| Use MNIST labels for prediction on the domino test set | MNIST-CIFAR-10 domino · MNIST-CIFAR-10 domino test | accuracy: around 99% accuracy | Helped |
| Use ResNet-18 features from later training stages (epoch 256 rather than epoch 64) with small-regularization logistic regression | Logistic regression with small regularization on ResNet-18 features · MNIST-CIFAR-10 domino OOD test, validation part | accuracy: approximately 8.13% ± 0.80% | Hurt |
| Use later-stage features with unregularized logistic regression | Unregularized logistic regression on ResNet-18 features · MNIST-CIFAR-10 domino OOD test, validation part | accuracy | No clear effect |
| Use later-stage features with unregularized logistic regression | Unregularized logistic regression on ResNet-18 features · MNIST-CIFAR-10 domino OOD test | quadratic mean of regression coefficients: approximately 50 × bigger | Hurt |
Surprising findings
- OOD accuracy of features from later training stages drops, suggesting the network forgets random initial features (extreme simplicity bias), while unregularized regression instead relies on much larger weights.
Limitations
- Theory for XOR-like case relies on an unproven assumption about convergence of a 4-neuron network, validated only empirically.
- Smooth ReLU approximation and odd input dimension d assumed in main theorems; extension to even d is only conjectured.
- Image experiments are a single domino dataset setting and a qualitative test of the hypothesis.
Experimental setup
Datasets: synthetic XOR-like dataset in R^2 and R^d MNIST-CIFAR-10 domino dataset
Benchmarks: MNIST-CIFAR-10 domino OOD test set
Metrics: accuracy of logistic regression on OOD test set quadratic mean of regression coefficients neuron angle alignment neuron scale
Baselines: ResNet-18 with CIFAR-10 training recipe
Builds on: Lyu et al. 2021 (max-margin implicit bias of homogeneous networks) Maennel et al. 2018 (gradient descent aligns neurons with few directions) Shah et al. 2020 (simplicity bias) Pezeshki et al. 2021 (gradient starvation)
Closest papers by meaning
- Simplicity Bias in Overparameterized Machine Learning (AAAI 2024 · similarity 0.90)
- Simplicity Bias and Optimization Threshold in Two-Layer ReLU Networks (ICML 2025 · similarity 0.90)
- Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wild (CVPR 2025 · similarity 0.89)
- Simplicity Bias via Global Convergence of Sharpness Minimization (ICML 2024 · similarity 0.89)
- The Rich and the Simple: On the Implicit Bias of Adam and SGD (NeurIPS 2025 · similarity 0.88)
- Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.