IM-Context: In-Context Learning for Imbalanced Regression Tasks

TL;DR

Empirical evaluations demonstrate that in-context learning substantially outperforms existing in-weight learning methods in scenarios with high levels of imbalance, particularly for addressing imbalanced regression.

Abstract

Regression models often fail to generalize effectively in regions characterized by highly imbalanced label distributions. Previous methods for deep imbalanced regression rely on gradient-based weight updates, which tend to overfit in underrepresented regions. This paper proposes a paradigm shift towards in-context learning as an effective alternative to conventional in-weight learning methods, particularly for addressing imbalanced regression. In-context learning refers to the ability of a model to condition itself, given a prompt sequence composed of in-context samples (input-label pairs) alongside a new query input to generate predictions, without requiring any parameter updates. In this paper, we study the impact of the prompt sequence on the model performance from both theoretical and empirical perspectives. We emphasize the importance of localized context in reducing bias within regions of high imbalance. Empirical evaluations across a variety of real-world datasets demonstrate that in-context learning substantially outperforms existing in-weight learning methods in scenarios with high levels of imbalance.

Figures (7)

• Figure 1: Training distribution of two datasets: Boston (a) and AgeDB (b). The empirical expectations of errors are computed assuming ideal retrieval of neighbors for a new query sample (i.e., retrieving samples that are semantically relevant in terms of the target variable, rather than just close in the input space). The behavior of examples from different shot regions varies distinctly with the number of context examples. In the many-shot regions, the error bound stabilizes as more examples are provided, whereas in the few-shot regions, additional context examples lead to an increase in the error bound.
• Figure 2: Empirical Error: Averaging vs In-context learning (using GPT-2 model from garg2022can).
• Figure 3: Impact of neighbors retrieval on performances (using GPT-2 model proposed in garg2022can).
• Figure 4: Localization vs all training set set as context
• Figure 5: An overview of the proposed approach for imbalanced regression. Rather than relying on in-weight learning, which trains models directly on the training data, we propose leveraging in-context learning. For each query sample, we retrieve $k = k_s' + \tilde{k}_s$ neighboring examples from both (1) the training set ($k_s'$ neighbors) and (2) and inverse density dataset ($\tilde{k}_s$ neighbors), where the number of samples in each region is inversely proportional to its representation in the original set, and feed these as context to the model. This serves a dual purpose: avoiding bias toward the mean of the training set, which is crucial for tail regions, and reducing the memory requirement of the transformer.

Theorems & Definitions (3)

• Proposition 3.1
• Theorem 3.2
• Proposition 3.3