Generalization Properties of Retrieval-based Models

TL;DR

This paper provides a theoretical analysis of retrieval-based models, exploring their generalization capabilities through local and global approaches, and demonstrating how local sub-tasks can enhance overall accuracy.

Contribution

It introduces a formal framework for understanding retrieval-based models, analyzing local empirical risk minimization and kernel-based global models for classification.

Findings

Local sub-tasks enable low complexity models to achieve high accuracy.

Explicit local risk minimization improves generalization in retrieval models.

Kernel methods provide a global approach to retrieval-based classification.

Abstract

Many modern high-performing machine learning models such as GPT-3 primarily rely on scaling up models, e.g., transformer networks. Simultaneously, a parallel line of work aims to improve the model performance by augmenting an input instance with other (labeled) instances during inference. Examples of such augmentations include task-specific prompts and similar examples retrieved from the training data by a nonparametric component. Remarkably, retrieval-based methods have enjoyed success on a wide range of problems, ranging from standard natural language processing and vision tasks to protein folding, as demonstrated by many recent efforts, including WebGPT and AlphaFold. Despite growing literature showcasing the promise of these models, the theoretical underpinning for such models remains underexplored. In this paper, we present a formal treatment of retrieval-based models to characterize their generalization ability. In particular, we focus on two classes of retrieval-based classification approaches: First, we analyze a local learning framework that employs an explicit local empirical risk minimization based on retrieved examples for each input instance. Interestingly, we show that breaking down the underlying learning task into local sub-tasks enables the model to employ a low complexity parametric component to ensure good overall accuracy. The second class of retrieval-based approaches we explore learns a global model using kernel methods to directly map an input instance and retrieved examples to a prediction, without explicitly solving a local learning task.