Out-of-Distribution Generalization in Kernel Regression

TL;DR

This paper develops an analytical framework using statistical physics to understand and optimize out-of-distribution generalization in kernel regression, applicable to various kernels and datasets.

Contribution

It introduces a replica-based analytical formula for out-of-distribution error in kernel regression and identifies distribution mismatch as a key factor affecting generalization.

Findings

Analytical expression for OOD generalization error derived

Mismatch quantified by an overlap matrix impacts performance

Optimization procedures for training and test distributions developed

Abstract

In real word applications, data generating process for training a machine learning model often differs from what the model encounters in the test stage. Understanding how and whether machine learning models generalize under such distributional shifts have been a theoretical challenge. Here, we study generalization in kernel regression when the training and test distributions are different using methods from statistical physics. Using the replica method, we derive an analytical formula for the out-of-distribution generalization error applicable to any kernel and real datasets. We identify an overlap matrix that quantifies the mismatch between distributions for a given kernel as a key determinant of generalization performance under distribution shift. Using our analytical expressions we elucidate various generalization phenomena including possible improvement in generalization when there is a mismatch. We develop procedures for optimizing training and test distributions for a given data budget to find best and worst case generalizations under the shift. We present applications of our theory to real and synthetic datasets and for many kernels. We compare results of our theory applied to Neural Tangent Kernel with simulations of wide networks and show agreement. We analyze linear regression in further depth.