Technical Summary: Canonical Regularisation of Wide Feature-Learning Neural Networks

Problem Statement

Wide neural networks operating in the feature-learning regime—where the network actively learns representations rather than behaving as a fixed kernel machine—are the engine of modern deep learning. However, theoretical understanding of these networks lags significantly behind the kernel regime (formalized by the Neural Tangent Kernel, NTK). A critical gap exists regarding the implicit regularisation and function-space prior induced by gradient flow training.

In the kernel regime, it is well-established that gradient flow selects the unique solution minimizing the squared parameter-space distance from initialization (anchored ridge, $\|\theta - \theta_0\|^2$). This property underpins the Neural Network–Gaussian Process (NN-GP) correspondence, allowing for principled noise modeling and uncertainty quantification. However, practitioners often apply standard weight decay ($\|\theta\|^2$) or anchored ridge to feature-learning networks by default. This paper argues that these regularisers are geometrically misaligned in the feature-learning regime, leading to a "silent" bias that distorts the network's inductive bias, particularly in transfer learning scenarios where pretrained features are informative.

Methodology and Theoretical Framework

The authors develop a geometric framework to characterize the parameter-space dynamics of overparametrised networks satisfying a Uniform Jacobian Conditioning assumption (Assumption 2.1). This assumption posits that the Jacobian $J(\theta)$ has singular values bounded away from zero and infinity in a neighborhood of the interpolating solutions.

1. Geometric Decomposition

The authors decompose the parameter space $\Theta$ around an interpolator into three key objects:

• Flow Orbit ($O_{flow}$): The set of all parameters reachable by gradient flow from an initialization $\theta_0$. In the feature-learning regime, this can be a high-dimensional sub-Riemannian set.
• Gauge-Fixed Flow Manifold ($M_{flow}(\theta_0)$): A specific $n$-dimensional submanifold within the flow orbit that contains the realized gradient flow trajectory. It serves as a parameter-space chart for the function space near the interpolator.
• Output Fibres ($\Theta(c)$): Directions transverse to the flow manifold where the training predictions $g(\theta)$ remain constant, but predictions on unseen inputs can change drastically.

2. The Pathology of Ridge Regularisation

The paper proves that in the feature-learning regime, the flow manifold $M_{flow}(\theta_0)$ is curved, unlike the affine subspace of the kernel regime.

• Anchored Ridge Failure: The gradient of anchored ridge, $-2(\theta - \theta_0)$, is a Euclidean chord. On a curved manifold, this gradient acquires a component along the output fibres (orthogonal to the flow direction).
• Consequence: As regularisation strength $\lambda \to 0$, the solution does not converge to the unregularised gradient flow limit $\theta^*$. Instead, the regularisation gradient pulls parameters along the fibres, distorting the implicit prior and potentially destroying useful inductive biases learned during pretraining. This is formalized in Theorem 3.1.

3. Axiomatic Derivation of the Canonical Regulariser

To resolve this, the authors axiomatise the canonical function-space energy as the minimum control-theoretic energy required to transport network outputs from initialization to a target under gradient flow dynamics.

• Uniqueness: By imposing axioms of isometry invariance and orthogonal additivity, they prove that in the kernel regime, this energy uniquely corresponds to the squared RKHS norm (recovering anchored ridge).
• Generalisation: Lifting this energy to parameter space via Definition 4.3 yields the Geodesic Ridge regulariser:
  $$d^2(\theta, \theta_0) = d^2_{flow}(\theta_0, g|^{-1}_{\theta_0}(g(\theta))) + d^2_{fibre}(\theta, g|^{-1}_{\theta_0}(g(\theta)))$$
  This measures the squared geodesic distance along the flow manifold plus the geodesic distance along the output fibre.
• Implicit Prior: The corresponding function-space prior is identified as a Riemannian Gibbs Process, which generalises the NTK Gaussian Process to curved geometries.

4. Arc Ridge: A Practical Surrogate

Since computing geodesic distances is intractable at scale, the authors propose Arc Ridge as a minimax-robust surrogate.

• Definition: Arc ridge is the squared total path length of the training trajectory: $L_2(\gamma) = (\int_0^T \|\dot{\gamma}(t)\| dt)^2$.
• Properties: It upper-bounds geodesic ridge, retains the "no-bias" property along the realized gradient flow direction, and is formally equivalent to early stopping under exact gradient flow. It requires no second-order information, accumulating only a single scalar (path length) during training.

Key Results

1. Bias of Standard Regularisers: The paper proves that both standard ridge and anchored ridge bias the gradient flow limit in the feature-learning regime (Theorem 3.1). This bias persists even as regularisation strength vanishes, causing the solution to drift along output fibres invisible to the training loss.
2. Geodesic Ridge as the Canonical Solution: The authors derive geodesic ridge as the unique regulariser that aligns with the geometry of feature-learning networks, ensuring the vanishing-regularisation limit recovers the unregularised gradient flow solution (Theorem 4.2).
3. Arc Ridge Equivalence: Arc ridge is shown to be a tight, tractable upper bound to geodesic ridge. Crucially, it is proven to be equivalent to early stopping in the limit of exact gradient flow (Proposition F.4), providing a theoretical justification for early stopping as a geometry-aware regularisation technique.
4. Empirical Validation: Experiments on transfer learning tasks (Age regression on UTKFace using ResNet18; Review score regression on Yelp using DistilBERT) demonstrate that:
   • At high regularisation strengths, standard and anchored ridge cause significant degradation in test performance due to the destruction of the pretrained prior.
   • Arc ridge degrades only mildly, consistent with benign undertraining.
   • At low regularisation strengths, all methods converge, as early stopping dominates the bias.

Significance and Claims

The paper claims to provide the first principled framework for understanding and correcting the implicit bias of gradient flow in wide feature-learning networks. Its primary contributions are:

• Geometric Insight: It identifies the curvature of the flow manifold as the root cause of the failure of Euclidean-based regularisers (like anchored ridge) in feature-learning regimes.
• Theoretical Unification: It generalises the NN-GP correspondence to the feature-learning regime by identifying the Riemannian Gibbs Process as the canonical prior and geodesic ridge as the canonical regulariser.
• Practical Resolution: It proposes arc ridge (squared path length) as a scalable, minimax-robust alternative that aligns with the geometry of training. This reveals a deep, previously unformalised relationship between early stopping and canonical regularisation, suggesting that early stopping is not merely a heuristic but a geometry-preserving regularisation strategy.

The authors emphasize that while their theory assumes uniform Jacobian conditioning (verified empirically in Appendix H), the geometric intuition explains phenomena such as the degradation of pretrained representations during fine-tuning and the "grokking" phenomenon, offering a path toward more robust regularisation strategies in deep learning.