🔬 applied physics

Adaptive inference and function vectors in deep transformers

This paper introduces a mean-field theory of deep transformers as distributed inference systems that utilize internal "function vectors" to hierarchically infer latent context variables, demonstrating that feedforward blocks and depth enable a richer class of adaptive in-context learning algorithms than previously understood.

Original authors: Ravin Raj, Gautam Reddy

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ravin Raj, Gautam Reddy

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 solve a mystery, but you only have a few scattered clues (the "context") and a specific question to answer (the "query"). You need to figure out the hidden rule that connects all the clues to the answer.

This paper proposes a new way to understand how Transformers (the AI brains behind tools like chatbots) solve these mysteries. Instead of just looking at the math, the authors treat the Transformer like a team of detectives working together to solve a case.

Here is the breakdown of their theory using simple analogies:

1. The Team of Detectives (The Tokens)

Imagine a room full of detectives (these are the "tokens" or pieces of data). They all know a little bit about the case, but no single detective has the whole picture. Their goal is to figure out the "hidden rule" (called the latent context) that explains everything.

In a standard AI, these detectives might just shout out their thoughts all at once. But this paper suggests that deep Transformers work like a relay race or a multi-stage investigation.

2. The "Function Vector" (The Shared Notebook)

The core idea is something called a Function Vector. Think of this as a shared notebook that gets passed around the room.

Round 1: Each detective looks at their clue and writes a short summary in the notebook.
Round 2: They read the notebook, realize they missed something, and add a new, more refined note.
Round 3: They read the updated notebook again and add even more detail.

By the end of the process, this notebook (the Function Vector) contains a compressed, perfect summary of everything the team learned about the hidden rule. When the final detective (the "query") asks for the answer, they just read this notebook.

3. The Two Types of Workers (MLP vs. Attention)

The paper explains that the Transformer has two types of workers that take turns:

The Local Thinker (MLP): This is the individual detective who looks at their own clue and the current notebook. They decide, "Based on what I know and what's in the notebook, what is the most important new thing I should add?" They act like a router, choosing which information is worth sharing.
The Group Mixer (Attention): This is the person who takes everyone's new notes and blends them together into a single, updated summary for the next round.

4. The Big Discovery: Why "Deep" Matters

The authors tested two scenarios to see if having many layers (a "deep" team) is better than having just one big layer.

Scenario A: The Smooth Hill (Gaussian Prior)
Imagine the hidden rule is on a smooth, flat hill. You can find the top just by looking at it from one angle.

Result: It doesn't matter if you have a team of 10 people passing a notebook back and forth, or one super-smart person doing it all at once. They get the same result. The "depth" doesn't help here.

Scenario B: The Maze (Tree Prior)
Now, imagine the hidden rule is inside a complex maze with many turns (a "tree" structure). To find the exit, you have to make a series of choices: Left or Right? Then Up or Down?

The "Shallow" Team (Non-Adaptive): This team tries to guess the whole path at once. They might get the first few turns right, but they get lost quickly because they can't adjust their plan based on new info.
The "Deep" Team (Adaptive): This team moves step-by-step.
- Step 1: They check the first intersection.
- Step 2: Based on that result, they decide which path to check next.
- Step 3: They adjust again.
- Result: Because they can adapt their strategy at every step (using the Function Vector to remember where they are), they navigate the maze much better than the shallow team.

5. The Proof: The "Key-Patching" Experiment

To prove this, the researchers did a "surgery" on the AI. They took the "keys" (the decision-making notes) from a detective who solved a different maze and forced them into the current detective's notebook.

What happened? If they swapped the notes early in the process, the detective got confused and failed. If they swapped them late, it didn't matter as much.
Why? This proved that the AI wasn't just memorizing a static answer. It was actively building a strategy layer by layer, just like the theory predicted.

The Bottom Line

This paper argues that deep Transformers aren't just fancy pattern matchers. They are adaptive inference machines.

When the problem is simple, a shallow approach works fine. But when the problem is complex and hierarchical (like a tree with many branches), the Transformer uses its depth and internal notes (Function Vectors) to act like a smart detective: it gathers clues, updates its theory, and decides what to look for next, all within a single pass of reading the data. This allows it to solve complex puzzles that simpler, "shallow" models cannot.

Problem Statement
Transformers have emerged as a general-purpose substrate for learning complex correlations in diverse domains, yet their internal mechanisms for in-context learning (ICL) remain mysterious. While empirical work has identified two primary motifs for ICL—attention circuits implementing specific algorithms (e.g., copying, gradient descent) and the formation of "function vectors" (internal representations of the task)—the theoretical understanding of the latter remains limited. Specifically, it is unclear how deep transformers utilize function vectors to perform inference, and what specific computational roles depth and feedforward (MLP) blocks play in this process. Existing theories often focus on attention-based circuits, leaving the mechanisms enabled by multi-layer interactions between attention and MLPs underexplored.

Methodology
The authors propose a theoretical framework modeling a deep transformer as a mean-field interacting system that implements distributed Bayesian inference. The model is subject to architectural constraints on communication (channel dimension $d$ ), locality, and depth ( $L$ ).

Theoretical Model: The framework posits that a transformer processes a set of tokens coupled by a latent context variable $\theta$ . Inference is mediated through "function vectors" ( $\Phi_\ell$ ), which act as compressed state variables.
- At each layer $\ell$ , a local nonlinearity (MLP) processes a token and the current state $\Phi_\ell$ to produce an embedding $\lambda_{\ell+1}$ .
- An attention-like operation pools these embeddings across tokens to form a new statistic $\phi_{\ell+1}$ .
- The state is updated via direct sum: $\Phi_{\ell+1} = \Phi_\ell \oplus \phi_{\ell+1}$ .
- The MLP blocks are hypothesized to act as "routers" (choosing which information to convey) and decoders, while attention pools information.
Optimization Strategy: The authors frame the selection of embeddings as a dynamic programming problem. The optimal strategy at layer $\ell$ minimizes the expected loss at layer $L$ , effectively treating the embedding choice as an experimental design strategy for adaptive inference.
Analytical Cases: The theory is analyzed under two prior distributions for the latent variable $\theta$ :
- Gaussian Prior: Where the latent variable follows a standard normal distribution.
- Tree Prior: A hierarchically structured prior where $\theta$ is constructed based on a path through a balanced binary tree. This structure tests the advantage of depth in resolving latent variables at increasingly finer scales.
Empirical Validation: The predictions are tested using a constrained linear attention transformer architecture with three configurations:
- Model A: A deep ( $L=M$ ) model with channel dimension $d=1$ and MLP blocks at every layer.
- Model B: A single-layer model with channel dimension $M$ and identical MLPs.
- Model C: A single-layer model with channel dimension $M$ and no MLPs (linear attention only).
- Causal Perturbation: A "key-patching" experiment is conducted where key matrices from a "donor" context are inserted into a "target" context at specific layers to test the routing mechanism.

Key Contributions

Theory of Function Vectors: The paper provides a formal theory defining function vectors as internal state representations that compress context information for prediction. It demonstrates that these vectors allow a model to infer latent context variables at increasingly finer scales across layers.
Role of Depth and MLPs: The theory predicts that for non-Gaussian, hierarchical priors, depth and MLP blocks are essential for implementing adaptive inference strategies that outperform non-adaptive, single-layer approaches. Conversely, for Gaussian priors, depth offers no advantage over a single layer with equivalent total communication budget.
Mechanistic Mapping: The work maps the theoretical components to transformer architecture, assigning the MLP the role of a "router" that selects statistics to measure and the attention mechanism the role of pooling these statistics.
Experimental Confirmation: Numerical experiments confirm that deep transformers with MLPs (Model A) can learn adaptive inference strategies on hierarchical tree priors, significantly outperforming optimized single-layer non-adaptive strategies (Model B) and linear attention models (Model C).

Results

Gaussian Prior: In the case of a Gaussian prior, the Mean Squared Error (MSE) loss is independent of the function vector's specific realization and depends only on the measurement strategy. The theory and experiments show that a deep model (Model A) performs no better than a single-layer model with equivalent channel capacity (Model B).
Tree Prior: For the hierarchical tree prior, the theory predicts that an adaptive strategy (measuring one coordinate at a time based on previous results) yields an MSE loss that scales exponentially with depth ( $\alpha^M$ $α^{M}$ ), whereas non-adaptive strategies scale as a power of $M$ $M$ .
- Experiments show Model A achieves a significantly lower MSE loss than Model B, confirming the model learns an adaptive inference strategy.
- Model C (no MLPs) performs similarly to the non-adaptive theoretical baseline, confirming the necessity of MLPs for the adaptive routing mechanism.
Causal Perturbation: Key-patching experiments reveal that inserting keys from a donor context (leaf $\sigma'$ ) into a target context (leaf $\sigma$ ) causes a large increase in loss if the paths diverge early in the tree and the intervention occurs early in the network. This confirms that the model's internal state (function vector) evolves adaptively based on the specific context path.

Significance
The paper claims to reconcile two perspectives on in-context learning: the view of ICL as Bayesian inference over a latent variable and the mechanistic view of ICL as algorithmic execution (e.g., gradient descent). By showing that deep transformers can implement Bayesian inference as a distributed, adaptive, layerwise procedure mediated by function vectors, the authors demonstrate that depth and MLP blocks enable a richer class of ICL algorithms than previously described (such as preconditioned gradient descent). The work places function vectors within a circuit-level framework, suggesting that the organization of billion-parameter models favors this structure because it allows for efficient, adaptive inference under communication constraints. The authors note that while their theory uses linear pooling, it respects exchangeability and may be sufficient for scenarios where the specific token-to-token retrieval of softmax attention is not the primary driver of ICL.