SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Avanika Narayan*
{kushb, avanika, chrismre} @cs.stanford.edu, cdesa@cs.cornell.edu
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).<sup>2</sup>
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 fine-tuning 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).
We first investigate why this quality gap exists. We decompose an LLM's in-context learning capability into two abilities: learning good representations for the task and performing probabilistic inference,
2Our code and model is available at https://github.com/HazyResearch/TART
Figure 1: Taxonomy of task adaptation strategies. (Left) Comparison of different adaptation strategies across three desiderata: task-agnostic, quality, scalability. (Right) Demonstration of parameter updates across adaptations strategies, colored regions represent parameter changes as a result of the adaptation strategy.
or reasoning, over these representations<sup>3</sup>. Is the gap because the data representations learned by the models lack task-specific information or is there some deficiency in the model's ability to use this information. 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
<sup>3</sup>In Section 2.3, we precisely define the this representation-reasoning decomposition in eq. (1)
Figure 2: TART. (Left) Inference module training procedure: The inference module is trained on sequences of synthetically generated logistic regression tasks. (Right) End-to-end framework: TART composes a pre-trained LLM with the inference module. TART uses the LLM to embed the input text. These embeddings, along with the train labels, are passed as a sequence to the inference module which generates a final prediction.
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).
- 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 in-context [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 increases 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  $(\mathrm{Acc}_{\mathrm{LR}})$ . Using this as an intermediate, we decompose the performance gap
$$
\Delta_ {\text {p e r f}} := \mathrm {A c c} _ {\mathrm {F T}} - \mathrm {A c c} _ {\mathrm {I C L}} = \underbrace {\mathrm {A c c} _ {\mathrm {F T}} - \mathrm {A c c} _ {\mathrm {L R}}} _ {\Delta_ {\text {r e p}}} + \underbrace {\mathrm {A c c} _ {\mathrm {L R}} - \mathrm {A c c} _ {\mathrm {I C L}}} _ {\Delta_ {\text {r e a s}}} \tag {1}
$$
where  $\Delta_{\mathrm{rep}}$  represents the gap in performance which can be attributed to insufficient representation capacity and  $\Delta_{\mathrm{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
(b)
Figure 3: All results for GPT-NEO (125M). (a) Accuracy of in-context learning vs. linear probing on model embeddings: representations have sufficient information. (b) Fine-tuning majorly improves task-specific reasoning across datasets. (c) Accuracy (averaged across 6 datasets) of task-specific fine-tuned model vs. accuracy of model fine-tuned on AGNews and evaluated on task. Fine-tuning hurts task-agnosticity.
and the accuracy of linear probing – the smaller the gap the better the representations. Similarly, the models reasoning ability, or 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  $\Delta_{\mathrm{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 task-agnostic, 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\mathrm{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}),\ldots ,(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
$$
\ell \left(T _ {\theta}\right): = \mathbb {E} _ {x, y} \left[ \frac {1}{k} \sum_ {i = 1} ^ {k} \ell_ {\mathrm {C E}} \left(T _ {\theta} \left(z _ {1: i - 1}, x _ {i}\right), y _ {i}\right) \right], \tag {2}
$$
where  $z_{1:i-1}$  corresponds to the first  $i-1$  examples and  $\ell_{\mathrm{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  $\theta$  comprises a different  $d$ -dimensional logistic regression problem, sampled as
$$
\text {S e q u e n c e} s _ {t}: w _ {t} \sim \mathcal {N} (0, I _ {d}), \quad x _ {i, t} \sim \mathcal {N} (0, I _ {d}), \quad y _ {i, t} \sim \sigma (\alpha \langle x _ {i, t}, w _ {t} \rangle) \quad \text {f o r} i \in [ k ], \tag {3}
$$
where  $\sigma$  represents the sigmoid function and the multiplier  $\alpha$  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  $\langle x_{i,t}, w_t \rangle$
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
(b)
Figure 4: Properties of TART's inference module. (a) Comparison with learned logistic function: inference module recovers underlying probabilities. (b) Variation in error with different noise levels for model trained on  $\alpha = 10$ . (c) Comparison of TART performance when using LOO embeddings and vanilla embeddings.
$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