W2: Transformers are Turing-complete and the theoretical result does not provide sufficient insights

We respectfully disagree with the reviewer on this. Previous works [4,5] have indeed established that Transformers, with access to infinite memory, are Turing complete. However, while Turing completeness implies broad theoretical capabilities, it does not address the depth or width required to achieve the specific result presented in our paper. In fact, based on the results of [4,5], the required depth and/or width would likely be impractically large, and adapting existing results to our specific setting may not be straightforward.

Given this context, we emphasize that the primary contribution of our work lies in the empirical observations that drive our study. The theoretical results are included to complement these observations, offering a plausible explanatory framework, but they are not the primary focus of our novelty.

Q1: Does the quantitative relation in Thm 1 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?

In our Section 7, we conduct experiments to show that larger models have better task superposition capability, that they can solve more tasks in parallel and better calibrate to ICL distribution (Finding 5).

We conduct additional experiments to investigate whether we need $K$ heads for $K$ tasks superposition. In particular, we trained a 1-head small model on addition tasks (the setting in Section 4) and tested on prompts with mixture of in-context examples of plus2 and plus5. In Figure V, we observe that while the 1-head model can in-context learn a single task (that corresponds to two end sides where $\lambda=0$ and $1$), it cannot perform task superposition when provided with examples of two tasks. For example, when $\lambda=0.5$ when we provide prompt with equal number of task examples from plus2 and plus5, half of the output probability mass lies on the plus2 task answer but the output probability for plus5 task answer is almost 0. This indicates that we need at least $2$ heads for $2$ tasks superposition.

Q2: Learnability of a Transformer that has the task superposition capability

Thanks for your suggestion. We think showing the learnability of a Transformer that has the task superposition capability is indeed a valuable and interesting direction for future work.

References

[1] Reynolds, L., & McDonell, K. (2021). Multiversal views on language models. arXiv preprint arXiv:2102.06391.

[2] moire. Language models are multiverse generators, January 2021. https://generative.ink/posts/language-models-are-multiverse-generators

[3] Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., ... & Stoica, I. (2023, October). Efficient memory management for large language model serving with pagedattention

[4]: Jorge Pérez, Pablo Barceló, and Javier Marinkovic. 2021. Attention is turing complete. J. Mach. Learn. Res. 22, 1, Article 75 (January 2021), 35 pages.

[5]: Giannou, Angeliki, et al. "Looped transformers as programmable computers." International Conference on Machine Learning. PMLR, 2023.

We thank the reviewer for their thoughtful feedback on our paper. Below, we address the specific questions raised by the reviewer.

W1&Q3: Task superposition in practice

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.

We would first like to clarify that task superposition is not just "choosing a random task if there are multiple different ICL tasks". Rather, task superposition is the phenomenon that when provided with in-context task examples from multiple tasks, the model will output non-negligible probabilities on the corresponding task answers (we provide a more detailed explanation of task superposition framework in the "A more explicit definition of task and task superposition" section of our general response). Below we further discuss the contribution of our work.

It seems unclear whether real-world tasks have the superposition phenomenon or how the superposition phenomenon affects LLMs in practice. Are there any insights on how the superposition phenomenon of the LLM influences real-world tasks? Is it desired or undesired? Can this be disadvantageous in some scenarios? Could it be possible to control which task the LLM will perform?

We appreciate reviewer's thoughtful questions regarding the practicality and implications of task superposition. We acknowledge that the immediate practical applications of task superposition may not be fully apparent. However, we believe that our findings contribute significantly to the fundamental understanding of LLMs and open avenues for future research in several important areas.

1. Task superposition could have practical benefits in scenarios where an LLM serves multiple users simultaneously. By leveraging superposition (e.g., using a better decoding strategy that can maintain the model’s multi-task state throughout the generation process), an LLM might efficiently handle diverse tasks in parallel, improving scalability and responsiveness in LLMs serving [3] under multi-user environment.
2. The phenomenon of task superposition has intriguing implications for AI alignment. It suggests that alignment interventions might influence the distribution over latent skills conditioned on the prompt (i.e., $\mathbb{P}(`skill` | `prompt`)$ ) without necessarily affecting the model's performance conditioned on a specific skill and input (i.e., $\mathbb{P}(`answer` | `skill`, `prompt`)$ ). This raises questions about the effectiveness of alignment strategies that focus solely on modifying model outputs without considering the underlying superposition of skills. While this is speculative, we believe it highlights a critical area for future research in understanding and improving alignment methodologies.
3. Our observations align with the "simulator-in-superposition" hypothesis [1,2] that emerged with the advent of GPT-3. This hypothesis suggests that LLMs can simulate multiple potential continuations or behaviors simultaneously, reflecting a superposition of different skills or tasks. By demonstrating that LLMs can internally represent and process multiple tasks in parallel when provided with mixed in-context examples, we provide empirical support for this theoretical framework.
4. We believe that explicitly characterizing the phenomenon of task superposition is valuable in its own right. Adding to our understanding of how LLMs work contributes to the broader field of AI research.

Our goal with this paper is to shed light on an emergent behavior of LLMs that has not been previously documented, and we think is surprising. We believe that understanding such phenomena is crucial for the research of LLMs.

We thank the reviewer for their thoughtful feedback on our paper. Below, we address the specific questions raised by the reviewer.

W1&Q1: Difference between our work and the work by Bai et al.

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.

We acknowledge the valuable insights from Bai et al. [1] and plan to expand our discussion of these connections in the revised paper. However, we respectfully disagree with the characterization that our work is equivalent to algorithm selection in Bai et al. While there are connections, the works differ in several aspects.

1. Settings:

• Bai et al. focuses on algorithm selection where during test time Bai et al. tests transformers using in-context examples from one task, and shows that the model selects the right algorithm to solve one target task.
• We test the model with multiple qualitatively different ICL tasks presented in a single prompt, and it solves multiple tasks in superposition.

2. Training Requirements:

• In Bai et al.'s empirical setting, for the algorithm that the model selects during test time, it requires explicit training on that algorithm.
• While it is true that superposition can be think of as an algorithm that outputs probability distribution that calibrates with in-context examples, applying Bai et al.'s setting here assume we have explicitly trained the model on all such algorithm that calibrates output distribution with respect to in-context examples, i.e., the model was trained on data where the $`prompt`$ mixes examples from qualitatively different tasks and labels being non one-hot to calibrate the in-context task examples.
  • This is likely not the case in the pre-trained model since the labels during LLMs training are in one-hot format (probability 1 on one token and 0 elsewhere); furthermore, prompts consisting of task examples from different tasks (within a single prompt) are unlikely to be part of some training data.
  • In fact, in Finding 2 we find that we don't need to train the model on such data, and the model can do task superposition even if we train the model to in-context learn one task at a time.

3. Scaling Experiment:

• To our best knowledge, our Finding 5 that studies model's task superposition capability as the model scales was also not observed in other works.

We would be happy to further clarify these distinctions if the reviewer finds the current explanation insufficient.

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

W2&Q2: More analysis on Findings

A similar “task vectors” analysis could also be applied to finding 2.

We extract task vectors from our small trained-from-scratch models using the same pipeline as the pretrained models. While our task vector extraction works well for large, pretrained models, we could not find task vectors that work well for our small models. For example, in Figure I we plot the accuracy on task ret2 and plus2 when using vectors extracted from different layers and observe that the maximum accuracy we can get is lower than 0.2 (while for large pretrained model such accuracy is usually near 1).

We plot the task vectors decomposition in Figure II and conduct the linear interpolation experiment in Figure III. We notice that task vectors of uniform mixtures for retrieval tasks do not lie in the middle of three individual task vector clusters and convex combinations in both cases do not produce task superposition. A possible explanation is that, while current task vectors are extracted from a specific layer, for the small models, the feature that represents a task is likely not localized to a specific layer (as indicated in Figure I), which means that we need to modify how we extract the task vectors for small models. We believe this is an important area for further research.

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?

There is a correlation between the inter-cluster distances of task vectors in section 6 and calibration of LLM output probabilities in Figure 2. To illustrate this, we plotted the task vector scatter plot in the style of Figure 4 alongside the distances between centroids of the clusters and the corresponding probabilities on task answers in Figure IV. We observe that in most cases, for task vectors of mixtures, a closer distance to task vectors of a task gives higher probability in the corresponding task answer. In particular,

• The model's preference of op1+op2 over copy(op1) or copy(op2) corresponds to smaller distance from mixture cluster centroid to the op1+op2 cluster centroid.
• The model's preference of capital over continent and capitalization corresponds to smaller distance from mixture cluster centroid to the capital cluster centroid.

That said, we do not believe this relationship is necessarily proportional or linear. For example, in Figure IV(c), for the 0.5 copy(op1), 0.5 op1+op2 mixture, it is significantly closer to op1+op2 cluster than to copy(op2) cluster, but the model only assigns slightly more probability to op1+op2 task answer than to copy(op2) task answer. So, while we show that the model composes task vectors internally, it is more than a convex combination of task vectors. It will be an interesting future research area to capture the mechanism of how model mixes the task vectors to produce superposed answers.

(We will also incorporate this analysis in our future revised manuscript).

Q1: How do you generate model output probabilities?

We first clarify that, given $K$ tasks $g_1,...,g_K$ and a $`prompt`$ consisting of $m$ task examples (from $K$ tasks) and a query $x^{(m+1)}$, we calculate the output probabilities of task answers $g_1(x^{(m+1)}),...,g_K(x^{(m+1)})$ by calculating $\mathbb{P}(g_1(x^{(m+1)})\mid`prompt`),...,\mathbb{P}(g_K(x^{(m+1)})\mid`prompt`)$.

If task answers are single tokens, then this can be calculated in a forward pass as $`LM`$ outputs a distribution over $\mathcal{V}$ given $`prompt`$. If task answers are multi-tokens, we calculate each $\mathbb{P}(g_i(x^{(m+1)})\mid`prompt`), i=1,...,K$ autoregressively. That is, if in the token space $`prompt`=[v_1,...,v_M]$ has $M$ tokens, $g_i(x^{(m+1)})=[u_1,...,u_N]$ has $N$ tokens, $v_j, u_k\in\mathcal{V}, j\in[M], k\in[N]$, then \begin{aligned} \mathbb{P}(g_i(x^{(m+1)})\midprompt)&=\mathbb{P}([u_1,...,u_N]\mid[v_1,...,v_M])\\ &=\mathbb{P}(u_1\mid[v_1,...,v_M])\prod_{j=2}^N\mathbb{P}(u_j\mid[v_1,...,v_M, u_1,...,u_{j-1}]). \end{aligned}

Our Appendix B formalizes on this, please let us know if you would like us to extend it further.

Q2: More explanation on "perform multiple tasks in a single call"

We provide an example in our response to W2, where we show that given $`prompt`\_{\text{mix}}$ with task examples from two tasks, the model puts significant portion of probability mass on two task answers respectively. We would happy to further explain if the reviewer finds this unclear.

We would also like to note that when task answers are multi-tokens, the model is still doing multiple tasks in parallel when provided with $`prompt`$ and predicting the first token by putting non-negligible probability mass on first tokens of each task answers. On the other hand, as we discussed in the "Limitation and Future Directions" section of paper, using existing decoding strategy, after the first token is generated, the model tends to converge on predicting tokens for a single task. We show that the model has the capability for multi-task execution and highlights a critical area for future research -- developing decoding strategies that can maintain the model’s multi-task state throughout the generation process (e.g., recent "superposed decoding" [1] offers some hope towards this direction).

[1] Shen, E., Fan, A., Pratt, S. M., Park, J. S., Wallingford, M., Kakade, S. M., ... & Kusupati, A. (2024). Superposed Decoding: Multiple Generations from a Single Autoregressive Inference Pass. arXiv preprint arXiv:2405.18400.

We thank the reviewer for their thoughtful feedback on our paper. Below, we address the specific questions raised by the reviewer.

W1: Clearer definition of task/task superposition and how it differs from existing understandings of ICL

We add a more explicit formulation of task and task superposition and explains how it differs from existing understandings of ICL in the "A more explicit definition of task and task superposition" section of our general response. We would happy to further explain if the reviewer finds this unclear.

W2: Distinction from a more complex single task

Under our new definition and framework in the general response, a more complex single task can be think of as a vector of multiple simple tasks. For example, let $g_1, g_2:\mathcal{X}\rightarrow \mathcal{Y}$ be two simple tasks, we have a $g_3=(g_1, g_2):\mathcal{X}\rightarrow \mathcal{Y}\times\mathcal{Y}$ as a more complex task.

Here is a concrete example. Let $\mathcal{X}, \mathcal{Y}$ be integers. Let $g_1$ be $g_1(x)=x+3$ (plus-three task) and $g_2(x)=x-1$ (minus-one task). In the existing in-context learning framework, we form prompts where task examples are from a single task:

$`prompt`_1=$1->4; 2->5; 3->6; 4->7; 5-> (task examples from $g_1$)

$`prompt`_2=$1->0; 2->1; 3->2; 4->3; 5-> (task examples from $g_2$)

and if we provide two prompts to language model $`LM`$, we can get $\mathbb{P}(8\mid`prompt`_1)=0.98$ and $\mathbb{P}(4\mid`prompt`_2)=0.99$. We can observe that $`LM`$ assigns most of the probability mass on a single task answer in both cases and this indicates $`LM`$ in-context learns $g_1$ and $g_2$ respectively.

In our setting, we form prompt where task examples are from different tasks

$`prompt`\_{\text{mix}}=$1->4; 2->1; 3->2; 4->7; 5-> (task examples from $g_1$ and $g_2$)

and $`LM`$ gives $\mathbb{P}(8\mid`prompt`\_{\text{mix}})=0.49$ and $\mathbb{P}(4\mid`prompt`\_{\text{mix}})=0.48$, putting two significant portions of probability mass on two task answers respectively. We call this phenomenon task superposition, and we can observe this from a single forward pass because we can get both $\mathbb{P}(8\mid`prompt`\_{\text{mix}})$ and $\mathbb{P}(4\mid`prompt`\_{\text{mix}})$ from $`LM`(`prompt`\_{\text{mix}})$, which is a probability distribution over $\mathcal{V}$. This indicates that $`LM`$ performs two tasks simultaneously when provided with $`prompt`\_{\text{mix}}$.

We can also consider a more complex single task $g_3=(g_1, g_2)$ and form the prompt

$`prompt`_3=$1->4,0; 2->5,1; 3->6,2; 4->7,3; 5-> (task examples from $g_3$)

and note that task examples are all from a single task $g_3$, we will get $\mathbb{P}((8,4)\mid `prompt`_3)=0.96$ from $`LM`$. Since $`LM`$ predicts next token autoregressively, the full process would be

1. $\mathbb{P}($8$\mid$1->4,0; 2->5,1; 3->6,2; 4->7,3; 5->$)=0.99$. Append 8 to the prompt.
2. $\mathbb{P}($,$\mid$1->4,0; 2->5,1; 3->6,2; 4->7,3; 5->8$)=0.99$. Append , to the prompt.
3. $\mathbb{P}($4$\mid$1->4,0; 2->5,1; 3->6,2; 4->7,3; 5->8,$)=0.98$, and we get model's output 8,4 with probability $0.99\times0.99\times 0.98=0.96$ (or an additional step where the model predicts an empty space indicating the end symbol).

Here the model is still performing a single complex task using multiple steps (each step performing a single simple sub-task) where the model puts most of its probability mass to a single token at each step. We highlight that this "more complex single task" setting is different from our task superposition setting as this setting first calculate $4$ and then calculate $8$ in different steps (in different model forward passes).

We will incorporate this discussion in our revised manuscript.

We thank the reviewer for their thoughtful feedback on our paper. We have updated our draft with a new Appendix E which contains additional experiments to address the reviewer's concerns. We break down the aforementioned weaknesses in two parts:

W1: Is superposition really happening on large pre-trained LLMs?

Yes. Based on your feedback, we conduct the following experiments that provide clear evidence that task superposition is occurring in current language models.

In Figure 8 we show the probability of the most likely answer other than the task answer of tasks specified in the prompt. For example, in Figure 8a, we observe that for GPT-3.5 and Llama-3, the highest probability for a non-task answer (an individual “other” task) is still significantly lower than the four addition tasks that were given in the prompt. Because the next most likely answer still has lower probability than the specified task answers, any individual “other” answer must also has much lower probability than the specified tasks. This shows that the models focus on performing superposition on the tasks specified in-context. We note here that, Qwen-1.5 has similar probability as some of the task answers. We attribute this to the model’s limited ability to perform some of the specified tasks, specifically adding in language other than English.

Our controlled ablation studies (Figure 7) show that task probabilities directly reflect in-context demonstrations: task outputs have negligible probability when their corresponding task is absent from the prompt, but gain substantial probability when the task is included. This systematic behavior cannot be explained simply by noise.

These findings demonstrate that task superposition is a real phenomenon that emerges specifically in response to multiple in-context tasks.

(Experimental setup details can be found in Appendix E1 and E2 of our updated draft.)

W2: Analysis on how task superposition impact task performance and usefulness of task superposition

In Table 2, we show the task accuracy with one task prompting $(K=1)$ and with multiple tasks prompting $(K>1)$, where we define correctness of a task as whether the correct task answer shows up in the top-$K$ output probabilities. We find that as we increase the number of tasks, all models exhibit an accuracy decrease in correctly performing each individual task. Notably, however, the accuracy degradation of Llama3-70B is less than that of Llama2-70B across most tasks. We believe this is indicative that improved model training techniques lead to better preservation of task superposition, and therefore, while task superposition can impact task performance, we believe such performance drop will be even lower for future models. Additionally, this also suggests that task superposition may be used to help benchmark future model training.

(Experimental setup details can be found in Appendix E3 of our updated draft.)