Multi-Task In-Context Learning
- Multi-Task In-Context Learning is a framework where multiple related tasks are used in-context to infer latent task identities without parameter updates.
- It employs latent-task inference, distributed task representations, and mixed curriculum strategies to enhance adaptability and overcome catastrophic forgetting.
- Empirical results demonstrate that careful prompt construction and curriculum design lead to improved task retrieval and performance across diverse domains.
Multi-task in-context learning denotes a family of settings in which in-context adaptation is shaped by multiple related tasks rather than a single homogeneous task. In recent work, this has meant training a transformer to perform in-context learning over a mixture of related latent task families, conditioning prediction on several prior-task datasets and one target dataset, configuring multimodal demonstration sequences in mixed-task regimes, or learning from repeated interaction histories across task families without parameter updates at test time (Bhasin et al., 2024, Zhu et al., 18 Jun 2026, Li et al., 21 May 2025, Lin et al., 3 Feb 2026). Across these variants, the common question is how context should represent task identity, shared structure, and task-specific evidence when the model must reuse or infer more than one task distribution.
1. Latent-task inference and the meaning of “multi-task” in context
A recurring formalization treats pretraining or meta-training as exposure to a distribution over tasks rather than a single task. In the PAC framework of latent-task pretraining, the pretraining distribution is a mixture
and a downstream prompt acts as evidence for which latent task is active rather than as a small supervised dataset from which a completely new rule is learned (Wies et al., 2023). This yields the paper’s central claim that, in that stylized regime, in-context learning is “more about identifying the task than about learning it.”
A related Bayesian account replaces a single prior with a Gaussian mixture over task groups,
and interprets the in-context predictor as the MMSE estimator under posterior inference over tasks (Lin et al., 2024). In that formulation, multi-task structure is explicit in both the pretraining prior and the inference procedure: context can reweight latent task groups and also shift the posterior within each group. This supports the paper’s distinction between task retrieval, in which context mainly selects a pretrained skill, and task learning, in which context mainly updates the operative task itself. The same work gives a theoretical explanation for the “early ascent” phenomenon: with a small number of examples, the posterior may initially retrieve an incorrect skill, raising risk, before later examples drive true task learning (Lin et al., 2024).
These perspectives suggest a broad taxonomy. In one regime, “multi-task” refers to latent task mixtures in pretraining. In another, it refers to multiple explicit tasks or datasets inside one context window. In a third, it refers to structured sequences of related episodes in which the same hidden task persists across episodes. A plausible implication is that the field’s central design problem is not merely how to concatenate more demonstrations, but how to make the model infer which shared structure should be reused and which task-specific variation should be updated.
2. Internal task representations: from single task vectors to distributed task codes
One influential mechanistic hypothesis is that ordinary in-context learning can often be approximated by
where the support set is compressed into a single latent task vector and later layers apply that representation to the query (Hendel et al., 2023). In the reported experiments, the patching-based “hypothesis” procedure retains about 80–90% of regular ICL accuracy, while the no-demonstration baseline is only around 10–20% on average; task vectors also cluster by task, and conflicting-task interventions show that an injected task vector can redirect the model away from the task specified by the visible demonstrations (Hendel et al., 2023). This establishes a compact single-task bottleneck as a useful starting point for multi-task analysis.
Subsequent large-scale evidence argues that this bottleneck is often insufficient. On QuiteAFew, a benchmark of 3,096 few-shot tasks with 30 unique input-output pairs per task, task-vector performance peaks around the 15th layer of Llama-3-8B, varies substantially by task type, and breaks down on complex tasks that require multiple sub-operations (Tikhonov et al., 29 May 2025). In the paper’s synthetic automobile-description-to-JSON task, different output tokens are recovered by injecting different hidden states from the few-shot example, leading to the explicit conclusion that one must speak not about one task vector for the entire task, but about “many task vectors for the task” (Tikhonov et al., 29 May 2025). This directly weakens any universal “one prompt, one vector” account of multi-task in-context learning.
A complementary formulation, Learnable Task Vectors, defines a layerwise task representation as
a weighted sum of attention-head outputs within each layer, with one learned scalar per head and frozen transformer weights (Saglam et al., 8 Feb 2025). The method is trained causally by impairing normal ICL and optimizing only the head weights; it is evaluated on both language tasks and function-regression tasks, and its ablations tie success to alignment of the last hidden-state distribution with that of an optimally performing ICL model (Saglam et al., 8 Feb 2025). On language tasks, it records the highest accuracy in 29/30 evaluations and is statistically superior in 19/30. Taken together with the single-vector and multi-vector results, this suggests that task representations are real, causally important, and often distributed across layers and heads rather than localized to one delimiter token.
Multimodal work reaches a similar conclusion from a different angle. Multimodal Task Vectors compress many-shot multimodal demonstrations into mean last-token activations at selected attention heads and reinsert them at inference time without finetuning (Huang et al., 2024). On ViLA-1.5-8B, this yields 55.2 on VizWiz and 40.6 on OK-VQA, exceeding 8-shot multimodal ICL, and transfer experiments show that head locations extracted on VizWiz can be reused on OK-VQA (Huang et al., 2024). This does not constitute full multi-task prompting, but it does show that compact task-structured internal states can be reused across related multimodal tasks.
3. Training multi-task in-context learners: curriculum, rehearsal, and robustness
The most controlled study of multi-task in-context training uses synthetic function classes of the form
with prompts consisting of interleaved examples and a query, and training minimizing average squared prediction error over context positions (Bhasin et al., 2024). Instead of natural-language tasks, the model is exposed to related linear, quadratic, and cubic Hermite-based function classes, allowing precise comparison of single-task training, random multi-task mixing, pure sequential easy-to-hard training, and a mixed curriculum that retains earlier tasks when harder tasks are introduced.
In that study, the transformer is a GPT-2 architecture from HuggingFace with 22.4M parameters, 12 heads, 8 layers, and embedding size 256, trained for 500,000 steps with batch size 64 and 100 pairs per batch (Bhasin et al., 2024). The primary metric is normalized MSE. The central empirical result is that mixed curriculum is the only curriculum that reliably learns all three tasks, whereas pure sequential curriculum performs poorly and random mixing, although better, still struggles on the quadratic class. The clearest convergence result is that 0% of the single-task quadratic models converged, whereas 60% of the mixed curriculum models converged. Supplementary plots indicate that mixed curriculum reaches accurate solutions after roughly 20/80/90-shot examples for linear/quadratic/cubic tasks, respectively, and a mixed-curriculum cubic model initialized from linear and quadratic training begins around 200 normalized MSE instead of 450, while having seen only 1/9 of the training examples of the single-task cubic model (Bhasin et al., 2024).
The same work frames these gains as a combination of optimization warm-up and rehearsal. Easier tasks teach the general structure of using context to infer a latent mapping, while retaining earlier tasks reduces catastrophic forgetting (Bhasin et al., 2024). The negative result is equally important: progression alone is not enough. The sequential curriculum, although easy-to-hard in the textbook sense, often performs worse than random mixing. The paper also reports that instruction prompting does not help in this synthetic setting: one-hot instruction vectors give little benefit, while preset random instruction vectors hurt, especially on quadratic and cubic tasks (Bhasin et al., 2024).
Mechanistic analysis reinforces the notion of shared circuits. The paper isolates attention from each token to its corresponding 0 token and labels heads with strong such attention retrospective heads. The same retrospective heads appear across task evaluations in curriculum-trained models, and masking 7 retrospective heads sharply increases normalized MSE, whereas masking 7 non-retrospective heads has much less effect (Bhasin et al., 2024). This suggests that multi-task exposure can stabilize reusable in-context machinery rather than merely produce separate task-specific solutions.
4. Prompt construction, multimodal task mapping, and compositional contexts
In multimodal settings, multi-task ICL depends strongly on how demonstrations are selected, ordered, and related to the query. TACO formalizes this with the notion of task mapping, distinguishing local mappings
1
from a global mapping
2
and defining prompt interpretation as composition
3
(Li et al., 21 May 2025). This framing is explicitly tested in mixed-task multimodal ICL through a Hybrid dataset formed by sampling 5,000 training instances from each of six datasets spanning VQA, captioning, and classification. TACO itself is a 140M-parameter sequence configurator with four transformer decoder blocks, trained to imitate Oracle sequence configurations and equipped with task-aware attention and a learned task guider (Li et al., 21 May 2025).
Across five LVLMs and nine datasets, TACO beats all baselines on all nine datasets. The paper highlights an average 3.26% improvement on VQA tasks and a 3.61% gain on Hybrid, and ablations show that removing the 4 token, disabling task-guider updates, or omitting task-mapping components consistently degrades performance, especially on the mixed-task Hybrid benchmark (Li et al., 21 May 2025). The conceptual claim is not simply that similar examples should be retrieved, but that demonstrations must induce a coherent latent reasoning structure. This is one of the clearest prompt-centric accounts of multi-task multimodal ICL.
A different line of work argues that cross-task examples are often most useful as scaffolding rather than as the final prompt. In the graph-based pseudo-labeling framework “From Cross-Task Examples to In-Task Prompts,” source-task examples are first used to pseudo-label a small target subset via cross-task ICL, and the remaining target examples are labeled by GLIP, a graph-based label propagation method that avoids additional LLM queries (Chen et al., 28 Oct 2025). The final ICL prompt is then constructed from the fully pseudo-labeled target-task dataset. The paper explicitly reports that cross-task methods can help at 1-shot or 4-shot, but can degrade as more cross-task examples are added; by contrast, in-task methods improve as the number of demonstrations grows (Chen et al., 28 Oct 2025). This supports a narrower interpretation of multi-task prompting: cross-task examples may be useful for bootstrapping target-task demonstration pools even when they are not ideal as permanent prompt contents.
Compositional visual ICL sharpens the same issue. In medical imaging, a single context element can be a full task sequence
5
and the model must generate a sequence of intermediate outputs
6
rather than a single target image (Reiß et al., 1 Jul 2025). The paper introduces a synthetic compositional task generation engine that builds valid pipelines from generative, geometric, and discriminative tasks, rather than random task bags. It also shows that codebooks are a hard bottleneck: off-the-shelf VQ-GAN codebooks reconstruct ordinary images well but degrade sharply on many-class segmentation maps, while task-aware codebook training and color-remapping augmentation materially improve representability (Reiß et al., 1 Jul 2025). Among masking objectives, sequence-token masking, which masks the entire target output sequence, performs best and reaches 61.21% segmentation IoU after 52 epochs in the reported setup (Reiß et al., 1 Jul 2025). The broader lesson is that multi-task ICL does not automatically yield compositional ICL; the episode distribution must preserve task order, intermediate dependencies, and output semantics.
5. Beyond static prompting: Bayesian prefixes, decision-making, online adaptation, and wireless systems
One generalization of multi-task ICL makes the prior itself an in-context object. In Multi-Task Bayesian In-Context Learning, the model receives several prior-task datasets
7
as a prefix
8
followed by a target-task context and query (Zhu et al., 18 Jun 2026). The target predictive distribution is then conditioned jointly on target data and the prior prefix, so that changing the prefix has the operational effect of changing the prior at test time. Across Gaussian, Student-9, and flow-based priors, the transformer matches hierarchical Bayesian references while being orders of magnitude faster, and on the ERA5 temperature benchmark the MT 0 model improves test NLL/MSE from 1 to 2 relative to MT 3 under the IID split (Zhu et al., 18 Jun 2026). This is a particularly explicit instance of multi-task ICL as hierarchical inference over related datasets.
In sequential decision-making, the same idea appears as learning from histories rather than static demonstrations. For structured bandits, a GPT-2-based decision transformer is pretrained to predict a reward vector over all actions from trajectories
4
and then chooses actions greedily or by softmax exploration without parameter updates at test time (Mukherjee et al., 2024). The key design shift is reward prediction rather than action imitation, which removes the need for privileged optimal-action labels and permits the model to outperform the demonstrator across linear, nonlinear, bilinear, latent, and changing-action structured bandits (Mukherjee et al., 2024). Here “multi-task” means a family of related bandit tasks sharing latent structure, and “in-context” means online adaptation from a growing interaction history.
ORBIT makes the same move in multi-episode online RL. It trains an LLM as a history-dependent policy
5
where the history includes all prior episodes on the same underlying task instance, and optimizes a cross-episode reward with GRPO over a task distribution of partially observable environments (Lin et al., 3 Feb 2026). Training uses five environments—RPS, Minesweeper, Hangman, Wordle, and Blackjack—and evaluation is on unseen Maze and Mastermind. The paper reports Episode 3 success on Maze of 0.55 for ORBIT versus 0.45 for the RL baseline, and on Mastermind of 0.59 versus 0.21; it also reports stronger conditional exploration after failure, with 6States (Ep27F1) of 4.69 for ORBIT versus 4.11 for RL and 1.64 for the base model (Lin et al., 3 Feb 2026). The claimed result is that a relatively small Qwen3-14B after ORBIT training matches GPT-5.2 on the paper’s unseen online-learning benchmarks.
Wireless physical-layer learning supplies a domain-specific but technically precise example. ICWLM is trained directly on mixed wireless datasets to solve three tasks—weighted sum-rate precoding, max-min SINR precoding, and channel prediction—inside one decoder-only transformer (Wen et al., 24 Jul 2025). Prompts are sequences of wireless-domain demonstration pairs,
8
with task identity inferred from the demonstrations rather than a separate task head (Wen et al., 24 Jul 2025). The implemented model has 8 transformer layers, 8 attention heads, and embedding dimension 1024; the reported results show that, without demonstration pairs, the model struggles to discern the intended task, while a small number of demonstrations yields performance comparable to WMMSE on sum-rate maximization and robust generalization to unseen SNR values and varying user velocities (Wen et al., 24 Jul 2025). This is a clean example of one shared architecture using prompt content to switch among several structured regression problems.
6. Limits, misconceptions, and open problems
A central misconception is that in-context learning is always online task induction from scratch. Latent-task PAC theory and Bayesian mixture theory both argue otherwise: in important regimes, a prompt primarily identifies a task or skill already encoded in pretraining, and only later examples drive substantial within-task adaptation (Wies et al., 2023, Lin et al., 2024). Another misconception is that one compact task code must suffice. The single-task-vector hypothesis is powerful, but large-scale and compositional evidence indicates that complex tasks often require multiple subtask-specific vectors or layerwise control signals rather than one monolithic embedding (Tikhonov et al., 29 May 2025, Saglam et al., 8 Feb 2025).
A second misconception is that easy-to-hard sequencing alone is enough for multi-task ICL. In the controlled function-class study, pure sequential curriculum often underperforms random mixing, while mixed curriculum with rehearsal materially improves convergence, data efficiency, and OOD robustness (Bhasin et al., 2024). The same paper shows that instruction tokens are not uniformly beneficial; in that synthetic continuous setting, one-hot vectors provide little benefit and preset random vectors hurt (Bhasin et al., 2024). This suggests that task identity cues can act as either structure or noise depending on how they interact with the underlying task family.
The empirical scope of much of the literature remains narrow. The curriculum study is largely confined to noiseless synthetic function classes and small transformers, and the authors explicitly state that extensibility to natural language is an open question (Bhasin et al., 2024). TACO’s mixed-task evidence is strongest on the Hybrid dataset and broad multi-benchmark generality, not on a dedicated same-task versus cross-task transfer protocol (Li et al., 21 May 2025). The graph-based pseudo-labeling framework assumes one suitable source task and leaves multi-source or mixed-task scenarios for future work (Chen et al., 28 Oct 2025). Multi-task Bayesian ICL shows strong behavior under several synthetic prior families, but also reports that real-world OOD seasonal shift in ERA5 can make auxiliary prior tasks misleading, with Set-MT more robust than the plain decoder-only model under the severe 2019 OOD split (Zhu et al., 18 Jun 2026). In visual compositional ICL, codebook granularity, long-sequence degradation, and semantic hallucination remain major bottlenecks (Reiß et al., 1 Jul 2025).
The open problems follow directly from these tensions. The field still lacks a general account of how to represent multiple simultaneously relevant tasks in one context, how to route or compose internal task representations, how to separate shared structure from task-specific semantics, and how to scale from controlled latent-task mixtures to heterogeneous naturalistic task ecologies. Current evidence supports several partial conclusions: multi-task exposure can stabilize in-context algorithms; prompt construction can matter as much as model architecture; task representations are real but often distributed; and cross-task information is often most useful when it induces a coherent prior, a coherent task mapping, or a coherent target-task demonstration pool rather than when it is naively mixed into the final prompt.