The Geometric Structure of Models Learning Sparse Data

TL;DR

This paper explores how models learn in sparse data regimes by leveraging local geometric structures, introducing the concept of normal alignment, and proposing regularization strategies to improve robustness and training efficiency.

Contribution

It formalizes the concept of normal alignment, proves its benefits for training and robustness, and introduces GrokAlign and RFAMs as practical methods based on these insights.

Findings

Normal alignment minimizes training objectives under norm constraints.

GrokAlign accelerates training dynamics in deep networks.

RFAMs show increased adversarial robustness over RFMs.

Abstract

The manifold hypothesis (MH) is often used to explain how machine learning can overcome the curse of dimensionality. However, the MH is only applicable in regimes where the training data provides a sufficiently dense sample of the underlying low-dimensional data manifold, or where such a low-dimensional manifold is conceivably present. We describe the regimes where the MH is not applicable as sparse. In this paper, we demonstrate that models succeed in the sparse regime by exploiting a highly structured local geometry, a property we formalize as normal alignment. We prove that normal-aligned classifiers -- whose input-output Jacobians are rank-one and align perfectly with the training data -- minimize the training objective under norm constraints and achieve maximal local robustness under a non-zero Jacobian constraint. For continuous piecewise-affine deep networks, normal alignment manifests geometrically as centroid alignment within the network's induced power diagram partition and results from the feature-learning regime. Motivated by these theoretical insights, we introduce GrokAlign, a regularization strategy that actively induces normal alignment. We demonstrate that GrokAlign significantly accelerates the training dynamics of deep networks relevant to the grokking phenomenon. Furthermore, we apply the principle of normal alignment to Recursive Feature Machines (RFMs) to introduce Recursive Feature Alignment Machines (RFAMs). We show that RFAMs exhibit greater adversarial robustness compared to RFMs when trained on tabular data.