TART: A plug-and-play Transformer module for task-agnostic reasoning
This addresses the problem of inefficient task adaptation in LLMs for AI researchers and practitioners by offering a task-agnostic reasoning enhancement, though it is incremental as it builds on existing Transformer architectures.
The paper tackles the performance gap between in-context learning and task-specific tuning in large language models by identifying a lack of probabilistic reasoning ability, and proposes TART, a plug-and-play Transformer module that improves reasoning across models, tasks, and modalities, achieving results where GPT-Neo (125M) outperforms BLOOM (176B) and nears GPT-3 (175B) within 4% on the RAFT Benchmark.
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 analysis 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 propose TART which generically improves an LLM's reasoning abilities using a synthetically trained Transformer-based reasoning module. TART trains this 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 binary classification tasks), and even across different modalities (audio and vision). Additionally, 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). Our code and models are available at https://github.com/HazyResearch/TART .