🤖 machine learning

Cross-Layer Interaction under Weight-Space Ablation: A Closed-Form Attention Jacobian Bound and a Test on a Real Pretrained Model

This paper extends theoretical bounds on cross-layer interactions in weight-space ablation by deriving a closed-form attention Jacobian bound verified on a real pretrained model and demonstrating mixed empirical results on an emergent indirect object identification circuit.

Original authors: Abdallah Khemais

Published 2026-08-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Abdallah Khemais

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 understand how a giant, super-smart robot brain works. This brain isn't a single lump of metal; it's built like a long assembly line with many stations, or "layers," where information gets passed down the line. At each station, the brain does two main things: it pays attention to specific clues (like a detective focusing on a witness) and it processes those clues through a thinking module (like a calculator crunching numbers). Scientists call the parts that carry these clues "carriers."

To figure out what each part does, researchers play a game of "what if." They try two different tricks: Activation Patching, where they swap a piece of the robot's current thought with a thought from a different situation to see if the robot gets confused, and Weight Ablation, where they effectively "turn off" a part of the brain's wiring to see if the robot stops working. Usually, these two tricks tell the same story. But sometimes, when the robot's stations are connected in tricky ways, the tricks give different answers. The big question is: Why do they disagree, and how much does one part of the brain mess up another part that is far away down the line? Understanding this is crucial because if we want to fix or control these giant brains, we need to know exactly how their internal gears interact, not just how they look on the surface.

This paper is like a detective story that takes a clue found in a tiny, simple model and tries to see if it holds up in a massive, real-world robot. The author starts with a discovery from a "companion paper" (a sibling study by the same writer) which proved that in a single, isolated station, turning off a specific type of thinking module (the MLP) doesn't mess up the attention part at all. It's like turning off the calculator in a factory station without affecting the detective's magnifying glass. However, that proof only worked for one station at a time. Real robots have dozens of stations, and the parts the scientists want to turn off are often far apart.

The author's first major move is to break down the chaos of a multi-station robot into a neat sum of parts. They show that the total confusion caused by turning off parts across many layers can be split into two things: a "same-station" effect (which they understand perfectly) and a "cross-layer" remainder (the messy stuff that happens when Station A changes the input for Station B). They prove that this messy remainder isn't just random noise; it has a precise mathematical shape, like a double integral (a fancy way of saying it's the accumulated effect of two changes happening together).

But here is the twist: while they can describe the shape of the mess, they can't yet put a hard number on how big it gets when you chain many stations together. To do that, they needed a new tool: a "Jacobian bound" for the attention part of the brain. Think of this as a speed limit sign for how fast a change in one clue can ripple through the attention mechanism. The author derived this speed limit in a closed-form formula (a neat, exact equation) and tested it against a real, pre-trained robot brain called Qwen2.5-1.5B-Instruct. They checked 12 specific spots in this real brain, and the speed limit held true every single time—zero violations. However, they are honest about the limit: they verified the speed limit for one station, but they haven't yet proven that if you chain 28 stations together, the errors don't explode into something useless. So, the math is exact, but the final "big number" for the whole chain remains an open question.

To see if their theories work on something real, the author went hunting for a hidden circuit inside the Qwen robot that helps it solve "Indirect Object Identification" (figuring out who gave an item to whom in a sentence like "The ball was given to Mary by John"). They didn't design this circuit; it emerged naturally as the robot learned to speak. They found a small, shared group of neurons that acts as the core of this skill, appearing in almost every test case, plus some extra "backup" neurons that are specific to each sentence.

When they tested the three-way pattern (does turning off the circuit collapse the robot's logic? do the two tricks disagree? is there a messy interaction?), the results were a genuine mix.

Collapse: On four out of five test sentences, turning off the circuit made the robot confused, just as predicted. But on one sentence ("Anna/Mike"), the robot barely cared, showing that these circuits aren't always 100% reliable.
Dissociation: The two tricks (patching vs. turning off) always disagreed, confirming that they are indeed measuring different things.
Interaction: On three of the five sentences, they measured a "non-zero interaction," meaning the parts did mess with each other. Crucially, these interactions happened between layers that were far apart, not right next to each other. This suggests that the messy "cross-layer remainder" the author identified is real and measurable in a real brain, even though the exact mathematical bound for it is still being worked out.

In short, the paper proves that the "messy interaction" between distant parts of a neural network is a real, measurable phenomenon that can be broken down into exact mathematical pieces. They found a new tool to measure one piece of that mess and verified it on a real robot, but the final step of chaining that tool across the entire robot's depth is still a work in progress. The takeaway isn't a solved problem, but a much clearer map of where the mystery lies.

Problem Statement
This paper addresses two specific limitations in a companion work (same author) regarding the causal analysis of neural networks via activation patching and weight-space ablation. The companion paper established an exact first-order interaction formula for two carriers (an attention head and its layer's MLP) within a single residual block of an idealized model. However, it left two critical gaps:

Multi-layer Scope: The companion theorem is silent on interactions between carriers separated by multiple layers, which is the generic case in real transformers. It does not account for how an ablation in an earlier layer propagates to affect the computation of a later layer.
Empirical Scope: The companion paper's results were verified only on small, synthetic transformers trained from scratch for a specific conditional task. It remains unknown whether these theoretical predictions hold for "emergent" circuits found within large, genuinely pretrained models trained on natural language.

Methodology
The paper extends the theoretical framework of the companion work using exact mathematical decompositions and closed-form bounds, then validates these findings on a real pretrained model.

Theoretical Extension (Multi-layer Decomposition): The author derives an exact identity (Proposition 1) showing that the interaction gap for an arbitrary subset of ablated carriers spanning multiple layers decomposes into a sum of "same-block" terms (one per touched layer) plus a "cross-layer remainder" ( $R_\times$ ).
Exact Cross-Layer Identity: For the case of two touched layers, the remainder $R_\times$ is isolated exactly as a double integral of a mixed second partial derivative (Proposition 2), applying the fundamental theorem of calculus one level up.
Closed-Form Bounds:
- Attention Jacobian: To bound the cross-layer remainder, the author derives a closed-form, pointwise bound on the Jacobian of a single attention head's output with respect to a perturbed key or value token (Proposition 3).
- Curvature Constant: The author computes the curvature constant $\Lambda$ required for the companion paper's second-order bound in closed form, using only the trained weights (Proposition 4), rather than assuming it.
Empirical Verification:
- Synthetic Check: The multi-layer decomposition is tested on the five synthetic transformer instances from the companion paper.
- Real Model Check: The author searches for an emergent circuit for Indirect Object Identification (IOI) in Qwen2.5-1.5B-Instruct using a greedy activation-patching search. They then measure the three-way pattern (collapse, dissociation, and interaction) on this circuit.

Key Contributions

Exact Multi-Layer Decomposition: The paper provides an exact decomposition of multi-layer interaction effects into same-block terms and a specific cross-layer remainder. It proves that same-block terms are zero if only MLPs are ablated and equal to the companion paper's bounded term if the layer feeds the readout directly.
Closed-Form Attention Jacobian Bound: The paper derives a closed-form bound for the operator norm of the attention sub-block's Jacobian. This bound is verified pointwise against the real weights of Qwen2.5-1.5B-Instruct with zero violations.
Explicit Curvature Constant: The paper supplies the previously unexhibited curvature constant $\Lambda$ in closed form, computed from the spectral norms of the trained MLP weights and normalization scales.
Empirical Evidence on Pretrained Models: The paper identifies an emergent IOI circuit in a 1.5B-parameter pretrained model and reports mixed empirical results regarding the theoretical predictions of collapse, dissociation, and interaction.

Results

Theoretical: The decomposition identity holds exactly. The derived attention Jacobian bound is verified pointwise on real weights (12 sites, 0 violations), though the author notes that chaining this bound across many layers remains unverified and likely loose due to compounding looseness.
Synthetic Experiments: On the synthetic models, the multi-layer decomposition identity holds to floating-point precision. However, the cross-layer remainder $R_\times$ is often of the same order of magnitude as the same-block sum, and the two terms frequently partially cancel each other out. Thus, the same-block sum alone is not a reliable proxy for the total interaction in multi-layer configurations.
Real Model (Qwen2.5-1.5B-Instruct):
- Circuit Discovery: A shared core (Layer 25/Head 9) and instance-specific support were found for the IOI task.
- Collapse: Observed clearly on 4 of 5 instances, but absent on one (Anna/Mike), where ablating the circuit barely narrowed the decision gap.
- Dissociation: A significant gap between patching and ablation effects was observed on all 5 instances, confirming they are distinct operations.
- Interaction: Measurable non-zero interaction was found on 3 of 5 instances. Crucially, the layer pairs involved in these interactions were not in the single-block configuration proven in the companion theorem. Therefore, these results are suggestive of the multi-layer case but do not constitute a direct test of the single-block theorem.

Significance and Claims
The paper claims to close the gap between single-block theoretical guarantees and the reality of multi-layer transformer architectures, but does so with explicit honesty about remaining limitations.

It establishes that multi-layer interactions can be exactly decomposed, isolating the "cross-layer remainder" as the specific quantity that remains unbounded.
It provides the necessary mathematical ingredients (the attention Jacobian bound and curvature constant) to potentially bound this remainder, verifying them pointwise on real weights.
It demonstrates that the qualitative patterns (collapse, dissociation, interaction) observed in synthetic models do appear in emergent circuits of real pretrained models, but with "genuinely mixed fidelity."
The author explicitly states that they do not claim to have bounded the cross-layer remainder across arbitrary depths, nor do they claim the interaction results on Qwen2.5 prove the single-block theorem applies to multi-layer cases. Instead, they present the mixed results as evidence that the theory is not merely an artifact of synthetic training, while highlighting that the sign and magnitude of multi-layer interactions remain an open, complex question.

More like this