Everything Everywhere All At Once: Llms Can In-context Learn Multiple Tasks In Superposition

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.

Steering models through in-context learning has been a growing area of interest. Recent work has hypothesized that in-context learning can be encapsulated by a high-dimensional description of a task, which can be used to replace, (Hendel et al., 2023) compose (Todd et al., 2024) or augment (Liu et al., 2024) the latent states of a model, in order to alter its default behavior. Task vectors can be combined via arithmetic operations to solve a variety of tasks (Ilharco et al., 2023). Prior work has also been investigating the power of tokens in defining a task (Bai et al., 2024).

Our findings on superposition are inspired by notions of language models as multiverse generators (Reynolds & McDonell, 2021; moire, 2021). One consequence of LLMs as a superposition of tasks is that the outputs may collapse to unintended simulacra, a behavior known as the “Waluigi effect” (Nardo, 2023).

Superposition has been defined in various related contexts of learning models. Feature superposition (Elhage et al., 2022) refers to a neural network’s ability to represent multiple learned concept in a single neuron. Though our discovery of task superposition describes the same abstract idea, we stress that it is distinct from feature superposition because task superposition is most apparent in the final output of a model. Feature superposition is a microscopic-level observation whereas task superposition is a macroscopic-level observation.

One limitation of our work is the current gap between the demonstrated capability of LLMs to perform task superposition and its practical application in real-world scenarios. While we have shown that LLMs possess the capacity to execute multiple tasks simultaneously, conventional decoding algorithms are not equipped to fully leverage this capability. This limitation stems from what we term ”generation collapse,” a phenomenon where, after the first token is generated, the model tends to converge on predicting tokens for a single task, effectively negating its ability for multi-task execution.

This collapse presents a substantial challenge in harnessing the full power of task superposition. It highlights a critical area for future research: developing decoding strategies that can maintain the model’s multi-task state throughout the generation process. Recent work by Shen et al. (2024b) offers some hope that this direction may be fruitful, by proposing a “superposed decoding” algorithm. Their method efficiently generates multiple streams of tokens from a single inference pass by utilizing superposed token embeddings. While this approach represents a significant step forward, it also highlights the potential for further innovation in this area.