SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\alpha$ .
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,  $\alpha$ , of the logistic regression problem. Recall that we trained our inference module by fixing the noise level  $\alpha = 10$ . At inference time, we vary the noise level to [0.5, 1, 10, 20], where lower values correspond to noisier problem. The reasoning module generalizes to easier problems without any drop in accuracy but as we make the problem harder ( $\alpha = [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  $\mathcal{T}$  represent the class of transformer models and  $T_{S} \in \mathcal{T}$  denote the trained reasoning module on set  $S$  of synthetic regression with  $n_{\mathrm{syn}}$  sequences sampled from distribution  $P_{\mathrm{syn}}$  in eq. (3). The error of the transformer  $T_{S}$  when evaluated on a distribution  $P_{\mathrm{NL}}$  over natural language sequences is
$$
\operatorname {e r r} _ {P _ {\mathrm {N L}}} \lesssim W _ {1} \left(P _ {\mathrm {N L}}, P _ {\mathrm {s y n}}\right) + \sqrt {\frac {\operatorname {C o m p} (\mathcal {T})}{n _ {\mathrm {s y n}}}} + \operatorname {e r r} _ {P _ {\mathrm {s y n}}} \left(T _ {S}\right), \tag {4}
$$
where  $W_{1}$  denotes the Wasserstein-1 metric,  $\mathrm{Comp}(\mathcal{T})$  represents the complexity of class  $\mathcal{T}$ , and  $\hat{\mathrm{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_{\mathrm{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_{\mathrm{NL}}$  corresponds to the distribution induced on the embeddings of the natural language data (post the PCA and renormalization) while the distribution  $P_{\mathrm{syn}}$  is the one corresponding to the sampling process in equation (3).
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.
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.
(b)
Figure 5: Effects of scale. (a) Effect of number of in-context examples on performance for different task adaptation strategies. (b) Effect of model size on the performance of different task adaptation strategies. (c) Beyond context length limitations, performance comparison with respect to number of in-context examples.
4.2 Natural 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 1000x larger in size.
Performance with number of in-context examples. Our results demonstrate that performance of TART scales with number of in-context examples (see Figure 5a). Across 14 tasks and 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 10x less samples, with the context window saturating at 24 examples only and lags TART by an average of  $19.1\%$ .
Extension to other architectures. We compare the performance of TART by changing the 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 in-context 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 fine-tuning 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.
(b) CIFAR-10
Figure 6: TART can generalize across domains using the same inference module that was used for language benchmarks: Performance across vision tasks (MNIST, CIFAR-10) and an audio task (Speech Commands).
(c) Speech Commands
4.3 Extensions to other modalities
We demonstrate that TART is not only agnostic to models and tasks, but also modalities. We extend TART to classification tasks on modalities beyond language: vision and audio. For vision tasks, we use representations from Google's 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.
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.
We are grateful to Simran Arora, Rishi Bommasani, Niladri Chatterji, Arjun Desai, Sabri Eyuboglu, Neha Gupta, Karan Goel, Erik Jones, Ananya Kumar, Cassidy Laidlaw, Megan Leszczyński, Piero Molino, Laurel Orr, Michael Poli, Dimitris Tsipras, Michael Wornow, Ce Zhang, and Michael Zhang for their helpful comments and feedback, and discussions which helped shape this project. We thank Neel Guha for helpful discussions in shaping the narrative of this paper.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML);
US DEVCOM ARL under No. W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. CDS was supported by a NSF CAREER (award 2046760).
The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
(b)
Figure 7: Fine-tuning with NL synthetic task. (Left) Averaged over 6 different tasks, fine-tuning with the NL synthetic task provides a lift over base in-context learning, and scales with number of examples. (Right) Dataset level comparisons between task-specific fine-tuning, in-context learning and synthetic fine-tuning: synthetic fine-tuning outperforms base in-context learning on 4 out of 6 datasets, but lags task-specific tuning.
A Fine-tuning model with NL-based Probabilistic Inference Tasks
As highlighted in Section 1, we describe the details for directly fine-tuning an LLM on 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), \ldots, (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:
$$
w _ {t} \sim \mathcal {N} (0, I _ {d}), \quad y _ {i} \sim \sigma (\alpha \langle x _ {i}, w \rangle), \text {f o r} i \in [ k ], \tag {5}
$$
where we set noise parameter  $\alpha = 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}, \ldots, x_{k}: y_{k}$  and the model is trained using gradient descent on the 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 {6}
$$
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$ .
More concretely, a sample input sample sequence  $s$  to be used for training looks like:
"sports love null car ... cat: positive, null love null car ... null: negative,
sports null hat null ... cat : positive"
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
Figure 8: Comparison of linear probing, in-context learning, and fine-tuning. Accuracy of in-context learning vs. linear probing on model embeddings across three model families: representations have sufficient information.
(c) BLOOM (560M)
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.
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
Figure 9: Linear probing vs. in-context learning. Scatter plot of accuracy of in-context learning vs. linear probing on model embeddings across model families and different number of in-context examples: linear probing consistently outperforms in-context learning indicating that the learned representations have sufficient information. Each point in the plot represents a dataset.
(b) PYTHIA (160M)
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 task-agnosticity 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 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.
Figure 10: Effects of fine-tuning on reasoning. Across datasets (each point in plot represents a dataset) and model families, fine-tuning improves task-specific reasoning which improves its performance over base in-context learning.
(a) GPT-NEO (125M)
(b) PYTHIA (160M)
Figure 12: Effects of fine-tuning on task-agnosticity (dataset level) Accuracy of task-specific fine-tuned model vs. accuracy of model fine-tuned on AG-News-0 and evaluated on task. Fine-tuning consistently hurts task-agnosticity across all three model families and datasets.
Figure 11: Effects of fine-tuning on task-agnosticity Accuracy of task-specific fine-tuned model vs. accuracy of model fine-tuned on AG-News-0 and evaluated on task. Fine-tuning hurts task-agnosticity across all three model families.
(c) BLOOM (560M)