Grokking in Linear Models for Logistic Regression

TL;DR

This paper demonstrates that grokking, or delayed generalization, can occur in simple linear models with logistic loss, driven by gradient descent dynamics and data asymmetries, without the need for deep neural networks.

Contribution

It provides a theoretical and empirical analysis of grokking in linear models, revealing the underlying phases and data conditions that lead to delayed generalization.

Findings

Grokking occurs in linear models under certain data distributions.

The implicit bias of gradient descent induces a three-phase learning process.

Grokking can happen even without depth or complex representations.

Abstract

Grokking, the phenomenon of delayed generalization, is often attributed to the depth and compositional structure of deep neural networks. We study grokking in one of the simplest possible settings: the learning of a linear model with logistic loss for binary classification on data that are linearly (and max margin) separable about the origin. We investigate three testing regimes: (1) test data drawn from the same distribution as the training data, in which case grokking is not observed; (2) test data concentrated around the margin, in which case grokking is observed; and (3) adversarial test data generated via projected gradient descent (PGD) attacks, in which case grokking is also observed. We theoretically show that the implicit bias of gradient descent induces a three-phase learning process-population-dominated, support-vector-dominated unlearning, and support-vector-dominated generalization-during which delayed generalization can arise. Our analysis further relates the emergence of grokking to asymmetries in the data, both in the number of examples per class and in the distribution of support vectors across classes, and yields a characterization of the grokking time. We experimentally validate our theory by planting different distributions of population points and support vectors, and by analyzing accuracy curves and hyperplane dynamics. Overall, our results demonstrate that grokking does not require depth or representation learning, and can emerge even in linear models through the dynamics of the bias term.