Generalization in Nonlinear Least Squares via Learned Feature Geometry

arXiv:2606.08799 · stat.ML, cs.LG · Submitted 2026-06-07 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Generalization in Nonlinear Least Squares via Learned Feature Geometry".

Jane: This paper investigates how generalization in ridge-regularized nonlinear least-squares models is governed by data-dependent geometry rather than worst-case complexity.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, we're talking about the paper titled "Generalization in Nonlinear Least Squares via Learned Feature Geometry," and honestly, the title sounds super technical. It basically suggests that how well a model generalizes isn't just about how many parameters it has or some fixed complexity measure we use, but more about the actual shape of the data geometry after training.

Jane: That makes sense when you think about neural networks; they can have millions of parameters, but maybe what matters is the specific region in that high-dimensional space where the model actually settled during training.

Lu: Exactly! This paper shifts our focus from a general hypothesis class complexity thing, like VC dimension, to something much more concrete: an effective dimension that's derived directly from how the gradient model looks at those trained parameters.

Meng: From my side, I wonder if this means we can stop worrying so much about the absolute size of the network and start looking at its learned structure instead.

Lalam: I think this is fascinating because it suggests that generalization isn't just about memorizing training data points; it's about finding a stable, geometrically simple solution within that data's own landscape.

The paper's summary: Tom: Okay, so what the paper actually does is introduce this concept of an on-average algorithmic stability measure that links generalization gaps to this "effective dimension" derived from the empirical Jacobian Gram matrix at the trained parameters. It's a big theoretical step for understanding why these overparameterized models seem to generalize so well in practice.

Jane: In simpler terms, they're saying that the gap between what we see on training data and how it looks on new data is controlled by this effective dimension, which is calculated using the residual–curvature term: ((two) deff(θb; λ):= tr(Hb−1λ Gb two).

Lu: That formula is really powerful because it captures the geometry of the gradient model precisely at the trained parameters, which differs from just looking at initialization as we sometimes do in other analyses. This makes the bounds much more relevant to how real-world training happens.

Meng: It sounds like they're providing a new way to predict stability without needing to analyze every single possible optimization path, which is helpful for practical engineering applications where we don't know the exact path an optimizer will take.

Lalam: And from an information perspective, this suggests that the relevant complexity isn't some abstract measure of the entire function space but something intrinsically tied to the data geometry itself, which could fundamentally change how we design robust AI systems for real-world deployment.

The paper's improvements: Tom: The main improvement they highlight is showing that this effective dimension can be controlled by the covering complexity of the trained Jacobian features, leading to a bound where the relevant complexity scales with intrinsic dimension when data lies on an m-dimensional manifold.

Jane: That's really neat because it connects back to manifold transfer ideas; it implies that if our data is inherently low-dimensional, we get better theoretical guarantees that scale with that intrinsic dimension instead of the ambient space size.

Lu: And for the specific case of one-hidden-layer ReLU networks, they make this mechanism explicit by defining an "activation-stable region" where the Jacobian map becomes linear, which lets them control complexity based on occupied pieces and their local Lipschitz constants.

Meng: That connection to manifold structure is important for practical use because it allows us to set expectations about generalization based on the data's underlying shape rather than just treating every dimension as equally complex.

Lalam: I see how this helps us improve our cultural impact, because if we can quantify exactly how much feature learning has happened—how much that Jacobian Gram matrix has been compressed—we get a clearer picture of whether our AI is truly discovering structure or just overfitting noise.

Conclusion: Tom: So, wrapping up the main points of "Generalization in Nonlinear Least Squares via Learned Feature Geometry," the authors argue that generalization gaps are controlled by this data-dependent effective dimension derived from feature compression and manifold structure, which offers bounds dependent on learned geometry rather than just parameter count.

Jane: It really boils down to saying that trained Jacobian features can be compressed, and that this compression is what gives us strong generalization guarantees tied directly to the intrinsic structure of the data.

Lu: This framework is important because it provides a way to prove stability for local minimizers in ridge-regularized nonlinear least squares models, even in complex settings like multilayer neural networks.

Meng: For my work, this means we can use this framework to implicitly select regularization levels that balance fitting the data well with keeping the local curvature favorable and geometrically simple.

Lalam: I think the most significant implication is that it gives us a mathematical tool to understand how AI systems learn structure from data, moving us toward AI that is inherently more robust and less sensitive to arbitrary parameter counts.

University of Oxford · Google DeepMind

stat.ML, cs.LG

Submitted: 2026-06-07

Updated: 2026-09-30

Importance score: 83/100

The gist: This paper investigates how generalization in ridge-regularized nonlinear least-squares models is governed by data-dependent geometry rather than worst-case complexity.

Key concepts

Effective Dimension (deff): This is a data-dependent measure of complexity calculated using the residual-curvature term. It captures how complex the local geometry of the model's gradient changes at a specific trained parameter point, effectively quantifying how much information is relevant to generalization.
Algorithmic Stability: This concept links prediction stability (how much predictions change when parameters are slightly perturbed) to the effective dimension. The paper shows that on average, this stability is bounded by the effective dimension, providing a new way to measure generalization gaps.
Covering Complexity: This bound relates the effective dimension to how many points are needed to cover the space of trained Jacobian features within a certain radius. It demonstrates that if these features can be compressed, the complexity is low, leading to better generalization guarantees.
Manifold Transfer: This technique connects the complexity measured in the high-dimensional parameter space to the intrinsic dimension of the data manifold. It shows that bounds scale with how many dimensions are actually needed to represent the data structure, rather than just ambient space.

Terminology

Generalization — 88× in this paper · explained in 13 episodes
Intrinsic Dimension — 18× in this paper · explained in 2 episodes
Regularization — 17× in this paper · explained in 3 episodes
Algorithmic Stability — 10× in this paper · explained in 1 episode
Benchmark — 9× in this paper · explained in 2 episodes
Core — 8× in this paper · explained in 2 episodes
Generalization Bounds — 8× in this paper · explained in 2 episodes
Trace — 6× in this paper · explained in 2 episodes
In International Conference on Machine Learning — 5× in this paper
Neural Tangent Kernel — 5× in this paper
Hilbert Space — 4× in this paper · explained in 2 episodes
Ambient Dimension — 3× in this paper · explained in 1 episode
Robustness — 3× in this paper · explained in 5 episodes
VC Dimension — 3× in this paper · explained in 2 episodes
Degeneracy — 2× in this paper · explained in 3 episodes
In Conference on Computational Learning Theory — 2× in this paper
Institutional Review Board — 2× in this paper
International Conference on Learning Representations — 2× in this paper
Loss Function — 2× in this paper · explained in 2 episodes
Activation Function — 1× in this paper · explained in 1 episode
Adam — 1× in this paper · explained in 2 episodes
Alignment — 1× in this paper · explained in 1 episode
Approximation Bound — 1× in this paper · explained in 1 episode
Clustering — 1× in this paper · explained in 2 episodes
Generative Model — 1× in this paper · explained in 2 episodes
Generative Models — 1× in this paper · explained in 2 episodes
Grounding — 1× in this paper · explained in 3 episodes
Homogeneity — 1× in this paper · explained in 2 episodes
Inverse Problem — 1× in this paper · explained in 2 episodes
Synthetic Data Generation — 1× in this paper · explained in 5 episodes
Theoretical Guarantees — 1× in this paper · explained in 2 episodes
Transformer — 1× in this paper · explained in 2 episodes

Summary

This paper investigates how generalization in ridge-regularized nonlinear least-squares models is governed by data-dependent geometry rather than worst-case complexity. It introduces an on-average algorithmic stability measure that links generalization gaps to an effective dimension derived from the empirical Jacobian Gram matrix at the trained parameters, providing a new theoretical framework for understanding why overparameterized models generalize well.

Theoretical Framework: Algorithmic Stability and Effective Dimension

The core contribution is deriving error bounds for local minimizers in terms of a data-dependent effective dimension, denoted as deff. This quantity captures the geometry of the gradient model at the trained parameters through the residual–curvature term:

((2) deff(θb; λ):= tr(Hb −1λ Gb 2)

where Gb is the empirical covariance of trained Jacobian features and Hbλ is the Hessian of the regularized empirical objective.

The paper establishes a prediction-stability bound (Theorem 3), showing that on average, prediction change is controlled by this effective dimension: 1/n Xn i=1 E hf(xi; θb) − f(xi; θb(i)) 2i ≤ 4 E[deff(θb; λ)] αn. This result is a generalization of classical bounds, as it shows that the relevant complexity is governed by learned geometry on the data, even in large ambient and parameter spaces.

Geometric Control: Covering Complexity of Trained Jacobian Features

The paper demonstrates that the effective dimension is small whenever the trained Jacobian features can be compressed. This leads to a covering complexity bound:

((7) deff(θb; λ) ≤ min p, n, inf ε>0 CJ (ε) + ε 2 λ − ρ)

where CJ(ε) is the empirical covering number of the trained Jacobian features at radius ε. The paper further connects this to data geometry via manifold transfer:

((8) dlin(G, t b) ≤ min(...) for appropriate radii)

This shows that the relevant complexity is governed by learned geometry on the data, allowing bounds to scale with intrinsic dimension when the data lies on an m-dimensional manifold.

One-Hidden-Layer ReLU Specifics

For one-hidden-layer ReLU networks, the geometric mechanism becomes explicit through activation regions. The paper defines an activation-stable region where the Jacobian map is linear:

((26) g(x) = TU x)

The complexity is then controlled by the number of occupied pieces (M) and their local Lipschitz constants (Lr). The final bound for this case shows that when the optimizer lies in a bounded-radius regime, deff(θb; λ) ≲m / CM X M r=1 Lr !2/(m+2) (λ − ρ) −m/(m+2)

Numerical Evidence and Interpretation

The theory is validated through extensive experiments on synthetic manifolds, clustered distributions, and benchmark datasets like California Housing and Wine Quality. Key findings include:

((6) Training compresses the relevant Jacobian geometry)

Figure 6 shows that the trained Jacobian Gram has a much faster spectral decay than the initialization Gram, with the median effective dimension dropping significantly after training.

The paper also illustrates the tightness of the residual-curvature linearization by comparing theoretical bounds to observed generalization gaps, showing that trained deff bound remains above but close to the observed gap.

Limitations and Future Directions

The authors note several limitations:

((1) The analysis focuses on fixed-design square-loss regression with strongly log-concave noise, explicit ridge regularization, and a local non-degeneracy condition.)

The framework relies on the selection assumption for the learning rule to ensure a nondegenerate local branch. Furthermore, extending the analysis to random design or other loss functions presents technical challenges. The current approach relies on characterizing a nondegenerate local minimizer via an explicit inverse-Hessian characterization, which may not capture solutions shaped by entire optimization trajectories.

Conclusion

In summary, the paper provides a novel way to quantify generalization in overparameterized models by shifting the focus from global hypothesis class complexity to the local geometry of the fitted solution. By linking algorithmic stability to a data-dependent effective dimension controlled by feature compression and manifold structure, it offers data-dependent bounds that are highly relevant for understanding how trained neural networks generalize. The results show that trained Jacobian features can be compressed, leading to generalization guarantees that depend on learned geometry rather than parameter count.

Improvements for AI systems

As a fastidious and diligent researcher, I have thoroughly analyzed this paper, Generalization in Nonlinear Least Squares via Learned Feature Geometry. The core contribution is shifting generalization theory from fixed function-class complexity (like VC dimension) to a data-dependent measure of learned geometry: the effective dimension of the trained Jacobian features.

Based on the findings—specifically that generalization gaps are controlled by this learned, compressed effective dimension—here are specific, actionable improvements for AI systems:

The improved AI system will be characterized by its ability to perform high-stakes inference and predictive modeling with provably tighter generalization guarantees than current methods, especially in overparameterized settings.

Here are the specific improvements and what the system can achieve:

(Geometric Feature Compression) The system will incorporate a mechanism to explicitly compute and utilize the low-rank structure of its trained Jacobian features.
(Data-Dependent Regularization) The regularization term will dynamically adjust based on the local curvature (inverse Hessian) of the fitted model, rather than relying on fixed priors like standard L2 regularization or NTK initialization geometry.
(Adaptive Complexity Control) The system's generalization bounds will be governed by the empirical covering number of its trained feature set, allowing it to scale complexity based on the actual data manifold geometry rather than worst-case ambient parameter counts.

The improved AI system can achieve the following specific capabilities:

(Significantly Tighter Generalization Bounds) The primary improvement is a generalization bound that scales with the intrinsic dimension of the data and the complexity of its learned representations (activation-stable regions), rather than scaling polynomially or exponentially with the total number of parameters.
(Superior Performance in Overparameterized Models) It will excel in highly overparameterized models (like deep neural networks) where classical complexity measures are often vacuous, as it focuses on the geometry of the specific local minimizer found by training.
(Improved Robustness to Data Geometry) The system will be inherently more robust when deployed on data lying on low-dimensional manifolds (e.g., in manifold learning tasks or high-dimensional sensor data), as its effective dimension is explicitly tied to the intrinsic dimension of the manifold, leading to better performance than models relying on ambient space complexity.
(Explicit Understanding of Feature Learning) The system will provide diagnostics that explicitly show how much feature learning (the compression of the Jacobian Gram matrix) has occurred during training, allowing researchers to quantify whether generalization is driven by simple interpolation or meaningful geometric structure discovery.
(Optimized Model Selection/Regularization) By utilizing the residual-curvature term, the system can implicitly select a regularization level that balances fitting the data (low residual error) against maintaining a favorable local curvature (low effective dimension), leading to solutions that are simultaneously accurate and geometrically simple.