Context Tuning for In-Context Optimization

Jack Lu, Ryan Teehan, Zhenbang Yang, and Mengye Ren

New York University

The 43rd International Conference on Machine Learning (ICML 2026)

arXiv PDF Code

TL;DR: Context Tuning directly optimizes an LLM's memory representation for efficient adaptation without updating model weights.

Abstract

We introduce Context Tuning, a simple and effective method to significantly enhance few-shot adaptation of large language models (LLMs) without weight updates. In-Context Learning (ICL) forms a memory representation of the demonstrations in a single forward pass but cannot refine it when insufficient. Prompt-based methods offer lightweight adaptation by optimizing a trainable prompt or prefix but initialize it independently of the demonstrations. In contrast, Context Tuning leverages the model's inherent ICL ability to initialize a trainable memory representation from demonstrations, then refines it through gradient-based optimization. Extensive evaluations on benchmarks such as CrossFit, UnifiedQA, MMLU, BIG-Bench Hard, and ARC demonstrate that Context Tuning outperforms both ICL and traditional prompt-based adaptation methods while achieving competitive accuracy with Test-Time Training at significantly higher training efficiency.

Overview

CT-KV is our strongest variant of Context Tuning:

Compared with in-context learning (ICL), CT-KV refines the model's initial memory representation of the provided input-output examples instead of using it directly for prediction, substantially improving accuracy.
Compared with Test-Time Training (TTT), CT-KV achieves competitive accuracy without updating model weights and in half the training time or less.
Combined with TTT, CT-KV achieves the highest accuracy, showing that KV cache tuning and model weight updates are complementary.

Accuracy vs. training time, averaged over 26 NLP tasks.

Context Tuning for In-Context Optimization

CT-KV keeps the LLM frozen and turns the key-value (KV) cache formed from the provided examples into a trainable memory representation. During optimization, Leave-One-Out Masking asks the model to predict each output from the other examples, while Token Dropout improves generalization. At inference, the model conditions on the full optimized cache. Our paper also presents CT-Prompt, a prompt embedding variant.

CT-KV initializes a key-value prefix from the provided examples and optimizes it with Leave-One-Out Masking (left). At generation time, the model conditions on the full optimized prefix to answer a new query (right).

Experiments

We evaluate Context Tuning on NLP-LR, MMLU, BBH, and ARC. The experiments span pretrained LLMs from 1B to 32B parameters.

Representative test examples from BBH, NLP-LR, and MMLU, followed by three input-output examples and a test example from ARC.

Comparing Context Tuning to Baselines

CT-KV outperforms in-context learning (ICL), Prompt Tuning, Prefix Tuning, LoRA, rank-stabilized LoRA, and DoRA across all four benchmarks. It achieves competitive accuracy with TTT without updating model weights and in half the training time or less, while TTT+CT-KV achieves the best accuracy on every benchmark. On NLP-LR, CT-KV's single-task adaptation surpasses MetaICL's multi-task meta-training under matched samples (44.2% vs. 43.3%).

Accuracy and training time per task in seconds. Means and standard deviations are computed over five sets of examples, except ARC, which has a fixed set. Bold and underlined values mark the best and second-best accuracy for each benchmark.

Robustness to Example Count and Quality

(a) CT-KV remains ahead of ICL and Prefix Tuning as more examples are provided.

(b) CT-KV performs best on both benchmarks even when up to 75% of example labels are corrupted.

NLP-LR and MMLU accuracy versus (a) the number of examples and (b) the label corruption probability.

Scaling Up the Pretrained Models

Across five pretrained models ranging from 12B to 32B parameters and spanning multiple architectures, CT-KV outperforms ICL and Prefix Tuning.

BBH accuracy across pretrained models of increasing size.

Ablating Our Design Choices

Leave-One-Out Masking and Token Dropout both improve CT-KV on three of the four benchmarks.

Ablations of Leave-One-Out Masking and Token Dropout across four benchmarks. Means and standard deviations are computed over five sets of examples, except ARC, which has a fixed set.

Qualitative Results

We show how CT-KV predictions evolve during optimization on two ARC tasks. Iteration 0 is equivalent to ICL. Green labels indicate correct predictions, and red labels indicate incorrect predictions.

Color Mapping

At iteration 0, the model fills the interior of every bordered square with yellow. During optimization, it gradually discovers the correct fill color for each one.

Cross Completion

At iteration 0, the model already identifies that red should be used to complete the cross shapes, but does not understand that it should avoid overwriting black squares. By iteration 200, the prediction becomes more consistent with the provided examples and solves the task.

BibTeX

@inproceedings{lu2026contexttuning,
  title     = {Context Tuning for In-Context Optimization},
  author    = {Lu, Jack and Teehan, Ryan and Yang, Zhenbang and Ren, Mengye},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2026}
}