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

At a certain scale, many LMs exhibit the ability to adapt to new tasks without updating their parameters by simply conditioning on input examples or instructions provided. Given a sequence of input-output pairs $(x_{1},y_{1}),\ldots,(x_{k},y_{k})$ and a new input $x_{k+1}$, an LM can generate the corresponding output $\hat{y}_{k+1}$ by sampling from:

While the possibility of in-context learning (ICL) as implicit machine learning simulation is discussed in previous work (Akyürek et al., 2023), empirical evidence shows that in-context learning with language models does not always resemble standard machine learning algorithms (Zhao et al., 2024; Min et al., 2022b). Furthermore, ICL often struggles with novel tasks “out-of-the-box.” For example, large language models exhibit poor performance on datasets like ARC (Opiełka et al., 2024; Bober-Irizar & Banerjee, 2024).

Test-time training (TTT) enables parametric models to adapt during inference through dynamic parameter updates in response to each test input. This approach remains relatively unexplored in the era of large language models. The general TTT process is as follows: starting with initial model parameters ${\bm{\theta}}_{0}$, for each test input (or batch of inputs) $d$, we generate a temporary training dataset $\mathcal{D}_{\mathrm{TTT}}$. We then optimize these parameters to minimize a loss function

resulting in temporarily updated parameters ${\bm{\theta}}_{d}$, which are subsequently used for prediction.$^{2}$

In previous work (Sun et al., 2020, e.g.,), $\mathcal{D}_{\mathrm{TTT}}$ is typically constructed by applying an unsupervised objective (e.g., masked autoencoding) to the input $x$ alone. In this paper, we extend TTT to the few-shot learning setting, treating it as a form of transductive learning by leveraging few-shot demonstration examples to improve predictions. Although TTT can also be applied to chain of thought (Wei et al., 2022, CoT;), we focus on direct transduction, where demonstrations consist of input-output pairs $(x,y)$ without intermediate reasoning steps or explicit function descriptions.

The few-shot learning setting we consider provides richer context in the form of demonstration pairs $(x_{1},y_{1}),\ldots,(x_{K},y_{K})$. One simple method for TTT is Direct I/O training, where we directly treat each input-output $(x_{k},y_{k})$ pair as training instances. Our key insight is that the few-shot examples can also be used to construct a more robust and expansive $\mathcal{D}_{\mathrm{TTT}}$ of synthetic in-context learning tasks, allowing for effective model adaptation during test time. Additionally, when task-specific knowledge is available, this structure can be leveraged to further expand the dataset, as demonstrated in our experiments on ARC (Section 4). We also explore the general case where no task-specific information is used, as tested on BBH (Section 5).

Our experiments in this paper characterize each component of the TTT pipeline, investigating different design choices across the following stages: (1) constructing an input-specific training dataset $\mathcal{D}_{\mathrm{TTT}}$ at test-time; (2) fine-tuning the LM by optimizing a loss function $\mathcal{L}$ over the dataset $\sum_{d\in\mathcal{D}_{\mathrm{TTT}}}\mathcal{L}(\mathrm{LM}(d;{\bm{\theta}}))$; and (3) sampling from the updated model with an augmented inference strategy based on self-consistency to obtain a final prediction.

Given a task with $K$ training input-output pairs $\{(x_{k},y_{k})\}_{k=1}^{K}$, a test-time training dataset $\mathcal{D}_{\textrm{TTT}}$ can be created by either following an in-context learning setup or a direct input-output (direct I/O) setup (top row in Figure 2):

We optimize the standard LM loss on $\mathcal{D}_{\textrm{TTT}}$. For the in-context leave-one-out setup, we experiment with $3$ different ways to take the loss (middle row in Figure 2):

We find in Sections 5.3 and 4.3 that the first method (taking the loss over both demonstration and test outputs) works best.

Once we have the test-time training dataset $\mathcal{D}_{\textrm{TTT}}$ (constructed via either the in-context or direct I/O approach), we perform a small number of gradient steps on task-specific LoRA adapters (Hu et al., 2022). This approach allows computationally efficient adaptation while maintaining the model’s general capabilities. By default, we learn task-specific LoRA adapters for each ARC or BBH task at test-time. That is, we obtain $K$ different LoRA adapters, where $K$ is the number of test tasks. We also experiment with using a single shared LoRA adapter from the aggregated dataset of few-shot examples drawn from multiple tasks (bottom row in Figure 2)—a test-time version of meta-ICL (Min et al., 2022a). We find that the shared adapter degrades performance on ARC, whereas it improves performance on BBH. We discuss this in more detail in Section 5.3.

The Abstraction and Reasoning Corpus (ARC) aims to evaluate the abstract reasoning capabilities of language models through their ability to solve visual puzzles. Each puzzle (henceforth referred to as a task) consists of input-output pairs of 2D grids (up to $30\times 30$ in size) containing shapes or patterns in up to $10$ different colors, as displayed in Figure 3. The output of each pair is obtained by applying an intuitive and shared transformation or rule $y=f(x)$. Each task has 2-7 demonstration examples and 1-3 test examples.

In this section, we compare the final implementation of our method with different design choices for TTT. FT serves as the baseline, using only the fine-tuned model with demonstrations in-context. No Transformations omits the augmentation step. Direct I/O Data replaces in-context tasks with the direct input-output task formulation (Section 3.1). Shared TTT uses a single LoRA adapter across all tasks instead of learning one per task. No Demonstration Loss removes the loss on demonstration outputs (Section 3.1).

Results are presented in Figure 5. Our TTT method is effective, improving fine-tuned model accuracy approximately 6$\times$ $\bf(5\%\to 29\%)$. In-context formatting is especially important; using the direct input-output data to construct $\mathcal{D}_{\textrm{TTT}}$ causes an $11$-task drop ($38\%$). Removing transformations causes a $16$ task drop ($55\%$). Regarding optimization, per-task LoRA adapters outperform a single shared adapter by $7$ tasks ($24\%$). Including losses on the demonstration outputs yields a modest but consistent gain ($26\%\to 29\%$).

We perform full fine-tuning of 1B and 3B Llama 3.2 (instruction-tuned) and 8B Llama 3 (instruction-tuned) using synthetically generated data, as detailed in Appendix B, and then use our default TTT implementation. We show results using different model sizes in Figure 6. Increasing the model size consistently improves FT performance, with the 8B model achieving the highest accuracy of $36\%$. At all model sizes, TTT leads to significant improvements in performance. We also observe that for smaller model sizes, TTT effectively closes the performance gap, with the 1B and 3B models achieving similar accuracy after TTT.

To analyze the impact of augmented inference and voting, we run several ablations: (1) Vanilla, which generates two predictions without transformations or advanced voting; (2) Transformed Inference, applying a single transformation (Rotate, Transpose, or Flip) to measure its isolated effect; (3) Hierarchical Voting, our full pipeline combining augmented inference and structured voting; (4) Flattened Voting, which selects the top-2 predictions from a single voting round over all generated outputs; and (5) Oracle, an upper bound that selects the correct answer if present.

As shown, individual transformations are modestly effective on their own (with Transpose performing worst), but their aggregation improves results markedly. Hierarchical voting further outperforms a flattened voting approach and closely approaches the oracle’s accuracy, suggesting that our two-stage aggregation effectively identifies the correct solution when it is present.

Following our experiments on $80$ tasks, we present comprehensive results on the full ARC public evaluation set, comparing our system against existing approaches. Our analysis focuses on three key aspects: the impact of our TTT methodology, the benefits of combining our approach with existing methods, and the differences between fully neural and program synthesis methods.

BIG-Bench Hard (Srivastava et al., 2023; Suzgun et al., 2023, BBH;) is a benchmark comprising $27$ challenging tasks across $23$ task types, designed to evaluate large language models on reasoning, compositionality, and generalization. Unlike ARC, BBH features a broader natural language structure and lacks a shared input format, making it unsuitable for invertible transformations. However, this broader scope offers a valuable testbed for evaluating TTT’s effectiveness in a more generalized setting. Despite the absence of invertible transformations—previously used in ARC to expand the TTT dataset and enhance inference—TTT still significantly improves performance on BBH.

In this section, we evaluate our method and its ablations, primarily comparing the zero-shot baseline, ICL, and TTT. No Example Permutation updates the model on a single in-context prompt instead of multiple shuffled versions. Direct I/O treats each input-output pair as separate training instances. Shared TTT uses a single adapter across tasks instead of task-specific adapters. No Demonstration Loss removes the loss applied to demonstration outputs. Loss on Inputs and Outputs extends the loss calculation to both inputs and outputs. These ablations are as detailed in Section 3. As these results are averages over $5$ runs, the standard errors of the mean for each method are given in Section F.1, averaging $0.4\%$.

The results in Figure 8 show that TTT achieves an overall accuracy of $57.8$%, outperforming standard ICL ($50.5$%) and Direct I/O learning ($51.5$%). This demonstrates that TTT’s capabilities extend beyond ARC to more diverse and complex reasoning tasks, proving its effectiveness in a broader range of natural language problem-solving scenarios.

We observe that TTT without example permutations—performing multiple gradient steps on a single in-context prompt before inference—reduces accuracy to a still-impressive $55.7\%$. Computing the loss only on the test output lowers accuracy to $54.4\%$, while applying it to both inputs and outputs achieves $55.9\%.$

Our task-specific results show that performance improvements from TTT are highly task-dependent. Among the $27$ tasks in BBH, TTT results in a performance decline of at least $2\%$ compared to ICL in only $2$ tasks. In contrast, $12$ tasks show an improvement of at least $2\%$, with $9$ of these showing improvements of at least $5\%$. The four tasks with the most significant performance boost from TTT over ICL or zero-shot and the task with the most significant performance decrease are shown in Figure 9. These tasks in order of TTT’s improvement over ICL are Dyck Languages (parentheses matching), Ruin Names (humorous name modifications), Movie Recommendation (choosing similar films), Hyperbaton (adjective ordering), and Boolean Expression (evaluating a boolean expression). Detailed results for every task are given in Section F.2.

We hypothesize that improvements from TTT may be driven by tasks involving distribution shifts and structured patterns. For example, tasks like Dyck Languages and Hyperbaton follow clear grammatical or programmatic rules, which could align well with TTT’s ability to adapt to latent structural regularities during test-time.

Conversely, tasks requiring explicit step-by-step computation show limited gains with TTT. For instance, Boolean Expressions declined from $85.7\%$ to $80.4\%$ under TTT. This task’s algorithmic nature—dependent on sequential reasoning rather than pattern-based transduction—and its likely pre-training exposure suggest TTT’s updates may not resolve its specific demands. While these particular observations align with our hypothesis, the reason certain tasks benefit more from TTT remains an open question.

The idea of updating model parameters at test-time using instance-specific data traces back to early work on local learning (Bottou & Vapnik, 1992). More recently, Sun et al. (2020) propose a simple test-time self-supervision scheme to adapt an image classifier when facing distribution shifts. In language modeling, Hardt & Sun (2024) fine-tune on retrieved neighbors at test-time for notable gains, while Hübotter et al. (2025) optimize retrieval via active data selection.

Abstraction and Reasoning Corpus (Chollet, 2019; Chollet et al., 2025, ARC;) is a collection of extremely challenging few-shot visual reasoning problems. Most approaches to ARC fall into two main categories: program synthesis and fully neural. Program synthesis approaches (Butt et al., 2024; Wang et al., 2024; Li et al., 2025; Greenblatt, 2024) first try to find the transformation function $f$, and then apply it to the test example. Fully neural approaches (Veldkamp et al., 2023; Bober-Irizar & Banerjee, 2024) try to directly predict the output $y^{\textrm{test}}$, only implicitly modeling $f$. In this work, we use a fully neural approach, using an LM to predict the test outputs. Recent work has explored hybrid methods, leveraging inference scaling and deep learning-guided program synthesis (Greenblatt, 2024; Li et al., 2025). Similarly, we find that integrating our neural model with program synthesis improves performance.

In development of ARC, we used a set of 80 tasks for validation/ablation experiments. Standard hyper-parameters (learning rate, epochs) were optimized using this set, which might have introduced some bias.

While the base Llama-3 performs poorly on the public validation set of ARC, the public availability of the dataset introduces the possibility that these models may have seen these examples during pre-training. Similarly, while the base model achieves reasonable performance on BBH, its public availability raises similar concerns.

We use numpy’s array printing format for all experiments as shown in Figure 10.

We use the following (Table 2) tasks validation tasks for our development.

We follow the competition rules that if any of the two pass@2 predictions of the system is correct, we consider that test correct. In the reported task-level accuracies, we did not give partial points if all tests are not solved, except the final table Table 1.

(Hodel, 2024) provides domain-specific language (DSL), ReARC, as well as the transformation $f_{i}$ that solves the task-$i$, and the data generation function $g_{i}$ that are implemented in this DSL for each training task in the $\mathcal{D_{\textrm{ARC}}^{\textrm{train}}}$ dataset. These functions enable sampling of new input-output pairs that maintains the same underlying transformation principle:

where $d$ represents a newly generated input-output pair that can be solved using the same transformation function $f_{i}$ as the original task-$i$ $^{4}$.

We perform full fine-tuning on LLama-3 family models by using the torchtune library. We train each model up to 16000 steps. We use 2xNVIDIA A100 GPU for 1B models, 4xNVIDIA A100 GPU for 3B and 8B models. We present hyperparameters in Table 4.

We perform full fine-tuning 1B, 3B Llama 3.2 instruction-tuned, and 8B Llama 3 instruction-tuned using augmented data. The format and training objective is same as the ones described for TTT in 3. Hyperparameter details are given in Section C.2. We do the following ablations for augmented data:

We show results using different model sizes in Figure 13. Increasing the model size consistently improves FT performance, with the 8B model achieving the highest accuracy of $36\%$. We also observe that TTT effectively closes the performance gap for smaller models, with the 1B and 3B models achieving similar accuracy after TTT.

We provide the augmentations used in TTT in Table 5, please refer to our code base for their implementations. After applying these augmentations, we additionally shuffle colors and shuffle training examples. Note that these transformations are applied to all input and output grids. The procedure for generating the dataset for TTT is shown in Figure 14.

We use the torchtune(torchtune Maintainers & Contributors, 2024) library to train LoRA adapters on Llama-3 family of models. We apply LoRA training to query and value projection weights of the self-attention layer, to the MLP weights and to the output projection layer (was only available for Llama-3 8B in torchtune). We present hyperparameters of this training in Table 6. We also found that using quantized LoRA adapters (Dettmers et al., 2023) instead of standard (full-precision) LoRA leads to only a small drop in performance ($29\to 26$ tasks solved with the 1B-parameter model), making it a viable option in memory-constrained settings.

We resort to the vLLM (Kwon et al., 2023) library for prediction as it provides fast kernels and batched inference for our models and LoRA inference. We just use greed decoding as we did not see improvements with temperature sampling in our early experiments. We use 90, 180 degree rotations, horizontal, vertical, and diagonal (transpose) flips as our invertible transformations.

With that, the whole TTT and inference process takes approximately $12$ hours for $100$ randomly sampled validation tasks when using an NVIDIA A100 GPU.

This procedure is shown in Figure 15. We initially described 10 training tasks with the hierarchical-style shown in Figure 11. Then, for other training tasks tasks, we obtained less quality crowd-worker annotations from LARC (Acquaviva et al., 2022) project. By using our high-quality seed annotations and their LARC version, we $10$-shot prompt and LM to produce high quality annotations for the other training tasks.

We use the following simple prompting template with k-shot prompting for all data generation procedures, where numbers filled with examples sampled from seed set. In simple few-shot generation, we exclude examples. We use GPT-4 and GPT-4o to generate the new scripts.

Recent work has shown that scaling test-time compute can significantly improve the performance of LMs. One of the most common techniques to do this is by sampling multiple responses, and then selecting the best response using a ranker. However, while sampling is very effective in domains with multiple possible solutions (programs in code) or multiple possible paths to the final answer (math), it can be detrimental when generating answers directly, as there is no way to directly enforce diversity across samples while ensuring coherence within samples. As an alternative inference-time scaling, we use an augmented inference strategy that generates multiple prediction candidates by using geometric transformations, combined with a greedy decoding scheme.

For a given task with training examples ${(x_{k},y_{k})}_{k=1}^{K}$ and test input $x_{\textrm{test}}$, we use invertible geometric transformations to produce equivalent transformed versions of the task, as shown in Figure 5. Let $\mathcal{T}$ be some set set of invertible geometric transformations (e.g., rotations and reflections). For each transformation $t\in\mathcal{T}$, we apply $t$ to all training demonstrations and the test input and run our model with these transformed inputs. We then apply the inverse transformation to obtain the final prediction for that transformation.

We further augment our predictions by permuting the order of training examples. For each transformation $g$, we sample $n=2$ different permutations of the demonstration sequence, resulting in $n\cdot|\mathcal{T}|$ total predictions per task. This is to mitigate any bias in the model’s processing of the demonstration sequence. (Bober-Irizar & Banerjee, 2024) also find transpose and rotation is helpful to produce extra prediction candidates.

We employ a hierarchical voting strategy to determine the final prediction from the set of candidates $\{y\}_{i=1}^{n\cdot|\mathcal{T}|}$. This approach involves two stages of voting to progressively narrow down the best candidates: first, by selecting the most frequent predictions within each transformation, and then by conducting an overall vote across transformation-specific candidates to identify the top-2 most frequent predictions. The details of each stage are as follows:

We write our own evaluation function for BIG-Bench Hard available in our codebase. We found that existing evaluation frameworks did not properly measure zero-shot performance due to insufficient answer-extraction parsing and answer-format prompting. We also wanted more control in splitting each individual task’s dataset into demonstration examples and evaluation sets. For all results, we average results over different selections of the $10$ few-shot examples with the following $5$ random seeds: $42,43,44,45,46$. The full TTT and inference process takes approximately 15 minutes on an NVIDIA A100 GPU.

The standard error of the mean for each method in Figure 8 over the $5$ seeds is given in Table 7.

We search over the following hyperparameters:

We similarly use the torchtune(torchtune Maintainers & Contributors, 2024) library for test-time training and the vLLM (Kwon et al., 2023) library for inference.

The full results for all tasks over all methods and ablations are shown in Figure 17.