💬 NLP

Hessian-Enhanced Token Attribution (HETA): Interpreting Autoregressive LLMs

This paper introduces Hessian-Enhanced Token Attribution (HETA), a novel framework designed to improve the interpretability of decoder-only autoregressive language models by combining semantic transition vectors, Hessian-based sensitivity scores, and KL divergence to produce more faithful and human-aligned token attributions, alongside a new benchmark for evaluating generative attribution quality.

Original authors: Vishal Pramanik, Maisha Maliha, Nathaniel D. Bastian, Sumit Kumar Jha

Published 2026-04-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Vishal Pramanik, Maisha Maliha, Nathaniel D. Bastian, Sumit Kumar Jha

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 watching a magician pull a rabbit out of a hat. You see the rabbit appear, but you have no idea how the magician did it. Did they use a hidden compartment? A sleight of hand? A secret signal?

In the world of Artificial Intelligence, Large Language Models (LLMs) are the magicians. They write poems, answer questions, and solve problems, but they are "black boxes." We know the input (the question) and the output (the answer), but we don't know which specific words in the question actually caused the model to generate that specific answer.

This paper introduces a new tool called HETA (Hessian-Enhanced Token Attribution) to solve this mystery. Think of HETA as a super-powered detective that can look inside the magician's hat and point exactly to the trick.

Here is how HETA works, broken down into three simple detective tools:

1. The "Causal Map" (Semantic Flow)

The Problem: In a normal conversation, if I say, "The cat sat on the mat," the word "cat" is important. But in a long story, the model might get distracted by irrelevant details. Old methods often just looked at where the model "looked" (attention), which is like guessing the magician's trick by watching where his eyes go. But magicians often look one way to trick you while doing the real work with their hands elsewhere.

The HETA Solution: HETA draws a Causal Map. It traces the actual path of information from the beginning of the sentence to the specific word being predicted. It asks: "Did this specific word actually travel through the model's 'brain' to help create the next word?" If the word didn't take a direct path to the answer, HETA ignores it. It filters out the noise and focuses only on the words that truly matter.

2. The "Curvature Sensor" (Hessian-Based Sensitivity)

The Problem: Imagine a hill. If you are standing on a flat part of the hill, a tiny step doesn't change your height much. But if you are on a steep cliff edge, a tiny step sends you plummeting.
Old AI tools mostly look at the "slope" (gradients). If the ground is flat, they say, "This word doesn't matter." But sometimes, even on flat ground, a small nudge can trigger a massive change because of the shape of the hill underneath.

The HETA Solution: HETA uses a Curvature Sensor. Instead of just looking at the slope, it feels the shape of the hill. It detects those hidden "cliff edges" where a tiny change in a word causes a huge shift in the answer. This helps it find important words that other tools miss because they are hiding in "flat" areas of the math.

3. The "What-If" Simulator (KL Divergence)

The Problem: Sometimes a word seems important, but maybe it's just a coincidence. How do we know for sure?

The HETA Solution: HETA runs a What-If Simulator. It takes the sentence, erases one word, and asks the model: "Okay, now that I removed this word, does the answer change?"

If the model suddenly gives a wrong answer or gets confused, HETA knows that word was critical.
If the model gives the same answer, that word was just decoration.
This is like removing a brick from a wall to see if the wall collapses.

Putting It All Together

HETA combines these three tools into one super-solution:

The Map ensures we only look at words that have a direct connection to the answer.
The Sensor finds the hidden, complex relationships that simple math misses.
The Simulator proves that the word actually changes the outcome.

Why Does This Matter?

Before HETA, trying to explain why an AI made a decision was like guessing the ingredients of a soup by tasting the final dish. You might get close, but you'd often be wrong.

With HETA, we can finally say: "The AI predicted 'slice' because of the words 'pizza,' 'knife,' and 'cut,' and it ignored the word 'red' even though it was in the sentence."

This makes AI:

Trustworthy: We know why it made a decision.
Safe: If an AI makes a mistake in a hospital or a bank, we can trace it back to the exact word that caused the error and fix it.
Better: By understanding how the model thinks, we can build smarter models in the future.

In short, HETA turns the "black box" of AI into a glass box, letting us see exactly how the magic happens.

1. Problem Statement

The paper addresses the critical challenge of interpreting autoregressive (decoder-only) Large Language Models (LLMs). While Explainable AI (XAI) methods like Integrated Gradients, LIME, and attention rollout exist, they face significant limitations when applied to generative models:

Architectural Mismatch: Most methods are designed for encoder-based models (bidirectional) and fail to account for the causal, autoregressive nature of decoder-only models where future tokens cannot influence past ones.
First-Order Limitations: Standard gradient-based methods rely on linear approximations. They often fail in "flat" regions of the loss landscape (e.g., saturated ReLU or GELU activations) where gradients vanish, yet the token still has a significant causal impact on the output.
Attention Misinterpretation: Attention weights indicate where a model "looks" but do not necessarily reflect causal influence, as they ignore the magnitude of value vectors and non-linear interactions.
Instability: Existing methods often produce inconsistent attributions when decoding hyperparameters (temperature, top-p) change or when inputs are syntactically rephrased.

2. Methodology: Hessian-Enhanced Token Attribution (HETA)

HETA is a novel framework designed specifically for decoder-only LLMs. It decomposes token attribution into three complementary components to capture causal directionality, non-linear sensitivity, and information-theoretic impact.

The final attribution score for a token $x_i$ contributing to a target token $x_T$ is defined as:
$\text{Attr}(x_i \to x_T) = M_T[i] \cdot \left( \beta S^{(T)}_i + \gamma I(x_i \to x_T) \right)$

Key Components:

Semantic Transition Influence (Causal Gate, $M_T[i]$ ):
- Goal: Enforce target-conditioned causality.
- Mechanism: It traces attention-weighted value flows across layers, but strictly aggregates only paths that terminate at the target token $T$ (respecting the causal mask).
- Formula: It computes a rollout of attention matrices and value vectors, normalized to form a simplex. This ensures that tokens without a causal path to the target receive zero attribution, filtering out spurious correlations.
Hessian-Based Sensitivity ( $S^{(T)}_i$ ):
- Goal: Capture second-order (non-linear) effects that gradients miss.
- Mechanism: It estimates the curvature of the log-likelihood surface with respect to input embeddings.
- Implementation: Since computing the full Hessian is infeasible, HETA uses Hessian-Vector Products (HVPs) via Hutchinson estimators (using Rademacher vectors) and Pearlmutter's trick. It optionally employs low-rank approximations for scalability.
- Significance: This captures influence in regions where gradients are zero but finite perturbations cause significant output changes (e.g., in saturated activation regions).
KL Divergence for Information Contribution ( $I(x_i \to x_T)$ ):
- Goal: Measure the information-theoretic impact of a token.
- Mechanism: It calculates the Kullback-Leibler (KL) divergence between the target token's probability distribution with the original context and the distribution when token $x_i$ is masked.
- Significance: This quantifies the "surprisal" or information loss caused by removing a token, providing a probabilistic grounding for its importance.

Optimization: The weights $\beta$ and $\gamma$ balance the curvature sensitivity and information gain. The method is implemented with windowing and low-rank approximations to handle long contexts efficiently.

3. Key Contributions

Novel Framework (HETA): The first attribution method to explicitly combine semantic flow tracing, second-order Hessian sensitivity, and KL-based information loss tailored for decoder-only architectures.
Curated Benchmark Dataset: Introduction of a new dataset (2,000 instances) combining NarrativeQA (distractor context) and SCIQ (answer-bearing context). This allows for systematic evaluation of whether attribution methods correctly identify predictive evidence while ignoring semantically rich but irrelevant distractors.
New Evaluation Metric (DSA): The Dependent Sentence Attribution (DSA) metric, which quantifies the degree to which attribution mass is concentrated on the correct evidence span versus distractor spans.
Comprehensive Empirical Validation: Extensive experiments across four diverse models (GPT-J, LLaMA-3.1, Phi-3, Qwen2.5) and multiple benchmarks, demonstrating superior performance over strong baselines.

4. Experimental Results

HETA was evaluated against state-of-the-art baselines including Integrated Gradients (IG), Attention Rollout, ContextCite, TDD, Peering, and ReAgent.

Faithfulness (Robustness): HETA achieved the highest Soft-NC and Soft-NS scores (measuring stability under input perturbations). For example, on GPT-J, HETA scored 10.3 on LongRA, outperforming the next best method (ReAgent) by over 2x.
Alignment (Human/Annotation): On the curated dataset, HETA achieved a DSA score of 4.80 (GPT-J) and 5.10 (LLaMA-3.1), significantly outperforming all baselines (which often scored negative values, indicating they focused on distractors).
Robustness to Hyperparameters: HETA showed remarkable stability (<1% variation) when decoding parameters (temperature, top-p, top-k) were varied. In contrast, baselines fluctuated by 2–5%.
Robustness to Syntax: HETA maintained high consistency under active/passive voice rephrasing and Gaussian noise perturbations.
Ablation Studies: Removing any single component (Semantic Flow, Hessian, or KL) resulted in significant performance drops, confirming that all three are necessary for faithful attribution.

5. Significance and Impact

Theoretical Advancement: HETA moves beyond linear approximations, providing a mathematically grounded approach that accounts for the non-linear, high-dimensional nature of LLM decision surfaces.
Practical Reliability: By decoupling attribution from stochastic decoding heuristics, HETA offers a more reliable tool for debugging, safety auditing, and understanding model behavior in high-stakes domains.
Scalability: Despite the computational cost of Hessian estimation, the authors demonstrate that low-rank approximations and windowing allow HETA to scale to long contexts (2k+ tokens) with manageable overhead (approx. 1.4x runtime compared to gradient methods, but with vastly superior quality).
Standard Setting: The paper establishes a new standard for interpretability in generative models, highlighting that future attribution methods must integrate causal gating and higher-order sensitivity to be truly faithful.

In conclusion, HETA represents a significant leap forward in LLM interpretability, solving the "black box" problem for autoregressive models by combining causal semantics, geometric curvature, and information theory into a unified, robust attribution framework.