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

Vishal Pramanik; Maisha Maliha; Nathaniel D. Bastian; Sumit Kumar Jha

cs.CLApr 14, 2026

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

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

Organizations: Department of Computer & Information Science & Engineering, University of Florida, Gainesville, FL 32611, USA · School of Computer Science, University of Oklahoma, Norman, OK 73019, USA · Department of Electrical Engineering and Computer Science, United States Military Academy, West Point, NY 10996, USA

Abstract

Attribution methods seek to explain language model predictions by quantifying the contribution of input tokens to generated outputs. However, most existing techniques are designed for encoder-based architectures and rely on linear approximations that fail to capture the causal and semantic complexities of autoregressive generation in decoder-only models. To address these limitations, we propose Hessian-Enhanced Token Attribution (HETA), a novel attribution framework tailored for decoder-only language models. HETA combines three complementary components: a semantic transition vector that captures token-to-token influence across layers, Hessian-based sensitivity scores that model second-order effects, and KL divergence to measure information loss when tokens are masked. This unified design produces context-aware, causally faithful, and semantically grounded attributions. Additionally, we introduce a curated benchmark dataset for systematically evaluating attribution quality in generative settings. Empirical evaluations across multiple models and datasets demonstrate that HETA consistently outperforms existing methods in attribution faithfulness and alignment with human annotations, establishing a new standard for interpretability in autoregressive language models.

Figures & tables

Figure 1: Overview of HETA. The pipeline (a) rolls out attention–value flows that end at the target token to form a causal gate over input tokens, (b) estimates token-level curvature via scalable Hessian–vector products to capture nonlinear interactions, and (c) measures KL-based information impact under token masking. The final attribution combines causal gating, curvature sensitivity, and information gain to produce target-conditioned, token-level explanations.

Method	GPT-J	LLaMA	Phi-3	Qwen2.5
IG	-0.34	-0.28	-0.41	-0.31
ContextCite	-0.12	-0.09	-0.18	-0.14
Peering (PML)	-0.25	-0.21	-0.30	-0.22
TDD-backward	-0.31	-0.27	-0.36	-0.29
Attention Rollout	-0.44	-0.39	-0.52	-0.41
fAML	2.10	2.30	2.05	2.20

Table 1: Attribution alignment on the curated dataset using the (Dependent Sentence Attribution) metric. We evaluate robustness across both model size and architecture by testing HETA on diverse decoder-only LMs, spanning different parameter scales and design choices. Higher scores indicate stronger alignment with human-annotated tokens. HETA outperforms all baselines.

Panel of Figure 2 — Figure 2: (a)-(c) Analysis of HETA components. Each bar plot shows the effect of ablating key components of HETA. The full HETA model achieves the highest attribution faithfulness and alignment across all metrics, while removing individual components consistently degrades performance.(d) Input importance distributions for a generative task using our proposed HETA method.

Panel of Figure 3 — Figure 3: (a)-(c) Analysis of robustness of HETA vs baseline methods (Left) Sensitivity under Gaussian perturbations (lower is better), where HETA maintains the lowest variance across input noise. (Center) Active/Passive robustness (higher is better), reflecting attribution consistency across syntactic rephrasings. (Right) Alignment F1 score against annotated tokens (higher is better). HETA outperforms all baselines, validating the complementary role of transition flow, curvature, and information gain.(d) Input importance distributions for a generative task using our proposed HETA method using Qwen 2.5 3B.

Model	Metric	HETA	ContextCite	IG	PML	TDD-bw	AttnRoll	fAML	ProgInf	SEA-CoT	ReAgent
Llama-3.1 70B	Soft-NC	0.6	2.6	2.7	2.4	3.1	2.3	2.5	2.4	2.5	2.3
	Soft-NS	0.8	3.1	3.0	2.8	3.4	2.7	2.9	2.8	3.0	2.6
	DSA	0.6	2.4	2.3	2.2	2.8	2.1	2.3	2.3	2.4	2.2

Table 2: Sensitivity of attribution metrics to decoding hyperparameters (max relative change

\Delta\%

; lower is better). HETA varies

<\!1\%

across all models/metrics; each baseline fluctuates

>\!2\%

. Grid: temperature

\in\{0.2,0.5,0.9\}

, top-

p\in\{0.8,0.9,0.95\}

, top-

k\in\{20,50,100\}

, repetition penalty

\in\{1.0,1.2\}

; 3 seeds.

Attribution Method	LongRA		TellMeWhy		WikiBio
	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$
GPT-J 6B
ContextCite	1.42	0.03	1.46	-0.22	0.49	-0.08
Integrated Gradients	1.87	0.45	1.54	0.04	1.38	0.77
Peering (PML)	2.05	0.50	1.68	0.06	1.50	0.83
TDD-backward	1.10	-0.12	1.89	-0.03	0.11	0.51

Table 3: Attribution faithfulness on benchmark datasets (LongRA, TellMeWhy, WikiBio) using GPT-J 6B, Phi-3-Medium-14B, and LLaMA-3.1 70B. Evaluated with Soft-NC and Soft-NS; higher is better. Mean of 3 runs; std

<\pm 0.06

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Attribution Method	LongRA		TellMeWhy		WikiBio
	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$	Soft-NC $\!\uparrow$	Soft-NS $\!\uparrow$
Qwen2.5 3B
Integrated Gradients	1.20	0.14	1.30	0.09	1.10	0.18
ContextCite	1.90	0.56	1.65	0.47	1.45	0.79
Peering (PML)	2.05	0.61	1.80	0.50	1.60	0.85
TDD-backward	1.05	-0.08	1.82	0.00	0.12	0.49

Appendix

Table 4: Attribution faithfulness on Qwen2.5 3B across LongRA, TellMeWhy, and WikiBio using Soft-NC and Soft-NS. Higher scores indicate stronger robustness and alignment. Mean of 3 runs; std

<\pm 0.06

Configuration	Soft-NC	Soft-NS	DSA
HETA (Full)	9.78	2.31	4.70
Transition Only	3.12	1.52	2.21
Hessian Only	2.89	1.45	2.97
KL Only	2.23	1.21	2.74
No Transition Gating	4.31	1.84	1.68
Uniform Transition	3.89	1.76	1.54

Appendix

Table 5: Ablation study of HETA components. Reported values are averaged across all datasets for GPT-J 6B. Mean over independent 3 runs and std

<\pm 0.2

$\beta$ (Hessian)	$\gamma$ (KL)	Faithfulness $\downarrow$	Sensitivity $\downarrow$	Robustness $\uparrow$	F1 (Alignment) $\uparrow$
0.0	1.0	0.179	0.022	0.70	0.82
0.2	0.8	0.143	0.027	0.78	0.84
0.5	0.5	0.108	0.025	0.91	0.89
0.8	0.2	0.124	0.041	0.86	0.81
1.0	0.0	0.254	0.087	0.54	0.68

Appendix

Table 6: Performance of HETA for different

(\beta,\gamma)

under the constraint

\beta+\gamma=1

. Lower is better for faithfulness/sensitivity; higher is better for robustness/F1. Mean over 3 runs; std

<\pm 0.05

Variant	Sensitivity $\downarrow$	Act./Pass. Robustness $\uparrow$	F1 (Alignment) $\uparrow$
Full HETA	0.025	0.91	0.89
Transition Only	0.031	0.72	0.76
Hessian Only	0.087	0.54	0.68
KL Only	0.022	0.70	0.82
No Transition Gating	0.049	0.68	0.79
Uniform Transition	0.038	0.61	0.74

Appendix

Table 7: Ablation of the HETA framework. Lower sensitivity and higher robustness/F1 indicate better attribution quality. Mean over 3 runs and std

<\pm 0.04

Variant	Runtime (s)	AOPC $\uparrow$	F1 $\uparrow$
HETA (Full)	455	0.61	0.89
HETA-LR (rank=64)	330	0.59	0.86
HETA-LS (6 layers)	305	0.57	0.84
HETA-GS (grad-squared only)	240	0.52	0.81

Appendix

Table 8: Runtime and attribution quality for Hessian approximations on 1,000 examples (sequence length = 512). AOPC and F1 represent attribution faithfulness and human alignment, respectively.

Variant	AOPC $\uparrow$	Runtime (s)
HETA (Full)	0.58	1,230
HETA-WIN (512-window)	0.54	690
HETA-LR+WIN	0.55	580
HETA-LS+WIN	0.52	525

Appendix

Table 9: Faithfulness and runtime for long-context attribution (sequence length = 2,048). Windowed methods use 512-token chunks with 50% overlap. Mean over 3 runs and std

<\pm 0.04

Model	Metric	HETA	ContextCite	IG	PML	TDD-bw	AttnRoll	fAML	ProgInf	SEA-CoT	ReAgent
GPT-J 6B	Soft-NC	0.8	3.0	3.2	2.6	3.8	2.5	2.9	2.7	2.8	2.4
	Soft-NS	0.9	3.6	3.4	3.0	4.2	2.9	3.1	3.0	3.2	2.7
	DSA	0.7	2.7	2.5	2.3	3.1	2.2	2.4	2.6	2.5	2.1
Llama-3.1 70B	Soft-NC	0.6	2.6	2.7	2.4	3.1	2.3	2.5	2.4	2.5	2.3
	Soft-NS	0.8	3.1	3.0	2.8	3.4	2.7	2.9	2.8	3.0	2.6
	DSA	0.6	2.4	2.3	2.2	2.8	2.1	2.3	2.3	2.4	2.2

Appendix

Table 10: Sensitivity of attribution metrics to decoding hyperparameters (max relative change

\Delta\%

; lower is better). HETA varies

<\!1\%

across all models/metrics; each baseline fluctuates

>\!2\%

. Grid: temperature

\in\{0.2,0.5,0.9\}

, top-

p\in\{0.8,0.9,0.95\}

, top-

k\in\{20,50,100\}

, repetition penalty

\in\{1.0,1.2\}

; 3 seeds.

Variant	AOPC $\uparrow$	DSA $\uparrow$	Runtime $\downarrow$	Memory $\downarrow$ $\widehat{\varepsilon}_{M}(W)\,\downarrow$
HETA (Full)	0.60	4.68	1,280	1.00 $\times$ 0.000
HETA-LR (rank=64)	0.58	4.55	850	0.78 $\times$ 0.006
HETA-WIN ( $W{=}512$ , 50% overlap)	0.56	4.40	690	0.72 $\times$ 0.021
HETA-LS (top 6 layers)	0.55	4.30	760	0.75 $\times$ 0.012
HETA-LR+WIN (r=64, $W{=}512$ )	0.59	4.58	610	0.66 $\times$ 0.010

Appendix

Table 11: Compute–quality ablation for HETA (long contexts). Backbone: GPT-J 6B . Results are averaged over LongRA , TellMeWhy , and WikiBio with sequence length 2,048 (batch size 8). “Runtime” is seconds per 1,000 examples; “Memory” is peak activation memory relative to Full;

\widehat{\varepsilon}_{M}(W)

is an empirical window-leakage proxy (lower is better). Mean of 3 runs; std

<\!\pm 0.02

for AOPC/DSA and

<\!\pm 20

s for runtime. The same ranking and gaps are observed on Llama-3.1 70B when restricting curvature to the last 6 layers (omitted for brevity).

GPT-J 6B
Method	Patch-Corr $\uparrow$	Intervention@5 $\uparrow$	Misattr. $\downarrow$
HETA (Full)	0.78 $\pm$ 0.03	0.143 $\pm$ 0.012	0.12 $\pm$ 0.02
HETA – No-Gate	0.53 $\pm$ 0.04	0.089 $\pm$ 0.010	0.27 $\pm$ 0.03
Attention Rollout	0.49 $\pm$ 0.04	0.081 $\pm$ 0.011	0.30 $\pm$ 0.03
Token Gradients	0.41 $\pm$ 0.05	0.067 $\pm$ 0.009	0.35 $\pm$ 0.04
Causal Tracing (baseline)	0.62 $\pm$ 0.03	0.105 $\pm$ 0.011	0.21 $\pm$ 0.03

Appendix

Table 12: Intervention-based causal sanity check. Mean over 2,000 examples;

\pm

denotes std across examples. Top-

k{=}5

. HETA’s target-conditioned gate yields the strongest alignment with interventional effects (highest Patch-Corr and Intervention@5, lowest Misattribution).

Method / Variant	Causal Pass@5 $\uparrow$	CF-AOPC $\uparrow$	Fact EM $\Delta$ $\uparrow$	Tool Hit@1 $\Delta$ $\uparrow$	Span-F1 $\uparrow$	Decoding $\Delta\%$ $\downarrow$
HETA (Full)	0.86	0.63	+3.7	+4.1	0.81	0.8
HETA (LR+WIN)	0.84	0.60	+3.3	+3.8	0.78	0.9
HETA (w/o Hessian)	0.77	0.53	+2.5	+2.7	0.72	1.6
HETA (w/o KL)	0.73	0.49	+2.0	+2.3	0.69	1.8
HETA (w/o Transition)	0.70	0.45	+1.6	+1.9	0.64	2.1
ContextCite	0.58	0.36	+1.2	+1.4	0.55	2.9

Appendix

Table 13: Causal and downstream robustness ablation. Higher is better for Causal Pass@5, Counterfactual AOPC, Fact-checking EM

\Delta

, Tool Hit@1

\Delta

, and Span-F1. Lower is better for Decoding Stability

\Delta\%

. Results averaged over four models (LLaMA-3.1 70B, Phi-3 14B, GPT-J 6B, Qwen2.5 3B); std

<\!\pm 0.05

for additive metrics and

<\!\pm 0.2

pp for

\Delta\%

Figure 4: Word-level attribution visualization for predicting the final word “slice.” Each word is shaded based on its importance score for predicting “slice.” Darker red indicates higher attribution.

Figure 5: Word-level attribution visualization for predicting the final word “friends.” Bounding boxes highlight influential context words (e.g., “shared,” “pictures,” “zoo”) contributing to the prediction of “friends.” Darker red denotes higher importance.

Figure 6: Word-level attribution visualization for predicting the final word “bush.” Attribution scores emphasize context words such as “hat,” “stuck,” and “park,” which strongly influence the prediction of “bush.”

Method	AOPC $\uparrow$	DSA $\uparrow$	Runtime (s) $\downarrow$	Peak Mem (GB) $\downarrow$
HETA (Full)	0.61	4.70	455	53.2
HETA–LR (rank=64)	0.59	4.55	330	41.6
HETA–LS (4 layers)	0.57	4.38	305	39.9
HETA–WIN (512/50%)	0.55	4.29	295	34.1
HETA–LR+WIN	0.58	4.52	245	31.7

Appendix

Table 14: GPT-J 6B @

L{=}1024

B{=}8

(1,000 ex). Mean over 3 runs; std

<\!\pm 0.04

Method	AOPC $\uparrow$	DSA $\uparrow$	Runtime (s) $\downarrow$	Peak Mem (GB) $\downarrow$
HETA (Full)	0.60	4.85	520	47.0
HETA–LR (rank=64)	0.58	4.68	375	39.0
HETA–LS (4 layers)	0.56	4.55	345	37.2
HETA–WIN (512/50%)	0.55	4.47	325	33.1
HETA–LR+WIN	0.57	4.63	275	30.2

Appendix

Table 15: Phi-3 14B @

L{=}1024

B{=}8

(1,000 ex). Mean over 3 runs; std

<\!\pm 0.05

Method	AOPC $\uparrow$	DSA $\uparrow$	Runtime (s) $\downarrow$	Peak Mem (GB) $\downarrow$
HETA (Full)	0.59	4.40	310	24.0
HETA–LR (rank=64)	0.58	4.30	235	20.1
HETA–LS (4 layers)	0.56	4.18	220	19.2
HETA–WIN (512/50%)	0.54	4.12	205	17.3
HETA–LR+WIN	0.56	4.26	180	16.1

Appendix

Table 16: Qwen2.5 3B @

L{=}1024

B{=}8

(1,000 ex). Mean over 3 runs; std

<\!\pm 0.05

Method	AOPC $\uparrow$	DSA $\uparrow$	Runtime (s) $\downarrow$	Peak Mem (GB) $\downarrow$
HETA (Full, 6 layers)	0.60	5.10	1180	74.5
HETA–LR (rank=64)	0.58	4.92	860	61.3
HETA–WIN (512/50%)	0.56	4.80	720	55.8
HETA–LR+WIN	0.57	4.96	620	49.7

Appendix

Table 17: Llama-3.1 70B @

L{=}2048

B{=}4

(500 ex). Curvature on last 6 layers. Mean over 3 runs; std

<\!\pm 0.05

Method	MoRF $\uparrow$	LeRF $\downarrow$	ABPC $\uparrow$	AOPC-Del $\uparrow$	AOPC-Ins $\uparrow$
Integrated Gradients	40.8	43.2	0.35	0.52	0.42
Attention Rollout	41.3	42.9	0.40	0.53	0.43
ContextCite	41.7	42.5	0.45	0.54	0.44
ReAGent	42.6	41.9	0.52	0.57	0.47
GiLOT ( Li et al., 2024 )	43.1	41.3	0.58	0.59	0.49
HETA (Ours)	44.0	40.2	0.65	0.60	0.51

Appendix

Table 18: Extended faithfulness comparison on GPT-J 6B / LongRA using classical perturbation metrics. HETA achieves the highest scores across all metrics. GiLOT is consistently the second-best method. Mean of 3 runs; std

<\pm 0.04

Feature Attribution

Feature attribution estimates which input features or internal components contribute to a model’s predictions.

See all 231 papers

Data Attribution in Large Language Models via Bidirectional Gradient Optimization
Jun 3, 2026Frédéric Berdoz, Luca A. Lanzendörfer, Kaan Bayraktar +1Gradient-Based Attribution LLM Interpretability
ETH Zurich
The Attribution Contract: Feature Attribution for Generative Language Models
May 21, 2026Giang NguyenFeature Attribution LLM Interpretability
Guide Labs
Probabilistic Attribution For Large Language Models
May 20, 2026Shilpika Shilpika, Carlo Graziani, Bethany Lusch +2LLM Interpretability LLM Uncertainty Estimation
Argonne Leadership Computing Facility Argonne National Laboratory Lemont, IL, USA · Mathematics and Computer Science Division Argonne National Laboratory Lemont, IL, USA · Department of Computer Science University of Illinois Chicago Chicago, IL, USA

Jun 3, 2026J/K move · Enter open · S save

Data Attribution in Large Language Models via Bidirectional Gradient Optimization

Frédéric Berdoz, Luca A. Lanzendörfer, Kaan Bayraktar, Roger Wattenhofer

ETH Zurich

Gradient-Based Attribution LLM Interpretability Training Data Attribution

Open paper AskPDF arXiv

Large Language Models (LLMs) are increasingly deployed across diverse applications, raising critical questions for governance, accountability, and data provenance. Understanding which training data most influenced a model's output remains a fundamental open problem. We address this challenge through training data attribution (TDA) for auto-regressive LLMs by expanding upon the inverse formulation: How would training data be affected if the model had seen the generated output during training? Our method perturbs the base model using bidirectional gradient optimization (gradient ascent and descent) on a generated text sample and measures the resulting change in loss across training samples. Our framework supports attribution at arbitrary data granularity, enabling both factual and stylistic attribution. We evaluate our method against baselines on pretrained models with known datasets, and show that it outperforms previous work on influence metrics, thereby enhancing model interpretability, an essential requirement for accountable AI systems.

Abstract

Feature Attribution

Explore similar work

Data Attribution in Large Language Models via Bidirectional Gradient Optimization

The Attribution Contract: Feature Attribution for Generative Language Models

Probabilistic Attribution For Large Language Models