Generalization vs. Specialization under Concept Shift

TL;DR

This paper analyzes how concept shift affects the generalization of machine learning models, revealing phase transitions and nonmonotonic performance changes, with theoretical insights and experiments on regression and classification tasks.

Contribution

It provides an exact theoretical analysis of ridge regression under concept shift, uncovering phase transitions and complex effects on generalization performance.

Findings

Concept shift causes a phase transition in prediction risk.

Longer context length can harm generalization under concept shift.

Behavior observed in regression also appears in classification tasks.

Abstract

Machine learning models are often brittle under distribution shift, i.e., when data distributions at test time differ from those during training. Understanding this failure mode is central to identifying and mitigating safety risks of mass adoption of machine learning. Here we analyze ridge regression under concept shift -- a form of distribution shift in which the input-label relationship changes at test time. We derive an exact expression for prediction risk in the thermodynamic limit. Our results reveal nontrivial effects of concept shift on generalization performance, including a phase transition between weak and strong concept shift regimes and nonmonotonic data dependence of test performance even when double descent is absent. Our theoretical results are in good agreement with experiments based on transformers pretrained to solve linear regression; under concept shift, too long context length can be detrimental to generalization performance of next token prediction. Finally, our experiments on MNIST and FashionMNIST suggest that this intriguing behavior is present also in classification problems.