💬 NLP

Shared Lexical Task Representations Explain Behavioral Variability In LLMs

The paper demonstrates that LLM prompt sensitivity can be explained by the activation levels of "lexical task heads"—shared internal mechanisms that represent the task regardless of whether the prompt uses instructions or examples—and that performance variability stems from how strongly these heads are triggered or diluted by competing representations.

Original authors: Zhuonan Yang, Jacob Xiaochen Li, Francisco Piedrahita Velez, Eric Todd, David Bau, Michael L. Littman, Stephen H. Bach, Ellie Pavlick

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhuonan Yang, Jacob Xiaochen Li, Francisco Piedrahita Velez, Eric Todd, David Bau, Michael L. Littman, Stephen H. Bach, Ellie Pavlick

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very talented, but slightly moody, assistant how to perform different tasks.

Sometimes, if you give them a clear instruction (like a recipe), they do a perfect job. Other times, if you just show them a few examples (like showing them pictures of finished dishes), they might get confused or perform inconsistently. You might think the assistant is just being unpredictable, but this paper reveals that there is actually a very logical "internal wiring" causing this behavior.

Here is the breakdown of the paper using a "Kitchen & Chef" analogy.

1. The Discovery: The "Recipe Labels" (Lexical Task Heads)

The researchers discovered that inside the AI's "brain," there are specific parts called Lexical Task Heads.

Think of these as sticky notes that the chef (the AI) slaps onto their tools. If the task is to translate words, the chef doesn't just start translating; first, a specific part of their brain grabs a sticky note that literally says "Translation Mode" or "Language Switch."

The researchers found that these "sticky notes" are written in actual words that humans can read. When they looked at the AI's internal signals, they saw the AI's brain literally shouting words like "Opposite," "Plural," or "Capital City" before it even gave the final answer.

2. The Shared Secret: One Brain, Many Styles

The big mystery was: Does the AI use different brain parts for a written instruction versus a visual example?

It turns out, no. Whether you give the AI a written command ("Translate this...") or a list of examples ("Apple $\rightarrow$ Manzana"), the AI eventually reaches for the exact same sticky notes.

It’s like a chef: whether you give them a written recipe or just show them a photo of a cake, they eventually reach for the same "Baking Mode" mental switch. This means the AI has a "core understanding" of a task that exists independently of how you ask it.

3. Why the AI Fails: The "Weak Signal" and "Distractions"

If the AI has these shared "sticky notes," why does it still mess up? The paper offers two main reasons:

The Faint Sticky Note (Insufficient Activation): Sometimes, the prompt you use is "weak." It doesn't trigger the "Baking Mode" sticky note strongly enough. The note is there, but it's a tiny, blurry scribble. Because the signal is faint, the chef doesn't quite "get into the zone," and they make mistakes. The researchers proved this by "turning up the volume" on these sticky notes, which actually fixed the AI's mistakes!
The Wrong Recipe (Competing Tasks): Sometimes, a prompt is ambiguous. Imagine you tell a chef, "Make something with flour." They might make bread, or they might make a cake. If the prompt is slightly confusing, the AI might accidentally grab the "Cake Mode" sticky note when you actually wanted "Bread Mode." The "wrong" signal drowns out the "right" one.

4. The "Many-Shot" Magic: Why more examples help

You’ve probably noticed that giving an AI 20 examples works better than giving it just one. The researchers found the reason: Accumulation.

Every time you provide an example, it’s like hitting the "Task Mode" button harder and harder. Each example adds more "ink" to that sticky note, making the signal louder and clearer until the chef is 100% certain which mode they are in.

Summary in a Nutshell

The paper tells us that AI "prompt sensitivity" isn't just random magic. It’s a mechanical process where different ways of asking a question trigger internal task-labels with different levels of strength.

The takeaway for humans: If your AI is acting up, it’s likely because your prompt isn't "loud" enough to trigger the right internal sticky note, or it's accidentally triggering a different one!

Technical Summary: Shared Lexical Task Representations Explain Behavioral Variability In LLMs

1. Problem Statement

Large Language Models (LLMs) exhibit significant prompt sensitivity, where minor linguistic variations in how a task is presented (e.g., changing the wording of an instruction or the number of few-shot examples) lead to unpredictable changes in accuracy. While previous research has suggested that instruction-based and example-based prompting might rely on different internal mechanisms (such as "function vectors"), it remains unclear whether there are shared underlying circuits that drive task execution across different prompting styles. Understanding this variance is critical for building reliable and predictable AI systems.

2. Methodology

The researchers investigated the internal mechanisms of 13 autoregressive transformer models (including Llama-3.1, Gemma-2, and Qwen families) across 17 simple tasks (e.g., translation, antonyms, country-capital).

Key Methodological Innovations:

Identification of Lexical Task Heads (LTHs): The authors identified a new class of attention heads that encode task information in a "lexical" format. By projecting the outputs of these heads into the vocabulary space using the Logit Lens, they found that these heads literally output words describing the task (e.g., for an antonym task, the heads decode to "opposite" or "reverse").
Circuit Analysis: They distinguished between Lexical Task Heads (which represent the task definition) and Retrieval Heads (which extract the correct answer, e.g., "Tokyo").
Causal Interventions: The study employed activation patching and activation scaling. They patched activations from successful prompts into failing prompts and scaled the magnitude of LTH activations to observe if they could "repair" incorrect model outputs.
Comparative Analysis: They measured the overlap of these heads across different prompting styles (instruction vs. example) and compared them against "function vectors" to determine if they were distinct or shared.

3. Key Contributions

Discovery of Lexical Task Representations: The paper introduces LTHs as a distinct, interpretable class of task representations that are shared across prompting styles, unlike previously identified function vectors which are often style-specific.
Mechanistic Explanation of Prompt Sensitivity: The authors provide a two-part explanation for why prompts fail:
- Insufficient Activation: Behavioral variation is driven by the strength of LTH activation.
- Competing Representations: Prompt ambiguity can trigger "off-target" task circuits that dilute the signal of the intended task.
Functional Linkage: They demonstrate a clear hierarchy: LTHs $\rightarrow$ Retrieval Heads $\rightarrow$ Correct Answer. LTHs act as a "directional filter" that modulates retrieval heads to ensure they extract the specific attribute required by the task.

4. Key Results

High Circuit Reuse: On average, 73% of lexical task heads are shared between instruction-based and example-based prompting styles for the same task.
Correlation with Accuracy: There is a strong positive correlation between the number/strength of LTH activations and model accuracy. This explains the "many-shot" effect: as more examples are provided, the activation of LTHs increases, leading to better performance.
Causal Repair: Scaling up the activation of LTHs in failing prompts can successfully "fix" the model, steering it toward the correct task and improving accuracy.
Ambiguity and Competition: In ambiguous prompts (e.g., a prompt that could be interpreted as both "Product-Producer" and "Product-Country"), the model shows diminished activation of the target LTHs and increased activation of competing task heads, leading to errors.
Generalization: The mechanism was found to hold for complex tasks, including compositional tasks (two-hop reasoning) and code generation (where LTHs for "Python" or "JavaScript" could steer the model to the correct language).

5. Significance

This work bridges the gap between high-level behavioral observations (prompt sensitivity) and low-level mechanistic interpretability (transformer circuits). By showing that LLMs use a shared, lexical-based internal language to represent tasks, the paper provides a roadmap for:

Improving Robustness: Designing prompts that maximize the activation of target LTHs.
Interpretability: Using the vocabulary space as a "universal language" to understand how different models represent tasks.
Model Steering: Using LTH activation as a precise tool to control model behavior (e.g., switching programming languages or task modes) without needing complex fine-tuning.