TART: A plug-and-play Transformer module for task-agnostic reasoning
Authors: Kush Bhatia
Avanika Narayan
Christopher De Sa
Christopher Re
Venue: NeurIPS 2023 poster (2023)
Keywords: In-context learning
task-agnostic methods
large language models
TL;DR: We study task-agnostic methods for improving language model performance, making it competitive with task-specific adaptation approaches.

Abstract: Large language models (LLMs) exhibit in-context learning abilities which enable the same model to perform several tasks without any task-specific training. In contrast, traditional adaptation approaches, such as fine-tuning, modify the underlying models for each specific task. In-context learning, however, consistently underperforms task-specific tuning approaches even when presented with the same examples. While most existing approaches (e.g., prompt engineering) focus on the LLM's learned representations to patch this performance gap, our experiments actually reveal that LLM representations contain sufficient information to make good predictions. As such, we focus on the LLM's reasoning abilities and demonstrate that this performance gap exists due to their inability to perform simple probabilistic reasoning tasks. This raises an intriguing question: Are LLMs actually capable of learning how to reason in a task-agnostic manner? We answer this in the affirmative and, as a proof of concept, propose TART which generically improves an LLM's reasoning abilities using a synthetically trained reasoning module. TART trains this Transformer-based reasoning module in a task-agnostic manner using only synthetic logistic regression tasks and composes it with an arbitrary real-world pre-trained model without any additional training. With a single inference module, TART improves performance across different model families (GPT-Neo, Pythia, Bloom), model sizes (100M - 6B), tasks (14 NLP classification tasks), and even across different modalities (audio and vision). On the RAFT Benchmark, TART improves GPT-Neo (125M)'s performance such that it outperforms Bloom (176B), and is within $4$% of GPT-3.

TaRT: A plug-and-play Transformer module for task-agnostic reasoning
Kush Bhatia"
Avanika Narayan"
Christopher De Sa
Christopher Ré
(kushb, avanika, chrismre | @ ecs.stanford.edu cess@ccs.commll.ed.
Abstract
Large language models (LLMs) exhibit in-context learning abilities which enable the same model perform several tasks without any task-specific training. In contrast, traditional adaptation approaches, such as fine-tuning, modify the under- lying models for each specific task. In-context learning, however, consistently underperforms task- specificc tuning approaches even when presented with the same examples. While most existing approaches (c.g., prompt engineering) focus the LLM's learned representations to patch this performance gap. our experiments ac- malll reveal that L.LM representations contain sufficient information to make good predictions. As such, we focus on the LLM's reasoning abilities and demonstrate that this performance gap exists due to their inability to perform simple proba- bilistic reasoning tasks. This raises intriguing question: Are LLMs actually capable of learning how to reason a task-agnostic manner? We answer this in the affirmative and, proof of concept, propose TART which generically improves an LLM's reasoning abilities using synthetically trained reasoning module. TART trains this Transformer based reasoning module in task -agnostc manner using only synthetic logistic regression tasks and composes with an arbitrary real-world pre-trained model without any additional training. With single inference module, TART improves performance across different model families (GPT-NER PYTHIA, BLOOM), model sizes (100M 6B), tasks (14 NLP classification tasks), and even across different modalities (audio and vision). On the RAFT Benchmark, TART improves GPT-NEO (125M)'s performance such that it outperforms BLOOM (176B), and is within 4% of GPT-3 (175B).²
Introduction
Large language models (LLMs) show in-context learning capabilities which enable them to perform
a task given only a few examples, without updating the model parameters [9, 7]. This task-agnostic
capability allows for single model to be applied wide range of tasks [1, 41, 28). In contrast, traditional task adaptation approaches, such as fine-tuning, update the model parameters for each specific task.
Despite being task agnostic, in-context learning is seldom the practitioner's 's method of choice since consistently underperforms task-specific adaptation approaches [20, 9]. Most existing works attribute this performance gap to the limited context window of LLMs which can only accommodate a few task examples [15, 14, 23). However, tuning approaches exists even when presented with the same task examples; in-context learning can underperform fine-tuning by up 30 points across binary classification benchmarks (see Figure 5a).
show that this gap between in-context learning and fine-
We first investigate why this quality gap exists. We decompose an LLM''s in-context learning capability into two abilities: learning good representations for the task and performing probabilistic inference,
Equal Contribution
Our code and model is available at httpps//ggtthha..cmmmmaa/eeeeerr
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Task- agnostic
Data- Scalable
Quality
X
In-context learning
Base
LM Head
Base
LM Head
Fine-tuning
In-context learning
Fine-tuning
Adapters
Inference Module
Base
Head
Base
TART
(Our method)
Adapters
TART
Figure Taxonomy of task adaptation strategies. (Left) Comparison of different adaptation strategies across three desiderata: task- agnostgc quality, scalability. (Right) Demonstration of parameter updates across adaptations strategies, colored regions represent parameter changes as a result the adaptation strategy.
or reasoning. over these representations Is the gap because the data representations learned by the models lack task-specific information is there deficiency in the the model's ability to use this information. We take first step towards understanding this hypotheses experimentally in Section 2 by measuring both the reasoning and the representation gaps across variety of LLM families (GPT-NEO [6], PYTHIA [5], BL.OMM [33]) over a suite of binary classification tasks. Our results suggest that LLMs indeed have the necessary information in the representations, and that the majority of the quality gap (up to 79%) can be attributed to their inability to perform simple forms of reasoning, e.g. linear classification. We further find that adaptation methods (such fine-tuning) improve the model along both representation and reasoning, but primarily improve the task-specific classification (or reasoning) ability which accounts for 72% the gained performance on the task.
Rather surprisingly, most existing techniques for improving the performance gap, such as prompt engineering or active example selection, focus entirely the LLM's Icarned representations. With these methods, the users are effectively searching the language space for a natural language template that best represents the task. This search enables the users optimize the representations for their specific task which makes the LLM more likely to solve it. In contrast, our work explores the orthogonal direction of improving the LLM's reasoning abilities.
As a first step towards improving this reasoning ability, we fine-tune LLMs using synthetically generated binary classification tasks (see Section A.1 for examples). While this approach provides an improvement over the model's base in-context learning performance (up to 17%, see Figure 7), this approach requires fine-tuning each LLM individually and might affect their generative abilities [18]. This brings about a question: can we improve the reasoning capabilities of LLMs without interfering with it pre-existing capabilities, that is,
make the process both task and model agnostic?
We show that it is indeed possible to improve the reasoning capabilities in completely agnostic manner. Rather than directly fine-tuning, our proposed algorithm TART, Task-Agnostic Reasoning Transformer (Figure 2), improves an LLM's reasoning abilities using an independent synthetically trained module. TART trains this additional Transformer module using only synthetically generated binary classification tasks from the logistic model independent the downstream task the base LLM. This inference module can be composed, without any additional training, with the embeddings of an arbitrary pre-trained LLM to improve upon its reasoning abilities.
TART is task, model, and domain agnostic. Using a single inference module trained on synthetic data, we exhibit that TART not only generalizes across three model families (GPT-NEO, PYTHIA, BLOOM) over 14 NLP classification tasks, but even across different domains (vision and speech; see Figure 6). In terms of quality, we show that TART'S performance 18.4% better than in-context learning, 3.4% better than task-specific adapters, and is within 3.1% of full task-specific fine-tuning across a suite NLP tasks. On the RAFT Benchmark (2), TART improves GPT-NEO (125M)'s performance such that it outperforms BLOOM (176B), and is within 4% of GPT-3 (175B). TART is
In Section 2.3, precisely define the this representation-rcasoiigg decomposition in eq. (1)

---

(1,y)
The movie is good Label: Positive
Sentence:
(Sentiment classification, News classrfication)
Synthetic logistic regression data
Embeddingg
:
Base Model
:
Inference Module
Pre-trained LLM
31
9
Exk
(GPT-N80, Pythia, BLOOM)
Inference Module
Averaging
Layer
Figure 2: TART. (Left) Inference module training procedure: The inference module is trained on sequences of synthetically generated logistic regression tasks. (Right) End-to-end framework: TART composes pre-trained LLM with the inference module. TAKT uses LLM to embed input text. These embeddings. along with the train labels, are passed as sequence to the inference module which generates final prediction.
data-scalable and overcomes the limited context length bottleneck of in-context learning. While each example spans multiple tokens in LLM, often spanning hundreds of tokens, TART'S reasoning module encodes each example using only two tokens one for the context and the other for the label. This data-scalability can lead to improvements of up to 6.8% (see Figure 5c).
From a theoretical standpoint, we show that the generalization abilitics of TART depends mainly on the distribution shift between the natural text embedding distribution produced by the LLM and the synthetic data distribution, measured in terms of the Wasserstein-1 metric (Theorem in Section 3.4).
To summarize, our main contributions are follows:
Study why in-context learning does not perform well as task-specific fine-tuning despite having access to the same information, via a ieapeenttttio-rrasonnnn decomposition.
Propose task aagk-ignoticc method, TART, which bridges the performance gap to task- methods and trained using only synthetic classification data.
Demonstrate that TART works across different NLP tasks for range of model families. The same inference module generalizes to vision and speech domains, suggesting that this capability is somehow uniformly lacking across range of foundation models.
Related work. Prompt engineering focuses on improving the in-context task adaptation abilities of LLMs by modifying prompts. A line work improves performance by carefully designing the natural language task specifications [4, 42] while others improve performance by optimizing the examples chosen for the prompt [10, 24J, encouraging the models to sequentially reason [16, 42, 45] and aggregating prompts [39, 38). Unfortunately, prompt-based task adaptation is noisy [26). Alternatively, prompt tuning improves the in-context abilities of models by training small amounts of learnable vectors [21, 20, 25] for specific tasks. While these methods have been shown to improve in-context learning performance, they require task-specific fine-tuning and are not task-agnostic.
Recent works seek understand the in-context learning property of LLMs by presenting mecha- nistic interpretations in-context learning [36] and performing exploratory analysis of in-context leamning behaviors [43]. Existing literature demonstrates that LLMs can learn simple function classes in-context [11] and propose that L.MMs are performing gradient descent when learning tasks in- context [36]. Complementary
these, work provides insights the mechanisms of in-context learning and its deficiencies. Furthermore, task transfer strategies adapt LLMS to a pre-specified target task. Strategies range from parameter efficient finetuning (PEFT) [12, 47] to Low-Rank adaptation (LoRA) [13] which introduces trainable rank decomposition matrices into each layer.
Task adaptation strategies: Taxonomy and evaluation
We begin by describing the problem of adapting pre-trained language models for a collection of downstream tasks while being task-agnostic, competent in performance, and data-scalable. Given

---

these criteria, we evaluate existing task adaptation approaches and propose represeenlationnrrcssnnng decomposition understand their relative performances.
2.1 Problem statement and evaluation criteria
Our focus is on methods for adapting pre-trained large language models (LLMs) for downstream tasks. Specifically, given an LLM and limited labeled data for task, how does one adapt the model to the task? In this work, our focus is on classification tasks. For ease of presentation, we consider binary classification tasks in the main paper and defer the extension to multi-class setup to Appendix D. When evaluating a task adaptation strategy, we care about the following properties:
Task-agnostic. Given the general capabilities of pre-trained LLMs, we strive to utilize the same model across different tasks without requiring any task-specific training. With the increase in model sizes, the cost of deploying cask-specific models increase both during training (expensive hyper- parameter search) well during inference (deploying several models). In general, task agnossto methods will scale better with increasing model sizes by sidc-stepping both these costs.
Performance quality. We would like the adaptation approaches to be competitive in performance when compared with task-specific approaches across a wide range of tasks. For the binary classifica- tion tasks, the method should have accuracy comparable with task-specific approaches.
Data-scalable. The task adaptation method should be scalable with the number of labeled task examples. In particular, the method should be capable of learning from large datasets, and continually improve its performance quality.
2.2 Taxonomy of task adaptation strategies
We can broadly taxonomize the existing task adaptation strategies for LLMs as in-context learning. fine-tuning the model, and training task-specific adapters (see Figure 1).
In-context learning. In-context learning allows for adapting the model without updating any model parameters, by simply providing few demonstrations of the task in the LLM prompt. In-context learning is completely task-agnostic since the same model can be used across tasks since no weights updated at inference time. However, its performance is usually not par with task-specific methods and it does not scale well with data since the number of examples that can be utilized is bottlenecked by the context length of the model.
Fine-tuming, This traditional class of methods update the model weights to adapt it specifically for the task, typically by performing gradient descent over the labeled dataset. Fine-tuning methods not task-agnostic since they change the underlying model significantly but usually achieve state-of-the-are performance for any given task and are data scalable.
Adapters. Adapters adapt the underlying L.MM to a specific task by composing the LLM base model with an additional set of parameters which are optimized for the task. In contrast to fine-tuning which performs updates to the base model, adapters keep the base model frozen and only update the additional parameters. Performance of adapters is usually competitive with full fine-tuning.
2.3 Understanding performance via Representation- decomposition
From the taxonomy of task adaptation approaches, only in-context learning satisfies the task-agnostic property but it consistently underperforms the task-specific tuning approaches. This section investi- gates why this performance gap exists. We hypothesize that it is either because (a) the representations learned by the L.LM are insufficient to learn good predictor for the specific task, or (b) the LLM lacks the capability reason over these representations to make good predictions for the task.
To understand whether the representations have sufficient information, we train task-specific linear classifier using these representations, also known as linear probing. and evaluate its accuracy (AcCLR). Using this as an intermediate, we decompose the performance gap
Apert = ACCFT AccicL
ACCFT
ACCUR
ACCLR
ACCICL
(1)
Arep
Arros
where Arep represents the gap in performance which can be attributed to insufficient representation capacity and Areas is the performance gap due to insufficient reasoning abilities. Mathematically, we define the quality of models representations to be the gap between the fine-tuned model's 's accuracy

---

0.80
Reasoning gaine
Fine-lune (lask)
0.9
Fine-lune (eg news)
Acourady
0.75
Incorteet leaming Uneer proting Fine-tuiing
Accuracy
0.70
0.65
0.7
0.60
-0.1 0.0 0.1 02 -0.1
0.5
50
20
Number examples
Representation gains
Number of examples
(a)
(b)
(c)
Figure 3: All results for GPT-NEO (125M). (a) Accuracy of in-context learning linear probing model embeddings: representations have sufficient information. (b) Fine-tuning majorly improves task-specific reasoning across datasets. (c) Accuracy (averaged across datasets) of task-specific fine-tume tne-tumed model vs. accuracy of model fine-tuned on AGNews and evaluated on task. Fine-tuning hurts task-agnosticity.
and the accuracy of linear probing the smaller the gap the better the representations. Similarly, the models reasoning ability, or ability to perform linear classification, measured by the gap between the in-context learning and linear probing accuracies if this gap is small, that means the model does not improve by learning any task-specific decision boundary and therefore already possesses the ability to reason.
Using the above decomposition, we consider the following hypotheses:
HI. LLM representations have enough information to perform the task in-context, but they lack the reasoning abilities to perform the task well.
H2. Fine-tuning affects both the representations and reasoning but the improvement in reasoning abilities primarily leads better performance.
H3. tine-tuning and adapters are not task-agnostic because the task-specific training hurts their ability transfer reasoning.
We now analyze each of the task adaptation approaches through the lens of the above hypotheses. We perform all experiments with three different classes of language models (GPT-NEo, PYTHIA, BLOOM) across collection of 6 binary classification tasks. See Appendix B for further details.
In-context learning: LLMs lack reasoning abilities. We begin by studying the representation and reasoning gaps, as defined in eq. (1), for in-context learning. In Figure 3a, we plot the average accuracy across datasets for in-context learning, task-specific fine-tuning, and linear probing. that across models and different numbers of in-context examples, the reasoning gap Areas accounts for up to 79.1 of the performance gap between in-context learning and fine-tuning. This indicates that the LLM representations have sufficient information but they lack the ability reason them.
see
Fine-tuning: Improves task-specific reasoning. We next investigate how fine-tuning for a specific task affects the performance of the base model. In Figure 3b, show scatter plot of the gains that can be attributed to improved representations against the reasoning gains. We see that, across models, reasoning improvements accounts for 73.06% of the improvements. This indicates that while fine-tuning improves both reasoning and representations of the L.M,, the gains are predominantly due improvements in task-specific reasoning. Furthermore, this task-specific fine-tuning of the LLM hurts its performance on other tasks. In Figure 3c, we show that the accuracy of a model fine-tuned on the AGNews dataset [48], leads an average decrease of 25.77% on other tasks. Furthermore, this drop in accuracy can be attributed to the drop in task-specific reasoning capabilities ttese account for 72.58% of the drop (see Appendix B for more details).
Adapters: Impairs aanossigtyy via reasoning. Task-specific adapters do not change the underlying representation ability of the model. To study their ability to generalize across tasks, we train an adapter for the AGNews dataset and evaluate it other tasks. In Appendix B, we show that the performance drops across tasks by an average of 19.8%, indicating that adapters only learn task-specific reasoning abilitics.
5

---

TART: Task-Agnostic Reasoning Transformers
The above analysis showed how is the effective reasoning capabilities of the LLMs which limits its performance when compared with task-specific adaptation approaches. Building on this insight, we propose TART, which learns general-purpose reasoning module completely agnostic to the underlying base LLM and when composed with any LLM via its embeddings, generically improves upon its reasoning abilities. TART is a completely task-agnostic method which works across a suite of tasks without any task-specific training.
3.1 Overview of algorithm
TART comprises two components: generic task-agnostic reasoning module, and embeddings from the base LLM. The reasoning module is trained using only synthetic data (Gaussian logistic regression problems), agnostic of the auto-regressively trained language model, with the objective of learning to perform probabilistic inference (Section 3.2). This learned transformer module is then composed with the base LLM, without any training, by simply aggregating the output embedding and using those as an input along with the class label (Section 3.3). Together, these components make TART task- agnostic, boost performance quality by improving reasoning, and make the approach data-scalable by aggregating input embeddings into single vector.
3.2 Reasoning module: Can transformers learn to do probabilistic inference?
TART's reasoning module is a Transformer-basee model which is trained to perform probabilistic inference in-context using only synthetically generated data.
3.2.1
Training the reasoning module
The reasoning module is Transformer model which is auto-regressively trained family of logistic regression tasks, with each input sequence corresponding to a different logistic regression problem. We next describe the model architecture and the training procedure.
Model architecture. The reasoning module is based on the standard decoder-only Transformer architecture from the GPT-2 family. The GPT-2 architecture comprises causal (unidirectional) transformer with 12 decoder layers and 8 attention heads per layer for a total of 22 M parameters. The model has absolute positional encodings for each position. Please see Appendix C. for details. The architecture takes as input sequence of vectors and is trained to predict the next vector in the sequence. The input sequence consists of le pairs of labeled examples (æ11 31): (I22, v2),..., (rt. y/s). with each example covariates x and the other for the label y. This in contrast standard LLMs where each example spread over multiple tokens which limits how many examples can be put in the context. For example, with a context window 2048, our module can support 1024 examples while the base model can support only 10 examples, assuming each demonstration comprises 200 natural language tokens.
(Ia y/) using only two input positions of the transformer one for the
Training procedure. This module trained using gradient descent to minimize the population loss
((Tb)
BB..
(cE(
(2)
where 21:1-1 corresponds the first i examples and fcE is the cross-entropy loss evaluated the transformer prediction and the true yt. Each training sequence s, used to update the parameters comprises a different d-dimensional logistic regression problem, sampled as
Sequence 8: we ~ (O, Ia), Fi, ~ (0, Ia), Vi ~ (a((2,, wx))
for i E [k],
(3)
where a represents the sigmoid function and the multiplier a determines the noise level of the problem.
The above sampling model comprises three components, the regressor wi, the covariates 24, and the labels yi, sampled as
the regressor w; is sampled from a standard Normal distribution
the covariates Fi, sampled from standard Normal distribution
the corresponding Vi, sampled based the sigmoid of the inner product (mite, wx)
See Figure (Left) for a demonstration of a 2d dataset sampled from this process. The process is sampling noisy- linearly separable data in d dimensions. We train model with d = 16 and

---

0.5
0.5
Leamed Logittc Funcion
0.5
Absolute enor
0.4
TART inference module
0.4
1-1
0.70
Accuracy
Absolutte
0.
0.1
0.1
0.72
TART-Van
TART LOG
0.0
0.0
200
Number examples
Number
Number examples
(a)
(b)
(c)
Figure 4: Properties of TART'S inference module. (a) Comparison with learned logistic function: inference module recovers underlying probabilities. (b) Variation in error with different noise levels for model trained on = 10. (c) Comparison of TART performance when using LOO embeddings and vanilla embeddings.
and the training procedure in more detail in Appendix C.1. Observe that the embedding dimension of the base LLM and the input dimension of the reasoning module might not always match. In order to compose these models together, we perform dimensionality reduction via PCA to reduce the emhedding dimension.
= 256. We describe the model
3.2.2 Properties of reasoning module
The task-agnostic reasoning module described above is trained perform well on a family of logistic regression tasks. We study some properties of the reasoning module, in particular how well it learns to perform the task at an instance level and how robust it to variations in the noise level a.
Accuracy of probabilistic inference. For understanding the instance level performance of our reasoning module, we evaluate it sample of 64 different logistic regression problems, sampled according to eq. (3). For each problem, we train task-specific linear classifiers using logistic regression and compare them with task-agnostic reasoning module. In Figure 4a we plot the deviation of the predicted probabilities (averaged over the 64 problems) from the true probabilities for our reasoning module and the task-specific logistic solvers as function of the number of examples used for predictions. We observe that the error for our reasoning module decreases as function of the number of in-context examples and is within 2% of the task-specific logistic function.
Robustness to noise level. We study the robustness of the learned module to the noise levels, a, of the logistic regression problem. Recall that we trained our inference module by fixing the noise level 10. At inference time, vary the noise level to [0.5, 1, 10, 20), where lower values corresponds to noisier problem. The reasoning module generalizes to easier problem without any drop in accuracy but we make the problem harder (a = [0.5, 1)) the error increases progressively (see Figure 4b).
3.3 Role of representations: How to embed training examples?
The reasoning module composes with a base LLM through its final layer embeddings. A natural way to produce these embeddings is to place all the train examples in-context and then average the embedding vectors corresponding to the particular example (see Figure 2). At inference time, we append the test example to the training set, and average the embeddings corresponding to this example. We call these vanilla embeddings. Our experiments reveal that these embeddings seem to saturate (or even hurt performance) beyond certain number of in-context examples (see Figure 4c). One reason can be that the causal nature the model causes these embeddings to have asymmetric information the embeddings of each example is influenced by its preceding examples.
To counter this asymmetry, we propose leave-one-out (LOO) embeddings where the embeddings for each training point is formed by placing all the other train examples before it in the prompt such that all the embedding are formed with the same information content. In Figure 4c, changing the embedding style from vanilla to LOO consistently improves performance across models and tasks. embed much larger number of
The L00-emmeddingg help TART be data-scalable by enabling points than the context window can support. To do so, we use only subset of the train examples as the in-context prompt. The reasoning module, by its architecture design, can already accommodate many more examples than supported by the context window of the base LLM.
7

---

3.4 Theoretical analysis: Generalization of TART to natural language tasks
We study the generalization properties of the proposed task-agnostic method TART. Note that that the inference module is trained completely on synthetic data while evaluation time, input is the embeddings from a natural language task. In Theorem I we show that its performance on the natural language task depends on the distribution shift from the synthetic the distribution (see Appendix C.3 for formal statement and proof).
Theorem 1 (Informal). Let T represent the class of transformer models and Ts E T denote the trained reasoning module on set S of synthetic regression with Rayn sequences sampled from distribution Psyn in eq. (3). The error of the transformer Is when evaluated on distribution PNL over natural language sequences is
Prr S Wi(PNL, Psyn) (ynn+ + Comp(T) + Par (Ts)
(4)
T1syn
where W denotes the Wasserstein-1 metric, Comp(T) represents the complexity of class T, and eîr represents the error the empirical distribution.
A few comments are in order: The first term represents the distribution shift error between the true natural language task and the synthetic task. The second term corresponds to the generalization error on the logistic regression task, which can be made arbitrarily small because scales with ayyn the number of synthetic datapoints which can be generated without any cost. The third term above is the optimization error indicating how well has the reasoning module Tş fit to the synthetic training set.
Observe that for our algorithm TART, the distribution Pui. corresponds to the distribution induced on the embeddings of the natural language data (post the PCA and renormalization) while the distribution Payn is the one corresponding to the sampling process equation (3).
Experimental evaluation
We evaluate TART on a wide range of binary classification tasks across three domains: language, vision and audio. We demonstrate that TART improves base in-context performance and closes the gap with standard task task-specific strategies. We also conduct ablations to demonstrate that TART scales with model size and can support 10x more samples than in-context learning.
4.1 Experimental setup
Datasets. We briefly describe the datasets used, with details available in Appendix E. We consider 14 different binary classification tasks ranging from sentiment classification, news article categorization spam detection. The evaluation datasets include: SST [35], Rotten Tomatoes [29], SMS Spam [3], IMDB [27]. Civil Comments [8], AGNews [48], DBPedia [48], and the Youtube dataset [46]. Since AGNews and DBPedial4 are multi-class datasets, construct 4 binary classifi-
cation tasks from each dataset respectively. For each dataset, we truncate the input text to be at most 100 characters to enable us to fit sufficient number samples in-context.
Model families. We evaluate our method across three different families of models: GPT-NEO [6], PYTHIA [5], and BLOOM [33]. For our evaluations across datasets, we use GPT-NEO (125M), PYTHIA (160M) and BLOOM (560M). For ablations larger models, evaluate models with 1B parameters across each of the model families (i.e., GPT-NEO (1.3B), PYTHIA (1.4B) ) and BLOOM (1.7B)) and models with 3B parameters (i.c., GPT-NEO (2. (2.7B), PYTHIA (2.8B) and BLOOM (3B)). We additionally evaluate on GPT-J (6B) [37].
Baselines. We evaluate our models against all types of task-adaptation strategies described in Section 2.2: 1) in-context learning. 2) full fine-tuning, 3) last layer fine- fine-tuning, 4) LM head fine-tunnn fine-tuning, and 5) adapters. For each baseline, we perform an extensive hyper-parameter search over number of epochs and learning rate for each dataset in order to optimize performance (see Appendix E. for hyperparameter details). For TART, we chose base default set of parameters and use the same inference module with the exact same weights for all the experiments in this section.
8

---

0.8
0.8
Accuracy
Accuracy
Acourase
0.7
0.7
FT
conicel
0.7
0.6
0.6
Plutl Acoptar
Inconteal komiing
0.6
TART
TART
In-context learning
0.5
20
40
0.5
0.125
1.3
0.8
2
2
2'
2'
Number examples
Model Size (Billions)
Number of examples
(a)
(b)
(c)
Figure 5: Effects of scale. (a) Effect of number of in-context examples on performance for different task adaptation strategies. (b) Effect of model size on the performance of different task adaptation strategies. (c) Beyond context length limitations, performance comparison with respect number of in-context examples.
4.2 Natual language benchmark evaluations
Performance with respect to baselines. As shown in Appendix E.2, averaged across all tasks and model families, TART improves upon the base in-context learning performance by an average of 18.4 points, improves upon adapter heads by 3.4 points, and is within 3.1 points of full finettuning We also observe that TART consistently outperforms the task specific strategies LM head fine-tuning and last layer fine-tuning.
Performance on RAFT Benchmark We evaluate TART on all binary classification tasks in the RAFT Benchmark, following the protocol used in HELM [22]. When applied with GPT-NEO (125M), TART [0.634] outperforms BLOOM (176B) [0.595], and is within 4% points of GPT-3 (175B) [0.673], both of which are 1000x larger in size.
Performance with number of in-context examples. Our results demonstrate that performance of TART scales with number of in-context examples (see Figure 5a). Across 14 tasks and model families, when scaling from 18 to 64 examples, TART improves performance by an average of 4.8%. Correspondingly, full fine-tuning improves performance by 9.0%.
Scaling with base model size. We analyze how different task-adaptation strategies scale with respect model size using the GPT-Neo family: GPT-NEO (125M), GPT-NEO (1.3B) and GPT-J (6B). Figure 5b shows that when scaling from 100M to 6B parameters, performance of task-specific methods and TART increases using the same inference module across model sizes. Furthermore, the difference in performance between TART and fine-tuning baseline reduces from 7.5% 2.2% from the 100M scale 6B scale.
function scale. For TART, the performance increases by 9.8% while
Beyond context length. We evaluate the data-scaling properties for both in-context learning and TART (Figure 5c). To demonstrate the scaling property, do not truncate the input text 100 characters and utilize the entire text sequences. For TART, we observe that accuracy continues to improve when scaling from 18 256 in-context examples with 6.8% lift in performance. In comparison, ICL, which is bottlenecked by context length, supports 10x less samples, with the context window saturating at 24 examples only and lags TART by an average of 19.1%.
Extension to other architectures. We compare the performance of TART by changing the underly- ing architecture used to train the reasoning module. Specifically. we train a reasoning module that has the Hyena architecture [30] using the same synthetic logistic regression tasks. In appendix E.3, we show that the accuracies obtained by swapping the Transformer with the Hyena module maintains the same accuracy as before. Thus, TART's s performance is robust to the underlying architecture.
Performance in small sample regime. While our focus has been on scaling up number of in- context examples (ranging from 18-64), in Appendix E.4 we compare the performance of TART with fine-tuning and in-context learning the small sample regime. Specifically, we look at few-shot learning with number of examples varying from 4-10. Across 6 datasets, TART outperforms fine- tuning by (with max. gains up to 20.6%) and in-context learning by 1 11.6% (with max, gains up to 35.7%). See Figure 26 for detailed results.

---

1,00
0.95
1.00
0.96
Accuracy
Accuracy
0.90
Accuracy
0.90
0.97
0.85
TART
TART Adaptor FT-tul
TART
0.80
Adapter
FT-full
0.75
FT-ful
2'
0.70
Number of examples
Number examples
Number examples
(a) MNIST
(b) CIFAR-10
Speech Commands
Figure 6: TART can generalize across domains using the same inference module that was used for language benchmarks: Performance across vision tasks (MNIST, CIFAR- 10) and an audio task (Speech Commands).
4.3 Extensions to other modalities
We demonstrate that TART is not only agnostic to models and tasks, but also modalities. We extend TART to classification tasks on modalities beyond language: vision and audio. For vision tasks, we use representations from Google's 's 307M parameter pretrained Vision Transformer (ViT) model [44]: VIT-LARGE-PAN PATCH 1 6-224. For audio tasks, we use representations from OpenAl's 1.5B parameter pretrained Whisper model [31]: WHISPER-LARGE. In applying TART the representations from these models, we provide way for performing in-context learning in modalities beyond text. We refer the reader to Appendix E.5 for further details on the experiment setup.
Vision application. We evaluate the performance of TART binary classification versions of MNIST [19] (classes and 8) and CIFAR-IO [17] (classes plane and bird). As shown in Figure 6a and 6b, performance of TART is competitive with task-specific adaptation approaches.
Audio application. We evaluate TART on a binary classification version of the Speech Commands dataset [40], where the task is to classify "stop" and "go" utterances. As shown in Figure 6c, performance of TART is competitive with aask-adaptation approaches.
4.4 Extension multi-class classification problems
While we focused on binary classification tasks in this paper, in Appendix D we demonstrate how to extend TART to multi-class problems. Instead of training module on synthetic binary classification problem, we modify the synthetics to sample from a multinomial extension of the logistic model. We evaluated this multi-class model on 4-class classification tasks (AGNews, Ecommerce) as well as a 14-class classification problem (DBPedia-14), Our evaluations show that TART can outperform in-context learning by up 45.9% on AGNews, 52.6% on Ecommerce and 59.3% on DBPedia- 14.
5 Discussion
We look at the problem of task-agnostic learning with LLMs. We show that LLMs lack the ability perform simple reasoning over their learned representations and introduce TART, task, model and domain agnostic method for improving their reasoning abilities. In this work, we focus on classification tasks, showing that synthetic. logistic regression task data can be used train generic reasoning module capable of completing this class of tasks. In future work, we seek to understand whether synthetic tasks exist for training other generic reasoning modules, capable of improving base LLM performance on tasks such as generation or summarization.
Acknowledgements
We are grateful to Simran Arora, Rishi Bommasani, Niladri Chatterji, Arjun Desai, Sabri Eyuboglu, Neha Gupta, Karan Goel, Erik Jones, Ananya Kumar, Cassidy Laidlaw, Megan Leszezynski, Piero Molino, Laurel Orr, Michael Poli, Dimitris Tsipras, Michael Wornow Ce Zhang, and Michael Zhang for their helpful comments and feedback, and discussions which helped shape this project. We thank Neel Guha for helpful discussions in shaping the narrative of this paper.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML);
10

---

US DEVCOM ARL under No. W911NF-21-2-022 (Interactive Human-Al AI Teaming); ONR under No. N000141712266 1226666 (Unifying Weak Supervision); ONR N00014-20-1-2488. Understanding and Applying Non-Fuclidean Geometry in Machine Learning: N000142012275, (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. CDS was supported by NSF CAREER (award 2046760).
The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions recom- mendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.

---

References
[1] Monica Agrawal, Stefan Hegselmann. Hunter Lang, Yoon Kim, and David Sontag. Large language models are zero-shot clinical information extractors. arXiv preprint 2Xiv:2205.12689 2022.
[2] Neel Alex, Eli Lifland, Lewis Tunstall, Abhishek Thakur, Pegah Maham, C Jess Riedel, Emmie Hine, Carolyn Ashurst, Paul Sedille, Alexis Carlier, al. Raft: A real-world few-shot text classification benchmark. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[3] Tiago A. Almeida, Jose Maria Gomez Hidalgo, and Akebo Yamakami. Contributions to the study of sms spam filtering: New collection and results. In Proceedings of the 2011 ACM Symposium on Document Engineering (DOCENG' II), 2011.
[4] Simran Arora, Avanika Narayan, Mayee F Chen, Laurel J Orr, Neel Guha, Kush Bhatia, Ines Chami, Frederic Sala, and Christopher Ré. Ask me anything: A simple strategy for prompting language models. ICLR 2023, 2023.
[5] Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: suite for analyzing large language models across training and scaling. arXiv preprint arXiv:2304.00 01373, 2023.
[6] Sid Black, Leo Gao, Phil Wang. Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, March 2021. If you use this software, please cite using these metadata.
[7] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2 07258, 2021.
[8] Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nu- anced metrics for measuring unintended bias with real data for text classification. CoRR, abs/1903.0456/1 2019.
[9] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33: 13:1877-1901, 2020.
[10] Shizhe Diao, Pengcheng Wang, Yong Lin, and Tong Zhang. Active prompting with chain-of- thought for large language models. arXiv preprint 2rXXiv:332..12444 2023.
[11] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598. 2022.
[12] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning, pages 2790-2799. PMLR, 2019.
[13] Edward Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
[14] Chip Huyen. Prompting vs. finetuning vs. alternatives, April 2023.
[15] Jan Kocoñ, Igor Cichecki, Oliwier Kaszyca, Mateusz Kochanek, Dominika Szydło. Joanna Baran, Julita Bielaniewicz, Marcin Gruza, Arkadiusz Janz, Kamil Kanclerz, et al. Chatgpt: Jack of all trades, master of none. arXiv preprint arXiv:2302. 10724, 2023.
[16] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. JCMI. 2022 Workshop on Knowledge Retrieval and Language Models, 2022.
12

---

[17] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[18] Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Fine- tuning can distort pretrained features and underperform out-of-distribution. arXiv preprint arXiv:2202. 10054, 2022.
[19] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann..e www.nn.wee......o. ceccu..com/ex///mmsss.. 2, 2010.
[20] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efin efffcieet prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
[21] Xiang Lisa Li and Percy Liang. Prefix-tuning Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pages 4582 4597, 2021.
[22] Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. arXiv preprint arXiv:221 1.0022211001100 1.09/ 10, 2022.
[23] Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. Advances Neural Information Processing Systems, 35:1 35:1950-1999 1965, 2022.
[24] Jiachang Liu, Dinghan Shen, Yizhe Zhang, William B Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? Proceedings of Deep Learning Inside Out (DeeLlO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pages 100-114, 2022.
[25] Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P- tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Pupers), pages 61-68, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[26] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8086-8099 8098, 2022.
[27] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Language Technologies. pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
[28] Avanika Narayan, Ines Chami, Laurel On, and Christopher Ré. Can foundation models wrangle your data? Proc. VLDB Endow., 16(4):738-746, dee 2022.
[29] Bo Pang, Lillian Lee, and Shivakumar Vaithyanathan. Thumbs up? sentiment classification using machine learning techniques. In Proceedings of EMNLP, pages 79-86, 2002.
[30] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866 2023.
[31] Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision, 2022.
[32] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training, arXiv preprint, 2018
13

---

[33] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, al. Bloom: I 176b-paramen open-access multilingual language model. arXiv preprint arXiv:221 1.05200 2022.
[34] Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[35] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631-1642, 2013.
[36] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 2022.
[37] Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https: //gpthub com com/kinggfllllz/meelee transformer jax, May 2021.
[38] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Lc, Ed Chi, and Denny Zhou. Rationale- augmented ensembles in language models. arXiv preprint arXiv:2207.00747 2022.
[39] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, and Denny Zhou. Self- consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171, 2022.
[40] Warden. Speech Commands: A Dataset for Limited- Limited-Vocabulary Vocabulary Speech Recognition. ArXiv é-prints, April 2018 8.
[41] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022. Survey Certification.
[42] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou, Chain of thought prompting elicits reasoning in large language models, arXiv preprint arXiv:2201. 1903, 2022.
[43] Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.00 230.00366, 2023.
[44] Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision, 2020.
[45] Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. In Advances in Neural Information Processing Systems, 2022.
[46] Jieyu Zhang, Yue Yu, Yinghao Li, Yujing Wang, Yaming Yang. Mao Yang, and Alexander Ratner. WRENCH: A comprehensive benchmark for weak supervision. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
[47] Renrui Zhang, Jiaming Han, Aojun Zhou, Xiangfei Hu, Shilin Yan, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao. Llama-adapter: Efficient fine-tuning of language models with zero-init attention. arXiv preprini arXiv:2303. 230..16999, 2023.
[48] Xiang Zhang. Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
14

---

Fire-turingg (beok)
In-corteat warning
tuningg (syntholic)
0.9
0.8
Accuracy
0.8
Accuracy
0.7
0.7
0.6
0.6
Fine luning (lask)
0.5
In-context learning
Youtube
AG-Newa-O
DBPedia-O
Rotten Tomatoss
SMS Spam
SST
mnn-ttnnng (synthetic)
0.5
20
30
40
50
60
Number of examples
(a)
(b)
Figure 7: Fine-tuning with NL synthetic task. (Left) Averaged over 6 different tasks, fine-tuning with the NL
synthetic task provides lift over base in-context learning, and scales with number of examples. (Right) Dataset level comparisons between task-specific fine-tuning, in-context learning and synthetic fine-tuning: synthetic fine-tuning outperforms base in-context learning on 4 out of datasets, but lags task-specific tuning.
A Fine-tuning model with NL-based Probabilistic Inference Tasks
As highlighted in Section 1, we describe the details for directly fine-tuning an LLM on syntheti- cally generated probabilistic inference tasks improve reasoning capabilities. For the following experiments, we use GPT-NEO (125M) as the base model.
A.1 Training Task
We fine-tune the base model using a sequence pairs of synthetically generated labeled natural language examples (x, y). Each example in the sequence = (æ1, y1 v1),... (Ix. Mk-) consists of a list of strings constructed from fixed V size of dimension d 30 We use the following fixed vocabulary: "ssports", "love", "hate", "car", "school", "family", "work", "sleep", "water", "tree", "fox", "train", "random", "movie", "music", "book", "play", "house", "spell", "bar", "jump", "park", "run", "hill", "fast", "slow", "talk", "wallet", "orange". "apple", "ball", "cat"
To generate particular example z4, we sample each coordinate z4,j uniformly from the set {-1,+11 +11. If the sampled value is +1, we set the value to be the corresponding word in the vocabulary, that is, Jijj Vy. Otherwise, the word Tıf is set "null". For given sequence s, we generate each of the labels { as:
iyy -, Ia): Yi ~ (a w)), for [k]
(5)
where we set noise parameter a If the sampled output is O, we set the y: "negative" and "positive" otherwise.
Ty ys and the Finally, the inputs are formatted with following template: "x1 : y1 22 y2 model is trained using gradient descent on the loss
((Ta) = Ez
(6)
where #1-i-1 corresponds to the first i examples and ICE is the cross-entropy loss evaluated on transformer prediction and the true ye.
More concretely, sample input sample sequence s to be used for training looks like:
"sports love null car cat: positive, mull love null car null: negative,
sports mull hat mull cat positive"
15

---

Accursay
0.7
Incontest leaming
Unear probing
Accuracy
0.7
Invormtett leaming
Unear prozing
Accuracy
0.7
In-comees leeming
Unear probing
Fine suningg
Fine-tuiing
Fine-tniing
0.6
80
examples
Number examples
examples
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
BLOOM (560M)
Figure 8: Comparison of linear probing, in-context learning, and fine-tuning. Accuracy of in-context learning linear probing model embeddings across three model families: representations have sufficient. information.
A.2 Training Parameters
We train GPT-NEo (125M) on this synthetic task with a learning rate of 0.0001 and batch size of For each sequence we sampled a total of 60 examples and trained the model for 10000 steps.
A.3
Evaluation
We evaluate on 6 datasets: AG News [48], DBPedia [48], SST [35], SMS Spam [3], Youtube [46] and Rotten Tomatoes [29]. We truncate the input texts to 100 characters to fit more in-context examples. We evaluate over a range of context sizes (k=[18, 32, 48, 60).. evaluation time, we use the same "sentence label" format that was used to train the model. We evaluate over 3 random seeds. In Figure 7, we compare the performance of the model fine-tuned on probabilistic inference tasks and the base in-context learning performance. While the performance of the fine-tuned model is better than the base in-context learning capabilities, task-specific fine-tuning still outperforms it by an average of 16.87% (see Figure 7).
B Details for Representation-Reasonnng decomposition evaluations
In this section, we provide details for the experimental evaluation and additional results for the representation-reasoiigg decomposition introduced in Section 2.3.
B.1 Experimental setup
For these experiments, we evaluate three different language models: GPT-NEO (125M), PYTHIA (160M), and BLOOM (560M) on a collection of 6 binary classification datasets: AG News [48], DBPedia [48]. SST [35], SMS Spam [3], Youtube [46] and Rotten Tomatoes [29]. For each model, we run evaluations for three different random seeds, where the randomness was the set datapoints chosen for the training task. For the hyperparameterss, models across datasets. For details on these hyperparameters and the adapter architecture we evaluate over, Appendix E.
performed an extensive search for all
To conduct linear probing over the embeddings, we perform logistic regression over the output embeddings of each model and the given labels in the training set using the built-in logistic regression solver from the scikit-learn python library, utilizing the Ibg/s solver.
B.2 Detailed results
For each class of methods the task- aak-aadpt taxonomy from Section 2.2, we now describe the details of the experimental evaluation and present additional results.
In-context learning. To understand the representation and reasoning gaps for in-context learning, we evaluated three accuracies: a) using in-context learning with base models, b) fine-tuning the model for the task, and c) linear probing the model specifically for the task. The gap due to representation
16

---

Linear probing
Linear probing
Linear probing
Uun probinng
0.6
Exampies (K) 18
Examples (K)
Examples (K) 48
Examples (k) 64
0.5
0.5
0.8
0.7
0.8
0.5
0.8
0.7
0.8
0.5
0.6
0.7
0.0
0.5
0.6
0.7
0.0
0.9
1.0
In-context learning
In-context learning
In-context leamiing
In-context learning
(a) GPT-NEo (1 (125M)
Linea probing
Linear probing
Linear probing
Linear probing
0.7
Examples
Examples
(k)
(k)
0.6
0.9
06
0.9
0.5
0.9
In-context learming
In-context learning
In-conlext learning
In-context learning
(b) PYTHIA (160M)
Linear probing
Linear probing
Linear probing
Lineat probing
Examples (k) 18
Examples (K) 32
Examples (K) = 48
Examples (k) 64
0.5
0.8
0.7
0.5
0.6
0.7
0.8
0.5
0.8 In-conlext learning
0.8
0.7
0.5 0.6
0.7 0.8 0,9
In-context leamming
In-context learning
In-context leaminn
(c) BLOOM (560M)
Figure 9: Linear probing vs. in-context learning. Scatter plot of accuracy of in-context learning vs. linear probing on model embeddings across model families and different number of in-context examples: linear probing consistently outperforms in-context learning indicating that the learned representations have sufficient information. Each point in the plot represents dataset.
was taken to be the difference between the fine-tuning and linear probing accuracies while the reasoning gap was the gap between linear probing and in-context learning, as described in eq. (1).
In Figure 8, we show the average accuracies of in-context learning, linear probing, and fine-tuning across the 6 tasks. Linear probing closes the gap between in-context learning and fine-tuning, while being task-specific. In Figure 9, we show scatter plot of the accuracies of linear probing vs. the accuracies of base in-context learning. Linear probing consistently out performs in-context learning showing that the learned representations across these models have sufficient information to complete the tasks but lack reasoning abilities.
Fine-tuning. For the fine-tuning approach, we are interested in understanding two hypotheses: a) how does fine-tuning improve the model performance, and b) whether fine-tuning hurts task- agnosticity of the base model and if yes, what is the underlying reason for it.
For the first hypothesis, representations of the the underlying model. This is computed as the difference in performance of linear probing over the base model and over the fine-tuned model this evaluates how much the representations have changed specifically for this task. The reasoning gains are then computed by subtracting the representation gains from total gain (fine-tuning accuracy minus in-context accuracy). Figure 10 shows scatter plot of these representation gains and reasoning gains, plotted across different datasets and number of examples (k). Most of the gains which are realized by fine-tuning are because of improved task-specific reasoning capabilities across the model families.
evaluate the proportion of gains that can be attributed improved
17

---

0.4
0.4
0.4
Reasoning gains
Reasoning gains
Reasoning gains
0.2
0.2
0.2
0.1
0.1
0.1
0.0
-0.1 0.1
-0.1
-0.1
0.2 Representation gains
-0.1
0.0
0.1
0.2
-0.1
0.1
0.2
Representation gains
Representation gains
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
(c) B1.00M (560M)
Figure 10: Effects of fine-tuning on reasoning. Acmoss datasets (each point in plot represents dataset) and model families, fine-tuning improves task-specific reasoning which improves it performance over base in-context learning.
1.0
1.0
1.0
Fine-tuning
Fine-luning
Fine-luning (task)
0.9
Fine-luning (AG News)
0.9
Fine-luning (AG News)
0.9
Fine-luring (AG News)
Acouracy
Accuracy
0.8
0.7
0.7
0.7
0.5
0.5
0.8
Number of examples
Number examples
Number of examples
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
(c) BLOOM (560M)
Figure 11: Effects of fine-tuning on task-agnosticity Accuracy of task-specific fine-tuned model vs. accuracy of model fine-tuned on AG-News-0 and evaluated task. Fine-tuning hurts task-agnosticity across all three model families.
Freeii (grasa)
Pienku op (944)
0.9
0.9
0.9
0.8
Accuracy
0.8
Accurary
0.8
0.6
0.8
0.6
0.5
0.5
0.5
04
0.4
Youlbe
DePedin-
Rotten Tomatono
BMS Spam
SST
Youlsbe
DiPeslod
Rotten Tometoid
GMS Spam
SST
Youdume
DBPediod
Rollen Tomatons
BMS Spum
SST
GEO-NEO (125M)
(b) PYTHIA (160M)
BLOOM (560M)
Figure 12: Effects of fine-tuning on tanntagnoy (dataset level) Accuracy of task-specific fine-tuned model vs. accuracy of model fine-tuned on AG-News-0 and evaluated on task. Fine-tuning consistently hurts lask-agnosticity across all three model families and datasets.
For the second hypothesis, we first evaluate whether fine-tuning hurts task-agnosticity. For this evaluate two sets of accuracies: accuracy of model fine-tuned for the specific task and the accuracy of model on the task but fine-tuned on the AG News dataset. From Figures 11 and 12, we see that there is drop in accuracy over 25.77 77% across models and datasets. For the second part, we again decompose drop in accuracy into a representation drop and a reasoning drop. The representation drop is computed by training a linear probe over the two models (task-specific fine-tuned and AG News fine-tuned) and looking at difference between them. The reasoning drop, as before, computed by subtracting this representation drop from the total drop. Figure 13 shows that most of
18

---

0.5
0.3
Reasoning loss
Reasoning loss
Reasoning loss
0.3
0.2
0.2
0.2
0.1
0.1
0.0
0.0
-0.1
0.0 0.1
0.2
0.3
0.0 0.1
0.2
0.3
0.0 0.1
0.2 0.3
Representation loss
Representation loss
Representation loss
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
(c) B1.00M (560M)
Figure 13: Effects of fine-tuning task agnosticity. Scatter plot of reasoning loss against representation loss when base model trained on AG-News-O and evaluated on other tasks. Across datasets (each point plot represents a dataset), fine-tuning majorly impairs reasoning when transferring to tasks outside the specific fine-tuned task.
Adspar
Adapar (aah)
0.9
0.9
Acquracy
0.8
Accuracy
Accuracy
0.8
0.7
0.7
0.7
0.6
0.6
0.5
04
Yeslibel
DiPediiv
Rotten Tomatore
BMS Span
SIST
Youlobe
DBPedii-
Retten Tomutous
BMB spam
Vinatsbl
DiPPddivd
Rotien Tomatoen
GMS Spam
SST
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
BLOOM (560M)
Figure 14: Effects of fine-tuning on task-agnosticity of adapters. Accuracy of task specifc fine-tuned adapter vs. accuracy of adapter fine-tuned AG-News and evaluated task. Fine-tuning consistently hurts the generalization ability of adapters across datasets.
this drop in task-agnosticity be attributed to over-fitting of the reasoning abilities over the task for which the models are fine-tuned.
Adapters. Since adapters do not modify the underlying representations of the model, we look how single adapter generalizes across tasks. For this we train adapter the AG News dataset and evaluate it on the other datasets. We compare this set of accuracies with those obtained by task-specific adapters in Figure 14. The main conclusion is that task-specific adapters are not agnostic learners and over-fit to the task for which they are fine-tuned.
C Details for TART implementation
This section contains the details on training TART'S reasoning module and extended results on the choice of embeddings from Section 3.
C.1 TART'S reasoning module
Architecture details. We use the standard GPT-2 architecture [32] for training our reasoning module. We set the embedding size 256, number of decoder layers 12, and number of heads to for total of 22 million parameters. Since the GPT-2 backbone outputs sequence of embeddings,
19

---

Linear Layer
GPT-2 Backbone (Untrained)
Linear Layer
9x
Input Dimension d
Embedding
Output Embeddings
Dimension 4emb
Figure 15: TART reasoning module architecture. The reasoning module takes as input sequences of y) pairs of dimension linear layer is used to project the hidden dimension size of the GPT-2 backbone. Finally, linear layer is applied to the outputs of the backbone to generate predictions for cach 2g in the input sequence.
0.7
0.6
Loss
0.5
0.4
0.3
0.2
5000
10000 15000 20000 25000 Number of steps
Figure 16: Training loss vs. number of steps. The plot shows the variation training loss as a function of the number of steps of gradient descent for TART's reasoning module.
we additionally add a linear layer in the end to convert the output to scalar values (see Figure 155). Additionally, the binary labels y encoded one-hot vector to match the input dimension d of the corresponding covariates x.
Training procedure. We trained TART'S reasoning module with a context length 258 (allowing for up to 129 in-context examples). The batch size set to 64, learning rate 0.0001 and the model was trained for a total 24000 epochs. Each batch training data consists of sampling a sequence of 258 examples using eq. (3). In addition to these hyperparameters, on the input dimensions and on the number examples in the sequence to train our module the input dimensions started from value of 4 and were incremented by 4 every 1000 epochs while the number examples started from 18 and were incremented by 30 every 1000 epochs.
used curriculum
Combining reasoning module with base LLM. We trained the reasoning module with input dimension set to 16. However, most base models produce representations which are much higher dimensional (ranging from 784 to 2048). order to reduce the dimensionality of these representations, perform PCA on the output embeddings of the base model, learning the components using only training points available for that specific task. The test examples are then projected onto these principal components produce 16 dimensional input representations.
C.2 Choice of representations
As discussed in Section 3.3 there are two possible options for forming the representations, the vanilla embeddings and leave-one-out (LOO) embeddings. Figure 18 shows the schematic differences between the two style of embedding. In Figure 17, we plot the average accuracies across different datasets for both vanilla and LOO embeddings, observing that the I.OO embeddings consistently perform better across the different model families.
20

---

0.80
Accurscy
Accuraoy
0.78
Accurscy
0.74
0.74
0.72
TART-Vanila
TART-Vanila
TART-Vanila
TART-LOO
0.74
TART-LOO
0.72
TART-LOD
0.70
Number examples
examples
Number examples
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
BLOOM (560M)
Figure 17: LOO embeddings vs. Vanilla Embeddings. Comparison of TART performance when using LOO embeddings and vanilla embeddings. Vanilla embeddings see performance collapse, but 1,00 embeddings do
Embeddings
Ex.
Base Model Pre-trained LLM
Inference Module
Ex. k
Averaging Layer
(a) Vanilla embeddings
Embeddinge
Ex. 2
Base Model Pre-trained LLM
Averaging Layer
Ex.
Inference
K
Embaddlings
Base Model Pre-trained LLM
Ex. k-1
k
Averaging
Layer
(b) LOO embeddings
Figure 18: TART Embedding Protocols. (a) For the vanilla embeddings, the test example is appended to the
training set and sequence passed the base model. The representation for each train example this sequence is taken as the average emhedding across all tokens. (b) the LOO emheddings, we generate embeddings for each train example separately by placing all the other train examples before it in the prompt and averaging the embeddings over the final example's tokens.
21

---

C.3 Proof of Theorem
In this section, we provide a formal statement of Theorem 1 from Section 3.4. Our theorem quantifies
the expected error of the Transformer, trained on synthetic data, on natural language tasks in terms of the change in the two input distributions.
We begin by introducing some notation. We denote the class of Transformer family by
To = (Ta (2x+1)xd R
E +
(7)
where k represents the maximum number of in-context examples the Transformer can support, represents the input dimensions, and θ represents the corresponding parameter class over which the Transformer family is defined.
Observe that the Transformer family Te takes input sequence of k train examples, each corre- sponding to two tokens of hidden dimension d: a covariate x E Rd and binary label, encoded as a one-hot vector in dimension. This sequence of train examples is followed by a test example, for which we only have the features 2k+1-
Given this background, let Psyn denote the synthetic distribution over sequences { 3/ ).., (Irk. 1)). Similarly, let PNu. denote the corresponding distribution over sequences derived from natural language tasks where x; denotes the LLM embeddings of the example. Recall from Section 3.2.1, the synthetic training distribution Psyn given by
Sequence 9t
la),
Iie N/0 Ia), 3t. ((1,,,,,,,))))
for i [k] (8)
for each training point (Ii,t, The test point is also sampled similarly from independent standard normal distribution. Let l IR x + R be the loss function used for evaluating the performance of the reasoning module. Further, let use denote the expected loss under distribution P
eTr) : [(s.)))(((((()),,
and the corresponding empirical distribution over samples by eirp, where the dependence on the samples implicit. Given these samples, we denote the empirical risk minimizer
Ts arg min T'E
((((()))))
(10)
In addition these notation, we make the following Lipschitz assumption on the the loss function l and the Transformer model T.
Assumption 1. [Lipschitz loss.] For any two output labels y, v/ the loss function é is Lipschitz with constant C1, that is,
(e(Tr(s), v) E(Tr(s); s')l s cily-v/
(11)
Assumption 2. [Lipschitz models. For any two imput sequences s', each any model T E Te is Lipschitz with constant L, that is,
-(3)-T(s)) - L||s s'Il.
(12)
Given this setup, we are now ready to state a formal version of Theorem
Theorem 2 (Formal version of Theorem 1), Let Ts Te denote the trained reasoning module on set of synthetic logistic regression tasks with nsyn sequences sampled from distribution Payn in eq. (3). Let the loss function & satisfy Assumption probability at least 8. we have
the model class Te satisfy Assumption 2. Then, with
2VC(Te) Inm
lnm
2 2m(41/)
err (Ts) S ci max(1, L)- (PNNLL, Pon) + +
+4V
eir (Ts) (13)
Paayn
Tesyn
where W1 denotes the Wasserstein-I metric and VC(Te) represents the VC dimension of class Te
Proof. We begin by decomposing the error eir (Ts) into three components as
err (Ts) err (Ts) - errp (Ts) + pmn. (Ts) - eir (Ts) terrp p.. (Ts)
(14)
(
We now upper bound each of the terms (I) and (II) separately.
22

---

Bound on Term (I), Let denote an arbitrary joint distribution over the distributions Psyn and PNL. Then, can bound the first term as
err PhL (T') errp,, (T's) EPu (PAl(((((ss))) Ep (
E,le(T(s), ((T(5'),
inf inf (((T(s))))) - ((T((s)) 7
(15)
where (i) follows from the independence of the two expectations and (ii) follows from that (i) holds for any arbitrary joint distribution 7. The final bound on this term now follows:
inf B, (e(T(s),)) -((T())(T)))))))))))) inf - ((T(()))))))))))))))))) ((T(s) -((()))))))((((())))) ((T(s'),)))))
ca inf / ll -'1 s'l 10-(-I((s') - d(s)
7
inf / by 11- 1| s||d7 7
S max (1, D) (1
c1 max(1, L) Wi Payn)
(16)
where the inequalities (i) follows from Assumption and (ii) follows from Assumption This completes the bound Term (I).
Bound on Term (I.. Using a standard generalization bound [34, see Theorem 26.5], we have with
probability least 1
(4/8))
err (Ts) eir eir (Ts) < R(E Te) + +4V
Rayn
2 21n(4/5)
$ c1 RRTT) + +
Nsyn
2VC(Te) lum
2 2l((4/))
C1
+ 4
(17)
Tayn
Rayn
where R(LoTe) denotes the Rademacher complexity of the class To composed with the loss function l and inequality (i) follows from Sauer's Lemma.
Combining the bounds in equations (16) and (17) completes the proof of the theorem.
D Extension to multi-class classification
In this section, we present extension of TART to the multi-class classification setup. Recall from Section 3.2 that the module was trained on synthetic binary tasks sampled from the logistic model as
Sequence 3; : we ~ La), T1,t ~ (O, La), Vi. ~ ((( wx)) for i E [k] (18)
Observe that each sequence corresponded to a single hyperplane w, which defined the binary classification problem. For the multi-class extension, suppose that there are L labels. For each sequence 8t, we sample one hyperplane per class d-dimensional space:
Wt,1 (0 Ia) for I [L]
(19)
Each covariate Fa, is sampled again from the standard Normal distribution before. The label yi sampled from the multinomial distribution over [7.] given by:
P( = 1) x exp((xi,,, wy,)))
(20)
Given this distribution over sequences 8t, we train the TART module as before, but with the linear layer with output dimension L instead of as with the binary classification setup. We evaluate our
23

---

Max char
Max token length
Avg. char length
Avg. token length
Dataset
Num Labels
Test Size
length
AGNews
4
2000
814
238
236.97
52.31
Ecommerce
4
2000
50328
10544
678.768
158.02
DBPedia- 14
14
1988
691
240
284.08
67.52
Table Dataset (test) statistics for multi-clase datasets.
Dataset
Method
k: 18
32
k:
64
TART
63.55 (+38.5)
64.45 (439.6)
71.15 (+45.9)
72.35 (+45.4)
AGNews
ICL
25.1
24.9
25.3
27.0
TART
78.2 (+52.3)
79.4 (+52.6)
75.4 (+47.1)
79.3 (+50.8)
Ecommerce
ICL
25.1
24.9
25.3
27.0
TART
52.4 (+45.1)
63.8 (+52.8)
63.1 (+49.5)
68.4 (+59.3)
DBPedia-14
ICL
7.3
11.0
13.6
9.1
Table 2: TART vs. Base ICL Performance multi-class classification tasks. Evaluations of GPT-Neo- 125M model the multi-class setup. k = Number of in-context examples.
multi-class variant of TART on AGNews-4, Ecommerce-4 and the DBPedia-14 datasets (see Table for dataset statistics).
Table 2 shows the evaluation our proposed multi-class extension of TART on the above three datasets. While the in-context learning performance of the base model (GPT-Nco- 125M) is close to random, we see that the TART substantially improves upon this performance. It provides improvement of up to 45.9% AGNews, 52.6% on Ecommerce, and 59.3% DBPedia- 14 datasets.
E Details for experimental evaluation
We describe supplementary experimental details from Scction 4 as well as additional results for the natural language benchmark evaluations (Section E.2) and results for other modalities (vision and audio) (Section E.5).
F.1 Experimental setup
We begin by providing dataset statistics and details of the baselines.
E.1.1 Dataset construction and statistics
Table and 3 provides a detailed breakdown of dataset statistics. For each dataset, we use the original test with the exception of Civil Comments [8], AG News [48] and DBPedia [48], For the multi-class datasets-AG News [48] and DBPedia [48] - we construct 4 binary classification tasks for each datasets. More concretely, AG News labels news articles into four categories: World, Sports, Business, and Scieecee/Technoloee We create separate binary classification task for each category. sampling negatives from the remaining classes. DBPedia is 14-way ontology classification dataset. We create 4 separate binary classification tasks for the educational institution, company, artist, and athlete ontologies, sampling negatives from the remaining classes. For the train set, we sample class-balanced set of 64 examples from the original dataset. For each dataset, we sample 5 separate training sets, using different random seeds. In evaluations, we evaluate TART and the baseline methods across each of these different training sets.
24

---

Max char
Max token length
Avg. char
Avg. token
Dataset
Test size
length
length
length
AG-News-0
3800
732
259
237.07
51.36
AG-News-1
3800
814
213
232.01
51.46
AG-News-2
3800
814
225
236.10
52.25
AG-News-3
3800
892
259
234.86
51.38
Civil Comments
11576
1000
634
272.72
61.73
DBPedia-O
10000
2081
629
300.94
65.83
DBPedia-1
10000
2081
629
298.81
66.91
DBPedia-2
10000
2081
883
286.85
66.53
DBPedia-3
10000
2081
629
275.81
63.88
IMDB
25000
12988
2972
1293.79
292.82
Rotten Tomatoes
1066
261
63
115.52
25.36
SMS Spam
4181
612
258
81.46
23.76
SST
2210
256
60
102.40
22.34
Youtube
250
1125
292
112.50
31.84
Table 3: Dataset (test) statisties for all NLP datasets.
Max char length
Max token length
Avg. char length
Avg. token length
Dataset
AG-News-0
701
256
236.15
51.53
AG-News-1
749
180
232.48
51.61
AG-News-2
735
256
241.70
53.91
AG-News-3
1002
258
241.21
53.21
Civil Comments
1000
347
280.97
63.44
DBPedia-O
707
207
300.48
65.79
DBPedia-1 1
1023
280
299.89
66.57
DBPedia-2
628
203
288.21
66.22
DBPedia-3
758
203
279.45
64.24
IMDB
7068
1630
1284.20
290.00
Rotten Tomatoes
260
62
112.46
24.82
SMS Spam
911
217
106.90
31.87
SST
248
56
101.97
22.34
Youtube
1089
767
90.12
29.98
Table Dataset (train) statistics all NLP datasets.
0.8
0.B
Accuracy
0.7
burning
Accuracy
0.7
Inconiidd mamo
Accuracy
0.7
Akg
0.6
0.6
0.6
0.5
0.5
0.5
20
40
30
60
20
50
Number examples
Number examples
Number examples
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
BLOOM (560M)
Figure 19: Comparison of all methods. TART significantly improves base in-context learning performance and is competitive with fill-finttunng across model families.
E.1.2 Baseline methods
For each dataset, we compare TART to 4 baseline task-adaptation methods: 1) in-context learning, 2) full fine-tuning, 3) last layer fine-tuning, and 4) adapters. The last layer fine-tuning and the adapters trained as follows:
25

---

1.0
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.8
0.8
0.6
0
0.5
0.5
0.5
0.5
0.4
0.4
0.4
0.6
09
In-context learning
FT-full
FT-layer
Adapter
(a) Number of examples = 18
0.8
TART
0.7
0.6
0.6
0.5
0.4
0.4
0.4
04
0.4
0.5 0.8
0.7
0.8
09
1.0
0.4
0.5
0.8
0.7
088
0.9
1.0
0.4
0.5
0.8
0.7 FT-layer
0.8
0.9
1.0
0.4
0.5
0.6 0.7 0.8
0.9
1.0
In-context learning learning
FT-full
Adapter
(b) Number of examples = 32
1.0
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.4
0.5 0.6
0.7 08
0.9
1.0
0.4
0.5
0.8 0.7
088
0.9
1.0
0.4
0.5
0.8
0.7 FT-layer
0.8
0.9
1.0
0.4
0.5
06
0.7 Adapter
08
0.9
1.0
In-context learning
FT-full
(c) Number of examples = 48
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.9
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.6
0.8
0.6
0.5
0.5
0.5
0.5
0.4
0.4
0.4
0.6
09
In-context learning
FT-full
FT-layer
Adapter
(d) Number of examples = 64
Figure 20: Comparison TART and task-adaptation approaches (GPT-NEO (125M)). We see that for GPT-NEO ( (125M), TART outperforms in-context learning and is competitive with full fine-tuning and adapters across all &.
Last layer fine-tuning: Freeze all layers of transformer but the final transformer block and the language modeling head.
26

---

1.0
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.8
0.8
0.6
0.5
0.5
0.5
0.5
0.4
0.4
In-context learning
FT-full
FT-layer
Adapter
(a) Number of examples 18
TART
0.7
0.6
0.6
0.4
044
0.4
04
0.4
0.5
08
0.7
0.8
09
1.0
0.4
0.5
0.8
0.7
088
0.9
1.0
0.4
0.5
0.8 0.7 0.8
0.9
1.0
0.4
0.5
08 0.7 0.8 0.9 Adapter
In-context learning
FT-full
FT-layer
(b) Number of examples = 32
1.0
1.0
1.0
0.9
0.9
0.9
0.8
0.8
0.8
TART
0.4
0.5 06
0.7 08
0.
1.0
0.4
0.5
088 0.7
08
0.9
1.0
0.4
0.5
0.8
0.7 FT-layer
08
0.9
1.0
0.4
0.
06
0.7 Adapter
08
0.9
1.0
In-context learning
FT-full
(c) Number of examples = 48
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.6
0.8
0.6
0.5
0.5
0.5
0.5
0.4
In-context learning
FT-full
FT-layer
Adapter
(d) Number of examples 64
Figure 21: Comparison TART and adaatttpnn approaches (PYTHIA (160M)). We see that for РYТНAА (160M), TART outperforms in-context learning and adapters and is competitive with full fine-tuning across all &.
Adapter: Combine frozen LILM base transformer model with trainable adapter head ML.P composed of single linear layer followed by non-lincarity.
27

---

1.0
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.8
0.8
0.6
0.5
0.5
0.5
0.5
0.4
0.4
0.4
0.6
09
In-context learning
FT-full
FT-layer
Adapter
(a) Number of examples = 18
0.8
TART
0.7
0.6
0.6
0.4
0.4
0.4
04
0.4
0.5
08
0.7
0.8
09
1.0
0.4
0.5
0.8
0.7
088
0.9
1.0
0.4
0.5
0.8 0.7
0.8 FT-layer
0.9
1.0
0.4
0.5
0.6 0.7 0.8
0.9 1.0
In-context learning
FT-full
Adapter
(b) Number of examples = 32
1.0
1.0
1.0
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.4
0.5 0.6
0.7 08
0.9
1.0
0.4
0.5
0.8 0.7
08
0.9
1.0
0.4
0.5
0.8
0.7 FT-layer
0.8
0.9
1.0
0.4
0.5
06 0.7 08
0.9
1.0
In-context learning
FT-full
Adapter
(c) Number of examples = 48
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8
0.8
0.8
TART
0.7
0.7
0.7
0.7
0.6
0.6
0.8
0.6
0.5
0.5
0.5
0.5
0.4
0.4
0.4
0.6
In-context learning
FT-full
FT-layer
Adapter
(d) Number of examples 64
Figure 22: Comparison of TART and task adaptation approaches (BLOOM (560M)). We see that for B1.0OM (560M), TART outperforms in-context learning and adapters and is competitive with full fine-tuning across all k.
Hyperparameter search. For each baseline, we perform an extensive hyperparameter search over number of epochs and learning rate for each dataset in order to optimize performance. We search over a range of learning rates (1e-3, le-4, 3e-5, le-5, 8e-6), and range of epochs (5, 10, 15, 20, 50). For all models < IB parameters, we use batch size of 1. For all models > 1B parameters, we use batch
28

---

Model
TART
GPT-J (6B)
OPT (175B)
BLOOM (176B) GPT-3 (175B)
Accuracy
0.634
0.608
0.637
0.595
0.673
Table 5: RAFT (HELM) Binary Classification Performance (Average Accuracy). TART used with GPT- NEO (125M) model which is 1000x smaller than the corresponding 175B parameter models. TART outperforms
BLOOM (176B) and is competitive with OPT (175B) and GPT-3 (175B).
0.8
Acouracy
0.7
Acouracy
Accuracy
0.7
0.6
TART
TART
In-contedd learnin
In-context learning
In-context leariing
0.5
2'
0.
2'
0.5
2'
Number of examples
Number of examples
Number of examples
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
(c) BL.OMM (560M)
Figure 23: Beyond context window constraints. Performance comparison with respect number of in-context examples. Base in-context learning is bound with respect total numbers examples and performance saturates. TART is not bound by context length, and performance continues to scale as number of examples increases.
size of 8. We use these same batch sizes at evaluation time. We perform our hyperparameter searches with fixed number of train samples (64). We run our hyperparameter searches over 3 random seeds.
E.2 NL benchmarks
In this section, we provide additional results deferred from Section on the NLP benchmark evaluations, RAFT evaluations and demonstration of TART's data-scalability.
E.2.1 Performance on benchmark datasets
Figure 19 shows the performance of the baseline methods with TART averaged across the suite of 14 datasets. TART, while being task-agnostic, shows similar performance quality to tase-specific approaches across the different model families, and consistently outperforms in-context learning.
Figures 20, 21, and 22 show the scatter plots of the accuracies of TART with the baseline methods across datasets and different values of in-context examples k. An interesting observation is that as the number of examples k increases from 18 to 64, the performance of fine-tuning improves at better rate than that of TART.
F.2.2 Real-world Annotated Few-shot Tasks (RAFT) evaluation
For our evaluations on the RAFT benchmark [2], we follow the protocol (same train and test sets) used in HELM benchmark. The HELM benchmark [22] contains the evaluation results for many open and closed models enabling us to accurately compare the performance of TART with other models. We evaluate TART on all RAFT binary classification datasets (twitte--complainss, neurips-impace statement-riske overulling, ade-corpusv2, tweet-eval-hate terms-of-service tai-safety- -rsearc))) with the exception of wewlemmtti--rviee-ttte inclusion which contains zero positive samples in the train set. TART requires at least one example of each class in the training set. Table 5 contains a detailed performance comparison of TART with respect to other models. TART when combined with GPT- NEo (125M) is able to outperform BLOOM (176B) and is competitive with OPT (175B) and GPT-3 (175B), all of which have 1000x more parameters.
E.2.3 Beyond context length: TART is data-scalable
Setup. For these evaluations, we use the a subset of 6 datasets: AG-News-O, DBPedia-O, SST, SMS Spam, Youtube and Rotten Tomatoes. We evaluate the performance of TART over k=[18, 32, 48, 64, 128, 192, 256] where & is the number of in-context examples. When evaluating base models,
29

---

Embeddings
Pre-trained Model
Inference Module
Ex.
Ex. Averaging
Layer
Raw Signal- Image, Audio, Text
Stream embeddings. Another protocol for generating representations for in-context examples where each example is embedded by the base model separately.
Figure 24:
evaluate over k=[8, 24]-values given that the maximum input sequence length in the training set for AG News 256. With such sequence length, the maximum number of in-context examples that fit in the context-window is 8, hence the lower bound.
values of & that maximize the context window. We use lower-bound of 8
Embeddings. For these evaluations, In this setup, we use the context window of the LLM to encode single example at a time. The final embeddings then averaged and used in-context with TART'S reasoning module. This is in contrast to the vanilla and LOO embeddings which use multiple examples in-context with the base LLM to obtain the embeddings.
what call "streaming" embeddings (see Figure 24).
Evaluation. Figure 23 shows the performance of base in-context learning with TART across the three different model families. Observe that while in-context learning is bottlenecked by the context window of the base LLM, TART is able to learn from 10x more examples and exhibits increasing trend in accuracy with number of examples across models.
E.3 Generalization across architectures
In this section, we demonstrate that is possible to train TART reasoning modules on architectures be- yond transformers. More concretely, we train reasoning module that has the Hyena architecture [30] using the same synthetic logistic regression tasks.
E.3.1 Setup and prrpee-prraeters
We instantiate reasoning module with 12 Hyena blocks, hidden dimension size of 256, and sequence length of 2050. We train with a batch size of 16, using learning rate of 5e-05. We sample data with a noise parameter (a) equal we use the final checkpoint (i.e., 5000) of the reasoning module.
We train the model for 5000 steps. For our evaluations,
E.3.2 Evaluation
We evaluate the Hyena-based reasoning module applied GPT-NEO (125M) 6 datasets: SMS Spam, SST, AG-News-O, DBpedia- 18pedia-14-0,, Youtube, and Rotten Tomatoes. As seen in Figure 25, the Hyena-hased reasoning module transfers to natural language tasks and performs competitively with the transformer-baue reasoning module
E.4 Performance in low sample regime
In Figure 26, we evaluate our models in the low-sample regime where the number of in-context examples range from 4 10. TART outperforms both in-context learning and task-specific fine-tuning baselines.
30

---

TART-Hyena TART
TART-Hyana
TART
0.9
Accuracy
Accuracy
0.8
0.8
0.7
0.7
0.6
0.6
0.5
0.5
Youtube
AG-News-0
DBPedie-O
Rotten Tomatoss
SMS Spam
SST
Youtube
AG AG-News-0
DBPedia-O
Rotten Tomatbes
SMS Spam
SST
(a) k= 18
(b) k 32
TART-ifyera
TART
TART
0.9
0.9
Accuracy
Accuracy
0.8
0.7
0.7
0.6
0.6
0.5
0.5
Youtube
AG- AG-News-0
DBPedia-O 0
Rotten Tomatoes
SMS Spam
SST
Youtube
AG-News-O 0
DBPedia-0
Ratlen Tomatoes
SMS Spam
SST
(c) 48
(d) 64
Figure 25: Comparing TART with Hyena version of TART. Instead of using Transformer module, a Hyena model to learn the probabilistic reasoning task. Hyena-TART is comparable TART across datasets and achieves better performance larger value & (= 64).
E.5 Extension to other modalities: TART is domain-agnostic!
We begin by providng a description of the datasets we used to evaluate TART on audio and vision tasks, and then provide additional results comparing algorithm with baselines.
E.5.1
Dataset details
the Speech Commands (Version 0.01) dataset [40]. Speech Commands a multi-class classification task where the task is to detect preregistered keywords by elassifying utterances into predefined of words. We construct binary classification task over the keywords "stop" and "go", "up" and "down", and "yes" and "no" (see Table for more details).
For audio classification,
use CIFAR-IO [17] and MNIST [19]. Both tasks are multi-class classification tasks. We create binary classification tasks for each of the datasets. For CIFAR-10 the tasks are: airplane vs. bird, bird vs. horse, and ship automobile. For MNIST the tasks are: 0 vs. 1 vs. 6 and 2 vs. See Table for more details.
For image classification,
For both the audio and image datasets, we sample a class-balanced set of 256 samples from the training set. For the test sets, we filter the original test sets to only include samples of the two classes are learning to predict for (i.e., airplane and bird for CIFAR10 and 0 and 8 for MNIST).
31

---

Fine-turing
Fine-tuning
in-contest TART
context leamiin
TART
0.9
0.9
Accuracy
0.8
Accuracy
0.7
0.7
0.6
0.5
0.5
Youtube
AG- Ro-Newa-O
DBPedia-0
Rotten Tomatoes
SMS Spam
SST
Youtube
AG 1i--News-O
DBPedia-O
Rotten Tomatoes
SMS Spam
SST
(a) k 4
(b) k=
Flne-tuning
In-contest learning
In-coniest leamning
TART
0.9
0.9
Accuracy
0.8
Accuracy
0.8
0.7
0.6
0.6
0.5
Youtube
AG-News-0
DBPedia-O
Rotten Tomaloes
SMS Spam
SST
Youtube
AG-News-O 0
DBPedia-U 0
Rotten Tomatoes
SMS Spam
SST
(c) k
(d) == 10
Figure 26: TART in low sample regime
Dataset
Modality
Train size
Test size
MNIST (0 vs. 8)
image
256
1954
MNIST (1 vs. 6)
image
256
1940
MNIST (2 vs. 4)
image
256
2014
Speech Commands (stop go)
audio
256
500
Speech Commands (up vs. down)
audio
256
508
Speech Commands (yes no)
audio
256
525
CIFAR- 10 (airplane vs. bird) CIFAR-10 (bird vs. horse)
image
256
2000
image
256
2000
CIFAR-10 10 (ship vs. automobile)
image
256
2000
Table 6: Dataset statistics for all audio and image evaluation datasets.
E.5.2 Algorithms for comparison
For these evaluations, we use the "streaming embeddings" described in Figure 24 to obtain the embedding for TART. We evaluate over k=(18, 32, 48, 64, 128, 256).
compare against two baseline task-adaptation methods: 1) full fine-tuning and 2) adapters. We the same architectures described in Appendix 1.2. For vision tasks, Google's 307M
32

---

1.00
0.99
Accuracy
Accuracy
0.98
0.95
0.97
TART
TART
Adapter
Adapter
0.96
FT-full
FT-full
0.95
0.90
2"
2
2
2
2
2
Number of examples
Number of examples
(a) MNIST (1 vs. 6)
(b) MNIST (2 vs. 4)
Figure 27: Additional MNIST binary classification tasks. TART is competitive with task iasc-specific full fine- tuning and adapters.
0.95
0.95
Accuracy
Accuracy
0.90
0.90
0.85
TART
TART
0.85
Adapter
Adapter
0.80
FT-full
FT-full
0.80
2"
2
2
2
2
2
2
2
Number of examples
Number of examples
(a) CIFAR-10 (ship vs. automobile)
(b) CIFAR-10 (bird vs. horse)
Figure 28: Additional CIEAR-10 binary classification tasks. TART is competitive with task-specific full fine-tuning and adapters.
parameter pretrained Vision Transformer (ViT) model [44]: VIT-LARGE-PATCHI6-222 For audio tasks, use OpenAl's 1.5B parameter pretrained Whisper model [31]: WHISPER-LARGE.
Hyperparameter search For each baseline, we perform an extensive hyperparameter search over number of epochs and learning for each dataset in order optimize performance. We search over range learning rates (1e-3, 5e-04, le-4, 5e-5, 1e-5, and 8e-6) and range of epochs (5, 10, 15 and 20). For all models we use batch size 1. We perform our hyperparameter searches for fixed number of train samples (128) and run our hyperparameter searches over 3 random seeds.
E.5.3 Evaluation
We plot the accuracy function of the number of examples for TART, fine-tuning and adapter in Figure 27 (MNIST), Figure 28 (CIFAR-10), and Figure 29 (Speech Commands). TART is competitive with both these baselines, showing how task-agnostic methods can compete with task-specific adaptation methods across different modalities.
Broader Impact
Our submission focuses on understanding in-context learning and how it compares to tasc-specific fine-tuning approaches. We view our work as general understanding paper and, to the best of our knowledge, see no negative consequences of our work. We hope that our work furthers the state of understanding of large language models (LLMs).. Furthermore, because TART is plug-and-play solution that doesn't 't require any fine-tuning, we hope that our method can help domain experts (c.g., lawyers, scientists) incorporate LLMs into their workflows. In this sense, view TART
33

---

0.95
0.95
Accuracy
Accuracy
0.90
0.90
0.85
TART
TART
Adapter
Adapter
0.80
FT-full
FT-full
0.80
2"
2
2
2
2
Number of examples
Number of examples
(a) Speech Commands (yes vs. no)
(b) Speech Commands (up vs. down)
Figure 29: Additional Speech Commands binary classification tasks. TART is competitive with task- full fine-tuning and adapters.
contribution that is democratizing ML by making them more accessible and usable to experts in variety of domains.
G Compute Resource Estimates
Resources for training TART We use single NVIDIA RTX A6000 GPU ( ($2/hr) for hours to train our TART inference head, costing total of S18 to train.
Resources for training task-specific baselines We use 4 NVIDIA A100 GPU's $3.5 per GPU/hr) for 100 hours for hyperparamter tuning, costing total of $1,400. We use 8 NVIDIA A100 GPU's ( (3.5 per GPU/hr) for 50 hours to fine-tune all task-adaptation baseline models. Total cost for fine-tuning amounted to $1,440.
Resources for inference We use 8 NVIDIA A100 GPU's $3.5 per GPU/hr) for 480 hours for all inference runs costing total of $13,440.
34

---

Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0.790 0.036
0.552 0.049
0.724 + 0.070
0.852 0.019
0.824 0.031
AG-News-1
0,828 + 0.021
0.541 0.037
0.779 = 0.137
0.903 + 0.021
0.868 + 0.020
AG-News-2
0.751 0.023
0.513 = 0.014
0.626 = 0.057
0.765 + 0.025
0.755 0.012
AG-News-3
0,743 + 0.031
=
0.736 = 0.035
0.786
0.727 + 0.066
Civil Comments
0.559 + 0.027
0.499
0.520 = 0.033
=
0.535 0.036
DBPedia-O
0,866 = 0.030
0.611
0.802 = 0,020
0.825 0.012
0.837 = 0.022
DBPedia-1
0.778 0.036
0.579
0.766 = 0.056
0.740 0.041
0.778 0.044
DBPedia-2
0.798 0.042
- 0.136
0.862 = 0,041
0.908 0.011
0,832 = 0.048
DBPedia-3
0.812 0.032
0.611 0.135
0.817 + 0.034
0.859 0.025
0.848 0.028
IMDB
0.537 + 0.022
0.507 0.007
0.625 = 0.013
0.560 + 0.021
0.556 0.014
Rotten Tomatoes
0.535 0.030
0.550 0.043
0.689 = 0.019
0.541 0.037
0.524 0.018
SMS Spam
0,869 + 0.063
0.736 = 0.099
0.925 0.011
0.833 0.015
0.886 0.023
SST
0.544 + 0.021
0.542 + 0.024
0.715 = 0.009
0.555 + 0.026
0.547 + 0.009
Youtube
0.784 + 0.047
0.658 + 0.089
0.833 + 0.089
0.768 + 0.100
0.715 0.136
Table 7: Standard deviation of accuracy, number of examples = 18, GPT-NEO (125M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-O
0,805 + 0.029
0.498 = 0.002
0.758 = 0.127
0.601 + 0.072
0.669 + 0.044
AG-News-1
0.825 0.023
0.517 = 0.019
0.916 = 0.015
0.702 = 0.097
0.690 0.034
AG-News-2
0.758 + 0.027
0.500 = 0.000
0.596 = 0.031
0.500 0.001
0.588 0.026
AG-News-3
0.754 + 0.033
0.501 + 0.003
0.661 + 0.093
0.552 + 0.041
0.613 0.022
Civil Comments
0.575 + 0.019
0.500 + 0.003
0.525 + 0.049
0.500 + 0.000
0.515 0.008
DBPedia-O
0.861 + 0.018
0.508 + 0.016
0.786 + 0.054
0.638 + 0.109
0.704 0.038
DBPedia-1
0.813 0.020
0.499
0.787 = 0.067
0.599 + 0.078
0.710 0.059
DBPedia-2
0,870 + 0.035
0.502
0.878 = 0,071
0.701 + 0.113
0.767 + 0,052
DBPedia-3
0.850 0.050
0.002
0.864 = 0.034
0.603 0.114
0.734 + 0.030
IMDB
0,550 0.027
0
0.590 = 0.046
0.500
0.526 + 0.012
Rotten Tomatoes
0.544 = 0.027
0.491
0.589 = 0.074
0.500 0.000
0.522 = 0.023
SMS Spam
0,901 + 0.038
0.867 = 0.030
0.892 = 0.052
0.867 0.004
0.851 = 0.042
SST
0,572 0.016
0.517 0.002
0.617 = 0.074
0.517 0.000
0.539 0.007
Youtube
0.847 + 0.047
0.611 + 0.084
0.810 + 0.060
0.528 + 0.000
0.598 0.103
Table 8: Standard deviation accuracy, number examples 18, PythIa (160M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0.813 + 0.021
0.511 + 0.013
0.702 + 0.102
0.790 + 0.033
0.689 0.060
AG-News-1
0.851 + 0.020
0.511 + 0.014
0.801 + 0.086
0.891 + 0.037
0.752 0.032
AG-News-2
0.744 + 0.034
0.509 + 0.009
0.622 = 0.083
0.711 + 0.070
0.652 + 0.047
AG-News-3
0.763 0.026
0.508 0.014
0.626 = 0.016
0.775 0.035
0.706 0.027
Civil Comments
0.561 + 0.029
0
0.491 = 0.032
0.540
0.533 + 0.030
DBPedia-O
0.851 0.009
0.531 0.044
0.811 = 0.116
0.813 0.057
0.812 0.043
DBPedia-1
0.760 + 0.037
0.546 0.088
0.750 = 0.129
0.718 + 0.067
0.754 + 0.041
DBPedia-2
0,800 0.032
0.567
0.850 = 0.086
0.904 = 0.018
0.851 0.055
DBPedia-3
0.848 + 0.025
+
0.739 + 0.133
0.785 + 0.132
0.849 + 0.011
IMDB
0.552 0.031
0.550 0.044
0.630 + 0.016
0.608 0.014
0.526 0.025
Rotten Tomatoes
0.574 + 0.029
0.539 0.031
0.638 + 0.037
0.618 + 0.049
0.507 + 0.010
SMS Spam
0.830 0.112
0.613 0.261
0.883 0.130
0.911 + 0.035
0.885 0.045
SST
0.574 + 0.019
0.561 + 0.062
0.707 = 0.024
0.636 0.025
0.531 0.015
Youtube
0.762 0.1
0.584 + 0.144
0.753 = 0.140
0.726 0.052
0.769 0.084
Table 9; Standard deviation of accuracy, number of examples 18, BLOOM (560M)
35

---

Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0,808 0.030
0.551 + 0.041
0.795 + 0.049
0.874 + 0.022
0.830 0.034
AG-News-1
0,883 + 0.014
0.577 = 0.055
0.852 = 0.070
0.911 + 0.021
0.902 + 0.011
AG-News-2
0.764 0.019
0.540 0.027
0.705 = 0.063
0.812 + 0.019
0.782 0.019
AG-News-3
0,798 + 0.025
=
0.762 = 0.046
0.806 + 0.028
Civil Comments
0.575 + 0.054
0.498 0.002
0.534 = 0.029
0.022
0.579 0.044
DBPedia-O
0,886 0.021
0.632 0.096
0.884 = 0.019
0.866
0.859 = 0.020
DBPedia-1
0,809 0.031
0.593 0.062
0.802 = 0.030
0.783
0.813 0.023
DBPedia-2
0.886 0.011
0.586 0.078
0.078
0.916 = 0.029
0.922
0,899 = 0.017
DBPedia-3
0.868 0.022
0.565 + 0.041
0.902 + 0.016
0.908 0.012
0.868 0.023
IMDB
0.537 + 0.040
0.543 0.029
0.609 = 0.029
0.556 + 0.024
0.551 0.036
Rotten Tomatoes
0.549 + 0.035
0.554 0.036
0.667 = 0.036
0.550 0.045
0.528 0.020
SMS Spam
0,903 + 0.027
0.768 + 0.105
0.931 = 0.014
0.861 0.027
0.920 0.011
SST
0,573 + 0.009
0.564 + 0.045
0.711 + 0.022
0.594 + 0.022
0.551 + 0.022
Youtube
0.810 + 0.072
0.759 + 0.074
0.854 + 0.044
0.773 + 0.052
0.722 0.113
Table 10: Standard deviation of accuracy, number of examples = 32, GPT-NEO (125M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-O
0,835 + 0.020
0.499 = 0.002
0.791 = 0.100
0.855 + 0.017
0.712 = 0.028
AG-News-1
0,900 0.012
0.504 = 0.004
0.911 = 0.022
0.926 = 0.009
0.738 0.028
AG-News-2
0.773 = 0.012
0.510 = 0.006
0.687 = 0.034
0.742 = 0.044
0.683 = 0.025
AG-News-3
0.823 0.024
0.514 + 0.017
0.792 + 0.022
0.833 + 0.011
0.697 0.023
Civil Comments
0.596 + 0.024
0.499 + 0.007
0.588 + 0.052
0.536 0.020
0.542 0.016
DBPedia-O
0.890 0.016
0.514 0.026
0.880 + 0.049
0.777 + 0.055
0.782 0.044
DBPedia-1
0.833 0.018
0.522 0.023
0.834 = 0.061
0.768 + 0.024
0.760 0.022
DBPedia-2
0,912 + 0.010
0.033
0.916 = 0.033
0.903 + 0.012
0.859 + 0,019
DBPedia-3
0,879 + 0.021
0.520
0.900 0.019
0.863 0.018
0.812 0.022
IMDB
0,541 + 0.041
+.000
= =.011
0.575 0.024
0.540 + 0.018
Rotten Tomatoes
0,576 0.042
0.495 0.007
0.659 = 0.064
0.568 0.027
0.542 0.019
SMS Spam
0,937 + 0.017
0,856 0.080
0.913 = 0.044
0.953 0.014
0.877 = 0.036
SST
0,602 0.015
0.509 0.013
0.705 = 0.021
0.605 0.043
0.542 0.003
Youtube
0.862 + 0.045
0.626 + 0.081
0.874 + 0.046
0.762 + 0.052
0.654 0.050
Table 11: Standard deviation of accuracy, number of examples PYTHIA (160M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0.825 + 0.031
0.503 + 0.003
0.802 + 0.041
0.818 +: 0.014
0.739 0.049
AG-News-1
0.880 + 0.012
0.516 + 0.026
0.898 + 0.043
0.912 + 0.034
0.829 0.010
AG-News-2
0.771 + 0.009
0.519 0.025
0.736 = 0.071
0.773 0.039
0.684 + 0.040
AG-News-3
0,810 + 0.016
0.509 0.013
0.753 = 0.058
0.805 0.031
0.742 0.030
Civil Comments
0,587 + 0.025
0.500
0.549 = 0.016
0.564 0.062
0.555 + 0,009
DBPedia-O
0.857 0.035
0.592
0.801 0.125
0.822 0.043
0.834 0.038
DBPedia-1
0,802 + 0.031
0.558 0.047
0.870 = 0.037
0.800 0.041
0.813 + 0.034
DBPedia-2
0.879 0.018
0.609 0.025
0.938 = 0.018
0.920 = 0.031
0.903 0.021
DBPedia-3
0.866 + 0.035
0.634 0.110
0.812 + 0.163
0.911 + 0.004
0.876 + 0.037
IMDB
0.553 0.046
0.541 0.024
0.636 + 0.016
0.600 0.018
0.536 0.024
Rotten Tomatoes
0.589 + 0.037
+
0.713 + 0.029
0.630 0.026
0.527 + 0.015
SMS Spam
0.933 0.017
0.762 0.033
0.956 0.026
0.905 0.056
0.917 0.017
SST
0.579 + 0.032
0.502 = 0.020
0.020
0.739 = 0.013
0.638 0.035
0.562 + 0.017
Youtube
0.799 0.117
0.710 + 0.182
0.887 = 0.033
0.756 0.115
0.850 0.024
Table 12: Standard deviation of accuracy, number of examples = 32, BLOOM (560M)
36

---

Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0,812 0.021
0.560 :: 0.046
0.832 + 0.023
0.876 0.012
0.833 0.016
AG-News-1
0,904 + 0.009
0.623 = 0.051
0.924 = 0.026
0.932 + 0.007
0.923 + 0.009
AG-News-2
0,786 0.008
0.566 0.023
0.791 = 0.016
0.824 + 0,010
0.797 0.010
AG-News-3
0,822 + 0.030
0.533 0.0299
0.800 = 0.042
0.840 0.015
0.825 + 0.034
Civil Comments
0.591 0.029
0.003
0.568 = 0.038
:.029
0.614 = 0.028
DBPedia-O
0,911 0.011
0.676
0.892 = 0.017
0.898 0.012
0.894 = 0.023
DBPedia-1
0,823 0.025
0.644 0.115
0.838 = 0.058
0.019
0.833 0.026
DBPedia-2
0,902 0002:0000 0.015
0.931 = 0.032
0.900
0.913 = 0.007
DBPedia-3
0.883 0.023
0.129
0.891 + 0.015
0.911 0.006
0.889 0.021
IMDB
0.557 0.033
0.522
0.641 = 0.008
0.564 0.017
0.577 0.020
Rotten Tomatoes
0.572 0.014
0.548 1:48000 0.045
0.710 = 0.010
0.575 + 0.048
0.572 0.029
SMS Spam
0,882 + 0.044
0.860 + 0.036
0.946 0.019
0.881 0.013
0.925 + 0.010
SST
0,570 + 0.019
0.563 + 0.032
0.705 + 0.019
0.575 + 0.028
0.542 - 0.035
Youtube
0.810 + 0.039
0.792 + 0.081
0.923 1 0.014
0.874 + 0.029
0.832 0.053
Table 13: Standard deviation of accuracy, number of examples = 48, GPT-NEO (125M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-O
0,846 + 0.018
0.496 = 0.007
0.871 = 0.015
0.871 + 0.005
0.732 + 0.036
AG-News-1
0.923 0.005
0.528 = 0.044
0.942 = 0.004
0.935 = 0.009
0.780 0.040
AG-News-2
0.798 = 0.009
0.509 = 0.004
0.731 = 0.040
0.782 0.018
0.722 0.019
AG-News-3
0.845 0.012
0.515 + 0.014
0.787 + 0.046
0.852 + 0.011
0.718 0.026
Civil Comments
0.605 + 0.037
0.512 0.009
0.607 + 0.039
0.556 + 0.014
0.556 0.018
DBPedia-O
0.905 0.020
0.549 0.057
0.924 + 0.020
0.830 + 0.034
0.798 0.021
DBPedia-1
0,840 0.021
0.592 0.062
0.867 = 0.024
0.808 0.031
0.780 0.011
DBPedia-2
0,916 + 0.009
0.636 0.081
0.937 + 0.017
0.923 + 0.013
0.870 + 0,022
DBPedia-3
0,888 + 0.022
0.618 0.096
0.927 0.011
0.885 : 0.007
0.827 0.026
IMDB
0,557 + 0.034
0.503 0.008
0.632 + 0,013
0.566 0.030
0.545 + 0.008
Rotten Tomatoes
0.601 0.026
0.480 0.014
0.683 = 0.033
0.556 0.033
0.562 0.018
SMS Spam
0,956 + 0.008
0.869 1 0.044
0.966 = 0.011
0.960 0.008
0.891 = 0.035
SST
0,612 = 0.024
0.500 0.025
0.702 = 0.046
0.594 0.042
0.560 0.023
Youtube
0.872 + 0.019
0.654 + 0.063
0.904 + 0.033
0.826 + 0.046
0.666 + 0.027
Table 14: Standard deviation of accuracy, number of examples PYTHIA (160M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0.829 + 0.029
0.507 + 0.005
0.829 + 0.071
0.847 + 0.022
0.788 0.028
AG-News-1
0.909 + 0.010
0.543 + 0.052
0.940 + 0.011
0.935 + 0.004
0.856 0.014
AG-News-2
0.789 + 0.016
0.511 = 0.012
0.715 = 0.057
0.708 + 0.011
AG-News-3
0.832 0.018
0.530 0.038
0.834 = 0.029
0.04:
0.763 0.015
Civil Comments
0,602 + 0.030
1.506+0.00
0.605 = 0.044
0.592
0.566 + 0,008
DBPedia-O
0.889 0.022
0.703 0.085
0.933 - 0.021
0.888 - 0.024
0.873 0.027
DBPedia-1
0.817 + 0.023
0.891 = 0.028
0.836 0.036
0.834 + 0.024
DBPedia-2
0,900 0.011
0
0.949 = 0.013
0.940 0.009
0.914 0.013
DBPedia-3
0.884 + 0.020
0.815 0.105
0.928 + 0.024
0.917 + 0.015
0.891 + 0.037
IMDB
0.565 0.033
0.545 0.033
0.641 + 0.020
0.604 + 0.022
0.542 0.015
Rotten Tomatoes
0.605 + 0.015
0.505 + 0.005
0.704 + 0.025
0.672 + 0.042
0.531 + 0.018
SMS Spam
0.933 0.009
0.636 + 0.130
0.929 + 0.058
0.919 0.032
0.914 0.027
SST
0,610 + 0.030
0.489 + 0.004
0.695 - 0.020
0.673 0.029
0.538 0.011
Youtube
0.834 0.049
0.797 + 0.090
0.805 = 0.095
0.851 0.020
0.865
Table 15: Standard deviation of accuracy, number of examples = 48, BLOOM (560M)
37

---

Dataset
TART
In-context learning
Fine-tuning full
Finc-tuning layer
Adapters
AG-News-0
0,817 0.024
0.539 0.044
0.855 + 0.030
0.877 0.013
0.830 0.023
AG-News-1
0,904 + 0.007
0.561 = 0.048
0.923 = 0.027
0.939 + 0.003
0,922 + 0.010
AG-News-2
0,790 0.026
0.576 + 0.022
0.814 = 0.009
0.839 + 0.009
0.816 0.014
AG-News-3
0.833 + 0.015
0.550 0.035
0.803 = 0.017
0
0.842 + 0.018
Civil Comments
0.581 + 0.039
=.499=0.000
0.387 = 0.018
0.576 0.039
0.605 0.036
DBPedia-O
0,915 = 0.005
0.652
0.917 = 0,011
0.908 0.019
0.909 = 0.025
DBPedia-1
0.825=000
0.633 0.104
0.838 = 0.032
=
0.852 0.011
DBPedia-2
0,887 = 0.020
0.606 0.088
0.950 = 0,007
0.952 0.011
0.916 0.014
DBPedia-3
0.873 0.033
0.611 0.136
0.898 + 0.034
0.927 + 0.007
0.909 0.008
IMDB
0.558 + 0.029
0.515 = 0.024
0.646 = 0.010
0.563 + 0.021
0.566 0.035
Rotten Tomatoes
0.601 0.020
0.533 + 0.053
0.708 = 0.015
0.579 + 0.049
0.556 0.033
SMS Spam
0,869 + 0.023
0.825 + 0.074
0.934 = 0.016
0.898 0.008
0.927 0.004
SST
0.570 + 0.038
0.554 + 0.041
0.712 = 0.033
0.609 #
0.567 - 0.016
Youtube
0.790 + 0.050
0.819 + 0.050
0.926 + 0.014
0.891 + 0.031
0.837 0.055
Table 16: Standard deviation of accuracy, number of examples 64, GPT-NEo (125M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-O
0,851 + 0.024
0.502 = 0.003
0.858 = 0.009
0.869 + 0.010
0.764 + 0.030
AG-News-1
0.925 0.002
0.529 = 0.030
0.937 = 0.011
0.938 0.004
0.820 0.021
AG-News-2
0,812 = 0.007
0.513 = 0.013
0.752 = 0.051
0.795 0.010
0.738 0.008
AG-News-3
0.851 + 0.007
0.503 + 0.004
0.820 + 0.020
0.855 + 0.009
0.741 0.019
Civil Comments
0.606 + 0.029
0.500 0.001
0.659 + 0.032
0.566 0.023
0.566 0.018
DBPedia-O
0.910 0.013
0.518 + 0.020
0.912 + 0.027
0.858 + 0.019
0.825 0.015
DBPedia-1
0.839 0.027
0.542 0.028
0.897 0.016
0.824 + 0.031
0.788 0.018
DBPedia-2
0,916 + 0.011
0.609 0.106 0.106
0.953 = 0.013
0.922 + 0.012
0.882 - 0,019
DBPedia-3
0.887 = 0.028
0
0.940 0.013
0.904 + 0.007
0.857 0.020
IMDB
0,556 + 0.024
0.619 + 0,030
4.574 0.017
0.552 + 0.009
Rotten Tomatoes
0.624 0.024
0.485
0.686 = 0.040
0.577 0.033
0.568 0.019
SMS Spam
0,937 + 0.018
0.960 = 0.021
0.961 0.006
0.899 = 0.014
SST
0,606 = 0.022
0.508 0.022
0.688 = 0.047
0.602 0.036
0.567 0.017
Youtube
0.888 + 0.028
0.715 + 0.097
0.897 + 0.046
0.878 + 0.050
0.675 0.019
Table 17: Standard deviation of accuracy, number of examples 64, PYTHIA (160M)
Dataset
TART
In-context learning
Fine-tuning full
Fine-tuning layer
Adapters
AG-News-0
0.836 + 0.018
0.509 + 0.008
0,850 + 0.027
0.856 + 0.008
0.799 0.029
AG-News-1
0.918 + 0.007
0.543 + 0.033
0.900 + 0.024
0.933 + 0.008
0.875 0.015
AG-News-2
0.799 + 0.012
0.515 + 0.019
0.784 = 0.018
0,831 + 0.022
0.732 + 0.017
AG-News-3
0.836 0.011
0.504 0.003
0.811 = 0.034
0.853 0.011
0.784 0.022
Civil Comments
0,602 + 0.030
0.510 0.012
0.573 = 0.035
0.611 0.025
0.572 + 0,013
DBPedia-O
0.905 0.015
0.667 0.052
0.936 = 0.018
0.902 0.022
0.882 0.020
DBPedia-1
0,809 + 0.022
0.687
0.887 = 0.039
0.852
0.853 + 0.008
DBPedia-2
0.881 0.026
0.799
0.955 = 0.022
0.947 0.009
0.922 0.014
DBPedia-3
0.877 +: 0.014
0.106
0.906 + 0.027
0.921 0.017
0.899 + 0.024
IMDB
0.571 0.033
0.539
0.621 + 0.039
0.618 0.013
0.542 0.012
Rotten Tomatoes
0.597 + 0.025
0.684 + 0.053
0.672 + 0.041
0.543 0.024
SMS Spam
0.907 0.045
0.659 0.133
0.942 0.024
0.931 0.030
0.909 0.033
SST
0,620 + 0.039
0 0.015
0.672 = 0.039
0.716 + 0.032
0.554 0.020 0.020
Youtube
0.844 + 0.059
0.765 - 0.137
0.879 + 0.058
0.865 - 0.033
0.883
Table 18: Standard deviation of accuracy, number of examples = 64, BLOOM (560M)
38


Reviews:
[{"review_id": "NIrbOEdJm8", "reviewer": "Reviewer_VGrP", "summary": "The paper presents an recipe for adapting an LLM to perform classification tasks in a task agnostic manner. They first try to tease apart if existing \"in-context\" methods which construct prompts to describe the task and then ask for the inference result are not achieving great performance because of information extraction or reasoning ability. By constructing a linear probe experiment they get indications that it may be the reasoning ability which hampers performance. The authors then propose to improve the task performance by training an inference module to solve a generic task, here few shot logistic regression, and then show that they can apply this to the outputs of the LLM to obtain good performance on desired tasks. The down stream tasks considered range from sentiment classification, news article categorization to spam detection. They also verify that the performance gains on text tasks carries over to image tasks and speech.", "questions": "* Th1 seems to say that the error is controlled by how similar P_NL and P_syn are but they are really different. In some sense they should always be different because it's a different language with a very different structure. In NL we have a lot more expressivity and less regularity ? Did you do any experiment that broadens the task space to see if this improves ?\n* I think 3.3 should be more appropriately named how to leverage LLM embeddings. The embeddings are always the last layer it's more how they are used that changes ?\n* Are all the tasks considered binary classifications ? I guess line 339 should be corrected MNIST and CIFAR are switched. We can't easily evaluate this work when this is happening. ", "limitations": "None.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* There is still a lot of scope and practical use in extracting information from a text model in a reliable way.\n* The method is novel in so far I can tell.\n* Contains reasonable baselines.", "weaknesses": "* Although I find it intriguing the paper leaves me more questions than it answers. I think it would be very beneficial to make the presentation much more concrete and map some of the down stream tasks clear and exemplified in the main text. They remain very abstract and lots of intuition that the readers should build up is lost that way. \n* Although I think the line of scientific inquiry is quite intriguing I don't find the paper mature enough to be published in the current form. I find the writing can be improved a lot and maybe the context can be expanded a bit.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments on our work and we are glad you found our proposed method novel. First, we address your major concern regarding the presentation of our work, and then respond to your questions.\n\n**Quality of presentation.** While other reviewers (tMpn, CKNq) found the quality of our presentation to be high, we welcome your feedback on how it can be further improved. We will make the following changes to the final camera-ready draft to improve its readability:\n- Update the introduction section to better reflect our objectives and motivate our proposed method. We will use the outline from the general response to communicate our message better.\n- Add an additional section in the main body (derived from content from Appendix A) on using synthetic finetuning to improve a pre-trained model's ICL ability.\n- We will add some more details of our downstream evaluation tasks in language, vision and speech (expanding on lines 67-73). Currently, these are detailed in Section 4 with details deferred to Appendix D due to lack of space.\n- For Section 3.3, our choice of name “Which embeddings to take” was to reflect the fact that one can use the LLM in different ways to embed the input sequences (i.e.,  Vanilla embeddings or leave-one-out (LOO) embeddings). We agree that it might be more appropriate to focus on “How to embed” as opposed to “Which embeddings to take”. We will make this change.\n- We will fix the typos in the revised draft (MNIST <-> CIFAR).\n\n**Closeness of distributions P_NL and P_SYN in Theorem 1.** We would like to clarify that P_{NL} is the distribution induced on the embeddings of the natural language data rather than the data itself (see Lines 619-622). Furthermore, when we combine the TART module with the base LLM, we perform a PCA renormalization step (Lines 221-224) which projects the embeddings into the same space as the synthetics. These two facts combined ensure that P_syn and P_NL are indeed very similar to each other in our experiments. It is indeed surprising that projection+renormalized NL embeddings are close to the synthetic Gaussian distribution. We will add this discussion to Section 3.4 to make it clear.\n\n**Extensions of TART beyond binary classification.** In this work, we focus on the set of binary classification tasks across language, vision and speech domains. That said, based on the feedback from the reviewers, we provide additional experiments (see Table 1 and Table 2 in attached PDF) demonstrating that our method can indeed be extended to multiclass tasks, improving base ICL performance by up to 59%."}, {"review_id": "CH4qtgJxeL", "reviewer": "Reviewer_hArx", "summary": "The paper studies why in-content learning achieves inferior performance compared with finetuning and adaptor, then proposes an LM-based inference module that learns to perform logistic regression based on the sample and previous (sample, label) sequence, where the sample and linear cutting-plane are sampled vectors. When testing, the base LLM encodes the in-content examples into single vectors, concatenates them with their corresponding labels to form the (sample, label) sequence and combines with the test sample to predict the test label. Experimental results show the proposed inference module can achieve comparable results without task-specific parameters.", "questions": "What do the \"parameters w\" in Line 217 and \"w_t\" in Equation 3 means? Are these two refer to different things?\n\nHow does the TART obtain multi-way classification labels as it is trained with two-way classification? Can this method generalize to text generation with a much larger vocabulary?\n\nAuthors may compare TART with linear probing besides FT-layer given the same amount of examples.", "limitations": "The authors addressed the limitations and broader impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a new way to classify a test sample based on training examples. The proposed inference module learns to predict the sample label given (x,y) examples, which share the same separation plane. By sampling enough various planes and input x, an LLM is possible to learn how to find the linear separation plane given (x,y) examples, thus correctly predicting the test label without further training. The proposed method may apply to meta-learning tasks as both settings are similar.", "weaknesses": "Although the proposed inference module is interesting, I can not see it has a significant advantage over linear probing.\nFirstly, TART requires base LLM to encode examples and then fed the average pooled vector into the proposed inference module. This process requires a similar amount of computation and memory as in-context learning, and this process must repeat k times in the proposed LOO embedding if given k examples. On the contrary, linear probing can encode each sample independently and then train one linear layer for prediction.\nSecondly, since the inference module is trained with linear logistic regression, the LLM representation must be linear separable for task-specific labels. Thus, it has the same upper bound as linear probing, while one linear layer should be much easier to train than a GPT-2 model.\n\n\nLine 10 claims \"performance gap exists due to their inability to perform simple probabilistic reasoning tasks\" and Line 51, \"...prompt engineering or active example selection, focus entirely on the LLM’s learned representations.\" The authors do not clearly explain what the reasoning ability is and why prompts improve learned representations. I assume prompt engineering is proposed due to the unavailability of hidden representation inside GPT-3 or GPT-4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the detailed comments and feedback. We are glad that you find TART interesting and see applications of our work to settings beyond ICL (i.e., meta-learning).\n\n**Comparison with linear probing.** As we highlighted in the general response above, our main objective with this work is to understand why a gap exists between ICL and finetuning, and provide an algorithmic way to fix this gap. The caveat is to devise algorithms that can obtain the performance of finetuning while maintaining the benefits of ICL (namely, no need for training at inference time, providing a natural language interface to the user). As a first step towards this bigger goal, we begin by studying this problem for the class of binary classification tasks. While linear probing (for each task separately) would get close to finetuning performance, one has to forgo the advantages of ICL to use it. Furthermore, it is not possible to extend linear probing beyond the realm of classification tasks.\n\nOn the other hand, TART allows users to benefit from the ICL framework and provides a “no training” way to match finetuning performance. Additionally, it opens up the possibility of improving the generative ability of LLMs using synthetic tasks. This is a promising line of work, which requires a new set of tools than those in the paper, and one which we are currently exploring.\n\nAdditionally, we would like to clarify that the adapter head baseline, which is included in the original manuscript, is actually a slightly stronger baseline than the linear probing method. The adapter head comprises a single linear layer composed with a sigmoid nonlinearity. TART performs 3.4% better than these adapters.\n\n**Generalizing TART.** In the current form, you could obtain multiclass labels using a one-vs-all approach by looking at the logit probabilities for each class. However, to show that our method generalizes beyond binary classification, we trained an extension of TART for multiclass problems using multinomial logistic regression as the synthetic data generating task. We find that our TART inference module improves the base ICL performance by up to 59%. Please see Tables 1 and 2 in the attached PDF for details.\n\n**Explaining reasoning ability.** Throughout the paper, whenever we say the ability to perform reasoning we mean the ability to solve simple probabilistic reasoning problems, such as the logistic regression model. We make this connection mathematically precise in Equation (1) in  Section 2.3 (Line 149-150). The equation precisely defines what the representation and reasoning abilities are. The representation-reasoning decomposition in Eq. (1) allows us to formally state the hypotheses – H1, H2, and H3 —that we evaluate next. \n\n**Connecting prompt engineering with improved representations.** We agree that the connection between prompt engineering and improving representations was not elaborated on in the current version. We plan to revise and add this explanation: “Via prompt engineering users are searching the language space for a natural language template that best represents their task. This search enables the users to find the best representations for their specific task which makes the LLM more likely to solve it.” Thank you for pointing this out.\n\n**Minor typo.** The parameter $w$ in Line 217 should be $\\theta$. We will fix this in the draft. \n"}, {"review_id": "gi22EuDYmw", "reviewer": "Reviewer_tMpn", "summary": "The authors study why in context learning doesn't perform as well as task specific fine tuning by decomposing the process of in context learning into representation and reasoning. They discover that the gap is due mostly to deficiencies in reasoning and propose a new method called TART to bridge the gap. They demonstrate that TART is task and model agnostic.", "questions": "In a real world application of this technique it would be inconvenient for the user to have to indicate which tokens belong to which data points/labels, which I think the embedding averaging method requires. For example if I understand right\nSentence: The movie is good\nLabel: Positive\nThe user would have to manually indicate which tokens belong to the data (sentence) and label so it can be averaged.\nI can imagine using the base model to separate them out for you. Is that how this would work?", "limitations": "Yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Originality:\nThe authors note and address the problem that in-context learning underperforms model fine tuning. The paper's primary contributions are disentangling representation strength and reasoning ability and finding a way to measure them, then using the insight from their measurements to improve the model's reasoning ability and thus close the gap between in context learning and fine tuning. Both are very original. \nI thought the concept of training a reasoning model to do probabilistic inference using synthetic data sets was clever. The same for averaging embeddings of multi-token data points in order to piece together the base model and the reasoning module. \nI also like thinking of LLM reasoning as probablistic inference in the first place. That seems insightful.\n\nQuality and Clarity:\nThe quality here is high. They provide extensive experimental results to validate their idea, comparing TART with baselines across modalities and models. Every concept is backed up by a full sweep of experiments, they answer questions before I have a chance to ask them. For example extending this to vision and speech went above and beyond. The writing is clear well structured and the methodology is clearly explained. Their figures and tables are clear and accessible. \n\nSignificance:\nI believe the significance is high. The authors address a limitation of LLMs with a novel solution that is broadly useful. It has potential to significantly improve transformer model's in context learning performance across the range of tasks. \n\nOverall, great work.", "weaknesses": "There's a typo on line 339, it should be CIFAR-10 (classes plane and bird) and MNIST (classes 0 and 8)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments. We are glad that you found our work interesting on both the problem and the solution front, and found the paper well-written and clear.\n \n**Demarcating examples**. This is an interesting challenge. Currently, our code takes as input a demarcating limiter (which is “Sentence: and Label:” in the example you provided) and uses that to identify which positions to take the average embeddings over. If you are working with a fixed prompting template, implementing this is pretty straightforward. Alternatively, you could use Streaming embeddings (Figure 24) which side-steps this issue entirely because it embeds one example at a time.\n\n**CIFAR <-> MNIST.** Thanks for pointing out the minor typo. We will correct it in the updated draft. \n"}, {"review_id": "LAFrTWZND8", "reviewer": "Reviewer_CKNq", "summary": "This paper proposes and evaluates a method (TART) for improving the in-context learning of large, pretrained base models applied to downstream binary classification tasks. TART stacks a second transformer on top of the pretrained base transformer. The TART transformer is trained to rely on in-context examples to solve a large number of synthetically generated binary classification problems. TART’s training is task-agnostic in the sense that its synthetic training data is unrelated to the domain of the downstream tasks on which the combined base+TART model is evaluated. The experiments demonstrate that TART improves the combined model’s in-context learning across a range of binary classification tasks, and a diverse set of base models pretrained on language, vision, and audio data.", "questions": "I assume that TART stands for Task-Agnostic Reasoning Transformer. Is this explained somewhere?\n\nLine 589  If the context length is 258, and each example is one pair, then there would be 129 in-context examples, not 256. Right?\n\nI’m not sure what’s meant by “the binary labels y are encoded as a one-hot vector to match the input dimension $d$ of the corresponding covariates $x$.” Does this mean that only 2 of the 16 dimensions are ever 1.0, while the other 14 dimensions are always 0.0?", "limitations": "No problems noted.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Given the profound importance of large pretrained transformer models today, and their remarkable ability to learn from in-context examples, general methods for further boosting the performance of in-context learning are of great interest. \n\nThe paper’s proposed method, TART, is a well-motivated and novel approach. And the paper’s experiments clearly demonstrate the degree to which TART boosts the in-context learning of large, pretrained base models applied to downstream binary classification tasks. Exciting results!\n\nFinally, as noted in section 5 of the paper, the TART approach may lead to similar methods that improve LLM performance on a wider range of tasks beyond binary classification.", "weaknesses": "The paper claims to make three main contributions, which I summarize here:\n\n1.\tWHAT CAUSES THE ICL GAP?  Study why in-context learning does not perform as well as task-specific fine-tuning despite having access to the same information, via a representation-reasoning decomposition.\n\n2.\tTART CLOSES THE ICL GAP  Propose a new task-agnostic method, TART, which bridges the performance gap to task-specific methods and is trained using only synthetic data.\n\n3.\tTART IS GENERAL  Demonstrate that TART works across different NLP tasks for a range of model families. The same inference module generalizes to vision and speech domains as well.\n\nI follow this outline of claims in describing what I see as the paper’s weaknesses.\n\nClaim 1. WHAT CAUSES THE ICL GAP?\n\nThe paper mounts an ambitious study of the gap in performance between in-context learning and task-specific fine-tuning, given a common set of examples. The experimental results are suggestive, providing ample motivation for the TART approach. However, the paper ventures far beyond the question of TART’s motivation when it attempts to draw broad conclusions about a clear representation-reasoning decomposition. The field has arrived at no settled definition of reasoning, so any distinction between representation and reasoning remains nebulous. Transformers themselves illustrate the complex, intertwined relationship between representation and reasoning. \n\nAs a result the paper makes extraordinary claims (requiring extraordinary evidence) when it equates reasoning with probabilistic inference, and when it makes statements like the following:\n\n•\t“LLMs lack reasoning abilities”\n\n•\t“LLMs lack the ability to perform simple reasoning over their learned representations” \n\n•\t“this performance gap exists due to their inability to perform simple probabilistic reasoning tasks.” \n\nI believe that too much of the paper is devoted to pursuing such a challenging case, which would require far more evidence than the experiments provide. And I see it as unnecessary, since the experimental results are more than sufficient to provide adequate motivation for the TART approach. Furthermore, the paper’s heavy emphasis on its reasoning-centric terminology (like “reasoning module” or “probabilistic inference module”) obscures the relatively simple and limited nature of TART. To be concrete, the paper doesn’t need to wade into the philosophy of reasoning vs. representation to make the case that TART (as one might expect) boosts in-context learning by training a transformer (on task-agnostic, synthetic data!) to perform well at in-context learning. \n\nIn summary, while I agree that the paper does cast some useful light on the performance gap between in-context learning and task-specific fine-tuning, much mystery remains there, so I don’t see this as a strong contribution of the paper, although it is sufficient to motivate the approach.\n\nClaim 2. TART CLOSES THE ICL GAP\n\nThe experiments demonstrate to me that TART does succeed in largely closing the in-context learning gap, at least when given enough examples. But what happens when only a few examples are available at test time? The paper presents no performance comparisons against baselines tested on fewer than 20 examples, as far as I can see. Since the TART transformer is a task-agnostic module sitting on top of the base model, one would expect its performance to be no better than random in the absence of task-specific examples at test time. So how many test-time examples does it take for TART to show good results? Below 20 examples is an important range to explore, since in-context learning is known to be able to derive great benefit from as few as one or two examples. If it turns out that TART requires at least, for example, 10 examples to deliver benefit, then the paper would be over-claiming when it says that TART, “when composed with any LLM via its embeddings, generically improves upon its reasoning abilities.” \n\nThe all-important technical details of TART’s training are poorly explained in the current version of the paper. Section 3.2.1 leaves too much to the reader’s imagination. Line 217 talks about updating the parameters $w$, which must therefore refer to the parameters of the TART transformer. But equation (3) also contains a symbol $w_t$ which cannot refer to the transformer’s parameters. Rather, the $w_t$ must refer to the randomly sampled weights for the linear layer used to define a binary classification problem that the transformer is trained on. At least this is my current understanding. The apparent conflict of notation needs to be sorted out, and eq. (3) should be explained much more carefully in the text. For instance, it’s worth mentioning that the angle brackets refer to a dot product, since this notation is rarely used in papers on transformers. \n\nThe paper would be easier to follow if it used the term logistic regression correctly. Specifically, logistic regression is not a problem or a task. Rather, logistic regression is a method that can be used in solving classification problems. The term is used correctly on the following lines:\n\nLine 231  “For each problem, we train task-specific linear classifiers using logistic regression”  \n\nLine 535  “To conduct linear probing over the embeddings, we perform logistic regression over the output embeddings of each model and the given labels in the training set using the built-in logistic regression solver from the scikit-learn python library, utilizing the lbgfs solver.”\n\nBut the rest of the paper uses the term “logistic regression problem” as if logistic regression were a problem type rather than a solution method. \n\nTerminology for eq. (3) is tricky, because it involves (I believe) a linear layer of random weights which could be (but are not) obtained by the logistic regression method. And these random weights are used, along with random covariates $x$, to generate synthetic classification problems. Because the structure of this generator is related to the method of logistic regression, it is tempting to refer to the generated classification problems as logistic regression problems. But that’s conflating problems with methods. I suggest spending a few sentences to explain eq. (3) in simple terms, and replacing the phrase *logistic regression problem/task* with *binary classification problem/task* throughout the paper. \n\nIt seems that positional encodings must be required in training the TART transformer, to establish the alternating positions of $x$ and $y$ vectors in the sequences. But I find no mention of positional encodings in the paper.\n\n\nClaim 3. TART IS GENERAL  \n\nIt is shown that TART works across different binary classification tasks and model families, including vision and speech domains. However, the results are all limited to binary classification, which is not a very general class. As explained in the discussion, “In future work, we seek to understand whether synthetic tasks exist for training other generic reasoning modules, capable of improving base LLM performance on tasks such as generation or summarization.” Such results would be highly significant. Until then, the current restriction to binary classification limits the significance of this contribution.\n\n\nMINOR SUGGESTIONS\n\nLine 130  “at par when compared with” should be changed to “on par with”.\n\nLine 191  “TART comprises of two components” should be changed to either “TART comprises two components” or “TART is comprised of two components”.\n\nLine 584   A transformer outputs a sequence of vectors, not embeddings.\n\n\nCONCLUSION\n\nBecause of the problems noted in this section, I would vote to reject the paper in its current form. However, I have confidence that the authors can address all of these issues in the next version of the paper, so for now I’m optimistically voting for acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed comments and helpful feedback on our work. We are glad you found our approach well-motivated, novel and exciting.\n\nWe completely agree with the three main contributions that you highlighted in your review. While we addressed some of these concerns in the general response (study on representation-reasoning trade-off, generalization to multiclass, and TART in low-data regime), we provide a detailed response below.\n\n**Reasoning vs representation.** We completely agree that there is no consensus in the field on a definition of a reasoning vs. representation, and much less so for in-context learning which is quite nascent. As we mentioned in the general response, this motivated us to make an attempt at explaining why ICL fails. We make a first step towards this understanding through our representation-reasoning decomposition which we mathematically characterize in Equation (1). We feel that this is an important and interesting contribution in and of itself, and while just a first step, we believe it will indeed be helpful and inspiring in spurring further research along these directions. \n\nOur representation-reasoning definitions as well as the analysis with Hypotheses H1-H3 is currently focussed on binary classification tasks in this draft. Our statements and claims are indeed supposed to be restricted to this setup. Following our new extensions to multiclass tasks (see attached PDF for experimental results), we plan to extend this analysis to the multiclass tasks as well.\n\n**TART in low in-context example regime.** For most plots in our paper, we study the performance of TART (and baselines) beginning with 16 in-context examples (k). In the attached pdf, we show that the TART works even in the extremely low-data regime (with k = 4, 6, 8, 10). TART can improve the base ICL performance by up to 22.2% (k=4), 30.1% (k=6), 24.1% (k=8) and 31.3% (k=10). See Table 3 in attached PDF.\n\nFurthermore, on 2 out of 3 datasets, we see that TART outperforms full model finetuning by up to 15.8%. While it is challenging to adapt the LLM for the task with such few data points, TART is able to better use the sample information. We will update our draft with a more extensive study across all benchmarks in this regime. \n\n**Beyond binary classification.** Thank you for this suggestion. In response to this feedback, we train a multiclass TART module which outperforms the ICL baseline by over 59 points (see attached PDF for more details). We realized that our training algorithm was not specific to binary classification tasks and could be easily extended to multiclass problems by using the multinomial logistic regression as the synthetic. We will update the draft to feature these extensions in Sec. 4 (Experimental evaluation).\n\n**Use of term logistic regression.** We have actually used the term logistic regression in an overloaded manner to refer to both the logistic regression generative model as well as the estimator obtained by solving the ERM with respect to this model. We will clarify this in the paper and make it clear which of these we are referring to at a given time. Note that the logistic regression model is a very specific generative model of binary classification which we use in our paper for the synthetics.\n\n**Minor Comments.** We address the some of the minor concerns here:\n- TART stands for Task-Agnostic Reasoning Transformer (mentioned in Sec. 3 header). We agree that the acronym could be explained earlier and we will reflect this change in our draft.\n- For context length = 258, the number of examples should be 129. We will correct this in the updated draft.\n- For the binary classification task, we only use the last dimension out of the 16 to set the labels – it is 1 when the label is positive and 0 when the label is negative. Such a mapping is quite common and has been previously used in [1].\n\n*References*\n\n[1]  Garg, Shivam, et al. \"What can transformers learn in-context? a case study of simple function classes.\" Advances in Neural Information Processing Systems (2022)\n"}]

Author rebuttal:
We thank all the reviewers for their thoughtful comments and feedback. We are glad that they found TART to be **novel** and **practical** (VGrP, CKNq, tMpn), **well-motivated** (CKNq, tMpn), and supported by **extensive experiments** (VGrP, tMpn, CKNq). 

Recall that our paper shows that in-context learning (ICL) is much more competitive with finetuning than previously thought. Further, we show that this is achievable by training an independent Transformer module entirely on synthetic logistic regression data.

Based on the feedback, we clarify two points on our presentation quality and scope of TART. We will make the following small, yet important, changes to our paper:
- **Presentation.** Our work presents two approaches to improving the ICL performance: (1) finetuning the entire model with synthetic binary classification data and (2) training an independent module (TART) on synthetic data. Following VGrP and hArx’s comments on presentation quality, we realized we had underemphasized point (1), which is a natural first step towards TART. Additionally, this point actually demonstrates that pretrained models do not learn how to classify in-context from standard pretraining objectives. The experiments for (1) were mentioned briefly in the main body (deferred to App. A) and we will update the draft with an additional section incorporating this.
- **Scope of TART.** As a result of the questions posed by tMpn and CKNq, we realized that our TART module could be generalized in two useful ways: (1) to multiclass classification problems, and (2) in the low-data regime (4-10 examples). We conducted additional experiments (see attached PDF) and show that for multiclass problems, TART outperforms base ICL performance by up to 59%, far exceeding the gains we saw on binary problems. Further, in the low-data regime, it can be up to 31% better than in-context LLM accuracy and outperforms finetuning on 2/3 datasets. 

---
**Detailed response**

We synthesize all major reviewer concerns and considerations, and show how we incorporate these in our presentation either through minor rewrites or with additional experiments. 

Our paper is about understanding why the ICL ability of pretrained LLMs has sub-par performance when compared with SGD-based finetuning. ICL has fundamentally changed how people use AI systems by providing a new natural language interface for specifying tasks. Yet, it still lags finetuning methods in performance quality (up to 30% on binary classification). Our objective is to understand whether this performance gap is an inherent limitation of ICL or not.
- hArx suggests linear probing as an alternative. While we agree that it might be competitive with finetuning in several cases, it doesn’t address the main question that we consider: why is ICL under-performing finetuning? Additionally, linear probing takes away ICL’s intuitive interface from a user’s interaction with an LLM.

We begin by asking why this quality gap exists: do the data representations learned by the models lack task-specific information? Or is there some deficiency in the model’s ability to use this information, i.e., reasoning ability? For classification tasks, our results suggest that pretrained LLMs indeed have the necessary information in the representations, but fail because of their inability to do simple forms of reasoning, e.g. linear classification.
- We agree with CKNq that understanding this question is a challenging study to undertake, with little agreement in the literature. Our representation-reasoning decomposition is a first step towards formalizing this question (Eq. 1 provides a mathematical description). While Sec. 2.3 currently focuses on binary tasks, we will perform a similar analysis for multiclass tasks.

As a first step towards improving ICL, we finetuned an LLM with synthetic binary classification problems – observe that we are updating the weights of the model here. While this procedure improves model performance on classification tasks by up to 17% (see App. A), such finetuned models might lose their generative capabilities. This brings about the question: is it possible to improve the ICL abilities of these models without interfering with its other pre-existing capabilities?
- To partially address VGrP and hArx’s concern on presentation, we will add a section on this synthetic finetuning between Sec. 3 & 4.

We undertook this challenging task of teaching a model a new generic skill without affecting existing capabilities. Rather than directly finetuning, we trained an independent Transformer module (TART) entirely on synthetic data to learn this classification skill. By applying this module to the output embeddings of an arbitrary pretrained LLM, we improved the model’s ICL ability to be competitive with finetuning (within 3% from a gap of over 20%). Additionally, this surprising finding demonstrates that pretrained models do not learn how to classify (or reason over embeddings) in the pretraining phase.

While our work focuses on binary classification, we recognized (following tMpn and CKNq’s suggestion) that our procedure is not restricted to binary tasks. We trained a multiclass version of the TART module and our preliminary results suggest that it can outperform base ICL by up to 59%, more than the improvements on binary tasks. On CKNq’s suggestion, we also explored TARTs performance in the low-data regime (4-10 examples). Excitingly, we found that it can be up to 31% better than base in-context LLM and that it even outperforms finetuning on 2/3 datasets in that data regime (see attached PDF).
- We will add detailed experiments for multiclass tasks and extend our current evaluations to the low-data regime as suggested by tMpn and CKNq.

Further, we go beyond text and show that the same TART module can be combined with vision and speech models to enable ICL for these modalities. This suggests that this capability is somehow uniformly lacking across a range of foundation models.

---


Meta-review: This paper studied the gap of performance between in context learning and the task-specific tuning.
The reviewers generally like the main insight of decoupling the representation and reasoning, but had concerns regarding the presentation, generality and the effectiveness of the proposed method. During the rebuttal, the authors have done a good job in addressing the most of the concerns, so I would recommend accepting the paper, with the condition that the authors would update the paper accordingly to reflect the changes in their camera ready version.
Decision: Accept (poster)