PaperlayerView exact revision on arXiv

Source-linked AI summary

The Surprising Effectiveness of Test-Time Training for Few-Shot Learning

Ekin Akyürek, Mehul Damani, Adam Zweiger, Linlu Qiu, Han Guo, Jyothish Pari, Yoon Kim, Jacob Andreas

arXiv:2411.07279v2cs.AIcs.CLcs.LG

TL;DR

Language models often struggle to learn structurally novel tasks from a few examples, leaving open how well they acquire untrained reasoning skills. This paper applies test-time training to in-context examples; it reports gains on ARC and BBH, including a 7.3-percentage-point improvement over few-shot prompting on BBH. These findings indicate potential for TTT on novel reasoning and few-shot learning tasks.

  • Problem

    Whether language models can acquire untrained reasoning, planning, and abstraction skills from a few examples remains an open question.

    Source: §1 · p. 1

  • Method

    The paper combines in-context learning with test-time training, temporarily updating model parameters using training data constructed from few-shot examples.

    Source: §1 · p. 2, §3.1 · p. 3, §3.3 · p. 4

  • Results

    TTT improves few-shot performance on ARC and BBH; on BBH it yields a 7.3% absolute improvement over few-shot prompting, achieving 57.8% accuracy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings indicate TTT’s potential for improving language-model performance on abstract reasoning and few-shot learning tasks.

    Source: §7 · p. 8

    The public availability of ARC and BBH raises the possibility that models saw examples from these datasets during pretraining.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Language models (LMs) have shown impressive performance on tasks within their training distribution, but often struggle with structurally novel tasks even when given a small number of in-context task examples. We investigate the effectiveness of test-time training (TTT) -- temporarily updating model parameters during inference using a loss derived from input data -- as a mechanism for improving LMs' reasoning and few-shot learning capabilities. On the Abstraction and Reasoning Corpus (ARC), performing TTT with in-context examples yields up to $6\times$ higher accuracy compared to fine-tuned baselines -- reaching $53.0\%$ on the public validation set with an 8B-parameter LM and $61.9\%$ when ensembled with program-synthesis methods, matching average human performance. On BIG-Bench Hard (BBH), TTT on in-context examples surpasses standard few-shot prompting in the $10$-shot setting by $7.3$ percentage points ($50.5\%$ to $57.8\%$). Our findings highlight the limitations of in-context learning for novel tasks and demonstrate the potential of test-time training to enhance language model adaptability.

1. Introduction

Language models can handle familiar few-shot tasks, but whether they acquire untrained reasoning, planning, and abstraction skills remains an open question. This paper combines in-context learning with test-time training and reports gains on ARC and BBH, including large improvements on structurally challenging tasks.

1. Introduction

  • Few-shot learning can work for tasks related to pretraining, but acquiring skills absent from training—especially non-trivial reasoning, planning, and abstraction—remains an open question.The paper frames this question as relevant to understanding flexible novel-skill acquisition.Source: §1 · p. 1
  • The paper investigates combining in-context learning with explicit gradient-based updates on test data to improve challenging reasoning and rule-based generalization tasks.Its contributions include analysis of training-data selection, objectives, and interactions with pretrained parameters and in-context learning.Source: §1 · p. 2
  • The evaluation applies TTT to ARC and BIG-Bench Hard, two challenging benchmark suites.ARC evaluates visual-puzzle reasoning; BBH is the second benchmark suite named in the paper.Source: §1 · p. 2
  • TTT reaches 53.0% accuracy on ARC with an 8B model and 61.9% when ensembled with program synthesis; on BBH, it improves over few-shot prompting by 7.3 percentage points to 57.8%.The ARC ensemble result is described as comparable to human performance.Source: Figure 1 · p. 1, §1 · p. 2
  • TTT’s gains are particularly large on tasks involving structural rules or distribution shifts, with 20–50 percentage points of improvement over standard in-context prompting.Examples include Dyck languages and Ruin names.Source: §1 · p. 2

2. Preliminaries

In-context learning conditions an LM on task examples without updating its parameters. Test-time training instead temporarily updates parameters using test-time data, and this paper extends it to use few-shot demonstrations for prediction.

2.1. In-context Learning

  • In-context learning uses input-output examples and a new input as context for generating the corresponding output, without parameter updates.The paper notes that ICL can struggle with novel tasks such as ARC.Source: §2.1 · p. 2

2.2. Test-Time Training

  • Test-time training temporarily updates model parameters using a dataset formed for each test input, then uses the updated model for prediction.The general process includes optimizing a loss on the constructed dataset and sampling with an augmented self-consistency strategy.Source: §2.2 · pp. 2–3
  • This paper extends TTT to few-shot learning by using demonstration input-output pairs as transductive data, focusing on direct input-output prediction rather than intermediate reasoning or explicit function descriptions.It distinguishes this from prior TTT objectives built from the input alone.Source: §2.2 · p. 2
  • Few-shot examples can form synthetic in-context training tasks, while task-specific information can further expand the dataset for ARC; BBH tests the general case without such information.The figure also distinguishes leave-one-out from direct I/O data generation, alternative losses, and task-specific versus shared adapters.Source: §2.2 · pp. 2–3, Figure 2 · p. 3

3. TTT Design

The TTT design constructs task-specific training data from demonstrations, chooses which tokens contribute to the loss, and adapts either task-specific or shared LoRA parameters. Across these choices, the paper reports that in-context demonstrations outperform direct I/O training, the loss over demonstration and test outputs works best, and shared adapters have benchmark-dependent effects.

3. TTT Design

  • The design explores how to construct test-time training data, choose an optimization objective, and efficiently update model parameters.The overview presents these as key design choices for applying TTT to large language models.Source: §3 · p. 3, §3.1 · p. 3

3.1. Data Generation

  • Direct I/O training is more computationally efficient, but underperforms methods that use in-context demonstrations.Direct I/O treats each input-output pair as an independent training instance; leave-one-out instead holds one pair out while using the rest as demonstrations.Source: §3.1 · p. 3
  • For structured tasks such as ARC, invertible transformations including flips, rotations, and color permutations can augment the TTT dataset while preserving core data relationships.The transformations expand the training signal.Source: §3.1 · p. 3

3.2. Loss Function

  • The loss over both demonstration and test outputs works best among the tested loss choices.Alternatives include loss on only the test output or on all input and output tokens.Source: §3.2 · pp. 3–4

3.3. Parametrization

  • A shared LoRA adapter degrades ARC performance but improves BBH performance, while the default is a task-specific adapter for each test task.The adapters receive a small number of gradient steps on the constructed TTT dataset.Source: §3.3 · p. 4

4. Abstraction and Reasoning Corpus

This section evaluates test-time training (TTT) for ARC across data and optimization choices, model sizes, augmented inference, and comparisons with other systems. The reported results show gains from TTT and its combination with program synthesis, alongside lower accuracy on a semi-private evaluation.

4.1. Background

  • ARC tasks pair 2D-grid inputs and outputs under a shared transformation rule, with 2–7 demonstrations and 1–3 test examples per task.Source: §4.1 · p. 4

4.2. Experimental Details

  • Evaluation uses pass@2 on 80 balanced ARC validation tasks; augmented inference transforms tasks, votes within each transformation, then globally selects two predictions.The 80-task subset is balanced across four difficulty levels.Source: §4.2 · pp. 4–5, Figure 4 · p. 5

4.3. Impact of TTT Design

  • TTT raises accuracy approximately 6× (5% →29%) over fine-tuning; removing transformations costs 16 tasks (55%), while direct I/O costs 11 tasks (38%).Source: Figure 5 · p. 5, §4.3 · p. 5
  • Per-task adapters outperform a shared adapter by 7 tasks (24%), while adding demonstration-output loss raises accuracy from 26% to 29%.Source: Figure 5 · p. 5, §4.3 · p. 5

4.4. Impact of Model Size

  • Fine-tuned accuracy rises with model size, reaching 36% for 8B; TTT improves every tested size, with 1B and 3B reaching similar accuracy afterward.Source: §4.4 · p. 5, Figure 6 · p. 5

4.5. Impact of Augmented Inference

  • Voting across transformed predictions improves on individual transformations; hierarchical voting beats flat voting and approaches oracle performance.Source: Figure 7 · p. 6, §4.5 · p. 6

4.6. Comparison to Other Systems

  • 61.9% accuracy with the TTT pipeline, BARC neural model, and synthesizer is comparable to average human performance of 60.2%; TTT on BARC reaches 53.0%.Applying TTT to the fine-tuned model raises accuracy from 18.3% to 47.1%.Source: Table 1 · p. 6, §4.6 · p. 6
  • 47.5% accuracy was observed on the ARC-AGI semi-private evaluation; the authors suggest greater distribution shift may explain the decline.Source: §4.6 · pp. 6–7

5. BIG-Bench Hard

BBH tests TTT across diverse natural-language reasoning tasks. The section reports gains over in-context learning and examines how TTT design and task characteristics relate to performance.

5.1. Background

  • BBH comprises 27 challenging tasks across 23 task types and tests reasoning, compositionality, and generalization.Its natural-language structure and lack of a shared input format distinguish it from ARC.Source: §5.1 · p. 7

5.2. Experimental Details

  • The evaluation uses Llama 3.1 8B with task-specific LoRA, no synthetic pre-finetuning, 10-shot demonstrations, five seeds, and greedy decoding.BBH does not use ARC-style invertible transformations for augmented inference.Source: §5.2 · p. 7

5.3. Impact of TTT Design

  • 57.8% overall accuracy with TTT exceeds standard ICL (50.5%) and Direct I/O learning (51.5%).The reported comparison is on BIG-Bench Hard.Source: Figure 8 · p. 7, §5.3 · p. 7
  • Without example permutations TTT reaches 55.7%; test-output-only loss reaches 54.4%, while loss on inputs and outputs reaches 55.9%.These are design ablations against the overall TTT result.Source: §5.3 · p. 7
  • Unlike on ARC, a shared adapter improves BBH performance; task instructions help distinguish tasks, and some task updates help others.The authors give Logical Deduction Five Objects aiding Logical Deduction Three Objects without hurting Word Sorting as an example.Source: §5.3 · p. 8

5.4. Task-Specific Analysis

  • Among 27 tasks, 12 improve by at least 2% under TTT, including nine with gains of at least 5%; only two decline by at least 2%.Task-specific results vary substantially.Source: §5.4 · p. 8
  • The authors hypothesize that structured patterns favor TTT, but Boolean Expressions falls from 85.7% to 80.4%, and why tasks differ remains open.The proposed explanation links rule-like tasks to adaptation to latent structural regularities.Source: §5.4 · p. 8

6. Related Work

The related work situates TTT among test-time adaptation methods and reviews approaches to ARC reasoning. It contrasts neural prediction with program synthesis while noting hybrid methods.

Test-time training

  • Test-time parameter updates connect to local learning and test-time self-supervision; language-model work includes fine-tuning on retrieved neighbors and active retrieval selection.The cited examples include Bottou & Vapnik (1992), Sun et al. (2020), Hardt & Sun (2024), and Hübötter et al. (2025).Source: §6 · p. 8

ARC challenge

  • ARC methods include program synthesis that finds a transformation function and neural methods that directly predict outputs; hybrid approaches combine deep learning with program synthesis.The paper uses a fully neural approach and reports that integrating it with program synthesis improves performance.Source: §6 · p. 8

7. Conclusion

The paper investigates TTT on ARC and BBH and reports improved language-model performance on abstract reasoning and few-shot learning tasks. It concludes that TTT may help address novel reasoning tasks.

7. Conclusion

  • The paper investigates TTT for language-model performance on abstract reasoning and few-shot learning tasks in ARC and BBH.The conclusion describes improvement on both benchmarks.Source: §7 · p. 8
  • Its contributions include a robust TTT framework using leave-one-out in-context task construction, an optimization setup, and a post-TTT inference strategy.Source: §7 · p. 8
  • The results suggest potential for TTT to address novel reasoning tasks and for test-time methods to advance future language models.Source: §7 · p. 8

Limitations

The limitations section flags possible bias from ARC development-set hyperparameter tuning. The included appendix passages describe the development tasks and data presentation.

Limitations

  • ARC hyperparameters were optimized on the same 80 tasks used for validation and ablations, which might have introduced bias.Source: Limitations · p. 9

A. ARC Dataset

  • The appendix presents the ARC development tasks, data format, and evaluation details.Source: A. ARC Dataset · p. 12

A.1. Data Format

  • ARC experiments use NumPy array-printing format for data presentation.Source: §A.1 · p. 12

A.2. List of 80 Tasks Used For Development

A.3. Evaluation

The appendix describes ARC evaluation conventions and the construction of fine-tuning data from existing and generated tasks. It also specifies initial fine-tuning and grid-string formatting.

A.3. Evaluation

  • Evaluation counts a test as correct if either of the system’s two pass@2 predictions is correct; task-level accuracy generally gives no partial credit when not all tests are solved.The final table in Section 4.6 is an exception to the no-partial-credit rule.Source: §A.3 · p. 12
  • ARC grids are represented as strings of numpy arrays containing digits from 0 to 10, with each digit denoting a color.Source: Figure 10 · p. 12

B.1. Preparing Fine-tuning Data (Hodel, 2024)

  • LLM-based generation uses few-shot prompts with existing generators, optionally paired with task descriptions, or a two-stage process that generates a description before a new generator.The two-stage approach conditions generator creation on existing task pairs and the new description; the approaches yielded 6,426 generators.Source: Figure 11 · p. 13, §B.1 · pp. 13–14
  • Fine-tuning data draws on REARC generators, which sample new input-output pairs preserving each task’s transformation, and on rule-based and LM-generated task sources.Generated pairs can be verified by checking that the task’s transformation maps the input to the output.Source: §B.1 · p. 13
  • Task descriptions were generated from 10 manually written training-task seeds, using hierarchical fields for category, summary, and description.Source: §B.1 · p. 14

B.2. ARC Initial Fine-tuning Hyperparameters

  • Llama 3 family models were fully fine-tuned for up to 16,000 steps, with GPU counts varying by model size.The setup used 2 A100 GPUs for 1B models, 4 for 3B models, and 8 for 8B models.Source: §B.2 · p. 15, Table 4 · p. 15

B.3. Results

The fine-tuning experiments compare data-source and model-size variants, while the ARC TTT procedure augments task examples before fitting task-specific adapters. The 8B model has the highest reported fine-tuning accuracy, and TTT brings the 1B and 3B models to similar accuracy.

B.3. Results

  • The fine-tuning comparisons include a full-data condition, removal of geometric transformations, and exclusion of LM-generated tasks.The experiments also include an unfine-tuned model and use augmented data to fine-tune 1B, 3B, and 8B Llama models.Source: §B.3 · p. 15
  • 36% is the highest reported fine-tuning accuracy, achieved by the 8B model; after TTT, the 1B and 3B models achieve similar accuracy.Source: §B.3 · p. 15

C.1. Transformations

  • For ARC TTT, augmentations are applied to all input and output grids, followed by color and training-example shuffling.The TTT dataset is formed from leave-one-out tasks augmented through rule-based transformations, then used to train task-specific LoRA adapters on the base fine-tuned model.Source: §C.1 · p. 16, Figure 14 · p. 16

C.2. Training Setup & Hyperparameters

The appendix details task-specific LoRA training, LM-generated ARC data and descriptions, and transformed-input inference with hierarchical voting. These procedures specify how training examples and multiple predictions are constructed and combined.

C.2. Training Setup & Hyperparameters

  • TTT trains LoRA adapters on attention query and value projections, MLP weights, and the output projection; quantized adapters reduce 1B performance from 29 to 26 tasks solved.Greedy decoding is used, and inference applies rotations and horizontal, vertical, and diagonal flips.Source: §C.2 · p. 16

D. LM Data Generation

D.1. Getting Descriptions for Tasks

  • Ten manually written seed descriptions and crowd-worker annotations were used to prompt descriptions for other training tasks.The figure notes crowd-worker annotations are available only for training tasks.Source: §D.1 · p. 18, Figure 15 · p. 19

D.2. Few-shot Prompting Details

  • The data-generation prompts provide few-shot examples to generate new scripts, using GPT-4 and GPT-4o.The described prompting procedure generates scripts from examples of task descriptions paired with generator functions.Source: §D.2 · p. 18, Figure 15 · p. 19

E.1. Augmented Inference

  • Augmented inference transforms the demonstrations and test input, inverse-transforms each prediction, and also permutes demonstration order to create additional candidates.Two demonstration-order permutations are sampled per transformation to mitigate possible bias in processing their sequence.Source: §E.1 · p. 20

E.2. Ensembling Predictions (Voting Strategy)

  • Hierarchical voting first selects frequent predictions within each transformation, then globally votes for the top two candidates, prioritizing identity-transformation predictions in a tie.Within a transformation, row- and column-based majorities can supplement candidates when fewer than three unique predictions exist.Source: §E.2 · p. 20

F.1. Further Experimental Details

The BBH evaluation averages results across five random seeds and reports standard errors for the methods. The appendix points to task-specific results and groups tasks by their TTT-minus-ICL accuracy difference.

F.1. Further Experimental Details

  • BBH results average performance over five selections of 10 few-shot examples, using seeds 42 through 46.The full TTT and inference process takes approximately 15 minutes on an NVIDIA A100 GPU.Source: §F.1 · p. 21
  • Table 7 reports each method’s standard error of the mean for Figure 8.Source: Table 7 · p. 21, §F.1 · p. 21

F.2. Task-specific Results

Loading 2411.07279v2…