Source-linked AI summary
Temporally-Resolved Token Attribution Reveals the Generation Dynamics of Diffusion Language Models
Darpan Aswal, Céline Hudelot
TL;DR
DLMs generate gradually, so explaining them requires locating prompt support across layers and denoising steps, a gap left open by existing attribution tools. DLIG applies Integrated Gradients across that layer–step grid; across three tasks, it reveals how models draw on inputs and provides an inexpensive first check before interventional validation.
Problem
DLM explanations need to locate support across depth and denoising time, but existing token-attribution tools lack denoising-step resolution and no concrete method instantiates the iterative-generation attribution contract.
Source: §1 · p. 1Method
Results
Across three tasks, DLIG finds similar attribution depths for Yes and No responses, preserved localization of the right fact chain in errors, and use of both sides of infilling gaps, with context reliance correlated with quality only late in denoising.
Source: §8 · p. 8Takeaways & Limitations
DLIG makes the evolution of output commitment and its input support observable across denoising, offering an inexpensive first check before costlier interventional validation.
Source: §8 · p. 8DLIG measures gradient sensitivity rather than causal circuits and is instantiated only on one 355M discrete masked-diffusion model; extension to other families and larger natively trained DLMs remains future work.
Source: §9 · p. 8
Abstract
from arXiv · showhide
This work presents Diffusion Layer Integrated Gradients (DLIG), a token attribution method for diffusion language models (DLMs) that extends Integrated Gradients (IG~\cite{sundararajan2017axiomatic}) to arbitrary layers and denoising steps. DLIG attributes a DLM's progressive commitment to a self-generated or fixed completion for an input prompt. We establish direct correspondences between DLIG and the IG axioms of completeness, implementation invariance, linearity, and symmetry preservation. As a lightweight complement to interventional analysis, DLIG provides an inexpensive first check of mechanistic hypotheses across the denoising trajectory. We demonstrate this on word-sense disambiguation, multi-hop graph reasoning, and sentence infilling, revealing how DLMs draw on inputs across positions, layers, and denoising steps.
1 Introduction
Because DLMs form outputs gradually, explanations must locate supporting computation across both network depth and denoising time. DLIG extends Integrated Gradients to score prompt support at arbitrary layers and denoising steps, and the paper demonstrates it on three tasks.
1 Introduction
- DLM explanations must locate support across network depth and denoising time because tokens form gradually over iterative denoising.The same output token is arrived at progressively, with supporting computation distributed over both dimensions.Source: §1 · p. 1
- Post-hoc DLM attribution remains largely unaddressed: existing token tools lack denoising-step resolution, while adapted interventional analyses are expensive.A recent taxonomy formalizes an attribution contract for iterative generation, but no concrete method instantiates it.Source: §1 · p. 1
- DLIG scores commitment to a self-generated completion at a chosen layer and step, integrating gradients from a prompt-free baseline to the prompt-present state.Across a layer–step grid, its IG correspondence preserves axioms at each coordinate while exposing trajectory structure; demonstrations cover word-sense disambiguation, multi-hop reasoning, and infilling.Source: §1 · p. 1
2 Background and Related Work
DLMs generate through iterative refinement and span continuous-space, discrete-space, and hybrid AR–diffusion families. Prior explainability work leaves token attribution across layers and denoising steps open, and DLIG complements rather than replaces causal interventions.
2 Background and Related Work
- DLMs generate through iterative refinement and fall into continuous-space, discrete-space, and hybrid AR–diffusion families.The paper describes DLMs as achieving competitive performance with substantially higher throughput than standard next-token prediction models.Source: §2 · p. 2
- Prior diffusion explainability focuses mainly on vision, while DLM token-level attribution across layers and denoising steps remains open.A taxonomy formalizes attribution contracts for iterative generation, but no concrete method instantiates them.Source: §2 · p. 2
- DLIG is a lightweight complementary first check that can identify or rule out candidate mechanisms before more costly causal validation.The paper distinguishes this role from interventional methods such as activation patching and causal tracing.Source: §2 · p. 2
3 Diffusion Layer Integrated Gradients (DLIG)
DLIG measures how prompt-conditioned activations support a completion at a chosen layer and denoising step, using Integrated Gradients along a path from prompt-free to prompt-present activations. Its construction carries IG’s stated axiomatic guarantees to each layer–step coordinate.
3 Diffusion Layer Integrated Gradients (DLIG)
- At a fixed step, DLIG scores the model’s commitment to target tokens using summed target log-probability, defaulting to its self-generated completion or allowing a fixed target.The self-generated framing targets the model’s actual behavior along its denoising trajectory.Source: §3 · pp. 2–3
- DLIG compares prompt-conditioned activations with activations under a null prompt while holding generated tokens fixed, then integrates score gradients along their straight-line path.The resulting per-position, per-hidden-dimension attributions are collapsed across hidden dimensions to yield a score per token.Source: §3 · p. 3
- A token’s positive attribution indicates that it drives the model toward the target, while a negative attribution indicates suppression.The summed per-token attribution decomposes the prompt’s causal effect on the scored function; the integral is approximated with an m-point Riemann sum.Source: §3 · p. 3
3.1 Extending Axiomatic Properties of IG [1] to DLIG
- For WiC and infilling, the score is length-normalized and mean-centered because target length grows; ProsQA uses the summed score because its target length is step-invariant.The paper also notes that completeness implies the Sensitivity axiom.Source: §3.1 · pp. 3–4
- At each layer–step coordinate, DLIG inherits IG’s completeness, implementation invariance, linearity, and symmetry preservation guarantees.Completeness holds up to the O(1/m) error from the m-point Riemann sum; implementation checks are also reported as prerequisites for the correspondence in practice.Source: §3.1 · pp. 3–4
4 Experimental Setup
The experiments use DiffuGPT-M across word-sense disambiguation, multi-hop reasoning, and sentence infilling. Task-specific configurations and denoising-step grids are provided in an appendix; the main text reports qualitative findings, with quantitative results and tests in Appendix B.
4 Experimental Setup
- All three experiments use DiffuGPT-M, a masked diffusion language model adapted from GPT-2-medium.Task-specific configurations, training details, and attribution step grids are in Appendix C.Source: §4 · p. 4
- WiC asks whether a target word has the same sense in two contexts, with Yes/No answers and a re-split training set because official test labels are private.The re-split contains 4,928 training, 638 validation, and 500 test examples.Source: §4 · p. 4
- ProsQA tests binary choices whose gold option is reachable through a multi-hop chain in ontology facts that include off-path distractors.DiffuGPT-M is finetuned on a split of 17,886 train, 300 validation, and 500 test examples.Source: §4 · p. 4
- ROCStories infilling evaluates sentence-three generation from two preceding and two following sentences across 1,000 samples.Bidirectional context lets DLIG resolve how attribution mass splits between left and right context over denoising steps.Source: §4 · p. 4
- The main text reports findings qualitatively, while Appendix B contains full quantitative results, effect sizes, and statistical significance tests.Source: §4 · p. 4
5 Word-Sense Disambiguation
In word-sense disambiguation, Yes- and No-responses draw on essentially the same attribution depth, but differ in when and how strongly they commit. A single correctly answered example shows attribution concentrated on sense-bearing words, with an early opposing signal from “France” that flips to support deeper in the network.
5 Word-Sense Disambiguation
- The analysis compares mass-weighted attribution depth with the earliest step when answer probability exceeds 0.5 and stays above it.Depth averages the mass-weighted layer index over steps; label separation is reported as a Mann–Whitney AUC.Source: §5 · pp. 4–5
- In one correctly answered example, attribution centers on the two occurrences of “channel” and surrounding content, while “France” shifts from shallow-layer opposition to mid-layer support.Attribution is largest in shallow layers and decreases with depth; function words and prompt scaffolding stay near zero.Source: §5 · p. 5, Figure 1 · p. 5
- Yes- and No-responses overlap almost entirely in attribution depth, centering just above layer 7 of the 24-layer stack.This provides inexpensive evidence against the hypothesis that the labels draw on prompt evidence at systematically different depths.Source: §5 · p. 5
- The illustrated attribution map is a hand-picked clean example, not a representative sample of cases.The figure also hides tokens below a relative attribution threshold, retaining at least three largest tokens.Source: §5 · p. 5
- Yes-responses are confident from the first denoising step, whereas No-responses begin barely above threshold and slightly degrade toward the end.The per-example commitment is markedly stronger for Yes than for No.Source: §5 · p. 6, Figure 2 · p. 6
6 Multi-Hop Graph Reasoning
On ProsQA, DLIG localizes answer-driving support along the reasoning graph, including when the model answers incorrectly. The analysis compares gold-path support across outcomes and tracks how attribution varies by layer, denoising step, and hop.
Setup
- The analysis contrasts gold and incorrect answer sentences, uses contrastive-score signs to identify support for each option, and measures positive support on gold-path facts against chance and a matched random-edge null.Successes (n=286) and failures (n=69) are analyzed after excluding off-manifold generations (n=145).Source: §6 · p. 6
Results
- Positive attribution concentrates on gold-path tokens above chance and a random-edge null, localizes by layer and denoising step, and rises toward the answer edge.The pattern rises toward later layers, decays over denoising, and increases from root fact to answer edge at every path length.Source: §6 · p. 6, Figure 3 · p. 7
- On failures, attribution concentrates on the model’s cited fact chain; where that chain diverges from the gold path, the two are indistinguishable.Source: §6 · p. 6
- Gold-path precision is nearly identical for successes and failures, while contrastive magnitude shows no significant outcome discrimination, including after difficulty controls.The joint logistic model finds no signal in gold-path precision.Source: §6 · p. 7, Figure 4 · p. 7
7 Sentence Infilling
In sentence infilling, context reliance becomes associated with final infill quality as denoising proceeds. Attribution draws on both sides of the span, with stronger pull from the left and peaks near its boundaries.
Results
- The association between context reliance and infill quality is absent early and strengthens through denoising; pooling steps hides this temporal pattern.Infill quality is measured by final-span ROUGE-1 against the gold span.Source: §7 · p. 7, Figure 5 · p. 8
- Early denoising steps have a reduced sample because steps with no scoreable committed span token are excluded.Source: §7 · p. 7
- Attribution draws on both sides of the span, peaks at each boundary, decays outward, and remains stronger on the left; context reliance grows as the model commits later.Neither context-mass amount nor direction tracks infill quality; the attribution shape is task-level, while late strengthening tracks commitment.Source: §7 · p. 8, Figure 5 · p. 8
8 Conclusion
The conclusion presents DLIG as a token-attribution method that resolves output commitment across network depth and denoising time. Across three tasks, it reveals task-specific patterns and serves as an inexpensive first check before more costly interventional validation.
Conclusion
- DLIG decomposes a model’s commitment to an output into per-token contributions at chosen layers and denoising steps, revealing how support evolves through the trajectory.Source: §8 · p. 8
- Across word-sense disambiguation, multi-hop reasoning, and sentence infilling, DLIG reveals attribution patterns that a single-pass analysis would miss.Findings include preserved reasoning-path localization on errors and a late-emerging context-reliance association with infill quality.Source: §8 · p. 8
- DLIG is an inexpensive first check for identifying or ruling out where and when output-driving evidence emerges before more expensive interventional validation.The conclusion frames DLM explanations around the full denoising trajectory, not only the final output.Source: §8 · p. 8
9 Limitations and Future Work
Implementation checks support DLIG’s numerical and computational correctness across the tested configurations, while its interpretation remains limited to prompt sensitivity rather than causal mechanisms. The appendix also documents task-performance reporting and ProsQA outcome categories.
9 Limitations and Future Work
- DLIG measures where and when committed outputs are gradient-sensitive to prompts, not the causal circuits producing them.The authors recommend confirming its hypotheses with interventional tests and leave continuous-space, hybrid AR–diffusion, and larger natively trained models for future work.Source: §9 · p. 8
A.1 Implementation Checks
- At m=1000, every tested layer-step cell meets the completeness tolerance, and partial-forward and full-path DLIG agree exactly at every checked layer.The checks also report hook transparency, monotonic error reduction from m=200 to m=1000, and passing bidirectionality and token-alignment gates.Source: §A.1 · p. 13, Figure 7 · p. 13
A.2 Task Performance
- Table 1 reports task-level performance for WiC, ProsQA, and ROCStories infilling, based on a single evaluation run with within-run bootstrap intervals.The stochastic sampler means scores may vary across runs.Source: §A.2 · p. 13, Table 1 · p. 14
A.3 ProsQA Bucket and Contrastive-Target Examples
- ProsQA examples distinguish exact or conceptual successes, valid-but-wrong options, and off-manifold generations such as wrong subjects or invalid concepts.The appendix illustrates the contrastive target construction using a wrong_valid example whose generated option matches the dataset-defined incorrect option.Source: Table 2 · p. 14, §A.3 · p. 14
A.4 ROCStories infilling: attribution profile and robustness
For self-generated ROCStories infilling, pooled context mass does not capture the quality association that appears when correlations are resolved by denoising step. Fixed gold targets produce a stronger, flatter association, while the descriptive attribution profile remains similar across target framings.
A.4 ROCStories infilling: attribution profile and robustness
- The pooled self-generated context-mass relationship with infill quality is weak, while per-step analysis reveals an association that emerges as the span converges toward its final output.The pooled binned mean is roughly flat; early-step statistics exclude degenerate cases without scoreable committed-span tokens.Source: §A.4 · p. 14, Figure 8 · p. 15
- Context mass does not distinguish high- from low-ROUGE stories at any denoising step, and its right-to-left balance is likewise quality-neutral.The per-step AUCs do not separate meaningfully from chance; only the late-emerging per-step reliance–quality correlation carries a quality signal.Source: §A.4 · pp. 14–15, Table 4 · p. 16
- Without per-token normalization, the self-generated per-step correlation still starts near zero and strengthens over denoising; normalization sharpens the late-step signal.The authors report this as evidence that the trend is not a normalization artifact.Source: §A.4 · p. 15, Figure 9 · p. 17
A.4.1 Fixed-target framing
- Under a fixed gold target, context mass correlates with infill quality much more strongly and remains high and roughly flat across denoising, unlike the rising self-generated pattern.The authors hypothesize that target proximity to high-ROUGE generations adds gradient signal independent of context reliance; the self-generated framing removes this component.Source: §A.4.1 · pp. 15–16, Figure 10 · p. 18
- The attribution profile remains bidirectional and left-skewed under the fixed target, with mass peaking on both sides and decaying across denoising despite different overall magnitudes.Directional balance is near-constant across quality; fixed-target attribution is mildly right-heavy, whereas self-generated attribution is mildly left-heavy.Source: §A.4.1 · p. 16, Figure 11 · p. 18
B Statistical Evaluation
- Statistical tests use Wilcoxon signed-rank tests for one-sample claims, Mann–Whitney AUCs for group comparisons, and Pearson correlations for continuous associations.Difficulty-adjusted analyses and structure-matched random-edge nulls are used where needed to address confounds or compare structure-sensitive masses.Source: B Statistical Evaluation · pp. 16–17
C Model, Datasets and Training Details
The appendix records the model and task setups, attribution-step grids, and supporting statistical tables for WiC, ProsQA, and ROCStories infilling. ProsQA uses a sparser grid because its per-step contrastive attribution is more expensive.
C Model, Datasets and Training Details
- WiC and ROCStories infilling use all 32 odd steps of a 64-step trajectory, whereas ProsQA uses 11 sampled steps because its contrastive scoring is more expensive.The ProsQA attribution grid is subsampled for tractability.Source: C Model, Datasets and Training Details · pp. 18–19
- WiC has zero prompt-token attribution mass before the self-generated answer span commits; such steps are excluded from attribution-depth centroids, and 22.3% of all steps have zero mass.The appendix lists full statistical results for WiC, ProsQA, and infilling, plus data and generation configurations.Source: Table 7 · p. 19, Table 6 · p. 19, Table 8 · p. 20
- WiC lacks public official test labels, so the study carves a label-stratified test set from training data, yielding 4,928/638/500 train-validation-test examples.The task asks whether a target word shares its sense across two sentences; its Yes/No answer tokens do not overlap with the disambiguating context vocabulary.Source: C Model, Datasets and Training Details · p. 19
- ProsQA is a synthetic multi-hop graph benchmark with distractor facts and a 17,886/300/500 train-validation-test split; ROCStories infilling generates sentence 3 from surrounding sentences 1–2 and 4–5 over 1,000 samples.The infilling span attends bidirectionally, enabling analysis of left- and right-context attribution across denoising steps.Source: C Model, Datasets and Training Details · pp. 19–20