Title: TART: A plug-and-play Transformer module for task-agnostic reasoning
Venue: NeurIPS 2023
Accepted: yes
Keywords: In-context learning

task-agnostic methods

large language models
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

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 (175B). 2

1 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 a single model to be applied to a 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 method of choice since it 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, we show that this gap between in-context learning and finetuning approaches exists even when presented with the same task examples; in-context learning can underperform fine-tuning by up to 30 points across binary classification benchmarks (see Figure 5a).

Wefirst 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,

Task-

Data-

Quality

agnostic

Scalable

X

In-context learning

Base

LM

Base

LM

Head

Head

Fine-tuning

In-context learning

Fine-tuning

Adapters

Inference

Base

Head

Base

TART

Module

(Our method)

Adapters

TART

or reasoning , over these representations 3 . Is the gap because the data representations learned by the models lack task-specific information or is there some deficiency in the the model's ability to use this information. We take a first step towards understanding this hypotheses experimentally in Section 2 by measuring both the reasoning and the representation gaps across a variety of LLM families (GPT-NEO [6], PYTHIA [5], BLOOM [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 as fine-tuning) improve the model along both representation and reasoning, but primarily improve the task-specific classification (or reasoning) ability which accounts for 72 %of 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 on the LLM's learned 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 to 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, can we make the process both task and model agnostic?

We show that it is indeed possible to improve the reasoning capabilities in a 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 of the downstream task or 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 is 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 of 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

(x,y)

W

Ex1=

Sentence:The movie is good

(Sentiment classification,

Label:Positive

News classification)

Synthetic

Embeddings

logisticregression

data

Ex 1

BaseModel

Inference

.

。

Pre-trained LLM

Module

x1

y1

Ex k

(GPT-Neo,Pythia,BLOOM)

y

Inference

。。

Module

xXk

Averaging

yk

Layer

data-scalable and overcomes the limited context length bottleneck of in-context learning. While each example spans multiple tokens in an 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 abilities 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 1 in Section 3.4).

To summarize, our main contributions are as follows:

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.

Propose a new task-agnostic method, TART, which bridges the performance gap to task-specific methods and is trained using only synthetic classification data.

Demonstrate that TART works across different NLP tasks for a range of model families. The same inference module generalizes to vision and speech domains, suggesting that this capability is somehow uniformly lacking across a range of foundation models.

Related work. Prompt engineering focuses on improving the in-context task adaptation abilities of LLMs by modifying prompts. A line of 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, 24], 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 a 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 to understand the in-context learning property of LLMs by presenting mechanistic interpretations of in-context learning [36] and performing exploratory analysis of in-context learning behaviors [43]. Existing literature demonstrates that LLMs can learn simple function classes in-context [11] and propose that LLMs are performing gradient descent when learning tasks incontext [36]. Complementary to these, our work provides insights on 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.

2 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 a representation-reasoning decomposition to 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 a 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 task-specific models increase both during training (expensive hyperparameter search) as well as during inference (deploying several models). In general, task-agnostic methods will scale better with increasing model sizes by side-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 classification 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 a 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 are updated at inference time. However, its performance is usually not on 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-tuning. 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 are not task-agnostic since they change the underlying model significantly but usually achieve state-of-the-art performance for any given task and are data scalable.

Adapters. Adapters adapt the underlying LLM 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-Reasoning 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 investigates why this performance gap exists. We hypothesize that it is either because (a) the representations learned by the LLM are insufficient to learn a good predictor for the specific task, or (b) the LLM lacks the capability to reason over these representations to make good predictions for the task.

To understand whether the representations have sufficient information, we train a task-specific linear classifier using these representations, also known as linear probing, and evaluate its accuracy (Acc LR ). Using this as an intermediate, we decompose the performance gap

where ∆ rep represents the gap in performance which can be attributed to insufficient representation capacity and ∆ reas is the performance gap due to insufficient reasoning abilities. Mathematically, we define the quality of a models representations to be the gap between the fine-tuned model's accuracy

(a)

(b)

(c)

0.3

S

0.2

9

asoning

6

0.1

0.0

-0.1

0.1

0.0

0.1

0.2

Representationgains

0.80

0.75

In-context learning

0.70

Linear probing

Fine-tuning

0.65

0.60

20

30

40

50

60

Numberofexamples

1.0

Fine-tune (task)

0.9

Fine-tune (ag_news)

0.8

0.7

A

0.6

0.5

20

30

40

50

60

Numberofexamples

and the accuracy of linear probing - the smaller the gap the better the representations. Similarly, the models reasoning ability, or the ability to perform linear classification, is 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:

H1. 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 to better performance.

H3. Fine-tuning and adapters are not task-agnostic because the task-specific training hurts their ability to 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 a 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. We see that across models and different numbers of in-context examples, the reasoning gap ∆ reas accounts for up to 79 . 11 %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 to reason over them.

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, we show a 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 LLM, the gains are predominantly due to 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 to 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-these account for 72 . 58 %of the drop (see Appendix B for more details).

Adapters: Impairs task-agnosticity 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 on 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 abilities.

3 TART: Task-Agnostic Reasoning Transformers

The above analysis showed how it 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 a 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: a 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 taskagnostic, boost performance quality by improving reasoning, and make the approach data-scalable by aggregating input embeddings into a single vector.

3.2 Reasoning module: Can transformers learn to do probabilistic inference?

TART's reasoning module is a Transformer-based 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 a Transformer model which is auto-regressively trained on a 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 a 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.1 for details. The architecture takes as input a sequence of vectors and is trained to predict the next vector in the sequence. The input sequence consists of k pairs of labeled examples ( x 1 , y 1 ) , ( x 2 , y 2 ) , . . . , ( x k , y k ) , with each example z i = ( x i , y i ) using only two input positions of the transformer - one for the covariates x and the other for the label y . This is in contrast to standard LLMs where each example is spread over multiple tokens which limits how many examples can be put in the context. For example, with a context window of 2048, our module can support 1024 examples while the base model can support only 10 examples, assuming each demonstration comprises 200 natural language tokens.

Training procedure. This module is trained using gradient descent to minimize the population loss

where z 1: i -1 corresponds to the first i -1 examples and ℓ CE is the cross-entropy loss evaluated on the transformer prediction and the true y i . Each training sequence s t used to update the parameters θ comprises a different d -dimensional logistic regression problem, sampled as

where σ represents the sigmoid function and the multiplier α determines the noise level of the problem. The above sampling model comprises three components, the regressor w t , the covariates x i,t and the labels y i,t sampled as

the regressor w t is sampled from a standard Normal distribution

the covariates x i,t are sampled from a standard Normal distribution

the corresponding y i,t are sampled based on the sigmoid of the inner product ⟨ x i,t , w t ⟩

See Figure 2 (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 our model with d = 16 and

(a)

(b)

(c)

0.5

0.5

Learned Logistic Function

α = 0.5

Absoluteerror

0.4

TARTinference module

Absolute error

0.4

α = 1

0.76

α =10

0.3

0.3

α =20

0.74

0.2

0.2

0.1

0.1

0.72

TART-Vanilla

TART-LOO

0.0

0.0

0.70

0

100

200

0

100

200

20

30

40

50

60

Numberofexamples

Numberofexamples

Numberofexamples

k = 256 . We describe the model hyper-parameters 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 a dimensionality reduction via PCA to reduce the embedding dimension.

3.2.2 Properties of reasoning module

The task-agnostic reasoning module described above is trained to 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 is it to variations in the noise level α .

Accuracy of probabilistic inference. For understanding the instance level performance of our reasoning module, we evaluate it on a 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 our 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 a function of the number of examples used for predictions. We observe that the error for our reasoning module decreases as a 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, α , of the logistic regression problem. Recall that we trained our inference module by fixing the noise level α = 10 . At inference time, we 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 as we make the problem harder ( α = [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 a certain number of in-context examples (see Figure 4c). One reason can be that the causal nature of 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. The LOO-embeddings help TART be data-scalable by enabling it to embed a much larger number of points than the context window can support. To do so, we use only a 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.

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 at evaluation time, our input is the embeddings from a natural language task. In Theorem 1 we show that its performance on the natural language task depends on the distribution shift from the synthetic to the true distribution (see Appendix C.3 for a formal statement and proof).

Theorem 1 (Informal) . Let T represent the class of transformer models and T S ∈ T denote the trained reasoning module on set S of synthetic regression with n syn sequences sampled from distribution P syn in eq. (3) . The error of the transformer T S when evaluated on a distribution P NL over natural language sequences is

where W 1 denotes the Wasserstein-1 metric, Comp ( T ) represents the complexity of class T , and ˆ err represents the error on 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 n syn , 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 S fit to the synthetic training set.

Observe that for our algorithm TART, the distribution P NL corresponds to the distribution induced on the embeddings of the natural language data (post the PCA and renormalization) while the distribution P syn is the one corresponding to the sampling process in equation (3).

4 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-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.1. We consider 14 different binary classification tasks ranging from sentiment classification, news article categorization to 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 DBPedia14 are multi-class datasets, we construct 4 binary classification 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 of 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 14 datasets, we use GPT-NEO (125M), PYTHIA (160M) and BLOOM (560M). For ablations on larger models, we 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.e., GPT-NEO (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-tuning, 4) LM head 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.1 for hyperparameter details). For TART, we chose a base default set of parameters and use the same inference module with the exact same weights for all the experiments in this section.

(a)

(b)

(c)

0.8

0.8

0.8

0.7

FT-full

In-context learning

0.7

Adapter

TART

0.6

0.6

0.6

Ft-full

In-context learning

TART

Adapter

TART

In-context learning

0.5

20

30

40

50

60

0.5

0.125

1.3

6

0.5

2

4

2

5

2

Numberofexamples

ModelSize(Billions)

Numberofexamples

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 fine-tuning. We also observe that TART consistently outperforms the task specific strategies of 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 1000 x 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 3 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 to 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 as a function scale. For TART, the performance increases by 9 . 8% while using the same inference module across model sizes. Furthermore, the difference in performance between TART and fine-tuning baseline reduces from 7 . 5% to 2 . 2% from the 100M scale to 6B scale.

Beyond context length. We evaluate the data-scaling properties for both in-context learning and TART (Figure 5c). To demonstrate the scaling property, we do not truncate the input text to 100 characters and utilize the entire text sequences. For TART, we observe that accuracy continues to improve when scaling from 18 to 256 in-context examples with 6 . 8 % lift in performance. In comparison, ICL, which is bottlenecked by context length, supports 10 x 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 underlying 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 performance is robust to the underlying architecture.

Performance in small sample regime. While our focus has been on scaling up number of incontext examples (ranging from 18-64), in Appendix E.4 we compare the performance of TART with fine-tuning and in-context learning in the small sample regime. Specifically, we look at few-shot learning with number of examples varying from 4-10. Across 6 datasets, TART outperforms finetuning by 4.1% (with max. gains up to 20.6%) and in-context learning by 11.6% (with max. gains up to 35.7%). See Figure 26 for detailed results.

#2

(a) MNIST

(b) CIFAR-10

(c) Speech Commands

1.00

0.95

1.00

0.95

0.90

0.90

ccurac

0.97

0.85

0.85

TART

TART

A

0.80

TART

Adapter

0.80

Adapter

Adapter

FT-full

FT-full

0.75

FT-full

0.94

2

2

5

2

0.70

2?

Numberofexamples

Numberofexamples

Numberofexamples

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 307M parameter pretrained Vision Transformer (ViT) model [44]: VIT-LARGE-PATCH16-224. For audio tasks, we use representations from OpenAI's 1.5B parameter pretrained Whisper model [31]: WHISPER-LARGE. In applying TART to the representations from these models, we provide a 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 on binary classification versions of MNIST [19] (classes 0 and 8) and CIFAR-10 [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 task-adaptation approaches.

4.4 Extension to 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 a 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 to 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 to perform simple reasoning over their learned representations and introduce TART, a 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 to train a 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.

References

Monica Agrawal, Stefan Hegselmann, Hunter Lang, Yoon Kim, and David Sontag. Large language models are zero-shot clinical information extractors. arXiv preprint arXiv:2205.12689 , 2022.

Neel Alex, Eli Lifland, Lewis Tunstall, Abhishek Thakur, Pegah Maham, C Jess Riedel, Emmie Hine, Carolyn Ashurst, Paul Sedille, Alexis Carlier, et 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.

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'11) , 2011.

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. In ICLR 2023 , 2023.

Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. arXiv preprint arXiv:2304.01373 , 2023.

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 it using these metadata.

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:2108.07258 , 2021.

Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. CoRR , abs/1903.04561, 2019.

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 are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Shizhe Diao, Pengcheng Wang, Yong Lin, and Tong Zhang. Active prompting with chain-ofthought for large language models. arXiv preprint arXiv:2302.12246 , 2023.

Shivam Garg, Dimitris Tsipras, Percy S 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.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De 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.

Edward J 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.

Chip Huyen. Prompting vs. finetuning vs. alternatives, April 2023.

Jan Koco´ n, 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.

Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In ICML 2022 Workshop on Knowledge Retrieval and Language Models , 2022.

Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.

Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Finetuning can distort pretrained features and underperform out-of-distribution. arXiv preprint arXiv:2202.10054 , 2022.

Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist , 2, 2010.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, 2021.

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 1: Long Papers) , pages 4582-4597, 2021.

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:2211.09110 , 2022.

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 in Neural Information Processing Systems , 35:1950-1965, 2022.

Jiachang Liu, Dinghan Shen, Yizhe Zhang, William B Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures , pages 100-114, 2022.

Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. Ptuning: 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 Papers) , pages 61-68, Dublin, Ireland, May 2022. Association for Computational Linguistics.

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-8098, 2022.

Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies , pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.

Avanika Narayan, Ines Chami, Laurel Orr, and Christopher Ré. Can foundation models wrangle your data? Proc. VLDB Endow. , 16(4):738-746, dec 2022.

Bo Pang, Lillian Lee, and Shivakumar Vaithyanathan. Thumbs up? sentiment classification using machine learning techniques. In Proceedings of EMNLP , pages 79-86, 2002.

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.

Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision, 2022.

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. arXiv preprint , 2018.

Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ili´ c, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 , 2022.

Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

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.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax , May 2021.

Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, and Denny Zhou. Rationaleaugmented ensembles in language models. arXiv preprint arXiv:2207.00747 , 2022.

Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, and Denny Zhou. Selfconsistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171 , 2022.

P. Warden. Speech Commands: A Dataset for Limited-Vocabulary Speech Recognition. ArXiv e-prints , April 2018.

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.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 , 2022.

Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846 , 2023.

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.

Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. In Advances in Neural Information Processing Systems , 2022.

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.

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 preprint arXiv:2303.16199 , 2023.

Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems , 28, 2015.

(a)

(b)

Fine-tuning (task)

In-contextlearning

Fine-tuning (synthetic)

0.9

0.8

0.8

0.7

0.7

0.6

0.6

Fine-tuning (task)

0.5

In-contextlearning

Fine-tuning (synthetic)

0.5

matoes

20

30

40

50

60

AG-News-0

DBPedia-0

SMS Spam

Number of examples

Rotten

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 synthetically generated probabilistic inference tasks to 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 of k pairs of synthetically generated labeled natural language examples ( x, y ) . Each example x in the sequence s = ( x 1 , y 1 ) , . . . , ( x k , y k ) consists of a list of strings constructed from a fixed V size of dimension d = 30 . We use the following fixed vocabulary: [ 'sports', '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 a particular example x i , we sample each coordinate x i,j uniformly from the set {-1 , +1 } . If the sampled value is +1 , we set the value to be the corresponding word in the vocabulary, that is, x i,j = V j . Otherwise, the word x i,j is set to 'null'. For a given sequence s , we generate each of the labels { y i } as:

where we set noise parameter α = 5 . If the sampled output is 0, we set the y i to 'negative' and 'positive' otherwise.

Finally, the inputs are formatted with following template: ' x 1 : y 1 , x 2 : y 2 , ... , x k : y k ' and the model is trained using gradient descent on the loss

where z 1: i -1 corresponds to the first i -1 examples and ℓ CE is the cross-entropy loss evaluated on the transformer prediction and the true y i .

More concretely, a sample input sample sequence s to be used for training looks like:

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.8

0.8

0.8

0.7

In-context learning

Linear probing

0.7

Linear probing

In-context learning

0.7

In-context learning

Linear probing

Fine-tuning

Fine-tuning

Fine-tuning

0.6

0.6

0.6

20

30

40

50

60

20

30

40

50

60

20

30

40

50

60

Numberofexamples

Numberofexamples

Numberofexamples

A.2 Training Parameters

We train GPT-NEO (125M) on this synthetic task with a learning rate of 0.0001 and a batch size of 4. For each sequence we sampled a total of k = 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]). At 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-Reasoning decomposition evaluations

In this section, we provide details for the experimental evaluation and additional results for the representation-reasoning 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 in the set of datapoints chosen for the training task. For the hyperparameters, we performed an extensive search for all models across datasets. For details on these hyperparameters and the adapter architecture we evaluate over, see Appendix E.1.

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.

B.2 Detailed results

For each class of methods in the task-adaptation 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

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.9

Linear probing

0.8

Linear probing

Linear probing

Linear probing

0.7

0.6

Examples (k)= 18

Examples (k)= 32

Examples (k)= 48

Examples (k) = 64

0.5

0.5

0.6

0.7

0.8

0.5

0.6

0.7

0.8

0.5

0.6

0.7

0.8

0.5

0.60.70.80.9

1.0

In-context learning

In-context learning

In-context learning

In-context learning

0.9

Linear probing

Linear probing

Linear probing

Linear probing

0.8

?

0.7

0.6

Examples (k)= 18

Examples (k)= 32

Examples (k) = 48

Examples (k) = 64

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.80.9

1.0

In-context learning

In-context learning

In-context learning

In-context learning

0.9

Linear probing

Linear probing

Linear probing

Linear probing

0.8

0.7

0.6

Examples (k)= 18

Examples (k)= 32

Examples (k) = 48

Examples (k) = 64

0.5

0.6

0.7

0.8

0.5

0.6

0.7

0.8

0.5

0.6

0.7

0.8

0.50.60.70.80.9

1.0

In-context learning

In-context learning

In-context learning

In-contextlearning

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 a 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 taskagnosticity of the base model and if yes, what is the underlying reason for it.

For the first hypothesis, we evaluate the proportion of gains that can be attributed to improved 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 the total gain (fine-tuning accuracy minus in-context accuracy). Figure 10 shows a 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.

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.4

0.4

0.4

gains

0.3

gains

0.3

gains

0.3

0.2

g

0.2

g

0.2

g

g

g

soning

soning

0.1

0.1

0.1

Reas

6

S

a

?

S

R

R

0.0

0.0

0.0

-0.1

-0.1

-0.1

-0.1

0.0

0.1

0.2

-0.1

0.0

0.1

0.2

-0.1

0.0

0.1

0.2

Representation gains

Representation gains

Representation gains

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

1.0

1.0

1.0

Fine-tuning (task)

Fine-tuning (task)

Fine-tuning (task)

0.9

Fine-tuning (AG News)

0.9

Fine-tuning (AG News)

0.9

Fine-tuning (AG News)

ccuracy

0.8

0.8

0.8

0.7

0.7

0.7

A

0.6

0.6

0.6

0.5

0.5

0.5

20

30

40

50

60

20

30

40

50

60

20

30

40

50

60

Numberofexamples

Numberofexamples

Numberofexamples

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

Fine-tuning (AG News)

Fine-tuning (AG News)

Fine-tuning (AG News)

Fine-tuning (task)

Fine-tuning (task)

Fine-tuning (task)

0.9

0.9

0.9

0.8

0.8

0.8

0.7

0.7

0.7

0.6

0.6

0.6

0.5

0.5

0.5

0.4

0.4

0.4

DBPedia-0

SMS Spam

DBPedia-0

SMS Spam

DBPedia-0

SMS Spam

For the second hypothesis, we first evaluate whether fine-tuning hurts task-agnosticity. For this we evaluate two sets of accuracies: accuracy of a model fine-tuned for the specific task and the accuracy of a model on the task but fine-tuned on the AG News dataset. From Figures 11 and 12, we see that there is a drop in accuracy-over 25 . 77% across models and datasets. For the second part, we again decompose the 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 the difference between them. The reasoning drop, as before, is computed by subtracting this representation drop from the total drop. Figure 13 shows that most of

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.5

0.4

0.4

0.4

0.3

0.3

loss

loss

SSOl

S

0.3

oning

Reasoning

0.2

asoning

0.2

Reason

0.2

Reas

0.1

0.1

R

0.1

R

0.0

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) BLOOM (560M)

Adapter (AG News)

Adapter (AG News)

Adapter (AG News)

Adapter (task)

Adapter (task)

Adapter (task)

0.9

0.9

0.9

0.8

0.8

0.8

0.7

0.7

0.7

0.6

0.6

0.6

0.5

0.5

0.5

0.4

0.4

0.4

SST

SS

DBPedia-0

SMS Spam

DBPedia-0

SMS Spam

DBPedia-0

SMS Spam

this drop in task-agnosticity can 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 at how a single adapter generalizes across tasks. For this we train an adapter on 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 to 256, number of decoder layers to 12, and number of heads to 8 for a total of 22 million parameters. Since the GPT-2 backbone outputs a sequence of embeddings,

x1

y1

Linear

GPT-2Backbone

Linear

Layer

(Untrained)

Layer

xk

yk

Input

Embedding

ndino

Dimensiond

Dimension demb

Embeddings

0.7

0.6

S

0.5

S

0.4

0.3

0.2

0

5000

10000150002000025000

Number of steps

we additionally add a linear layer in the end to convert the output to scalar values (see Figure 15). Additionally, the binary labels y are encoded as a 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 of 258 (allowing for up to 129 in-context examples). The batch size was set to 64, learning rate to 0.0001 and the model was trained for a total of 24000 epochs. Each batch of training data consists of sampling a sequence of 258 examples using eq. (3). In addition to these hyperparameters, we used a curriculum on the input dimensions and on the number of examples in the sequence to train our module-the input dimensions started from a value of 4 and were incremented by 4 every 1000 epochs while the number of examples started from 18 and were incremented by 30 every 1000 epochs.

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). In order to reduce the dimensionality of these representations, we perform PCA on the output embeddings of the base model, learning the components using only the training points available for that specific task. The test examples are then projected onto these principal components to 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 the 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 LOO embeddings consistently perform better across the different model families.

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.80

0.78

0.76

0.78

0.76

0.74

0.76

0.74

0.72

TART-Vanilla

TART-Vanilla

TART-Vanilla

TART-LOO

0.74

TART-LOO

0.72

TART-LOO

0.70

20

30

40

50

60

20

30

40

50

60

20

30

40

50

60

Numberofexamples

Numberofexamples

Numberofexamples

(a) Vanilla embeddings

(b) LOO embeddings

Embeddings

Ex.

Inference

Pre-trained LLM

Module

Ex. k

Averaging

Layer

Embeddings

Ex. 2

Ex. k

Pre-trained LLM

Ex. 1

Averaging

Layer

Ex. 1

Inference

Module

Ex. k

yk

Embeddings

Ex. 1

·

Ex. k-1

Pre-trained LLM

Ex.k

Averaging

Layer

C.3 Proof of Theorem 1

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

where k represents the maximum number of in-context examples the Transformer can support, d represents the input dimensions, and Θ represents the corresponding parameter class over which the Transformer family is defined.

Observe that the Transformer family T Θ takes as input a sequence of k train examples, each corresponding to two tokens of hidden dimension d : a covariate x ∈ R d and a binary label, encoded as a one-hot vector in d dimension. This sequence of train examples is followed by a test example, for which we only have the features x k +1 .

Given this background, let P syn denote the synthetic distribution over sequences { ( x 1 , y 1 ) , . . . , ( x k , y x ) , ( x k +1 ) } . Similarly, let P NL denote the corresponding distribution over sequences derived from natural language tasks where x i denotes the LLM embeddings of the example. Recall from Section 3.2.1, the synthetic training distribution P syn is given by

for each training point ( x i,t , y i,t ) . The test point is also sampled similarly from an independent standard normal distribution. Let ℓ : R × R ↦→ R be the loss function used for evaluating the performance of the reasoning module. Further, let use denote the expected loss under a distribution P

and the corresponding empirical distribution over samples S by ˆ err P , where the dependence on the samples is implicit. Given these samples, we denote the empirical risk minimizer

In addition to these notation, we make the following Lipschitz assumption on the the loss function ℓ and the Transformer model T .

Assumption 1. [Lipschitz loss.] For any two output labels y, y ′ , the loss function ℓ is Lipschitz with constant c 1 , that is,

Assumption 2. [Lipschitz models.] For any two input sequences s, s ′ , each any model T ∈ T Θ is Lipschitz with constant L , that is,

Given this setup, we are now ready to state a formal version of Theorem 1.

Theorem 2 (Formal version of Theorem 1) . Let T S ∈ T Θ denote the trained reasoning module on set S of synthetic logistic regression tasks with n syn sequences sampled from distribution P syn in eq. (3) . Let the loss function ℓ satisfy Assumption 1 and the model class T Θ satisfy Assumption 2. Then, with probability at least 1 -δ , we have

where W 1 denotes the Wasserstein-1 metric and VC ( T Θ ) represents the VC dimension of class T Θ

Proof. We begin by decomposing the error err P NL ( T S ) into three components as

We now upper bound each of the terms (I) and (II) separately.

Bound on Term (I). Let γ denote an arbitrary joint distribution over the distributions P syn and P NL . Then, we can bound the first term as

where ( i ) follows from the independence of the two expectations and ( ii ) follows from that ( i ) holds for any arbitrary joint distribution γ . The final bound on this term now follows:

where the inequalities ( i ) follows from Assumption 1 and ( ii ) follows from Assumption 2. This completes the bound on Term (I).

Bound on Term (II). Using a standard generalization bound [34, see Theorem 26.5], we have with probability at least 1 -δ

where R ( ℓ ◦T Θ ) denotes the Rademacher complexity of the class T Θ composed with the loss function ℓ 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 an 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

Observe that each sequence corresponded to a single hyperplane w t which defined the binary classification problem. For the multi-class extension, suppose that there are L labels. For each sequence s t , we sample one hyperplane per class in d -dimensional space:

Each covariate x i,t is sampled again from the standard Normal distribution as before. The label y i,t is sampled from the multinomial distribution over [ L ] given by:

Given this distribution over sequences s t , we train the TART module as before, but with the linear layer with output dimension L instead of 1 as with the binary classification setup. We evaluate our

multi-class variant of TART on AGNews-4, Ecommerce-4 and the DBPedia-14 datasets (see Table 1 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-Neo-125M) is close to random, we see that the TART substantially improves upon this performance. It provides improvement of up to 45.9% on AGNews, 52.6% on Ecommerce, and 59.3% on DBPedia-14 datasets.

E Details for experimental evaluation

We describe supplementary experimental details from Section 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).

E.1 Experimental setup

We begin by providing dataset statistics and details of the baselines.

E.1.1 Dataset construction and statistics

Table 4 and 3 provides a detailed breakdown of dataset statistics. For each dataset, we use the original test sets 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 Science/Technology. We create a separate binary classification task for each category, sampling negatives from the remaining classes. DBPedia is a 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 a class-balanced set of 64 examples from the original dataset. For each dataset, we sample 5 separate training sets, using 5 different random seeds. In evaluations, we evaluate TART and the baseline methods across each of these 5 different training sets.

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.8

0.8

0.8

0.7

FT-full

In-context learning

0.7

FT-full

In-context learning

0.7

FT-full

In-context learning

FT-layer

TART

FT-layer

TART

FT-layer

TART

Adapter

Adapter

Adapter

0.6

0.6

0.6

0.5

0.5

0.5

20

30

40

50

60

20

30

40

50

60

20

30

40

50

60

Numberofexamples

Numberofexamples

Numberofexamples

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 are trained as follows:

(a) Number of examples = 18

(b) Number of examples = 32

(c) Number of examples = 48

(d) Number of examples = 64

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

RT

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

80.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.70.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

Last layer fine-tuning: Freeze all layers of transformer but the final transformer block and the language modeling head.

(a) Number of examples = 18

(b) Number of examples = 32

(c) Number of examples = 48

(d) Number of examples = 64

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

RT

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

80.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.70.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

Adapter: Combine a frozen LLM base transformer model with a trainable adapter head-an MLP composed of a single linear layer followed by non-linearity.

(a) Number of examples = 18

(b) Number of examples = 32

(c) Number of examples = 48

(d) Number of examples = 64

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

RT

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.50.60.70.8

80.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.70.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

1.0

1.0

1.0

1.0

0.9

0.9

0.9

0.9

0.8

0.8

0.8

0.8

ART

0.7

0.7

0.7

0.7

0.6

0.6

0.6

0.6

0.5

0.5

0.5

0.5

0.4

0.4

0.4

0.4

0.4

0.5(

0.60.70.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

In-context learning

FT-full

FT-layer

Adapter

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, 1e-4, 3e-5, 1e-5, 8e-6), and range of epochs (5, 10, 15, 20, 50). For all models < 1B parameters, we use a batch size of 1. For all models > 1B parameters, we use a batch

(a) GPT-NEO (125M)

(b) PYTHIA (160M)

(c) BLOOM (560M)

0.8

0.8

0.8

0.7

0.7

0.7

0.6

0.6

0.6

TART

TART

TART

In-context learning

In-context learning

In-context learning

0.5

2

4

2

5

2

2

8

0.5

2

4

2

5

2

2

8

0.5

2

4

2

2?

2

8

Numberofexamples

Numberofexamples

Numberofexamples

size of 8. We use these same batch sizes at evaluation time. We perform our hyperparameter searches with a 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 4 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 task-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 a better rate than that of TART.

E.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 (twitter-complaints, neurips-impactstatement-risks, overulling, ade-corpusv2, tweet-eval-hate, terms-of-service, tai-safety-research) with the exception of systematic-review-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 GPTNEO (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-0, DBPedia-0, SST, SMS Spam, Youtube and Rotten Tomatoes. We evaluate the performance of TART over k =[18, 32, 48, 64, 128, 192, 256] where k is the number of in-context examples. When evaluating our base models, we

Embeddings

Ex. 1

Pre-trained

Inference

Model

Module

Ex. K

Ex.k

Ex. 1

Ex.1 Averaging

Layer

Image, Audio, Text

evaluate over k =[8, 24]-values of k that maximize the context window. We use a lower-bound of 8 given that the maximum input sequence length in the training set for AG News is 256. With such a sequence length, the maximum number of in-context examples that fit in the context-window is 8, hence the lower bound.

Embeddings. For these evaluations, we use what we call 'streaming' embeddings (see Figure 24). In this setup, we use the context window of the LLM to encode a single example at a time. The final embeddings are 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.

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 an increasing trend in accuracy with number of examples across models.

E.3 Generalization across architectures

In this section, we demonstrate that it is possible to train TART reasoning modules on architectures beyond transformers. More concretely, we train a reasoning module that has the Hyena architecture [30] using the same synthetic logistic regression tasks.

E.3.1 Setup and Hyper-parameters

We instantiate a reasoning module with 12 Hyena blocks, a hidden dimension size of 256, and a sequence length of 2050. We train with a batch size of 16, using a learning rate of 5e-05. We sample data with a noise parameter ( α ) equal to 1. We train the model for 5000 steps. For our evaluations, we use the final checkpoint (i.e., 5000) of the reasoning module.

E.3.2 Evaluation

We evaluate the Hyena-based reasoning module applied to GPT-NEO (125M) on 6 datasets: SMS Spam, SST, AG-News-0, DBpedia-14-0, Youtube, and Rotten Tomatoes. As seen in Figure 25, the Hyena-based reasoning module transfers to natural language tasks and performs competitively with the transformer-based 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 k = 4 -10 . TART outperforms both in-context learning and task-specific fine-tuning baselines.

(a) k = 18

(b) k = 32

(c) k = 48

(d) k = 64

TART-Hyena

TART-Hyena

TART

TART

0.9

0.9

0.8

0.8

0.7

0.7

0.6

0.6

0.5

0.5

SST

SST

AG-News-0

DBPedia-0

SMS Spam

AG-News-0

DBPedia-0

SMS Spam

TART-Hyena

TART-Hyena

TART

TART

0.9

0.9

0.8

0.8

0.7

0.7

0.6

0.6

0.5

0.5

SST

SST

AG-News-0

DBPedia-0

SMS Spam

AG-News-0

DBPedia-0

SMS Spam

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 our algorithm with baselines.

E.5.1 Dataset details

For audio classification, we use the Speech Commands (Version 0.01) dataset [40]. Speech Commands is a multi-class classification task where the task is to detect preregistered keywords by classifying utterances into a predefined set of words. We construct a 3 binary classification task over the keywords 'stop' and 'go', 'up' and 'down', and 'yes' and 'no' (see Table 6 for more details).

For image classification, we use CIFAR-10 [17] and MNIST [19]. Both tasks are multi-class classification tasks. We create 3 binary classification tasks for each of the datasets. For CIFAR-10 the tasks are: airplane vs. bird, bird vs. horse, and ship vs. automobile. For MNIST the tasks are: 0 vs. 8, 1 vs. 6 and 2 vs. 4. See Table 6 for more details.

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 we are learning to predict for (i.e., airplane and bird for CIFAR10 and 0 and 8 for MNIST).

(a) k = 4

(b) k = 6

Fine-tuning

Fine-tuning

In-context learning

In-context learning

TART

TART

0.9

0.9

0.8

0.7

0.7

0.6

0.6

0.5

0.5

SST

SST

AG-News-0

DBPedia-0

SMS Spam

AG-News-0

DBPedia-0

SMS Spam

Fine-tuning

■

Fine-tuning

In-context learning

In-context learning

TART

TART

0.9

0.9

0.8

0.7

0.7

0.6

0.6

0.5

0.5

SST

SST

AG-News-0

DBPedia-0

SMS Spam

AG-News-0

DBPedia-0

SMS Spam

(d) k = 10

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].

We compare against two baseline task-adaptation methods: 1) full fine-tuning and 2) adapters. We use the same architectures as described in Appendix E.1.2. For vision tasks, we use Google's 307M

(c) k = 8

(a) MNIST (1 vs. 6)

(b) MNIST (2 vs. 4)

1.00

0.99

0.98

0.95

0.97

TART

TART

Adapter

Adapter

0.96

FT-full

FT-full

0.95

12

0.90

2

2

4

2

5

2

6

2

8

2

4

2

5

2

6

2

8

Number of examples

Number of examples

(a) CIFAR-10 (ship vs. automobile)

(b) CIFAR-10 (bird vs. horse)

0.95

0.95

0.90

0.90

0.85

TART

TART

0.85

Adapter

Adapter

0.80

FT-full

FT-full

2

4

#2

2

6

2?

2

8

0.80

2

4

2

5

2

6

2?

2

8

Number of examples

Number of examples

parameter pretrained Vision Transformer (ViT) model [44]: VIT-LARGE-PATCH16-224. For audio tasks, we use OpenAI'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 rate for each dataset in order to optimize performance. We search over a range of learning rates (1e-3, 5e-04, 1e-4, 5e-5, 1e-5, and 8e-6) and a range of epochs (5, 10, 15 and 20). For all models we use a batch size of 1. We perform our hyperparameter searches for a fixed number of train samples (128) and run our hyperparameter searches over 3 random seeds.

E.5.3 Evaluation

We plot the accuracy as a 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.

F Broader Impact

Our submission focuses on understanding in-context learning and how it compares to task-specific fine-tuning approaches. We view our work as a 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 a plug-and-play solution that doesn't require any fine-tuning, we hope that our method can help domain experts (e.g., lawyers, scientists) incorporate LLMs into their workflows. In this sense, we view TART as a

0.95

0.95

0.90

0.90

0.85

TART

TART

Adapter

0.85

Adapter

0.80

FT-full

FT-full

2

4

2

5

2

6

12

2

8

0.80

2

4

2

5

2

6

2?

2

8

Number of examples

Number of examples

contribution that is democratizing ML by making them more accessible and usable to experts in a variety of domains.

G Compute Resource Estimates

Resources for training TART We use a single NVIDIA RTX A6000 GPU ( $2/hr) for 6 hours to train our TART inference head, costing a total of $18 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 a 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 a total of $13,440.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nGiven 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nNo problems noted.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality:\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.\n\nWEAKNESSES:\nThere's a typo on line 339, it should be CIFAR-10 (classes plane and bird) and MNIST (classes 0 and 8)\n\nQUESTIONS:\nIn 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\nWhat 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.\n\nLIMITATIONS:\nThe authors addressed the limitations and broader impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}