Context is Environment

TL;DR

This paper proposes that viewing context as environment and leveraging in-context learning can significantly improve out-of-distribution generalization in AI models, supported by theory and experiments.

Contribution

It introduces In-Context Risk Minimization (ICRM), a novel algorithm that uses in-context learning to better adapt to test environments and improve domain generalization.

Findings

ICRM outperforms baseline methods in out-of-distribution tests

Paying attention to unlabeled context improves environment adaptation

Theoretical analysis supports the effectiveness of in-context learning for generalization

Abstract

Two lines of work are taking the central stage in AI research. On the one hand, the community is making increasing efforts to build models that discard spurious correlations and generalize better in novel test environments. Unfortunately, the bitter lesson so far is that no proposal convincingly outperforms a simple empirical risk minimization baseline. On the other hand, large language models (LLMs) have erupted as algorithms able to learn in-context, generalizing on-the-fly to eclectic contextual circumstances that users enforce by means of prompting. In this paper, we argue that context is environment, and posit that in-context learning holds the key to better domain generalization. Via extensive theory and experiments, we show that paying attention to context$\unicode{x2013}\unicode{x2013}$unlabeled examples as they arrive$\unicode{x2013}\unicode{x2013}$allows our proposed In-Context Risk Minimization (ICRM) algorithm to zoom-in on the test environment risk minimizer, leading to significant out-of-distribution performance improvements. From all of this, two messages are worth taking home. Researchers in domain generalization should consider environment as context, and harness the adaptive power of in-context learning. Researchers in LLMs should consider context as environment, to better structure data towards generalization.