In-context Learning in Presence of Spurious Correlations (2410.03140v1)

Abstract: LLMs exhibit a remarkable capacity for in-context learning, where they learn to solve tasks given a few examples. Recent work has shown that transformers can be trained to perform simple regression tasks in-context. This work explores the possibility of training an in-context learner for classification tasks involving spurious features. We find that the conventional approach of training in-context learners is susceptible to spurious features. Moreover, when the meta-training dataset includes instances of only one task, the conventional approach leads to task memorization and fails to produce a model that leverages context for predictions. Based on these observations, we propose a novel technique to train such a learner for a given classification task. Remarkably, this in-context learner matches and sometimes outperforms strong methods like ERM and GroupDRO. However, unlike these algorithms, it does not generalize well to other tasks. We show that it is possible to obtain an in-context learner that generalizes to unseen tasks by training on a diverse dataset of synthetic in-context learning instances.

An Examination of In-context Learning with Spurious Correlations

The paper "In-context Learning in Presence of Spurious Correlations" investigates the challenges and methodologies involved in training in-context learners, specifically LLMs, on tasks that include spurious features. Traditionally, LLMs like transformers have shown significant capability in in-context learning (ICL), where they adaptively solve tasks using a few examples without updating the model parameters. However, these learners can unfortunately be influenced by irrelevant or spurious correlations, which this paper addresses through a well-structured experimental approach.

Overview of Contributions

This work presents three primary contributions to the understanding and improvement of in-context learning in presence of spurious correlations:

1. Identification of Issues with Conventional Approaches: The paper articulates how traditional training of in-context learners on tasks containing spurious features leads to task memorization and reliance on these spurious features. These models fail to generalize well, especially when such correlations do not hold in unseen tasks.
2. Proposed Novel Techniques: To counteract the deficiencies observed, the authors propose innovative training techniques, namely:
   • Random Permutation of Input Embeddings: By randomly permuting input features, the researchers mitigate task memorization issues.
   • Revision of ICL Instance Formation: They devise methods to create ICL sequences that simulate distributional shifts, rendering the learner less sensitive to spurious correlations.
3. Diverse Task Generalization: Lastly, the paper demonstrates that training in-context learners on a synthetic and diverse set of tasks can yield models that generalize better across tasks with varying spurious correlations. This result is achieved by introducing architectural modifications and training procedures such as passing spurious feature annotations and promoting the emergence of induction heads.

Empirical Evaluation

The paper rigorously examines the efficacy of these strategies using several benchmark datasets, most notably Waterbirds and CelebA. On Waterbirds, they exhibit how the naive approach, which allows for task memorization, significantly underperforms compared to their proposed method. Their approach, augmented with permutation techniques and sequence simulations, achieves metrics that outperform well-known robust methods like ERM and GroupDRO.

For additional comparisons, the authors employ synthesized tasks using the iNaturalist dataset, evaluating their approach against traditional methods. The results affirm their hypothesis, revealing competitive or superior performance by the in-context learner, particularly in scenarios with unseen tasks where different spurious features are present.

Implications and Future Directions

The implications of this research are notable both theoretically and practically. From a theoretical standpoint, it advances our understanding of how spurious correlations impact LLMs in an ICL setup and provides mechanisms to mitigate these effects. Practically, the paper suggests pathways for creating more robust models capable of performing reliably under distribution shifts, a concern notably relevant for applications in domains like medical diagnostics where model errors could lead to significant repercussions.

Future avenues for research highlighted by the authors include delving deeper into understanding what specific algorithms in-context learners might be executing when trained as per their proposed methods. Additionally, exploring how these learners can handle multiple spurious features of varying severity and examining the trade-offs between in-weights learning and in-context learning in real-world, dynamic settings presents further opportunities for exploration.

In conclusion, this paper brings crucial insights into tackling the complexities arising from spurious correlations in in-context learning, potentially setting the stage for the development of more adaptive and resilient AI systems.