Diffusion Integrated Gradients: Controllable Path Generation for Flexible Feature Attribution¶

Conference: ECCV 2026
arXiv: 2606.22314
Code: None
Area: Diffusion Models / XAI Explainability
Keywords: Integrated Gradients, Diffusion Models, Path Attribution, Explainable AI, Guided Sampling

TL;DR¶

DiffIG reformulates attribution path generation in Integrated Gradients (IG) as a conditional generative modeling problem. It uses a diffusion model to learn the path distribution generated by the Stick-Breaking Process, and performs dual-guided sampling based on faithfulness and complexity during inference to generate adaptive non-linear integration paths. This approach significantly outperforms existing path attribution methods on the DiffID metric across Oxford-IIIT Pet and Mini-ImageNet.

Background & Motivation¶

Integrated Gradients (IG), due to its axiomatic properties such as completeness, has become the standard feature attribution method for DNNs. IG assigns feature importance scores by accumulating gradients along an integration path from a baseline to the input. However, the attribution quality of IG highly depends on the choice of the integration path: standard straight-line paths are optimal under conservative gradient field assumptions, but the piecewise linear structure of modern DNNs leads to highly discontinuous gradient fields. Straight paths accumulate a large amount of noisy gradients when traversing decision boundaries, resulting in misleading attributions.

Existing improvement methods fall into a dilemma: methods like Guided IG, IG2, and AGI formulate path search as an optimization problem, but the objective function is difficult to optimize globally on a continuous function space, relying instead on greedy or local heuristic search, which fails to guarantee globally consistent paths. Methods like SPI model the path distribution using stochastic processes but completely ignore the geometry of the model's decision boundary. Methods like EIG and MIG introduce data manifold information but are constrained by a fixed generation process, lacking flexibility.

Key Challenge: Path quality determines attribution quality, but existing methods either use greedy local search (suboptimal), rely on model-agnostic stochastic processes (blind), or are constrained by fixed manifolds (inflexible). This work's Key Insight is to shift path search from "explicit optimization" to "conditional generation"—using a diffusion model to learn the distribution of high-quality paths, and using guided sampling to flexibly control path properties at inference time. Core Idea: Use a diffusion generative model and guided sampling to replace hand-crafted path design, making path attribution a controllable generation process during inference.

Method¶

Overall Architecture¶

The core problem DiffIG aims to solve is: given a model \(f\), an input \(x\), and a baseline \(x'\), how to generate an integration path \(\gamma\) from \(x'\) to \(x\) such that the calculated attribution is both faithful to the model's decision logic and sufficiently concise and interpretable. The overall pipeline consists of four stages: (1) use the Stick-Breaking Process (SBP) to generate a large and diverse set of non-linear paths as training data; (2) train an unconditional diffusion model to learn the path distribution, while simultaneously training two independent guidance networks to predict the faithfulness score and complexity score of the paths; (3) inject faithfulness and complexity gradients into the reverse diffusion process via classifier guidance at inference time to control the path generation direction; (4) perform path generation in the latent space of a \(\beta\)-VAE to reduce computational overhead, and aggregate (using mean/median/variance-weighted) \(N\) sampled candidate paths to obtain the final attribution map.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["SBP Stochastic Process<br/>Generates Diverse Paths"] --> B["Train Unconditional Diffusion Model<br/>to Learn Path Distribution p_theta(gamma)"]
    A --> C["Train Guidance Networks<br/>J_phi1(Faithfulness) + J_phi2(Complexity)"]
    B --> D["Guided Reverse Diffusion Sampling<br/>g = lambda_faith * grad(J_phi1) + lambda_comp * grad(J_phi2)"]
    C --> D
    D --> E["Decode Latent Space to Pixel Space<br/>Enforce Endpoints gamma(0)=x', gamma(1)=x"]
    E --> F["Multi-Path Sampling of N Candidates<br/>Mean/Median/SPI-P Aggregation"]
    F --> G["Compute IG Attribution<br/>Along Aggregated Path"]

Key Designs¶

1. Conditional generative modeling reformulates path search: from explicit optimization to learning to generate

Traditional path attribution methods (GIG, IG2, AGI) formulate path search as \(\gamma^* = \arg\max \mathcal{J}(\gamma)\), finding a path that maximizes an objective function (e.g., negative noise cost). However, since \(\mathcal{J}\) is defined over a continuous function space, under the complex decision boundary topology of DNNs, such optimization problems are typically intractable and degenerate into greedy local searches. DiffIG's insight comes from trajectory planning in offline reinforcement learning: instead of explicitly optimizing a single path, it is better to learn to generate trajectories that "look like good paths". Specifically, the path search is reformulated as maximizing the conditional log-likelihood:

\[\theta^* = \arg\max_\theta \mathbb{E}_{\gamma\sim\mathcal{D}}[\log p_\theta(\gamma|y=\mathcal{J}(\gamma))]\]

where \(\mathcal{D}\) is the pre-collected path dataset, and the conditional label \(y\) is the objective quality score \(\mathcal{J}(\gamma)\) of the path. Instead of training a conditional model directly, DiffIG employs an unconditional diffusion + classifier guidance strategy: it first trains an unconditional diffusion model to learn \(p_\theta(\gamma)\), and then performs energy-based guided sampling to generate paths from the perturbed distribution \(\tilde{p}_\theta(\gamma^0) \propto p_\theta(\gamma^0)\exp(\mathcal{J}(\gamma^0))\). Guidance is implemented by injecting the gradient of the regression network \(\mathcal{J}_\phi\): \(\tilde{p}_\theta(\gamma^{\tau-1}|\gamma^\tau) = \mathcal{N}(\gamma^{\tau-1}|\mu_\theta(\gamma^\tau) + \omega\Sigma^\tau g, \Sigma^\tau)\), where \(g = \nabla_\gamma\mathcal{J}_\phi(\gamma)|_{\gamma=\mu_\theta(\gamma^\tau)}\). This design achieves three benefits simultaneously: adaptability (the diffusion model learns to avoid high-curvature noisy areas), global optimization (the generation process considers the entire path rather than local greedy steps), and controllability (varying the guidance objective during inference allows flexible customization of path behavior).

2. SBP synthetic path data generation: providing training signals for the diffusion model

Training a diffusion model requires a large number of diverse paths. DiffIG utilizes the Stick-Breaking Process (SBP) to generate them. The core idea of SBP is to incrementally and randomly break a "stick" of unit length, where the breaking points and proportions determine a non-linear interpolation curve from 0 to 1. For each feature dimension \(i\), a random measure \(G_i(t) = \sum_{k=1}^\infty \pi_k \delta_{t_k}(t)\) is sampled, where the stick-breaking weights are \(\pi_k = \beta_k\prod_{j=1}^{k-1}(1-\beta_j)\) with \(\beta_k \sim \text{Beta}(1,\alpha)\), and the breakpoints \(t_k \sim U(0,1)\). Then, the non-linear interpolation is defined using the CDF \(F_{G_i}(t)\): \(\gamma_i(t) = x'_i + F_{G_i}(t)(x_i - x'_i)\). Since \(F_{G_i}\) is monotonically non-decreasing and \(F_{G_i}(0)=0, F_{G_i}(1)=1\), the generated path naturally satisfies the endpoint constraints. The concentration parameter \(\alpha\) controls path diversity: a larger \(\alpha\) makes the path closer to a straight line (approaching standard IG), while a smaller \(\alpha\) leads to high non-linearity. DiffIG samples and generates the training set in the range of \(\alpha \in [1.0, 20.0]\), ensuring coverage of all path shapes from approximate straight lines to highly complex curves. Compared to manually designing path families, SBP provides a natural parameterization covering the entire path space, with excellent mathematical properties (automatic satisfaction of endpoint constraints, independent generation per dimension).

3. Dual-guidance networks: joint control of faithfulness and complexity

DiffIG employs two independent regression networks, \(\mathcal{J}_{\phi_1}\) and \(\mathcal{J}_{\phi_2}\), to predict the faithfulness score and complexity score of paths, respectively, and jointly guides the diffusion sampling during inference. The definition of both metrics directly quantifies attribution quality:

Faithfulness measures whether the attribution truly reflects the reasoning of the model: \(\text{Faithfulness}(\gamma) = \mathbb{E}_{r\in\mathcal{R}}[\text{Conf}^{\text{ins}}(r) - \text{Conf}^{\text{del}}(r)]\), where \(r\) is the perturbation ratio, \(\text{Conf}^{\text{ins}}\) represents the confidence of the correct class after inserting pixels in descending order of importance, and \(\text{Conf}^{\text{del}}\) represents the confidence after deleting pixels. High faithfulness means the attribution successfully identifies the characteristics driving the model's decision—deleting them drops the confidence rapidly, while inserting them restores it quickly.

Complexity controls the sparsity and readability of the attribution, defined by the entropy of the attribution weights: \(\text{Comp}(\gamma) = -\sum_i p_i \log(p_i + \varepsilon)\), where \(p_i = |\mathcal{A}_i| / \sum_j |\mathcal{A}_j|\). Low complexity indicates that attribution is concentrated on a few core regions, making it easy for humans to interpret; high complexity indicates that the attribution is scattered across too many pixels, making it difficult to explain.

The two guidance networks are trained using MSE objectives to regress the target scores of the noise-free path \(\gamma^0\) from the noisy path \(\gamma^\tau\). The joint gradient during inference is \(g = (\lambda_{\text{faith}}\nabla_\gamma\mathcal{J}_{\phi_1} + \lambda_{\text{comp}}\nabla_\gamma\mathcal{J}_{\phi_2})|_{\gamma=\mu_\theta}\). Here, \(\lambda_{\text{faith}}\) and \(\lambda_{\text{comp}}\) are hyperparameter knobs adjustable at inference—increasing \(\lambda_{\text{faith}}\) yields a more faithful attribution, and increasing \(\lambda_{\text{comp}}\) (to a more negative value) achieves a sparser attribution, enabling unprecedented test-time controllability. Notably, the guidance networks are proxy regressors trained on noisy latent paths rather than optimizing evaluation metrics directly in the pixel space, thus preventing "test-on-train" data leakage.

4. Latent space path generation and multi-path aggregation: practical designs for efficiency

Directly generating diffusion paths in high-dimensional pixel space (e.g., 256x256x3) is computationally expensive. Borrowing from Latent Diffusion Models, DiffIG uses a pre-trained \(\beta\)-VAE encoder \(E_{\text{enc}}\) to map the input \(x\) and baseline \(x'\) into a 4096-dimensional latent space, where the diffusion model and guidance networks are trained. At inference, the trajectory \(\gamma_z(t)\) is generated in the latent space and reconstructed back to the pixel space via decoder \(D_{\text{dec}}\), while the endpoint conditions \(\gamma(0)=x', \gamma(1)=x\) are explicitly enforced to strictly guarantee the completeness axiom (restoration errors in the decoder do not affect completeness).

Furthermore, the stochastic nature of diffusion sampling inherently supports generating multiple competitor paths. DiffIG offers two utilization strategies: (i) Best-of-\(N\) search—selecting the path with the highest joint target score among \(N\) candidates to compute attribution; (ii) Multi-path aggregation—computing the attribution maps for all \(N\) paths separately, and then aggregating them using mean, median, variance-weighted average, or the SPI-P operator to obtain the final attribution. Experiments show that median and mean aggregation perform significantly better than the Best-of-\(N\) strategy, as aggregation reduces the random variance of individual paths, producing more stable and consistent attributions. The completeness axiom still holds under mean aggregation (due to linearity preservation), though median aggregation does not guarantee completeness (which the paper deems acceptable when robustness is prioritized).

Loss & Training¶

Diffusion model training: DiT1D backbone is used, with \(M=100\) diffusion steps and DDPM solver. The loss function is \(\mathcal{L}(\theta) = \mathbb{E}_{\tau,\epsilon,\gamma^0}[\|\epsilon - \epsilon_\theta(\gamma^\tau)\|_2^2]\), which is the standard noise prediction MSE. Trained for 1 million steps, learning rate 2e-4, weight decay 1e-5, batch size 64.

Guidance network training: Two independent DiT1D models predict the faithfulness and complexity scores. The training objective is \(\min_\phi \mathbb{E}_{\tau,\epsilon,\gamma^0}[\|\mathcal{J}_\phi(\gamma^\tau) - \mathcal{J}(\gamma^0)\|_2^2]\). Also trained for 1 million steps.

Inference: Global guidance scale \(\omega=1\), faithfulness weight \(\lambda_{\text{faith}}\) is chosen from \(\{0, 1, 10, 100, 1000\}\), complexity weight \(\lambda_{\text{comp}}\) is chosen from \(\{-100, -10, -1, 0, 1, 10, 100\}\), and the number of paths \(N\) is chosen from \(\{1, 10, 30, 50\}\).

Key Experimental Results¶

Main Results¶

Evaluated on Oxford-IIIT Pet (37 pet categories, 370 validation images) and Mini-ImageNet (100 categories, 500 sampled images), across three architectures: VGG16, ResNet18, and InceptionV1. Metrics used are DiffID (Insertion AUC minus Deletion AUC, higher is better), Insertion AUC (higher is better), and Deletion AUC (lower is better).

DiffID Comparison on Oxford-IIIT Pet:

Method	ResNet18	VGG16	InceptionV1
IG	0.3213	0.4654	0.3051
GIG	0.3486	0.4859	0.3085
IG2	0.2767	0.3016	0.2228
AGI	0.3242	0.4930	0.4092
EIG	0.2680	0.3950	0.2698
MIG	0.2552	0.3681	0.2602
SPI	0.2738	0.4202	0.2292
DiffIG	0.5067	0.6368	0.4817

DiffIG consistently and significantly leads on Oxford-IIIT Pet: on ResNet18, the DiffID increases from 0.3486 (best of GIG) to 0.5067 (+45%); on VGG16, it increases from 0.4930 (best of AGI) to 0.6368 (+29%); on InceptionV1, it increases from 0.4092 (best of AGI) to 0.4817 (+18%). It also leads comprehensively on Mini-ImageNet (see Table 1 in the original paper), and this advantage remains true on ViT-B/16 (DiffID 0.5711 vs. AGI 0.4033).

Ablation Study¶

Ablation on Multi-Path Sampling Number \(N\) (Oxford-IIIT Pet / ResNet18):

Configuration	DiffID	Insertion	Deletion
DiffIG (N=1)	0.3791	0.5201	0.1410
DiffIG (N=10)	0.4898	0.6143	0.1245
DiffIG (N=30)	0.4940	0.6222	0.1254
DiffIG (N=50)	0.5067	0.6300	0.1233

The effect of multi-path sampling is significant and monotonic: as \(N\) increases from 1 to 50, DiffID increases from 0.3791 to 0.5067 (+33.7%). Performance is close to saturation at \(N=10\), while \(N=50\) achieves the best performance. Regarding aggregation strategies, median, mean, and variance-weighted mean deliver comparable performances, markedly outperforming the Best-of-\(N\) strategy (Best-of-\(N\) forms a distinct, inferior cluster in most settings).

Guidance Effectiveness: Heatmap analysis indicates that without any guidance (\(\lambda_{\text{faith}}=0, \lambda_{\text{comp}}=0\)), the diffusion model degenerates into unconditional path generation, yielding the lowest DiffID. Applying at least one guidance term significantly boosts performance, with optimal configurations varying depending on the dataset and model architecture. The effect of complexity guidance is visually intuitive: as \(\lambda_{\text{comp}}\) shifts from 0 to -100, the attribution maps become noticeably sparser and more focused on the core object area. Compared to DDPath (a contemporaneous diffusion attribution method that only reconstructs paths via unconditional diffusion inversion), DiffIG achieves an average DiffID improvement of 0.042-0.112 across three models on Mini-ImageNet, validating the critical role of guided sampling.

Key Findings¶

Multi-path aggregation is one of the largest contributing factors: \(N=50\) compared to \(N=1\) improves DiffID by approximately 33%, indicating that a large pool of high-quality candidates exists in the diffusion-generated paths, and aggregation effectively reduces the random variance of individual paths.
Faithfulness guidance is the primary driving force in most settings: The optimal configuration on Oxford-IIIT Pet consistently includes high \(\lambda_{\text{faith}}\) (1000), while complexity guidance contributes more on Mini-ImageNet (optimal layouts include higher negative \(\lambda_{\text{comp}}\)), indicating that guidance strategies require configuration-level adaptation.
Robustness across different VAE backbones: DiffIG's performance remains stable across three VAE backbones (MAR, Stable Diffusion 2.1, Kandinsky 2.1) and is insensitive to the SBP concentration parameter \(\alpha\).
Black-white mean baseline effectively mitigates dark region issues: Averaging the attributions under black and white baselines improves DiffID in 5 out of 6 settings (up to +0.115), making it a simple and effective solution.

Highlights & Insights¶

Paradigm shift from "hand-crafting paths" to "learning to generate paths": This is the most central insight. Path attribution methods have long been trapped in the loop of "designing an objective function -> greedy local search." DiffIG breaks this cycle by employing generative models—rather than struggling to optimize an intractable objective, it learns to generate samples that "look like good paths". This idea can be transferred to any XAI problem requiring continuous trajectory search (e.g., counterfactual explanation path generation, adversarial sample trajectory design).
Guidance networks perform proxy prediction on the noisy latent space rather than directly optimizing evaluation metrics: This represents an elegant separation—the guidance network predicts the quality of the path itself, whereas the evaluation metric is the faithfulness of the attribution map in the pixel space. Although there is a latent-to-pixel domain gap between the two spaces, experiments prove this gap is benign in practice (also validated in a controlled 3D-helix toy experiment that guidance can optimize pixel-space objectives without leaving the manifold). More importantly, this separation avoids circular reasoning ("test-on-train"), allowing DiffIG to rank first even on the independent Quantus Faithfulness Correlation benchmark.
Multiple lambda hyperparameters offer flexibility rather than tuning burdens: \(\lambda_{\text{faith}}\) and \(\lambda_{\text{comp}}\) control faithfulness and sparsity separately, which is impossible in traditional methods—where changing sparsity required changing the entire method or redesigning the objective function. DiffIG achieves zero-cost switching of attribution styles at inference time, which is a direct dividend of "turning paths into a controllable generation process".
The ingenuity of using SBP as a synthetic path data source: SBP has few parameters (only \(\alpha\)), good mathematical properties (automatic satisfaction of boundary constraints, independent dimension-wise generation), and wide coverage (a single formula covering lines to highly non-linear paths). Compared to generating paths using real model gradient fields (which require repeated forward-backward propagation), SBP is a purely statistical process with negligible generation cost.

Limitations & Future Work¶

Limitations of fixed baselines: DiffIG still employs a fixed black baseline, resulting in under-attributed dark visual areas (such as black eyes). While the proposed black-white mean scheme is effective, it is merely a workaround. Learning adaptive baselines remains the fundamental solution.
Need to retrain for each (dataset, model) pair: The diffusion model and guidance networks must be trained specifically for a given model \(f\) and data distribution, tallying about 0.9 GPU hours (on a single H100) of one-time training overhead per setting. For practical applications with diverse models and datasets, this accumulated computational cost is non-trivial.
Hyperparameter sensitivity: The optimal values of \(\lambda_{\text{faith}}\) and \(\lambda_{\text{comp}}\) differ significantly between datasets and models, requiring per-scenario tuning. Although performance degradation without guidance is predictable (always the worst), finding a "good-enough" default remains an open question.
Validation limited to CNNs and ViT image classifiers: The applicability of this paper to other modalities like text, tabular, and graph data remains open. Since the latent space solution relies heavily on VAE encoder quality, other modalities require different latent representation strategies.
100 diffusion steps + N paths: Running 100 steps for \(N=50\) paths takes 4.5 seconds per image (on H100), which is far slower than standard IG (0.09s). Future directions include more efficient sampling strategies (e.g., consistency models, distillation) and smart pruning of multiple paths.

vs Guided IG (GIG): GIG performs greedy path updates using local gradients of the objective function, while DiffIG performs global path generation using a diffusion model. GIG is trapped in local optima, whereas DiffIG implicitly achieves global search by generating complete paths. The core distinction lies in "optimizing a path" vs. "generating and selecting from a path distribution".
vs Stick-breaking Path Integration (SPI): SPI also uses SBP to generate path distributions but aggregates them directly via sampling, completely ignoring the geometry of the decision boundary. DiffIG treats SBP paths only as unlabeled training data and injects model-aware information via guidance networks—making model-adaptive paths using SPI's "raw materials."
vs DDPath: DDPath uses the reverse denoising trajectory of a pre-trained diffusion model as the integration path, assuming denoising trajectories naturally lie close to the data manifold. However, it relies entirely on unconditional diffusion inversion and cannot control path quality. DiffIG adds guided sampling on top of DDPath, upgrading "manifold constraint" to "manifold constraint + quality guidance," consistently outperforming DDPath.
vs Diffuser / Decision Diffuser in Offline RL: The design of DiffIG is directly inspired by trajectory planning works using diffusion models in offline RL. Treating path attribution as "generating a high-reward trajectory" is an insightful cross-domain transfer—Diffuser uses a reward predictor to guide trajectory generation, whereas DiffIG uses faithfulness/complexity predictors to guide path generation, showing isomorphic structures.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ Reformulating path attribution into conditional generative modeling is a brand-new perspective. No prior work has introduced diffusion generative models to feature attribution path design. The cross-domain transfer (offline RL planning -> XAI path generation) is natural and deep.
Experimental Thoroughness: ⭐⭐⭐⭐ Two datasets, three CNN architectures + ViT, seven baselines, multi-dimensional ablation studies (N, guidance, VAE backbones, SBP alpha, baseline choices), validation via independent metrics (Quantus), runtime analysis, and failure cases are thoroughly detailed in the appendix. One star is deducted because the optimal hyperparameter configurations vary substantially across different setups, and the unified configuration (Table 1) is not optimal for every sub-scenario.
Writing Quality: ⭐⭐⭐⭐⭐ Well-structured (motivation -> method -> experiments -> analysis), mathematically rigorous, and rich in polished charts (Figure 1 and Figure 2 are highly informational). The appendix contains full axiomatic proofs and hyperparameter analysis, ensuring high reproducibility.
Value: ⭐⭐⭐⭐ It introduces a promising generative framework to the XAI domain, and the concept of inference-time controllability has wide transfer potential. However, the training overhead (re-training per dataset-model pair) and inference speed (50x slower than IG at \(N=50\)) limit immediate practicality, making it more suitable for offline analysis requiring high-quality attribution rather than real-time deployment.