Concept-aware Data Construction Improves In-context Learning of Language Models (2024)

Abstract: Many recent language models (LMs) are capable of in-context learning (ICL), manifested in the LMs’ ability to perform a new task solely from natural-language instruction. Previous work curating in-context learners assumes that ICL emerges from a vast over-parametrization or the scale of multi-task training. However, recent theoretical work attributes the ICL ability to concept-dependent training data and creates functional in-context learners even in small-scale, synthetic settings.In this work, we practically explore this newly identified axis of ICL quality. We propose Concept-aware Training (CoAT), a framework for constructing training scenarios that make it beneficial for the LM to learn to utilize the analogical reasoning concepts from demonstrations. We find that by using CoAT, pre-trained transformers can learn to better utilise new latent concepts from demonstrations and that such ability makes ICL more robust to the functional deficiencies of the previous models. Finally, we show that concept-aware in-context learners are much more effective in in-context learning a majority of unseen tasks compared to traditional instruction tuning, and fare comparably also to previous in-context learners trained in large-scale multitask learning requiring magnitudes of more training data.

## introduction the in-context learning (icl), as initially uncovered by @xcite , is a setting requiring language models (lms) to infer and apply correct functional relationships from the pairs of inputs and outputs (i.e. demonstrations) presented in userprovided input prompt @xcite . given that a small set of demonstrations can be obtained for any machine learning task, in-context learning presents a much more versatile and practical alternative to training task-specific models. on the contrary, theoretical studies uncover different determinants of icl quality than the model or data scale, relating icl to specific data qualities, such as the occurrence of cases that can not be explained by mere statistical co-occurrence of tokens. notably, @xcite specify this as the occurrence of training exemplars that can only be resolved by identifying latent concepts, i.e. underlying functional relations that explain the correct prediction. in this and other work surveyed in sec-tion 2, authors prove that icl can also emerge with both small data and small models. our work explores the practical potential of concept-dependent data on the quality and robustness of in-context learning. in section 3, we propose and implement a data construction framework that encourages the occurrence of concept dependencies in training data, and hence, requires models to learn to utilise latent concepts that explain these irregularities (fig. 1 ). we call this framework concept-aware training (coat). in sections 4, we explore the impact of coat in controlled settings. we find that (i) it is possible to train language models for in-context learning of unseen latent concepts and (ii) that such conceptaware in-context learning is more robust to the functional deficiencies of existing in-context learners. finally, on a set of over 70 tasks of superglue and natural-instructions, we find that coat can also improve practical in-context learning performance over traditional instruction tuning approach; in many cases, coat enables icl of otherwise not learnable tasks, and with only two training tasks reaches icl performance comparable to in-context learners of similar or larger size trained on massive collections of over 1,600 tasks. not long after, a pivotal work of schick and schütze (2020) on a pattern-exploiting training (pet) has shown that even much smaller (110m) models like bert @xcite can be finetuned using self-training in a similarly small data regime, first disputing the assumption on the necessity of the scale in rapidly learning new tasks. a line of generation models further undermined the assumption of the size conditioning of icl. among the first, @xcite fine-tune smaller models (<1b parameters) on a large mixture of tasks in the few-shot instructional format and find that such models can perform previously unseen tasks. following approaches @xcite also train smaller models for instruction following on large mixtures of tasks, assuming that the model's ability to in-context learn an unseen task emerges from a large diversity of instructions and task types. a recently popularised reinforcement learning approach of instructgpt @xcite also presents an adaptation of an instruction-following objective, training on mixtures of instructions with automatic feedback. recently, the instruction following was extended by joint training on programming code generation @xcite and by chain-of-thought (cot) targets @xcite , where the model is trained to respond with a sequence of naturallanguage steps deducing the answer @xcite . these extensions were empirically shown to enhance icl ability @xcite and were adopted by flan models @xcite . analyses of icl recent studies shed some light on the functioning of icl in lms through controlled experimentation, finding that the lms' decision-making in icl does not align with humans. notably, @xcite report on the sensitivity of lms to the specific formulation of the instructions in the prompt, while @xcite measures sensitivity to the ordering of in-context demonstrations. further, we find that lms perform icl comparably well when the labels of the demonstrations are randomly shuffled @xcite or when the presented cot sequences do not make sense @xcite . we note that such behaviours differ from learning a functional relation from demonstrations that we expect from in-context learners @xcite and can be exploited to lead models to incorrect predictions. nevertheless, other studies report that under the right conditions, lms are able to learn functional relationships solely from the input prompt; for instance, @xcite al. (2023); li et al. (2023c) show that transformers can be trained to accurately learn regression functions solely from the prompt. xie et al. (2022) identify the key covariate of icl quality in the occurrence of training examples where correct predictions are conditioned by latent concepts. consider a pre-training example 'albert einstein was [mask]'; the correct prediction for [mask] can be resolved if the model can extract and apply a latent reasoning concept from context, such as that the context exhibits a concept of nationalities and hence, [mask] best substitutes 'german'. such concept dependencies occur @xmath0.76 p( = 0.43 p( = 0.89 p( = 0.62 <@xmath1. informativeness condition 2. non-triviality condition already picked predicted sample to be picked predicted sample predicted sample candidates language sparsely but naturally, allowing the emergence of a certain icl quality in lms @xcite . later work attributes icl to similar data properties labelled under the terms of statistical burstiness @xcite or compositionality @xcite . our work builds upon this theory, but compared to previous studies limited to in-silico experiments, we elevate the idea of concept-aware training to real-world settings, with publicly available datasets and pre-trained language models. we are first to measure the impact of concept-aware data construction in extrinsic evaluation over 70 diverse tasks and show its potential to substantially enhance data efficiency and robustness in training in-context learners, compared to previous work using magnitudes of more data and compute. ## concept-aware training aiming to create language models able to learn a new latent reasoning concept in-context, we propose a concept-aware training (coat) as an instruction-tuning framework specifying conditions for a selection of few-shot demonstrations for the training instructions (figure 2 ). we assume the format of training prompts widely used in the previous work training incontext few-shot learners, constructing training instructions from k demonstrations consisting of the input texts x with labels y followed by the predicted sample's input text x pred : in this setting, coat proposes to filter in-context demonstrations sequentially by two conditions. the main condition, denoted as informativeness condition, assures to pick demonstrations exhibiting a specific reasoning concept c that is shared between a picked demonstration (x i , y i ) and the predicted example (x pred , y pred ), thus picking only the demonstrations whose reasoning pattern is informative for the correct prediction. such settings make it beneficial for the trained model to learn to extract and apply concepts presented in demonstrations. however, as the sole informativeness condition may easily pick demonstrations very similar or identical to the predicted sample, we propose a second, non-triviality condition. this condition chooses from the informative demonstrations the ones with which it is 'difficult' for the model to respond correctly. this condition avoids the occurrence of in-context demonstrations identical to the predicted sample and may also increase the heterogeneity of different concepts that co-occur among the demonstrations, avoiding the over-reliance on the presence of a small set of specific concepts in small-data settings. we note that a body of previous work proposes better ways for picking in-context demonstrations during the inference (without updating the model) @xcite @xcite . while some of these strategies are applicable in picking informative demonstrations in coat, note that this line of work is complementary to ours in assuming an existing in-context learner. more importantly, our motivation is substantially different; coat uses demonstrations of instruction training as a vehicle for creating training cases conditioned on latent concepts. this idea is not restricted to instruction tuning, but we note that instruction tuning presents an opportunity to implement it easily. ## experiments our experiments provide empirical evidence towards answering three research questions (rqs): ## conclusion inspired by data-centric theories on emergence of in-context learning, we propose and implement a concept-aware training framework for constructing training scenarios that challenge language models to learn to utilise latent concepts from in-context prompts. we show that language models can be trained to benefit from unseen concepts (rq1), and that such icl is more robust in learning functional relations of a new task from demonstrations (rq2). finally, in extrinsic evaluation over 70 tasks, we demonstrate the practical efficiency of concept-dependent training data, with coat mod-els bringing significant improvements on 41 and 45 out of 60 natural-instructions tasks, or 6 and 5 of 10 superglue tasks (rq3), while reaching a performance comparable to multitask learning requiring magnitudes of more data. more broadly, our work pioneers an alternative direction for scaling the quality of in-context learning to the previously explored model and data scale axes. we wish towards inspire future work to a more proactive approach to refining training data properties so that fitting such data necessitates the emergence of the specific, robust abilities of the models, such as the concept-learning ability. specifically, future work can build upon our findings in researching ways to upscale conceptdependent data in unsupervised settings, allowing for pre-training more robust language models with a fraction of data and computing budget.