Everything Everywhere All at Once: LLMs can In-Context Learn Multiple Tasks in Superposition

Authors: Zheyang Xiong, Ziyang Cai, John Cooper, Albert Ge, Vasilis Papageorgiou, Zack Sifakis, Angeliki Giannou, Ziqian Lin, Liu Yang, Saurabh Agarwal, Grigorios Chrysos, Samet Oymak, Kangwook Lee, Dimitris Papailiopoulos
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Model, In-Context Learning, Task Superposition

Abstract:
Large Language Models (LLMs) have demonstrated remarkable in-context learning (ICL) capabilities. In this study, we explore a surprising phenomenon related to ICL: LLMs can perform multiple, computationally distinct ICL tasks simultaneously, during a single inference call, a capability we term "task superposition". We provide empirical evidence of this phenomenon across various LLM families and scales and show that this phenomenon emerges even if we train the model to in-context learn one task at a time. We offer theoretical explanations that this capability is well within the expressive power of transformers. We also explore how LLMs internally compose task vectors during superposition. Furthermore, we show that larger models can solve more ICL tasks in parallel, and better calibrate their output distribution. Our findings offer insights into the latent capabilities of LLMs, further substantiate the perspective of "LLMs as superposition of simulators", and raise questions about the mechanisms enabling simultaneous task execution.

Decision: reject
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730037552193

thread_type:
official_review

<thread_object>
creator: reviewer_q86s

<content>
summary:
This paper studies the "task superposition" phenomenon in ICL, i.e., 
LLMs can perform multiple different ICL tasks simultaneously. 
Empirically, it is observed that the phenomenon is ubiquitous across different LLMs and larger models can handle more tasks simultaneously 
and better represent the distribution of in-context tasks.
Theoretically, the authors show Transformers are expressive enough for task superposition.
They also relate task superposition to internal combinations of task vectors.

strengths:
1. The paper introduces the interesting "task superposition" phenomenon, which provides insights on the ICL mechanism of LLMs.

2. The paper reveals an impressive correspondence between task superposition and internal combinations of task vectors.

3. The paper theoretically proves Transformers are expressive enough for task superposition.

weaknesses:
1. The paper only illustrates the superposition phenomenon in the synthetic tasks. It seems unclear whether real-world tasks have the superposition phenomenon
or how the superposition phenomenon affects LLMs in practice.

2. The theoretical explanation only considers the expressive power of Transformers. Since it has been proved that Transformers are Turing-complete, 
the theoretical result does not provide sufficient insights.

questions:
1. Theorem 1 states a quantitative relation between the embedding dimension, the numbers of heads, tasks, in-context examples, and example length.
Does this quantitative relation reflect how the model hyperparameters influence the task superposition capability in practice? For example, 
to perform $K$ tasks, does the model need at least $K$ heads? 

2. Could it be possible to prove the learnability of a Transformer that has the task superposition capability? Some works [1] [2] prove the learnability of ICL
under the Bayesian model of ICL.

3. Are there any insights on how the superposition phenomenon of the LLM influences real-world tasks? Is it desired or undesired? 
In the examples of the paper, an LLM with the superposition phenomenon seems to choose a random task if there are multiple different ICL tasks simultaneously.
Can this be disadvantageous in some scenarios? Could it be possible to control which task the LLM will perform?

[1] Zhang, Yufeng, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. "What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization." arXiv preprint arXiv:2305.19420 (2023).

[2] Wies, Noam, Yoav Levine, and Amnon Shashua. "The learnability of in-context learning." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677331803

thread_type:
official_review

<thread_object>
creator: reviewer_ncnf

<content>
summary:
The paper studies the mechanism of in-context learning with multiple tasks. It discovers the “task superposition” of LLMs. Using both theory and pretraining experiments on toy tasks, it demonstrates the ability of LLMs to implement the “task superposition”. They also explore the composition of different task vectors and how the model size affects the “task superposition” ability.

strengths:
1. The paper designs tasks to study the ICL with multiple tasks in LLMs. 
2. The composition of task vectors is interesting.

weaknesses:
Overall, the paper provides multiple findings, but removing redundant findings and implementing more analysis of interesting findings would be better.

1. Some findings are redundant. Findings 1 and 5 are not surprising given previous works such as Bai et al. [2023].
2. Findings 2-4 are interesting. However, more work is required on these findings. A similar “task vectors” analysis could also be applied to finding 2. Findings 3 and 4 can also relate to the non-clibrated results in Figure 2. It would be better to emphasize these findings.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023b.

questions:
1. Can authors demonstrate their novelty? Bai et al. [2023] have shown that “A single transformer can adaptively select different base ICL algorithms—or even perform qualitatively different tasks—on different input sequences, without any explicit prompting of the right algorithm or task.” This paper considers the mixture of different tasks with varying parameters of mixture. It is equivalent to selecting the correct algorithm that fits the mixture probability in the in-context examples.
2. Figure 2 presents that the output of LLMs is not calibrated with the in-context examples. How does this finding interplay with the results in Section 6, where the composition of task vectors is not proportional to the composition of in-context examples?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730678601781

thread_type:
official_review

<thread_object>
creator: reviewer_p2uz

<content>
summary:
The paper studied a specific type of LLMs in-context learning task, which the authors termed “task superposition”.  The study is comprehensive and extensive, including theoretical analysis and in-depth empirical investigation.

strengths:
1. The paper is clearly written and well presented.
2. The investigation towards the proposed ability is comprehensive and in-depth.

weaknesses:
Due to the inherit hierarchical nature of the concept “task”, I think the definition of “task” is at the core of this line of work. For example, one can argue Fig 1(a) essentially only contains a single task, which is addition regardless of language, or one can also argue that it contains only two tasks, which are numerical addition and multi-lingual translation, or four tasks as defined in the paper.

Hence, with a slight change of perspective (i.e., one abstract level up), the so-called superposition of tasks ICL itself basically defines a single new task. Since we know that LLMs can do ICL for a single task, it is not surprising that LLMs can also do the (single) task of “superposition of tasks” through ICL. Therefore, the term “tasks superposition” is misleading and unclear.

Some suggestions:
- Provide a clearer definition of what constitutes a "task" in the framework.
- More detailed discussion on how the notion of task superposition differs from or adds to existing understandings of in-context learning.
- Address how task superposition is distinct from a more complex single task.

questions:
1. How do you generate the output probabilities for the answers (e.g., in Fig 1(a))? Do you generate multiple outputs for a fixed prompt, then filter out the wrong ones, and count the correct ones over multiple prompts? 
2. What does it mean by "perform multiple tasks in a single inference call”, and “solve tasks in parallel”? What exactly is this “a single inference call”, is there an example?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730758265802

thread_type:
official_review

<thread_object>
creator: reviewer_ssa4

<content>
summary:
This paper demonstrates evidence of task “superposition” in LLMs when doing in-context learning, i.e. that LLMs can perform, or predict tokens for, multiple tasks simultaneously when shown those tasks in-context.

The first major argument for this is an empirical study of both pre-trained and from-scratch transformers. The authors first show that, across 3 pre-trained LLMs (GPT-3.5, Llama-3 70B, and Qwen-1.5 72B), when a model is shown an equal mixture of in-context samples across 3 tasks, the output distribution of the model is often a simplex over the correct outputs for each task, i.e. the top 3 next-token predictions for each input are the correct next-tokens for each of the 3 tasks, and different models exhibit biases towards different tasks. This is in contrast to the model either failing to learn the task to perform due to conflicting information in the context, or simply outputting the answer for one of the tasks that appears in the context.

Additionally, the authors show that, when training models to perform simple ICL tasks from scratch, e.g. simple arithmetic, the probability of each task’s output is proportional to the percentage at which that task’s examples appear in the context during test-time, even after the models are trained to perform ICL on a single-task at a time.

The authors then demonstrate theoretically, through the construction of a transformer model, that a transformer with 7 layers is capable of performing task superposition over a set of in-context examples.

Finally, the authors demonstrate that “task vectors”, defined as the model representation of the final token of the ICL prompt, reflects superposition. The authors show that the task vector created by a mixture of ICL tasks appears as an interpolation of each task’s task vector in T-SNE embeddings. Additionally, they show that creating new “task vectors” by interpolating the task vectors of distinct tasks can result in the model exhibiting task superposition as well.

The authors close by analyzing the relationship between task superposition and model size. They show that larger models have output probabilities that are more calibrated with the task probabilities of the context, i.e. the output for a task which appears in 50% of the in-context samples has a 50% probability in the model output.

strengths:
The paper presents a relatively strong argument that LLMs can perform multiple tasks simultaneously when identifying the task from context. This, in turn, provides more evidence towards the notion that ICL operates as a “mixture of experts”, identifying a distribution of tasks from context and conditioning on that distribution when outputting an answer. As a result, this paper makes some important in-roads towards improving our understanding of ICL.

The “from-scratch” experiments are particularly strong in this setting, as they show a relatively tight relationship between the distribution of tasks in-context and the predicted distribution of the model over task outputs. Additionally, the analysis of how superposition occurs in task vectors is particularly convincing in demonstrating that there does indeed seem to be some kind of “task mixture” being identified by the model, and the analysis of how model scale impacts this seems to indicate that this task mixture capability becomes more accurate as model size increases, which is particularly useful in helping us understand how LLMs do ICL.

weaknesses:
The “core” empirical results on large, pre-trained LLMs do not leave a clear take-away that superposition is occuring. This may, in part, be due to how the results are presented, i.e. the take-away that the model spreads it’s output across the 3 task outputs relies on the probabilities for “other” outputs to be low, which it is not always. This seems to only really be true for all models in one of the four settings considered. Thus, while the paper makes a strong case that LLMs can perform multiple tasks in superposition, it does not necessary make a strong case that current LLMs will perform multiple tasks in superposition. 

Finally, the usefulness of task superposition is still up for debate - the authors note in their limitations section that it’s usefulness is limited due to the lack of decoding tools that can take advantage of it. However, there is also not an analysis of how task superposition can impact task performance, and that trade-off, if it exists, is not discussed in this work. It seems intuitive that task superposition may impact the performance of the tasks being superimposed, which may also have an effect on its usefulness. In fact, it is somewhat surprising to me that task performance is not a major consideration when making the claim that models can perform many tasks simultaneously.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732295626750

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response by Authors

comment:
# General Response

We sincerely appreciate the thoughtful and constructive feedback from the reviewers.

We include general update / clarification here in the general response. Other individual questions and comments from the reviewers are addressed separately in each reply.

(In some of our responses, we provide additional Figure / Table with anonymous link embedded for visualization. In case the link fails, Figure 7-8 & Table 2 can be found in Appendix E and Figure I-V can be found in Appendix F of our updated draft. We also fixed some small typos in our original submission. We will incorporate our discussion with the reviewers in our future revised manuscript and open source the code for running experiments.)

## A more explicit definition of task and task superposition
Here we define "task" formally in the setting of in-context learning.

Let $\mathcal{X}$ be the input space, $x\in\mathcal{X}$, and $\mathcal{Y}$ be the output space. We define a task as a function $g: \mathcal{X}\rightarrow \mathcal{Y}$; we denote a task answer as $g(x)$, and $(x, g(x))$ as a task example of the task $g$. Let $\texttt{prompt}=s_m\oplus x^{(m+1)}=[x^{(1)}, g(x^{(1)}),...,x^{(m)},g(x^{(m)}),x^{(m+1)}]$ be a prompt consisting of $m$ task examples and a query $x^{(m+1)}$. $M$ is a model that predicts $M(\texttt{prompt})\in\mathcal{Y}$. We say $M$ learns the task $g$ in-context if $M(\texttt{prompt})=g(x^{(m+1)})$, i.e., $M$ learns $g$ and performs the mapping on the query $x^{(m+1)}$.

Let $\texttt{LM}$ be a language model and let $\mathcal{V}$ be the vocabulary space (token space). Given a $\texttt{prompt}$, instead of outputing a $y\in\mathcal{Y}$ directly, $\texttt{LM}$ outputs a probability distribution over $\mathcal{Y}$, i.e., $\texttt{LM}(\texttt{prompt}) = [\mathbb{P}(y_1\mid \texttt{prompt}), ..., \mathbb{P}(y_{|\mathcal{Y}|}\mid \texttt{prompt})]$ for $y_j\in\mathcal{Y}, j=1,...,|\mathcal{Y}|$. (In practice $\texttt{LM}$ outputs a distribution over $\mathcal{V}$ but for simplicity here we assume inputs and outputs are single-tokens, i.e., $\mathcal{X}=\mathcal{V}=\mathcal{Y}$, and it naturally generalizes to the case when outputs are multi-tokens as we can predict the probability of multi-token output given prompt using a chain rule.) In previous works, that study in-context learning for language models [1,2], all task examples in a $\texttt{prompt}$ are from a **single** task $g$; $\texttt{LM}$ would put most of the probability mass on a single task answer, while the output is selected using decoding strategy like greedy decoding $M(\texttt{prompt})=\arg\max_{y}\mathbb{P}(y\mid \texttt{prompt})\in \mathcal{Y}$.

In our work, we are the first to study a setting that task examples in a single $\texttt{prompt}$ are from $K$ different tasks, i.e., we have tasks $g_1,...,g_{K}$ which form $\texttt{prompt}=[x^{(1)}, g^{(1)}(x^{(1)}),...,x^{(m)},g^{(m)}(x^{(m)}), x^{(m+1)}]$ where $g^{(j)}\in\{g_1,...,g_K\}$. Since $\texttt{LM}$ outputs a probability distribution over $\mathcal{Y}$, we calculate the probability mass it assigns to each of the task answers $g_1(x^{(m+1)}), ..., g_K(x^{(m+1)})$. In Figure 1 we depict $[\mathbb{P}(g_1(x^{(m+1)})\mid \texttt{prompt}),..., \mathbb{P}(g_K(x^{(m+1)})\mid \texttt{prompt})]$ and find that $\texttt{LM}$ assigns non-negligible probability on multiple task answers. In Section 3 we keep investigating this phenomenon and find that under various settings and for different models, given that the task examples in the $\texttt{prompt}$ are from multiple tasks, $\texttt{LM}$ will output non-negligible probabilities on the corresponding task answers. We term this phenomenon as *task superposition*.

We will incorporate this update in our revised manuscript.

[1] Brown, T. B. (2020). Language models are few-shot learners. arXiv preprint arXiv:2005.14165.

[2] Xie, S. M., Raghunathan, A., Liang, P., & Ma, T. (2021). An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732463088958

thread_type:
official_comment

<thread_object>
creator: area_chair_pdqx

<content>
title:
Reminder: Author-Reviewer Discussion Period Closing Soon

comment:
Dear Reviewers

This is a reminder that the author-reviewer discussion period will end on Nov 26 AoE.

Your engagement during this phase is critical for providing valuable feedback and clarifications. If you have any remaining questions or comments, please take a moment to participate before the deadline.

Thank you for your contributions to this important process.

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) Summary of Scientific Claims and Findings

This paper investigates a unique capability in Large Language Models (LLMs), referred to as “task superposition.” Through experiments, the authors demonstrate that LLMs can perform distinct in-context learning (ICL) tasks simultaneously within a single inference.

(b) Strengths of the Paper

The paper is the first to claim and explore the “task superposition” capability in LLMs. The authors empirically validate this phenomenon across various architectures and scenarios, showing its robustness.

(c) Weaknesses of the Paper and Missing Elements

The experiments primarily focus on synthetic or toy problems, with limited exploration of task superposition in real-world contexts.
The practical utility of task superposition in real-world applications requires further discussion and analysis.

(d) Decision and Rationale

While the paper introduces the intriguing concept of “task superposition,” it would benefit from further investigation into its implications for real-world tasks and performance. Addressing these gaps would significantly strengthen the contribution.