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 a 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 smallscale, 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, pretrained 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.

text: Findings of the Association for Computational Linguistics ACL 2024, pages 12335–12352 August 11-16, 2024 ©2024 Association for Computational Linguistics
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, pretrained 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.

heading: 1 Introduction
text: The in-context learning (ICL), as initially uncovered by Brown et al. (2020), 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 (Li et al., 2023b). 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.
*Corresponding author: stefanik.m@mail.muni.cz
Modern in-context learners may perform ICL with quality comparable to task-specialized models (Zhao et al., 2023; Štefánik et al., 2023). However, it remains unclear why some LMs are able of ICL in such quality while others are not; Initial work introducing GPT3 (Brown et al., 2020) followed by Thoppilan et al. (2022); Chowdhery et al. (2022); inter alia explains ICL as an emergent consequence of models’ scale. But more recent LMs (Sanh et al., 2022; Wang et al., 2022; Wei et al., 2021; Ouyang et al., 2022) are based on 10 to 100 times smaller models and reach comparable ICL quality, instead attributing the ICL ability to a vast volume and diversity of pre-training tasks and instructions.
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, Xie et al. (2022) 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-
12335
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.

heading: 2 Background
text: Methods for training in-context learners Incontext learning ability, including few-shot ICL, was first uncovered in GPT3 (Brown et al., 2020) trained unsupervisedly for causal language modelling. With no other substantial differences to previous GPT models, the emergence of ICL was attributed to GPT3’s scale, having grown to over 170- billion parameters since GPT2 (≈800M params).
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 (Devlin et al., 2019) 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, Min et al. (2022a) 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 (Sanh et al., 2022; Wang et al., 2022) 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 (Ouyang et al., 2022) 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 (Chen et al., 2021) and by Chain-of-Thought (CoT) targets (Wei et al., 2022), where the model is trained to respond with a sequence of naturallanguage steps deducing the answer (Zhao et al., 2023; Kadlčík et al., 2023). These extensions were empirically shown to enhance ICL ability (Fu and Khot, 2022) and were adopted by Flan models (Chung et al., 2022).
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, Lu et al. (2022) report on the sensitivity of LMs to the specific formulation of the instructions in the prompt, while Liu et al. (2022) 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 (Min et al., 2022b) or when the presented CoT sequences do not make sense (Wang et al., 2023). We note that such behaviours differ from learning a functional relation from demonstrations that we expect from in-context learners (Li et al., 2023b) 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, Akyürek et 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 in
language sparsely but naturally, allowing the emergence of a certain ICL quality in LMs (Wies et al., 2023). Later work attributes ICL to similar data properties labelled under the terms of statistical burstiness (Chan et al., 2022) or compositionality (Hahn and Goyal, 2023).
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.

heading: 3 Concept-Aware Training
text: 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 xpred:
[x1, y1, ⟨sep⟩, . . . , xk, yk, ⟨sep⟩, xpred]→ ypred 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 (xi, yi) and the predicted example (xpred, ypred), 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) (Li et al., 2023a; Gupta et al., 2023; Luo et al., 2024). 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.
What constitutes a concept applicable in CoAT? We broadly define the term concept as an arbitrary functional relation of input and prediction (Xie et al., 2022) that holds robustly for any sample of a given task. Hypothetically, once the model learns to model a specific concept perfectly, it will never produce a prediction that violates this concept (Štefánik and Kadlčík, 2023). In practice, attempts to clearly present useful concepts to the model can be obstructed by other covariates (Mikula et al., 2024), such as frequent predictive co-occurrences of tokens. Thus, we propose to pick CoAT’s concepts among features that are unlikely to be substitutable by non-robust covariates, presenting as best candidates to be features conditioned on a deep or holistic decomposition of the input.
Note that the goal of CoAT is not to represent chosen training concepts in the model’s weights explicitly but to improve the model’s ability to extract and apply available concepts from demonstrations. Towards this goal, CoAT fundamentally assumes that the ability to extract and apply one concept transfers to other concepts beyond the training distribution. We verify this ability later in Section 4.4.

heading: 3.1 Proposed Implementation
text: In our experiments, we implement the proposed CoAT framework in two training stages: First, we train LM on a scalable synthetic QA dataset, which, contrary to traditional QA datasets, contains annotations of reasoning concepts. Second, we refresh the LM’s ability to work with natural language prompts by further training on a QA dataset with only natural language inputs. Hence, contrary to previous work utilising massive multitask training, in total, we only use two QA datasets.
Informativeness condition We find a large collection of annotated reasoning concepts in a TeaBReaC dataset of Trivedi et al. (2022), containing more than 900 unique explanations over a large set of synthetic QA contexts. Each TeaBReaC’s explanation maps a natural question to the answer span through a sequence of declarative reasoning steps, such as “select→group→project”. Within CoAT, we use these explanations as the shared concepts C (Fig. 1); In the training prompts, all demonstrations exhibit the same reasoning chain as the predicted sample.
To restore the model’s ability to work with a natural language, in the second step, we fit the resulting model to natural inputs by further fine-tuning on AdversarialQA dataset (Bartolo et al., 2021); As the annotations of reasoning concepts in general QA datasets are scarce, in this case, we naively use the initial word of the question (“Who”, “Where”, . . . ) as the shared concept, aware that such-grouped samples are not always mutually informative.
Non-triviality condition In both training stages, we implement the non-triviality condition in the following steps. (i) We select a random subset of 20 samples that passed the informativeness condition (denoted Xinfo). (ii) From Xinfo, we iteratively pick a sequence of i ∈ 1..k demonstrations (with k : 2 ≤ k ≤ 8) as follows:
1. For each sample (x j, y j) ∈ Xinfo, we use the training model to compute a likelihood of generating the correct prediction ypred if a given sample (x j, y j) is included among demonstrations. Whenever ypred contains more than one token, we compute the likelihood as the average of the likelihoods of all ypred’s tokens in the teacher-forced generation.
2. In each step i, we pick among the demonstrations a sample with which the likelihood of generating correct prediction is minimal.
An overview of this process is depicted in Figure 2, with a schematic example of a training prompt in Figure 1. Full training prompts that our implementation of CoAT constructs in training on each dataset can be found in Table 2 in the Appendix.

heading: 4 Experiments
text: Our experiments provide empirical evidence towards answering three research questions (RQs):
1. Can we improve models’ ability to benefit from new reasoning concepts in-context?
2. Are the concept-aware in-context learners more robust to known functional artifacts?
3. Can concept-aware in-context learning also improve performance in real-world tasks?
The first two RQs validate our assumptions on concept-aware training: that (1) the implementation of CoAT indeed improves models’ utilisation of both seen and out-of-distribution concepts from demonstrations, and that (2) such an ability can
make the in-context learning of a CoAT-trained language model more robust to artefacts revealed in previous in-context learners (Wei et al., 2023). Finally, in (3), we assess whether the enhanced models’ ability to rely more on latent concepts can also improve the practical quality of low-resource in-context learning.

heading: 4.1 Training and Evaluation
text: To maximise comparability with the previous work, we fine-tune our models from mT5 pre-trained models of Xue et al. (2021). In both training stages (Sec. 3.1), we fine-tune all model parameters in a teacher-forced next-token prediction until convergence of evaluation loss.* We further detail the training parameters in Appendix A.
In all experiments, we construct evaluation prompts from k = 3 demonstrations chosen consistently for all models, with prompts including the options for expected labels. We complement all our evaluations with confidence intervals from the bootstrapped evaluation (population n = 100, repeats r = 200). We specify evaluation setup separately for each experiment (§4.4–4.6) with further details and examples in Appendix B.

heading: 4.2 Baselines
text: We assess the impact CoAT’s main design choices against two baselines, allowing us to measure the impact of both its data construction conditions.
Random demonstrations selection (Tk-random) We evaluate the impact of all CoAT’s components against a baseline trained in identical settings but picking the in-context demonstrations randomly with uniform probability over the whole training set. This baseline reproduces the methodology of a majority of the referenced work on instruction tuning, including Tk-Instruct (Wang et al., 2022) and Flan (Chung et al., 2022). Apart from the demonstration selection, all other settings, including training data, remain identical (§4.1) to assure comparability with CoAT models.
Demonstrations passing only informativeness condition (Tk-info) In this baseline, we perform ablation of CoAT’s non-triviality condition (Sec. 3) by picking the demonstrations passing only the informativeness condition. Hence, such-picked demonstrations in the training instructions are informative for the prediction but can exhibit cases
*Implementations of CoAT training and concept-learning evaluations are on https://github.com/MIR-MU/CoAT
where some of the demonstrations are similar or even identical to the predicted sample, making it trivial for the model to perform correct prediction. All other training settings are unchanged (§4.1).

heading: 4.3 Other evaluated models
text: We also evaluate three recent in-context learners for which we can assess which datasets were used in their training mix: (1) T0 of Sanh et al. (2022) trained on a mixture of 35 datasets of different tasks in zero-shot settings, mostly of QA type, mapped into a self-containing human-understandable interaction format; (2) Tk-Instruct of Wang et al. (2022) pre-trained in a few-shot format similar to ours, on a mixture of 1,616 diverse tasks, and (3) Flan models of Chung et al. (2022) that further extend data settings of Tk-Instruct to a total of 1,836 tasks, including chain-of-thought labels, i.e. a step-by-step reasoning chain mapping input prompt to a label.
All these models are based on the same pretrained model (T5), making the results comparable to the level of fine-tuning methodology. TkInstruct and Flan use the data construction reproduced in our Tk-random baseline, but applied in vastly larger data settings.

heading: 4.4 CoAT’s ability to improve models utilisation of latent concepts (RQ1)
text: We pose that if the model can utilize a new reasoning concept C from demonstrations, it will be able to improve the prediction in cases where the demonstrations use the same C as the predicted sample. Thus, to evaluate if training with CoAT improves models’ utilisation of concepts, we evaluate models’ performance in a few-shot setting where we ensure that the demonstrations share a specific latent concept with the predicted sample. Then, we quantify models’ ability to improve from the concept by computing the difference in accuracy between such concept-sharing evaluation and conventional evaluation using randomly chosen demonstrations.
We perform the first analysis on TeaBReaC with annotated reasoning chains as concepts C shared between demonstrations and predicted sample (Fig. 1), but to evaluate generalization to unseen concepts, we filter out all samples with reasoning chains that were present in training. This results in 316 evaluation scenarios presenting models with 14 previously unseen reasoning patterns. In this setting, we compare the concept-improving ability of CoAT models with Tk-random baseline.
The important limitation of evaluation with TeaBReaC’s concepts is that it remains unclear whether evaluation with synthetic samples is representative for concept learning in a natural language. Hence, we also perform this analysis with samples and concepts of natural-language tasks.
Previous work of Štefánik and Kadlčík (2023) evaluated ICL ability over four different functional concepts, all extracted from explanations of naturallanguage datasets. We adopt the concepts of this work and evaluate models for in-context learning of the following concepts: (i) reasoning logic of NLI samples of GLUE-Diagnostic dataset (Wang et al., 2018), (ii) entity relations annotated in human explanations (Inoue et al., 2020) in the HotpotQA dataset (Yang et al., 2018), (iii) functional operations annotated in general elementary-grade tests of OpenBookQA (Mihaylov et al., 2018), and (iv) shared facts in science exams of WorldTree dataset (Jansen et al., 2018; Xie et al., 2020). Examples of prompts with concept-sharing demonstrations for these datasets are shown in Table 3.
Identically to the case of synthetic concepts, we evaluate the ability of CoAT models to benefit from these concepts when exhibited in demonstrations and compare to uncontrolled demonstrations’ selection (Tk-random) used in previous work.
Results
Concept-aware training improves the ability to benefit from unseen concepts Figure 3 evaluates models’ ability to improve from presented concepts as the relative difference in performance between
random and concept-sharing demonstration selection. First, evaluation with unseen TeaBReaC concepts (left) assesses models’ ability to extrapolate the utilisation of latent concepts to 14 previously unseen reasoning chains. Both CoAT and randomdemonstration models (§4.2) can improve from concepts presented in demonstrations. However, the improvement of CoAT-trained models is significantly larger and exceeds gains of Tk-random by 2-fold and 4-fold with the smaller and larger model, respectively. This comparison verifies that CoAT’s data construction really improves our targeted skill of better utilizing concepts of demonstrations.
CoAT applied with synthetic data also improves the use of natural concepts Evaluation of improvements on selected natural concepts (Figure 3; right) shows that concept-learning ability obtained with synthetic TeaBReaC concepts transfers to natural-language settings, as the CoAT-trained models can benefit from concepts significantly more than models trained without concept-aware data construction (Tk-random).
Despite that, evaluations over the individual reasoning concepts (Figure 7 in Appendix C.3) reveal that even CoAT models can not benefit robustly from all concepts. Nevertheless, we note that in the cases where CoAT models do not improve, also none of the baselines benefit from presented concepts. This might be attributed to several reasons: (i) the presented concepts are not really informative for prediction, (ii) our training data allowed the models to memorize relevant knowledge and, hence, do not need (and benefit from) the concepts’ exposition, or (iii) our training concepts were simply not sufficient to generalize over these new concepts.

heading: 4.5 Robustness of concept-aware in-context learners (RQ2)
text: As we overviewed in Section 2, other work reports functional deficiencies of previous in-context learners, including surprising insensitivity of incontext learners to the assigned demonstrations’ labels (Min et al., 2022b). Wei et al. (2023) attribute this to models’ over-reliance on the semantic priors obtained in pre-training, which may override learning of the functional concepts. Such behaviour is defective, because the ability to learn functional relations is necessary for robust and interpretable in-context learning of truly unseen tasks.
To evaluate the impact of concept-aware training on models’ reliance on their semantic priors,
we follow the setup of Wei et al. (2023) and assess reliance on labels’ semantics in a standard few-shot evaluation (§4.1), with one of the two modifications; (i) We change the labels to tokens with irrelevant meaning for the prediction task, such as ‘Foo’, ‘Bar’ etc. (ii) We shuffle the labels so that semantically incorrect labels are assigned in the demonstrations, but the input-label mapping remains consistent. In both settings, the task’s functional relation can still be recovered from demonstrations, but the sole reliance on semantics will either not help or will mislead the model.
In this setting, we evaluate three model types: (i) CoAT-trained models, (ii) models with uncontrolled data construction (Tk-random & previous work), and (iii) models with uncontrolled data construction, but fine-tuned only on a natural QA dataset (denoted Tk-QA). We perform the evaluation over 8 SuperGLUE tasks with discrete labels.
Results Figure 4 shows the results. Evaluation with non-sensical labels (left) reveals that all models pre-trained on a synthetic TeaBReaC dataset (Tk-random, and Tk-CoAT) are more robust to the labels’ semantics than our natural-language baseline (Tk-QA). However, a comparison of Tkrandom and Tk-CoAT suggests that Tk-CoAT’s preference for learning functional relations is a
composition of both using a synthetic dataset in pre-training and CoAT’s data construction.
A comparison to previous models reveals that all multitask models experience substantially larger decay in performance than our models. We suspect this feature could be a bias specific to massive multi-task learning emerging when label semantics can explain a large portion of training data. This result is consistent with Wei et al. (2023), but contrary to their conclusions, we show that ICL robust to semantic distractions does not emerge exclusively with very large (≥ 100B) model scale.
Nevertheless, we note that the smaller CoAT model still relies on labels’ semantics when recognizable (Flipped labels; Fig. 4 right), less than previous models, but comparable to our baselines.

heading: 4.6 Practical effectiveness of concept-aware in-context learners (RQ3)
text: Finally, we assess the practical quality of conceptaware ICL on previously unseen tasks in a simulated low-resource application with only three randomly-chosen demonstrations. We evaluate on samples from two collections of tasks: (i) SuperGLUE (Wang et al., 2019) consisting of 10 tasks requiring a variety of reasoning skills, and (ii) a test split of Natural-Instructions (Wang et al., 2022) from which we pick 60 extractive tasks. For SuperGLUE tasks, we verbalize both the demonstrations and predicted sample using all available templates within PromptSource library (Bach et al., 2022) and report results for the best-performing template
for each model. For Natural-Instructions tasks, we prefix the demonstrations with the instruction provided with each task. To maximise evaluation reliability over all models, we analyse the error cases and choose to report the results in ROUGE-L for SuperGLUE, and in a standard accuracy for Natural-Instructions. We specify the metrics selection analysis and other evaluation details in Appendix B, with prompt examples in Table 4.
As a primary reference point, we again compare the results of CoAT-trained models to Tk-random, where we can make sure that all other training configurations except for the data construction method are identical. Further, we compare to Tk-info (without Non-triviality condition; §4.2) to also evaluate the importance of the non-triviality condition. Finally, to provide additional context to our results, we also compare the performance of CoAT-trained models to previous in-context learners (§4.3).
Results Figure 5 compares the accuracy of CoATtrained models to our baselines: (i) without systematic demonstrations selection (Tk-Random) and (ii) without the non-triviality condition (Tk-Info), over 60 tasks of NaturalInstructions collection. In comparison to Tk-Random, CoAT models reach significantly higher accuracy on 41 and 45 of 60 tasks, with comparable performance on a majority (13 and 14) of other tasks. The difference is further magnified on reasoning tasks, which we argue might better evaluate models’ ability to in-context learn a functional relation of the new task. A comparison of Tk-Info with Tk-Random shows that the performance on reasoning tasks is mainly fostered by the CoAT’s informativeness condition, but in a full task collection, Tk-CoAT still outperforms Tk-Info in 19 out of 60 tasks. Evaluations on other task segments can be found in Appendix C.2.
In the evaluation over the tasks of SuperGLUE collection (Table 1), we additionally report the specific values of ROUGE-L that our baselines
and CoAT models achieve. With a single exception, models utilising a concept-based selection of demonstrations (Tk-CoAT and Tk-Info) consistently reach higher scores than Tk-Random. Our analyses of models’ predictions reveal that in 7 out of 20 evaluations, Tk-Randommodels fail to follow the task’s instruction, consequentially responding out of valid label space. Tk-CoAT is shown to mitigate this issue in all cases except for a smaller CoAT-trained model on MultiRC. A comparison of Tk-CoAT with Tk-Info shows that non-triviality condition is more substantial for a smaller model, but the models of both sizes benefit similarly from the concept-sharing selection of demonstrations.
Comparison to multitask learners Figure 6 contextualizes the performance of CoAT models trained on two datasets of a single (QA) task with existing instructional models trained on massive mixtures of 35–1,836 tasks. Over all the NI tasks (Fig. 6; top), CoAT models outperform multitask learners on a majority of tasks in 3 of 6 competitions. CoAT models are outperformed by Flan models but perform at least similarly for the majority of the tasks in 5 out of 6 competitions. The evaluation on reasoning tasks (Fig. 6; middle) supports our hypothesis that CoAT particularly promotes improvements in in-context learning of new reasoning abilities, winning on reasoning tasks over Flan and Tk-Instruct in a comparable number of cases than the opponents.
Finally, we look at a few tasks with unseen labels for both Tk-Instruct and Flan models (Fig. 6; bottom) where multitask learners can not rely on shortcuts based on unseen tasks’ labels. Here, the results of competition between CoAT with Flan models turns over, with CoAT models performing significantly better on 4 out of 6 tasks. While this sole segment is not big enough for robust conclusions, the results further support our claim that
concept-based ICL is more robust to semantic distractions (RQ2).
Table 5 in Appendix C details models’ scores on SuperGLUE tasks, providing further evidence on overall comparability of CoAT models to multitask learners. For instance, a comparison with Tk-Instruct reveals that CoAT’s 1B and 3B models reach higher absolute scores on 3 and 5 out of the 7 Tk-Instruct’s unseen tasks.

heading: 5 Conclusion
text: 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.
Limitations
Although our main objective is to assess the efficiency of concept-aware training, we acknowledge the limitations of our comparison to the previous work, where several aspects convolute the representative comparison of different in-context learners: (i) each of the multitask learners was trained on a different, yet massive set of tasks, making it difficult to find a broader collection that is new for multiple models; For this purpose, we surveyed three standard collections used for few-shot evaluation: CLUES (Mukherjee et al., 2021), RAFT (Alex et al., 2021) and FLEX (Bragg et al., 2021), but found in total only three tasks unseen by the multitask learners of previous work, all of the same type (classification). Therefore, in our evaluations, we use (a) Tk-Instruct’s own evaluation set and (b) SuperGLUE, which significantly overlaps with the training tasks of previous work. (ii) many aspects make it “easier” for the model to improve, including the domain of labels or prompt format matching the training distribution (relevant to Tk-Instruct and Flan evaluated on Natural-Instructions).
Another aspect that we neglect in our experiments in favour of more in-depth analyses is the impact of pretraining projected into the properties of the foundation model that we use. We pick T5 as a base model to maximise comparability with previous work. While we do not identify any concrete reason to assume that CoAT would perform worse with other base models, one should note that our results do not provide any evidence in this respect.
Finally, we note that the applicability of CoAT is conditioned by the availability of the annotated concepts C in the training datasets, which might be difficult to obtain for natural-language datasets. Our implementation circumvents this issue by using a synthetically curated dataset. Hence, we simultaneously show that concept-aware abilities can also be obtained in the restrictive settings of synthetic-dataset pre-training, where we note that the volume and variability of the synthetic dataset can be scaled further much easier than the natural dataset(s) (Trivedi et al., 2022). Nevertheless, our experiments do not provide any empirical evidence for answering to what extent could further extension of synthetically-generated datasets, possibly covering even more complex concepts, scale to further performance gains.
Ethical Considerations & Broader Impact
The primary motivation of our work is to minimise the computing demands for the creation of accurate in-context learners by deepening our understanding of the covariates of the resulting quality. We believe that our presented method, as well as the future data-efficient methods improving our understanding of in-context learning, will enable the democratization of the creation of robust and accurate in-context learning models for both research and industry.
Finally, we note that data-efficient methods for training ICLs (as opposed to multitask training) might open possibilities for creating more accurate ICLs specialized to languages outside English, where training datasets are scarce. We look forward for the future work that will explore the potential of data-efficient instruction tuning specifically on the target-language datasets, creating in-context learners specially tailored for target languages outside English.

heading: A Training details
text: Table 2 shows a full training example for each stage of training: (1) TeaBReaC with synthetic contexts (top) and (2) AdversarialQA with natural-language contexts (bottom). In all our training setups, we fine-tune all model parameters for teacher-forced next-token prediction, conventionally used in training sequence-to-sequence language models. In the two training stages (TeaBReaC and AdversarialQA), we use a learning rate of 5e−5 and 2e−5, respectively. Other parameters remain identical between stages: effective batch size = 30 samples and early stopping with the patience of 2,000 updates based on evaluation loss on a standardized validation set of each dataset. We do not report the absolute values of evaluation loss as these are not directly comparable. In CoAT training, we use a random subsample of 20 informative examples as a candidate set for a selection of non-trivial demonstrations.
Other parameters of training configuration default to Training Arguments of Transformers library (Wolf et al., 2020) in version 4.19.1. For readability, we implement the relatively complex demonstrations’ selection as a new objective of the Adaptor library (Štefánik et al., 2022). The picked demonstrations are encoded into a format consistent with the evaluation.

heading: B Evaluation details
text: Tables 3 shows an example of an instruction for each evaluation that we perform within the conceptlearning evaluation. For readability, we only shorten the examples of HotpotQA, where we omit some sources of data available for the model. In the case of TeaBReaC not shown in this table, the evaluation prompt format is the same as in training (Table 2), whereas we make sure that the reasoning chains of evaluation samples differ from the training.
SuperGLUE Evaluation format As mentioned in Section 4.1, we verbalize both the demonstrations and predicted sample using all available templates of PromptSource library (Bach et al., 2022), obtaining prompts for each demonstration prompt xi and its label yi in a free-text form. The prompts commonly contain the full-text match of the possible labels as options for the model.
Following the example of Wang et al. (2022), we additionally prepend the demonstrations and
labels with keywords “Input” and “Prediction” and separate demonstrations with new lines. Thus, the resulting input→output pairs in evaluation take this format:
“Input: x1 Prediction: y1 <newline> Input: x2 Prediction: y2 <newline> Input: x3 Prediction: y3 <newline> Input: xpred Prediction: ” → “ypred”
where demonstrations (xi, yi) are picked randomly but consistently between all evaluated models.
Natural-Instructions Evaluation format In the evaluations on Natural-Instructions, we closely follow the example of Wang et al. (2022) and additionally prepend the sequence of demonstrations with an instruction provided for each task:
“<task instruction> <newline> Input: x1 Prediction: y1 <newline> Input: x2 Prediction: y2 <newline> Input: x3 Prediction: y3 <newline> Input: xpred Prediction: ” → “ypred”
where the <task instruction> contains the instruction as would be given to the annotators of the evaluation task, usually spanning between 3–6 longer sentences. The demonstrations are again picked randomly but consistently between models.
Examples of evaluation prompts for both SuperGLUE and Natural-Instructions can be found in Table 4.
Evaluation metrics selection Previous work training in-context few-shot learners is not consistent in the use of evaluation metrics, and the choice usually boils down to either using the exact-match accuracy (Sanh et al., 2022; Chung et al., 2022) or ROUGE-L of Lin (2004) (Wang et al., 2022), evaluating the longest common sequence of tokens. We investigate these two options with the aim of not penalising the models for minor discrepancies in the output format (in the accuracy case) but avoiding false positive evaluations in predictions that are obviously incorrect (in the ROUGE case).
Investigation of the models’ predictions reveals that the selection of the metric makes a large difference only in the case of Tk-Instruct models, where the situation differs between SuperGLUE and Natural-Instructions, likely due to the character of the evaluation prompts.
(1) On SuperGlue, e.g. on MultiRC task, for the evaluation prompt: "Does answer sound like a valid
Dataset Concept Training instruction Target
TeaBReaC Exactly-matching reasoning chain ["select"→ "maximum"→ "list"→ "maximum"→ "sum"] “Input: how many points did the Monte Vesuvio" score in their two highest scoring matches? Context: scores of games of Pentagon". 99 scores of games of monte vesuvio". 67 scores of games of Pentagon". 6 scores of games of monte vesuvio". 76 scores of games of Pentagon". 37 scores of games of monte vesuvio". 56 scores of games of Pentagon". 8 scores of games of Pentagon". 90 scores of games of Pentagon". 20 Answer: Prediction: 143 [2 more examples] Input: how many points did the Bell 212 score in their two highest scoring games? Context: scores of games of bell 212. 90 scores of games of S-50. 54 scores of games of bell 212. 41 scores of games of bell 212. 36 scores of games of S-50. 23 scores of games of bell 212. 6 scores of games of bell 212. 2 scores of games of S-50. Prediction: ” “131”
AdversarialQA Matching question-word “Who”
“Input: Who was the Speaker in 1909? Context: Second, Democrats have always elevated their minority floor leader to the speakership upon reclaiming majority status. Republicans have not always followed this leadership succession pattern. In 1919, for instance, Republicans bypassed James R. Mann, R-IL, who had been minority leader for eight years, and elected Frederick Gillett, R-MA, to be Speaker. Mann "had angered many Republicans by objecting to their private bills on the floor;" also he was a protégé of autocratic Speaker Joseph Cannon, R-IL (1903–1911), and many Members "suspected that he would try to re-centralize power in his hands if elected Speaker." More recently, although Robert H. Michel was the Minority Leader in 1994 when the Republicans regained control of the House in the 1994 midterm elections, he had already announced his retirement and had little or no involvement in the campaign, including the Contract with America which was unveiled six weeks before voting day. Prediction: Joseph Cannon, R-IL. [2 more examples] Input: Who created the legal system still in use in Florida? Context: As a result of these initiatives northeastern Florida prospered economically in a way it never did under Spanish rule. Furthermore, the British governors were directed to call general assemblies as soon as possible in order to make laws for the Floridas and in the meantime they were, with the advice of councils, to establish courts. This would be the first introduction of much of the English-derived legal system which Florida still has today including trial by jury, habeas corpus and county-based government. Neither East Florida nor West Florida would send any representatives to Philadelphia to draft the Declaration of Independence. Florida would remain a Loyalist stronghold for the duration of the American Revolution. Prediction: ” “British”
Table 2: Examples of training instructions with expected outputs, for both our datasets applied in training. Note that the shared reasoning concept is not a part of the model’s input.
answer to the question: question", Tk-Instruct-3B in our evaluation predicts "Yes." or "Yes it is" (instead of "Yes"), or "No not at all" (instead of "No"), likely due to the resemblance with the format of training outputs. As we do not wish to penalize these cases, we use ROUGE-L over all SuperGLUE evaluations.
(2) In Natural-Instructions evaluation, we find that Tk-Instruct often predicts longer extracts from the input prompt. This is problematic with ROUGE-L in the cases where the extract contains all possible answers, such as in the Tk-Instruct1B’s prediction: “yes or no” to the prompt whose instruction ends with “Please answer in the form of yes or no.”. As we encounter this behaviour in a large portion of Natural-Instructions tasks, we evaluate all models on Natural-Instructions for exactmatch accuracy after the normalization of the casing and the removal of non-alphabetic symbols. To make sure that the model is presented with the exact-matching answer option, we exclude from evaluation the tasks where the correct answer is not presented in the task’s instruction. The reference to the list of Natural-Instructions evaluation tasks can be found in Appendix C.4.
For the reported evaluations of the Reasoning tasks, we pick from the list of evaluation tasks the
ones concerned with the reasoning task by simply matching the tasks with ‘reasoning’ in their name, resulting in the collection of 20 evaluation tasks.

heading: C Further evaluations
text: C.1 SuperGLUE evaluations of other models
Table 5 compares the performance over the tasks of SuperGLUE collection (Wang et al., 2019) for CoAT models trained on two tasks of the same (QA) type with in-context learners trained on 35– 1,836 tasks of the comparable size. Despite the
Dataset Concept Model instruction Expected output
GLUE NLI Diag. Double negation “Input: I will say that she stole my money. Question: I won’t say that she didn’t steal my money. True, False, or Neither? Prediction: Neither Input: I won’t say that she didn’t steal my money. Question: I will say that she stole my money. True, False, or Neither? Prediction: Neither Input: A rabbi is at this wedding, standing right there standing behind that tree. Question: It’s not the case that there is no rabbi at this wedding; he is right there standing behind that tree. True, False, or Neither? Prediction: True Input: Even after now finding out that it’s animal feed, I won’t ever stop being addicted to Flamin’ Hot Cheetos. Question: Even after now finding out that it’s animal feed, I will never stop being addicted to Flamin’ Hot Cheetos. True, False, or Neither? Prediction: ” “True”
OpenBookQA Shared facts: {"Earth is greater in mass than Mars", "gravity means gravitational pull; gravitational force; gravitational attraction", "as the force of gravity increases, the weight of objects will increase."} “Facts: a decrease is a kind of change. increase means more. as mass of a planet; of a celestial body increases, the force of gravity on that planet will increase. to change means to become different. an animal is a kind of living thing. the gravitational force of a planet; of a celestial object does not change the mass of an object on that planet or celestial body. an increase is the opposite of a decrease. an astronaut is a kind of human. massive means great in mass. the Mars Rover is a kind of vehicle. a living thing is a kind of object. Earth is greater in mass than Mars. gravity means gravitational pull; gravitational energy; gravitational force; gravitational attraction. greater means higher; more in value. stay the same means not changing. a moon is a kind of celestial object; body. an increase is a kind of change. Earth is a kind of planet. as the force of gravity increases, the weight of objects will increase. less is similar to decrease. Mars is a kind of planet. Input: An object has a weight of 10 kg on the surface of Earth. If the same object were transported to the surface of Mars, the object would have a weight of 3.8 kg. Which best explains why the weight of the object changed when transported from Earth to Mars? (A) The density of the object is greater on Earth than it is on Mars. (B) The volume of the object is greater on Earth than it is on Mars. (C) Gravitational force is greater on Earth than it is on Mars. (D) Atmospheric pressure is less on Earth than it is on Mars. Prediction: Gravitational force is greater on Earth than it is on Mars [two more examples] Input: When astronauts walked on the Moon, they used weighted boots to help them walk due to the lower gravitational pull. What difference between Earth and the Moon accounts for the difference in gravity? (A) density (B) diameter (C) mass (D) volume. Prediction: ” “mass”
HotpotQA Shared relation in reasoning: “X is a genus” “Input: Are Broughtonia and Laeliocattleya both orchids? Hint: use the information from the paragraphs below to answer the question. Otaara, abbreviated Otr. in the horticultural trade, is an intergeneric hybrid of orchids, with "Brassavola", "Broughtonia", "Cattleya", "Laelia" and "Sophronitis" as parent genera. Paracaleana commonly known as duck orchids, is a genus of flowering plants in the orchid family, Orchidaceae that is found in Australia and New Zealand. Duck orchids have a single leaf and one or a few, dull-coloured, inconspicuous flowers. (...) Prediction: yes [two more examples] Input: Are both Parodia and Thalictrum flowering plants? Hint: use the information from the paragraphs below to answer the question. - Thalictrum ( ) is a genus of 120-200 species of herbaceous perennial flowering plants in the Ranunculaceae (buttercup) family native mostly to temperate regions. Meadow-rue is a common name for plants in this genus. - Parodia is a genus of flowering plants in the cactus family Cactaceae, native to the uplands of Argentina, Peru, Bolivia, Brazil, Colombia and Uruguay. This genus has about 50 species, many of which have been transferred from "Eriocactus", "Notocactus" and "Wigginsia". They range from small globose plants to 1 m tall columnar cacti. All are deeply ribbed and spiny, with single flowers at or near the crown. Some species produce offsets at the base. They are popular in cultivation, but must be grown indoors where temperatures fall below 10 degrees. Prediction: ” “yes”
Dataset Concept Model instruction Expected output
SuperGLUE - “Input: The soldiers were concealed in the brush. Select the most plausible cause: - They were armed with rifles. - They wore camouflage uniforms. Prediction: They wore camouflage uniforms. Input: The print on the brochure was tiny. Select the most plausible effect: - The man put his glasses on. - The man retrieved a pen from his pocket. Prediction: The man put his glasses on. Input: I excused myself from the group. Select the most plausible cause: - I turned off my phone. - My phone rang. Prediction: My phone rang. Input: My body cast a shadow over the grass. Select the most plausible cause: - The sun was rising. - The grass was cut. Prediction:” “The sun was rising.”
Natural-Instructions - “Indicate with ‘Yes‘ if the given question involves the provided reasoning ‘Category‘. Indicate with ‘No‘, otherwise. We define five categories of temporal reasoning. First: "event duration" which is defined as the understanding of how long events last. For example, "brushing teeth", usually takes few minutes. Second: "transient v. stationary" events. This category is based on the understanding of whether an event will change over time or not. For example, the sentence "he was born in the U.S." contains a stationary event since it will last forever; however, "he is hungry" contains a transient event since it will remain true for a short period of time. Third: "event ordering" which is the understanding of how events are usually ordered in nature. For example, "earning money" usually comes before "spending money". The fourth one is "absolute timepoint". This category deals with the understanding of when events usually happen. For example, "going to school" usually happens during the day (not at 2 A.M). The last category is "frequency" which refers to how often an event is likely to be repeated. For example, "taking showers" typically occurs 5 times a week, "going to Saturday market" usually happens every few weeks/months, etc. Input: Sentence: Jack played basketball after school, after which he was very tired. Question: How long did Jack play basketball? Category: Event Duration. Prediction: Yes Input: Sentence: He was born in China, so he went to the Embassy to apply for a U.S. Visa. Question: How often does he apply a Visa? Category: Frequency. Prediction: Yes Input: Sentence: Jack played basketball after school, after which he was very tired. Question: Was Jack still tired the next day? Category: Event Duration. Prediction: No Input: Sentence: It refers to a woman who is dangerously attractive, and lures men to their downfall with her sexual attractiveness. Question: How long does it take to lure men to their downfall? Category: Event Duration. Prediction: ” “Yes”
Table 4: Examples of evaluation instructions with expected outputs, for selected tasks of SuperGLUE and NaturalInstructions (RQ3). Displayed samples are from CoPA and MCTato Temporal Reasoning tasks, respectively. Note that in these evaluations, demonstrations are picked randomly, regardless of their concepts.
significantly smaller volumes and complexity of the training dataset, CoAT-trained models show competitive results to similar-size or even larger incontext learners of previous work. For instance, the 1-billion-parameter Tk-CoAT performs better than the 3-billion T0 in 3 cases (Ax-b, RTE, COPA) and comparably in another 3 cases (WSC, CB, WiC). In comparison with Tk-instruct of the same size, Tk-CoAT-1B outperforms Tk-instruct in 3 out of 7 unseen tasks (WSC, CB, ReCoRD), and reaches similar scores in most other cases, even in 2 out of 3 tasks that were included in Tk-instruct’s training mix. Similarly, larger Tk-CoAT-3B outperforms Tk-instruct on 4 of 7 new tasks (Ax-b, WSC, WiC, ReCoRD), but with larger gaps on the others.
C.2 Natural-Instructions: other task types
Figure 8 evaluates the impact of CoAT’s mechanism on the quality of in-context learning separately on the English and non-English tasks. The figure reveals that CoAT works particularly well for non-English tasks. Our analyses found this is mainly due to the low performance of the baseline on the non-English tasks. We speculate that this can be a consequence of the higher reliance of the
baseline on token semantics (Section 4.6, RQ2); As our models are fine-tuned on an English-only QA model, such learnt reliance is not applicable in multilingual settings.
Figure 9 compares the performance of CoAT models against the models of previous work, separately on the English and non-English tasks. We can see that CoAT is slightly better at the multilingual portion of Natural-Instructions, but the difference is not principal.
C.3 Per-concept evaluations Figure 7 evaluates the performance gains of the baseline models (§4.2) and CoAT-trained models individually per each of the concepts of the natural datasets. While the CoAT models are able to benefit from concepts the largest in the relative change of quality, they are also not consistent in the ability to benefit from all the concepts. However, as discussed in Section 4.4, this does not imply that CoAT is unable to utilize these concepts.
C.4 Evaluation tasks and other configurations SuperGLUE (Wang et al., 2019) consists of the following tasks (as ordered in our Results, §4.6): Winogender Schema Diagnostics (AxG) (Rudinger
et al., 2018), Broadcoverage Diagnostics (CB), The Winograd Schema Challenge, CommitmentBank (CB), Recognizing Textual Entailment (RTE), ContextWords in Context (WiC) (Pilehvar and Camacho-Collados, 2019), Reading Comprehension with Commonsense Reasoning (ReCoRD) (Zhang et al., 2018), BoolQ (Clark et al., 2019), Choice of Plausible Alternatives (COPA), MultiSentence Reading Comprehension (MultiRC).
Natural-Instructions consists of a larger mixture of tasks, which we do not enumerate here to maintain readability; the full list of evaluation tasks can be found in the original work of Wang et al. (2022) in Figures 11 and 12.
To maintain comparability of evaluations among models, we deterministically fix the demonstration selection procedure so that only the full prediction prompts for all the models are the same. In the analyses comparing the differences in performance
(§4.4; RQ1+2), we fixed the prediction samples (xpred) between different demonstrations’ sampling strategies to avoid perplexing our comparison with possible data selection biases. Further details can be found in the referenced implementation.

heading: D Computational Requirements
text: We run both training and evaluation experiments on a machine with dedicated single NVIDIA A100SXM-80GB, 40 GB of RAM and a single CPU core. Hence, all our reproduction scripts can run on this or a similar configuration. Two stages of training in total take at most 6,600 updates and at most 117 h of training for Tk-CoAT to converge.

heading: Future Work
text: 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. Limitations
Although our main objective is to assess the efficiency of concept-aware training, we acknowledge the limitations of our comparison to the previous work, where several aspects convolute the representative comparison of different in-context learners: (i) each of the multitask learners was trained on a different, yet massive set of tasks, making it difficult to find a broader collection that is new for multiple models; For this purpose, we surveyed three standard collections used for few-shot evaluation: CLUES (Mukherjee et al. , 2021), RAFT (Alex et al. , 2021) and FLEX (Bragg et al. , 2021), but found in total only three tasks unseen by the multitask learners of previous work, all of the same type (classification). Therefore, in our evaluations, we use (a) Tk-Instruct’s own evaluation set and (b) SuperGLUE, which significantly overlaps with the training tasks of previous work. (ii) many aspects make it “easier” for the model to improve, including the domain of labels or prompt format matching the training distribution (relevant to Tk-Instruct and Flan evaluated on Natural-Instructions). Another aspect that we neglect in our experiments in favour of more in-depth analyses is the impact of pretraining projected into the properties of the foundation model that we use. We pick T5 as a base model to maximise comparability with previous work. While we do not identify any concrete reason to assume that CoAT would perform worse with other base models, one should note that our results do not provide any evidence in this respect. Finally, we note that the applicability of CoAT is conditioned by the availability of the annotated concepts C in the training datasets, which might be difficult to obtain for natural-language datasets. Our implementation circumvents this issue by using a synthetically curated dataset. Hence, we simultaneously show that concept-aware abilities can also be obtained in the restrictive settings of synthetic-dataset pre-training, where we note that the volume and variability of the synthetic dataset can be scaled further much easier than the natural dataset(s) (Trivedi et al. , 2022). Nevertheless, our experiments do not provide any empirical evidence for answering to what extent could further extension of synthetically-generated datasets, possibly covering even more complex concepts, scale to further performance gains