🤖 machine learning

IG-Lens: Exact Additive Probability Attribution Across Transformer Layers via Telescoping Integrated Gradients

This paper introduces IG-Lens, a novel method that achieves exact, additive probability attribution across transformer layers by applying telescoping Integrated Gradients along a single path, thereby overcoming the nonlinearity and bias limitations of existing logit-based or non-additive readout tools while ensuring completeness without discretization error.

Original authors: Duc Anh Nguyen

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Duc Anh Nguyen

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 a large language model (like the one powering this chat) as a multi-story factory where a raw idea enters at the bottom and a finished product (a predicted word) comes out at the top.

For a long time, researchers have tried to answer a simple question: "At exactly which floor of this factory does the decision to say 'Hanoi' actually get made?"

Existing tools have tried to answer this, but they've been like bad accountants:

Some tools look at each floor and guess how "confident" the factory is, but their guesses don't add up to 100%. It's like saying Floor 1 is 40% sure, Floor 2 is 60% sure, and Floor 3 is 80% sure—totaling 180%, which makes no sense.
Other tools look at the raw numbers (logits) before they are turned into percentages. But turning raw numbers into probabilities is like baking a cake; you can't just add up the flour and eggs separately and expect to get the weight of the finished cake. The "baking" (the math called softmax) changes everything.
A third group of tools tries to measure the work done on each floor, but they measure each floor against a different starting point, so their numbers don't stack up to tell you the total work done.

Enter IG-Lens: The Perfect Accountant

The paper introduces IG-Lens, a new method that acts like a perfect accountant for this factory. It solves the problem by using a "telescoping" trick (think of a collapsible telescope that extends and retracts).

Here is how it works in simple terms:

1. The "One-Time Snapshot" Rule
Instead of letting the word travel all the way through the whole factory and getting mixed up with other words at every turn, IG-Lens takes a "snapshot" of the word's state at specific floors. It then asks: "If we took the state from Floor 10 and ran it through just the final 'finishing' machine (the part that turns raw data into a probability), what would the result be?"

2. The Telescoping Sum
IG-Lens breaks the journey into segments (e.g., Floor 10 to 11, 11 to 12, etc.).

It calculates the probability of the word at Floor 10.
It calculates the probability at Floor 11.
The difference between them is the exact amount of work done on that specific segment.
Because it measures the change at every step, when you add up all the changes from the bottom to the top, they perfectly equal the total change from start to finish. There is no missing math, no "baking error," and no rounding issues.

3. The "Prediction-Aware" Filter
Most methods look at how much the numbers wiggle (gradients) to guess importance. But sometimes numbers wiggle a lot without actually changing the final decision.
IG-Lens is smarter. It only credits a floor with work if the actual probability of the word changed. If the numbers wiggle but the prediction stays the same, IG-Lens ignores that wiggle. It's like a manager who only pays workers for the days they actually finished a task, not just for days they were busy moving boxes around.

Why This Matters (According to the Paper)

It's Exact: Unlike previous tools that give you an approximation, IG-Lens guarantees that if you sum up the contributions of every layer, you get exactly the final probability. It's mathematically perfect (down to the limits of computer math).
It Finds the "Onset": It can tell you the exact layer where the model "decided" on a word. For example, it might show that for the word "capital," the decision was mostly made by layer 12, but for "Hanoi," the decision happened much later, around layer 15 or 16.
It's Honest About Limits: The paper admits a trade-off. Because IG-Lens looks at the "snapshot" of a word at a specific layer, it doesn't count the work that layer might do later as the word travels up to the top. It measures: "How much did this layer contribute to the final answer, given where we started?" rather than "What was the total effect of this layer on the whole system?"

The Bottom Line

IG-Lens is a new way to look inside AI models that finally answers the question, "When did the model decide?" with a mathematically perfect "receipt" that adds up to 100%. It doesn't guess; it calculates the exact change in probability layer by layer, filtering out the noise to show you exactly where the decision happened.

The authors plan to test this by "breaking" the model at the layers IG-Lens identifies to see if the model actually stops working correctly there, which would prove the method is finding the real decision points.

Technical Summary: IG-Lens

Problem Statement

Decoder-only transformers are often conceptualized as iterative inference processes where the residual stream refines a running estimate of the next token. A critical interpretability question arises: at which specific layer is the probability of a predicted token actually decided?

Existing tools for layer-wise readout fail to answer this precisely in probability space:

Logit Lens & Tuned Lens: These report the probability level at each layer but do not provide an additive decomposition. Their estimates are biased and non-monotone across depth, meaning the sum of per-layer probabilities has no mathematical meaning.
Direct Logit Attribution (DLA) & Embedding Methods: These are additive but operate in logit space. Because the softmax function is highly nonlinear, an additive decomposition of logits does not translate to an additive decomposition of probabilities, which is the quantity of practical interest.
Layer Conductance & Layer Integrated Gradients (IG): These methods integrate gradients per layer but attribute each layer to its own independent baseline. Consequently, the per-layer terms do not "telescope" (sum up) to the total change in prediction probability.

Methodology: IG-Lens

The authors introduce IG-Lens, a method that provides an exact additive decomposition of probability across transformer layers.

Core Concept

IG-Lens applies a telescoping application of Integrated Gradients (IG) along a single path through the hidden states from a baseline to the final layer.

Path Definition: The method selects a sequence of layers $0 = L_0 < L_1 < \dots < L_k = n$ and defines a path of hidden states $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$ .
Head-Once Readout: Crucially, the readout function $f(h)$ is defined as the final unembedding applied once to a single hidden vector: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$ . This function contains no attention or MLP; all cross-token mixing is assumed to have occurred in the forward pass that produced the detached hidden states.
Telescoping Attribution: The contribution of layer $L_j$ is defined as the segment IG of $f$ along the straight line from $h_{L_{j-1}}$ to $h_{L_j}$ . By the gradient theorem, the integral of the gradient over a segment equals the difference of the endpoint values. Summing these segments yields a telescoping sum:
$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$
This ensures the sum of attributions is exactly the change in target probability, with the softmax and LayerNorm integrated through rather than linearized or frozen.

Prediction-Aware Estimator

While the integral is exact, finite-step estimators usually introduce discretization error. IG-Lens adopts a prediction-aware estimator (inspired by IDGI) as its default:

Instead of crediting steps based on raw gradients, it credits each step by its observed change in target probability ( $f(h^{(s)}) - f(h^{(s-1)})$ ).
For a scalar probability readout, this collapses each segment into a telescoping sum of endpoint values.
Result: Completeness holds to floating-point precision at any step count ( $m \ge 1$ ). Steps where the model shows gradient sensitivity but no output change are naturally filtered out, eliminating Riemann discretization error.

Implementation

The method uses a single-pass batched algorithm that computes the full token-by-layer map without requiring a backward pass (autograd).

By stacking all interpolation points into a single tensor and evaluating the pointwise function $f$ in one forward pass, the method avoids the overhead of gradient computation.
This efficiency is possible specifically because of the "head-once" design, which ensures independence between token points in the batch.

Key Contributions

Exact Additive Decomposition: IG-Lens is the first method to provide an exact additive decomposition of probability (not logits) across layers, where the sum of attributions equals $\Delta p$ .
Telescoping Identity: The paper provides a formal proof (Proposition 1) that the segment integrals telescope to the total change in prediction, keeping the softmax inside the integration path.
Prediction-Aware Estimation: The default estimator ensures machine-precision completeness regardless of step count and filters out "sensitivity without response" steps.
Efficient Implementation: A single-pass forward evaluation computes the full attribution map without backward calls, making it computationally efficient.
Positioning: The authors clearly delineate why existing methods fail in probability space (logit-space additivity, bias, or lack of telescoping) and position IG-Lens as the probability-space counterpart to logit-space methods.

Results and Evaluation Plan

Numerical Verification: On controlled tests and a real model (Llama-3.2-1B-Instruct), the default estimator reproduces per-segment endpoint differences to $\sim 10^{-7}$ (float32 precision). The telescoping sum holds exactly to floating point, independent of the number of steps $m$ .
Illustrative Output: In experiments with the prompt "What is the capital of Vietnam?", the method successfully concentrates attribution mass in the specific layer segments where the logit-lens value transitions from low to high. The sum of attributions for each token matches $\Delta p$ to floating-point precision.
Evaluation Strategy: The authors acknowledge that "onset" has no intrinsic ground truth. Therefore, they propose validating the method via causal ablation: if IG-Lens identifies layer $L$ as the onset, ablating $L$ should depress the target probability more than ablating layers with lower onset scores. This test is designed to be cheap to run and capable of falsifying the claim.

Significance and Claims

The paper positions IG-Lens as a quantitative refinement of the iterative-inference view. Unlike the Logit Lens, which asks "what is the running prediction at this layer?", IG-Lens asks "how much of the final probability is each layer responsible for?" and guarantees these responsibilities sum to the whole.

Modest Claims and Limitations:

Conditional Marginal: The authors explicitly state that IG-Lens measures a conditional marginal: the additional probability that the norm+head can read out of $h_L$ given $h_{L-1}$ . It discards the effect of $h_L$ routed through upper attention and MLP blocks.
Trade-off: The method chooses "sum = $\Delta p$ " over "total causal effect via the full forward pass." For total causal effects, component decomposition (like DLA) remains the appropriate tool.
Non-Monotonicity: The method handles non-monotone probability trajectories (where probability is built up and partially undone) by using signed attributions, offering a more faithful view of layer contributions than absolute-value-based methods.

The central open question remains empirical: whether the IG-Lens onset predicts the layer at which causal ablation most depresses the prediction, particularly for tokens with non-monotone trajectories.