📊 statistics

Generalization in Nonlinear Least Squares via Learned Feature Geometry

This paper establishes generalization error bounds for ridge-regularized nonlinear least-squares models by leveraging algorithmic stability and the Brascamp-Lieb inequality to derive data-dependent guarantees based on the learned gradient geometry and effective dimension, rather than parameter count or initialization.

Original authors: Ayub Kharel, Ilja Kuzborski, Patrick Rebeschini, Yasin Abbasi-Yadkori

Published 2026-06-09

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Ayub Kharel, Ilja Kuzborski, Patrick Rebeschini, Yasin Abbasi-Yadkori

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to predict the weather. You give it a massive notebook with millions of pages (parameters) and a huge stack of historical weather data.

In the past, scientists worried that if the notebook was too big, the robot would just "memorize" the specific weather patterns in the training data and fail miserably when faced with a new day. This is called overfitting. The old rule of thumb was: "If the notebook is bigger than the number of days you studied, the robot will fail."

But modern AI breaks this rule. We have robots with notebooks billions of times bigger than our data, and they still predict the weather perfectly. Why?

This paper, "Generalization in Nonlinear Least Squares via Learned Feature Geometry," tries to answer that question. It suggests that the robot doesn't just memorize; it learns a specific shape or geometry of the data that makes it smart.

Here is the breakdown using simple analogies:

1. The Problem: The "Perfect Memorizer" vs. The "Smart Learner"

Imagine you have a student taking a test.

The Old View: If the student has a huge memory (many parameters) but only studied 10 questions, they will just memorize the answers to those 10 questions. If you ask question #11, they will guess randomly.
The New Reality: The student has a huge memory, but when they studied, they didn't just memorize. They figured out the underlying pattern of the questions. Even though they have a huge brain, they only "use" a tiny, specific part of it to solve the problem.

The paper asks: How do we measure that "tiny, specific part" the student actually used?

2. The Solution: Measuring the "Used" Notebook

The authors propose a new way to measure complexity. Instead of counting how many pages are in the student's notebook (the total number of parameters), they count how many pages are actually relevant to the specific questions asked.

They call this the "Effective Dimension."

The Analogy: Imagine a library with 1,000,000 books.
- Old Measure: "This library is huge! It has 1,000,000 books, so it's too complex to learn from."
- New Measure: "Wait, for this specific topic (e.g., baking), only 50 books are actually relevant. The other 999,950 books are about space travel or cooking, which don't matter here. So, the effective size of the library for this task is just 50."

The paper proves that if this "effective size" is small, the model will generalize well (predict new data accurately), even if the total model is massive.

3. How They Found the "Used" Pages: The "Jacobian" Map

How do you know which 50 books are relevant? The authors look at the gradient (how the model changes its mind when you tweak the data).

The Metaphor: Imagine the model is a hiker on a mountain. The "gradient" is the direction the hiker looks to see which way is up.
When the model is first initialized (randomly), the hiker is looking in all directions at once (chaos).
After training, the hiker has found a path. They are only looking in a few specific directions that lead to the peak.
The paper measures the "Jacobian Geometry." This is a map of the directions the model actually cares about after training. If this map is simple (low-dimensional), the model is safe from overfitting.

4. The "Residual" Twist: Accounting for Curvature

The paper adds a clever twist. In simple math problems, the path to the peak is a straight line. But in complex AI, the path is curved.

The Analogy: Imagine walking on a flat field vs. walking on a curved hill.
The authors realized that the "curvature" of the hill (how much the model's predictions change non-linearly) matters. They created a formula that subtracts this "curvature" from the complexity count.
If the model fits the data perfectly (zero error), the curvature term vanishes, and the math simplifies to the classic "linear" version. But for real-world messy data, this extra term ensures the measurement is accurate.

5. The "Activation Regions" (The ReLU Network)

The paper specifically looks at neural networks with "ReLU" activation functions (a common type of AI neuron that acts like a switch: on or off).

The Metaphor: Imagine a city divided into neighborhoods. In some neighborhoods, the rules are simple (linear). In others, they are complex.
A neural network divides the world into many small "activation regions" (neighborhoods).
The paper shows that even though a network could create millions of neighborhoods, training usually only uses a tiny handful of them for the actual data it sees.
Key Finding: The complexity of the model isn't determined by how many neighborhoods exist, but by how many are occupied by the data. If your data only lives in 5 neighborhoods, the model only "sees" 5, regardless of how many empty neighborhoods exist in the background.

6. The Proof: Stability

How do they know this works? They use a concept called Algorithmic Stability.

The Analogy: Imagine you are teaching a class.
- Stable: If you remove one student from the class, the lesson plan doesn't change much. The teacher is stable.
- Unstable: If you remove one student, the teacher completely changes the lesson plan.
The paper proves that if a model has a low "Effective Dimension" (it's using a simple, learned geometry), it is stable. It won't panic if you swap out one data point. And because it is stable, it will perform well on new, unseen data.

Summary of the Paper's Claims

Size doesn't matter as much as we thought: A massive AI model can be simple if it only uses a small "effective" part of its capacity.
It's about the "Learned Geometry": The complexity depends on the shape of the solution the model found after training, not the shape of the model before training.
Data Compression: The model compresses the data into a lower-dimensional "manifold" (a smooth surface). The paper provides a formula to measure how small this surface is.
Verification: They tested this on synthetic data (made-up math problems) and real-world data (like house prices and wine quality). In every case, their new "Effective Dimension" formula predicted the model's success much better than old methods.

In short: The paper gives us a new ruler. Instead of measuring the size of the entire AI brain, we now measure how much of that brain is actually being used to solve the specific problem at hand. And it turns out, for good AI, that number is surprisingly small.

Technical Summary: Generalization in Nonlinear Least Squares via Learned Feature Geometry

Problem Statement
Modern machine learning models, particularly deep neural networks, often operate in overparameterized regimes where they can perfectly interpolate training data yet generalize well. Classical generalization theory, relying on uniform convergence bounds (e.g., VC dimension, Rademacher complexity) or fixed-kernel analyses (e.g., Neural Tangent Kernel, NTK), often fails to explain this phenomenon. These traditional measures either yield vacuous bounds for overparameterized models or fail to account for the specific geometry of the solution found by training, focusing instead on worst-case complexity or the geometry at random initialization.

This paper addresses the question: Which aspects of a fitted nonlinear model determine its ability to generalize? Specifically, it seeks to identify quantities that depend on the learned representation and reflect the data-induced geometry, rather than global class complexity.

Methodology
The authors study ridge-regularized nonlinear least-squares regression under a fixed design, where inputs are deterministic and randomness arises solely from response noise (assumed to be independent and strongly log-concave, e.g., Gaussian). The objective is to analyze the generalization of a nondegenerate local minimizer $\hat{\theta}$ of the regularized empirical risk.

The core methodology relies on algorithmic stability, specifically "on-average" prediction stability. The authors analyze the sensitivity of the fitted predictor when a single training response $Y_i$ is replaced by an independent copy $Y'_i$ .

Key technical components include:

Implicit Function Sensitivity: Using the implicit function theorem, the authors derive the derivative of the local minimizer $\hat{\theta}$ with respect to the response vector $Y$ . This reveals that the sensitivity is governed by the inverse Hessian of the regularized objective and the Jacobian of the predictor.
Brascamp–Lieb Inequality: To bound the variance of the prediction under strongly log-concave noise, the authors employ the multivariate Brascamp–Lieb inequality. This generalizes the standard Gaussian Poincaré inequality used in linear settings.
Learned Effective Dimension: The stability bound is expressed in terms of a data-dependent effective dimension, $d_{\text{eff}}(\hat{\theta}; \lambda)$ , defined at the trained parameters:
$d_{\text{eff}}(\hat{\theta}; \lambda) := \text{tr}\left( (\hat{H}_\lambda^{-1} \hat{G})^2 \right)$
Here, $\hat{G}$ is the empirical covariance of the trained Jacobian features ( $\nabla_\theta f(x_i; \hat{\theta})$ ), and $\hat{H}_\lambda$ is the Hessian of the regularized empirical objective at $\hat{\theta}$ . Unlike NTK theory, which uses the Jacobian at initialization, this metric captures the geometry after feature learning. The Hessian includes a residual-curvature term ( $\hat{\Delta}$ ) that vanishes in linear models but is present in nonlinear ones.
Covering Complexity: To bound $d_{\text{eff}}$ , the authors relate it to the covering number of the trained Jacobian features. They show that if the trained features can be covered by a small number of balls (i.e., they are compressible), the effective dimension is small. This is further specialized for data lying on low-dimensional manifolds and for ReLU networks, where the covering number is linked to the number of activation-stable regions occupied by the data.

Key Contributions

Stability Bound via Trained Effective Dimension: The paper proves that the average prediction stability (and consequently the generalization gap) is bounded by the expected effective dimension evaluated at the trained model. Theorem 3 establishes:
$\frac{1}{n} \sum_{i=1}^n \mathbb{E}[(f(x_i; \hat{\theta}) - f(x_i; \hat{\theta}^{(i)}))^2] \leq \frac{4 \mathbb{E}[d_{\text{eff}}(\hat{\theta}; \lambda)]}{\alpha n}$
where $\alpha$ is the strong log-concavity parameter of the noise. Corollary 4 translates this into a generalization gap bound of order $\sqrt{\mathbb{E}[d_{\text{eff}}]/n}$ .
Generalization of Linear Bounds: In the linear case (or when residuals are zero), the residual-curvature term vanishes, and $d_{\text{eff}}$ reduces to the classical ridge effective dimension of the Jacobian covariance. Thus, the result is a strict generalization of linear theory to the nonlinear regime.
Geometric Interpretation and Covering Bounds: The authors provide explicit bounds on $d_{\text{eff}}$ $d_{eff}$ based on the geometry of the data and the model:
- Manifold Support: If data lies on an $m$ -dimensional manifold and the Jacobian is piecewise Lipschitz, the bound scales with the intrinsic dimension $m$ rather than the ambient dimension.
- ReLU Networks: For one-hidden-layer ReLU networks, the bound is controlled by the number of activation-stable regions occupied by the data and the local feature contraction within those regions. This connects the generalization gap to the complexity of the learned partition of the input space.
Empirical Validation: The paper presents experiments on synthetic manifolds, clustered distributions, and benchmark datasets (California Housing, Wine Quality, Superconductivity). These experiments demonstrate:
- Jacobian Compression: The effective dimension of the trained Jacobian is significantly lower than that of the initialization (NTK regime), even in overparameterized settings.
- Tightness: The residual-curvature linearization used in the theory closely tracks the directly estimated nonlinear effective dimension.
- Agreement: The theoretical stability bounds are tight and closely track the observed generalization gaps.

Significance and Claims
The paper claims to provide a simple, first-principles derivation of generalization bounds for nonlinear least squares that depend on the learned geometry of the solution rather than the parameter count or initialization.

Shift from Fixed to Learned Geometry: Unlike NTK analyses that linearize around random initialization, this work analyzes the geometry at the fitted point. It demonstrates that feature learning can "compress" the relevant directions, leading to a lower effective dimension and better generalization.
Role of Activation Regions: The work makes an explicit theoretical connection between the number of activation-stable regions in ReLU networks and generalization bounds, supporting empirical findings that deep networks often utilize far fewer active regions than the worst-case combinatorial bounds suggest.
Limitations: The authors are modest about the scope, noting that the analysis is currently restricted to fixed-design square-loss regression with strongly log-concave noise and nondegenerate local minimizers. Extending these results to random designs, other loss functions, or capturing the implicit regularization of the full optimization trajectory (rather than just the local minimizer) remains an open challenge.

In summary, the paper argues that the generalization capability of nonlinear least-squares models is governed by the effective dimension of the trained Jacobian features, a quantity that reflects how the model has adapted its geometry to the data, and that this dimension can be significantly smaller than the model's parameter count or the sample size.