Meta-Learning Transformers to Improve In-Context Generalization

TL;DR

This paper introduces a meta-learning approach to train transformers on diverse small datasets, enhancing in-context generalization beyond training domains while addressing privacy and data quality issues.

Contribution

It proposes a novel training paradigm using multiple small, domain-specific datasets combined with meta-learning to improve in-context learning and generalization.

Findings

Models trained on curated datasets outperform single large-scale dataset models in generalization.

Meta-learning enhances robustness to forgetting in continual learning scenarios.

Transformers trained this way maintain performance in unsupervised and out-of-domain settings.

Abstract

In-context learning enables transformer models to generalize to new tasks based solely on input prompts, without any need for weight updates. However, existing training paradigms typically rely on large, unstructured datasets that are costly to store, difficult to evaluate for quality and balance, and pose privacy and ethical concerns due to the inclusion of sensitive information. Motivated by these limitations and risks, we propose an alternative training strategy where we leverage a collection of multiple, small-scale, and domain-specific datasets. We empirically demonstrate that the increased quality and diversity of such data improve the generalization abilities of in-context learners beyond their training domain, while achieving comparable performance with models trained on a single large-scale dataset. We investigate this paradigm by leveraging meta-learning to train an in-context learner on the Meta-Album collection under several settings. Firstly, we show the performance in a controlled environment, where the test domain is completely excluded from the training knowledge. Secondly, we explore the robustness of these models to forgetting in a continual scenario where the information is accessible for a limited time. Finally, we explore the more challenging unsupervised scenario. Our findings demonstrate that transformers still generalize for in-context prediction when trained on a curated dataset collection while offering advantages in modularity and replaceability.