Canonical Regularisation of Wide Feature-Learning Neural Networks

Abstract: Wide neural networks in the feature-learning regime drive modern deep learning, and yet they remain far less studied than their kernel-regime counterparts. We consider a critical yet under-explored difference between these two regimes: the regulariser and prior implied by gradient flow training. This canonical regularisation property is well-studied in kernel regime networks -- of all the infinite global minima, gradient flow selects exactly the vanishing ridge solution -- and underpins the celebrated NN-GP correspondence, precisely allowing the modelling of noise during training. However, we prove ridge regularisation biases gradient flow in feature-learning regime networks, even in the infinitesimal limit of vanishing regularisation. Over training, ridge distorts the inductive bias of the network, with a particular damage done to pretrained networks where the implicit prior is informative. We resolve this by axiomatising the canonical regulariser as a regime-agnostic function-space energy and lift, which uniquely identifies ridge in the kernel regime, and crucially generalises to the feature-learning regime. By studying the Riemannian geometry of feature-learning networks, we derive geodesic ridge from our framework, generalising ridge to the feature-learning regime. Correspondingly, we prove the canonical function-space prior is a Riemannian Gibbs Process, generalising the more familiar Gaussian Process. As a practical contribution, we propose arc ridge as a minimax-robust, scalable surrogate to geodesic ridge, revealing a deep relationship between early stopping and canonical regularisation across learning regimes. Finally, we demonstrate the consequences of our theory empirically on both image processing and NLP transfer-learning problems.

• The paper demonstrates that standard and anchored ridge regularisers degrade the inductive prior in feature-learning regimes by inducing fibre drift.
• It develops a geometric framework that recasts canonical regulariser as function-space energy and introduces geodesic ridge with its scalable surrogate, arc ridge.
• Empirical results on image and NLP regression validate that arc ridge maintains generalisation better than traditional methods under high regularisation.

Canonical Regularisation in Wide Feature-Learning Neural Networks

Summary and Motivation

Wide neural networks are central to contemporary deep learning advances, yet there remains a notable theoretical gap in understanding their behaviour outside the so-called "kernel regime". This paper interrogates the regularisation implicit in gradient flow training across two fundamentally distinct regimes: the kernel regime, where features are fixed, and the feature-learning regime, where features evolve during training. Canonical regularisation in the kernel regime is linked to anchored ridge, which is well-aligned with gradient flow and forms the basis of the NN-GP correspondence. However, the paper demonstrates that both standard and anchored ridge regularisers are insufficient and, in fact, provably degrade the inductive prior in the feature-learning regime—even with vanishing regularisation strength. The work recasts the canonical regulariser as a regime-agnostic function-space energy, generalises it to the feature-learning regime via a unique geometric lift, and introduces geodesic ridge and its scalable surrogate, arc ridge.

Geometric Structure and Regime Distinctions

The theoretical analysis exploits differential geometric constructs to describe the parameter space explored by wide, overparametrised networks under gradient flow. Three principal objects are identified: the flow orbit (the full set of parameters reachable via gradient flow), the gauge-fixed flow manifold (an $n$-dimensional chart mapping outputs to parameters), and output fibres (directions in parameter space unresponsive to training loss but impactful on generalisation).

In the kernel regime, these objects collapse to a single affine subspace, ensuring alignment between training dynamics and anchored ridge regularisation.

Figure 1: Left: flow manifold curvature in feature-learning. Right: affine manifold in kernel regime.

In contrast, the feature-learning regime exhibits manifold curvature, causing regularisation gradients of anchored (and standard) ridge to develop components along output fibres. This fibre drift is silent at the loss level but can severely degrade generalisation, particularly in transfer learning scenarios with pretrained representations.

Regularisation Pathology: Rigorous Characterisation

The paper proves that both standard and anchored ridge regularisers bias the gradient flow limit in feature-learning networks—even as the regularisation parameter $\lambda \to 0$. The bias arises from geometric misalignment: the Euclidean chord from initialisation to the interpolant is no longer confined to the flow manifold due to curvature, inducing parameter movement in directions that are invisible to the training loss but affect the inductive prior in function space. This is formalised via a theorem establishing inequivalence between the ridge-regularised minima and the pure gradient flow solution.

Figure 2: Trajectories of gradient flow under various regularisers for minimal overparametrised models; canonical regularisation uniquely avoids bias in feature-learning.

Canonical Regularisation Framework and Generalisation

To resolve the pathology, the framework identifies the function-space energy as the minimum control-theoretic energy required to transport network outputs via gradient flow. Two axioms—invariance and orthogonal additivity—uniquely single out the canonical regulariser in the kernel regime and generalise it to the feature-learning regime. The canonical regulariser, geodesic ridge, is defined as the squared geodesic distance from initialisation along the flow manifold, supplemented by a fibre term. Its gradient remains tangent to the flow manifold, eliminating fibre drift. The corresponding canonical function-space prior emerges as a Riemannian Gibbs Process, generalising the NTK Gaussian Process to curved geometry.

Practical Regularisation: Arc Ridge and Early Stopping

Computing geodesic ridge is infeasible at scale; the practical alternative, arc ridge, defined as the squared arc length of the training path, is introduced. Arc ridge upper-bounds the canonical regulariser and sandwiches it between anchored ridge and the training path length, and is shown to be minimax-optimal within this interval. Importantly, arc ridge connects with early stopping: regularisation by cumulative path length is formally equivalent to time-based stopping of gradient flow, requiring only noise estimates and not a validation set.

Figure 3: Empirical comparison on transfer-learning tasks: arc ridge preserves inductive prior at high regularisation, while standard and anchored ridge degrade test MSE sharply.

Empirical Results

Experiments validate the theoretical predictions on both image regression (UTKFace with ResNet18) and NLP regression (Yelp reviews with DistilBERT). At high regularisation, standard and anchored ridge regularisers induce catastrophic performance degradation, aligning with theory regarding prior destruction via fibre drift. Arc ridge, however, exhibits only mild performance losses consistent with benign undertraining—its generalisation properties are preserved across regularisation strengths, particularly in the regime where tasks benefit from pretrained features.

Implications and Future Directions

This work challenges prevailing regularisation practices inherited from kernel regime theory. It demonstrates that in the feature-learning regime, only regularisers respecting the geometry of the training manifold (geodesic or arc ridge) preserve the implicit prior and generalisation. Practically, arc ridge provides an unbiased, scalable approach aligning regularisation with the actual behaviour of wide networks under gradient flow, and elucidates the intimate relationship between regularisation and early stopping.

Figure 4: Empirical verification of uniform Jacobian conditioning along gradient flow trajectories, supporting theoretical assumptions.

Several research directions are proposed: improving scalable approximations to geodesic ridge, validating the framework across architectures and tasks, rigorously verifying uniform Jacobian conditioning, and analysing the Riemannian Gibbs Process from a Bayesian perspective. Extensions to non-interpolating regimes and geometry-aware optimisers also remain open for future investigation.

Conclusion

The paper supplies a rigorous, geometric characterisation of regularisation in wide feature-learning neural networks, linking gradient flow dynamics to the correct inductive prior and resolving the mismatch of traditional regularisers. It offers practical and theoretical tools for unbiased regularisation, advancing principled uncertainty modelling and transfer behaviour in modern deep learning systems.