From Baselines to Transport Geodesics: Axiomatic Attribution via Optimal Generative Flows

Cenwei Zhang; Lin Zhu; Manxi Lin; Lei You

cs.LGMar 5, 2026

From Baselines to Transport Geodesics: Axiomatic Attribution via Optimal Generative Flows

Authors: Cenwei Zhang, Lin Zhu, Manxi Lin, Lei You

Organizations: Aalto University, Espoo, Finland · Alibaba, Hangzhou, China · Technical University of Denmark, Copenhagen, Denmark

Abstract

Feature attributions often hide a critical modeling choice: they explain a prediction along a counterfactual path from a reference state to an input. Different baselines, interpolations, and generative trajectories define different paths and can therefor produce different explanations. We study this path ambiguity as a modeling problem. Our central question is whether the path can be chosen by the data-generating transport process, rather than by a hand-designed interpolation or by the sensitivity geometry of the model being explained. We separate attribution into fixed-path credit allocation and path selection. For a fixed path, we prove that the Aumann-Shapley line integral is the unique attribution rule under standard fixed-path axioms and explicit coordinate-trace regularity. For path selection, we minimize kinetic action over flows that transport a reference distribution to the data distribution, yielding a transport-geodesic attribution principle. We approximate this ideal with Rectified Flow and Reflow and derive stability bounds linking vector-field error to attribution error. Experiments show that lower-action, transport-consistent paths produce more stable and structured explanations, preserving competitive deletion faithfulness, without claiming data-manifold membership. Our code is available at https://github.com/cenweizhang/OTFlowSHAP.

Figures & tables

Steps $K$	MAE $\downarrow$	Std. dev.	SEM	Relative error $\downarrow$
10	1.483	1.346	0.177	19.30%
20	0.895	0.778	0.102	11.65%
50	0.411	0.317	0.042	5.34%
100	0.229	0.181	0.024	2.98%
200	0.101	0.079	0.010	1.34%

Table 1: Completeness residual for the discrete estimator in Eq. ( 11 ). The residual decreases as the integration grid becomes finer.

Method	Action $\widehat{\mathcal{A}}$ $\downarrow$	Pixel var. $\downarrow$	SSIM $\uparrow$	Rank corr. $\uparrow$
1-RF	3179.2 $\pm$ 295.8	0.0032 $\pm$ 0.0021	0.716 $\pm$ 0.080	0.662 $\pm$ 0.087
2-RF	3006.9 $\pm$ 340.4	0.0010 $\pm$ 0.0008	0.911 $\pm$ 0.051	0.882 $\pm$ 0.061

Table 2: Path action and attribution stability across seeds. Lower-action reflowed paths produce more stable maps and more consistent feature rankings.

Data	Method	GPS	FCE $\downarrow$	SATV $\downarrow$	EAS $\uparrow$	Del. zero $\downarrow$	Del. blur $\downarrow$
CIFAR-10	SmoothGrad	-	-	0.022 $\pm$ 0.008	0.363 $\pm$ 0.117	0.525 $\pm$ 0.923	1.216 $\pm$ 1.202
	GuidedBackprop	-	-	0.040 $\pm$ 0.018	0.439 $\pm$ 0.128	0.966 $\pm$ 0.978	1.627 $\pm$ 1.064
	GradientSHAP	-	-	0.669 $\pm$ 0.408	0.116 $\pm$ 0.146	0.691 $\pm$ 0.975	1.978 $\pm$ 1.050
	Integrated Gradients	1.000 $\pm$ 0.000	$(1.0\pm 0.1)\times 10^{4}$	0.643 $\pm$ 0.404	0.114 $\pm$ 0.134	0.670 $\pm$ 0.965	1.931 $\pm$ 1.105
	DDIM	1.059 $\pm$ 0.017	$(3.1\pm 0.5)\times 10^{3}$	0.073 $\pm$ 0.062	0.119 $\pm$ 0.160	0.444 $\pm$ 0.940	1.360 $\pm$ 1.143
	Transport Flow	1.023 $\pm$ 0.008	0.005 $\pm$ 0.003	0.062 $\pm$ 0.048	0.119 $\pm$ 0.138	0.456 $\pm$ 0.967	1.311 $\pm$ 1.177

Table 3: CIFAR-10 and CelebA-HQ benchmark results. We report path diagnostics only for methods that define a continuous path. Our method has much lower FCE than IG and DDIM, while deletion faithfulness remains competitive rather than uniformly best.

Figure 1: Qualitative visualization on CIFAR-10 and CelebA-HQ. Transport Flow explanations are spatially coherent and avoid some high-frequency artifacts of straight-line Integrated Gradients. The figure should be read together with Table 3 ; visual clarity alone is not a complete faithfulness metric.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Synthetic additive sanity check. The path-integral estimator recovers the analytical additive Shapley values, and the error decreases as the quadrature grid becomes finer.

Figure 3: Controlled two-dimensional Gaussian transport. The oracle displacement interpolation gives the reference geometry. One-step Rectified Flow can be visibly more curved, while Reflow straightens the trajectories and makes the intermediate densities closer to the oracle interpolation.

Metric	OT oracle	1-RF	2-RF	3-RF
$W_{2}^{2}(p_{0},p_{1})$	27.8976	-	-	-
Action gap $\Delta A$	0	$0.2205\pm 0.0033$	$0.00189\pm 0.00144$	$0.00199\pm 0.00141$
Relative field error	0	$0.2933\pm 0.0050$	$0.00284\pm 0.00014$	$0.00276\pm 0.00008$
Curvature proxy	$3.45{\times}10^{-5}\pm 2.63{\times}10^{-7}$	$0.0763\pm 0.00077$	$4.85{\times}10^{-4}\pm 1.50{\times}10^{-5}$	$4.35{\times}10^{-4}\pm 1.31{\times}10^{-5}$

Appendix

Table 4: Controlled Gaussian transport in

d=10

, reported as mean

\pm

standard deviation over five seeds. Reflow reduces the relative action gap, relative field error, and curvature proxy compared with one-step Rectified Flow.

Figure 4: High-dimensional controlled diagnostic. The left panel relates relative attribution error to relative field mismatch. The right panel reports neighborhood preservation along intermediate transport states. The result supports the stability view: better field alignment gives smaller attribution discrepancy.

Figure 5: Attribution convergence across flow checkpoints. The trend is consistent with the stability theorem: as the learned vector field approaches the reference checkpoint, the attribution gap decreases.

Figure 6: Representative action diagnostic for sampled paths. The plot illustrates how path movement and action are measured along a numerical trajectory.

Figure 7: Additional CelebA-HQ examples, part I.

Figure 8: Additional CelebA-HQ examples, part II.

Figure 9: Additional CIFAR-10 examples, part I.

Figure 10: Additional CIFAR-10 examples, part II.

Feature Attribution

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

See all 231 papers

Aumann-SHAP: The Geometry of Counterfactual Interaction Explanations in Machine Learning
Mar 14, 2026Adam Belahcen, Stéphane MussardFeature Attribution Shapley Value Attribution
AIRESS & UM6P
A Polynomial-Time Axiomatic Alternative to SHAP for Feature Attribution
Feb 28, 2026Kazuhiro Hiraki, Shinichi Ishihara, Takumi Kongo +1Feature Attribution Shapley Value Attribution
International Monetary Fund · Independent Researcher · Fukuoka University +1
Local Additive Feature Attribution: A Mathematical Taxonomy and Reporting Checklist
Jul 15, 2026Rebecca Afriyie Sarpong, Daniel CommeyGradient-Based Attribution Feature Attribution

Mar 14, 2026J/K move · Enter open · S save

Aumann-SHAP: The Geometry of Counterfactual Interaction Explanations in Machine Learning

Adam Belahcen, Stéphane Mussard

AIRESS & UM6P

Feature Attribution Shapley Value Attribution

Open paper AskPDF arXiv

We introduce Aumann-SHAP, an interaction-aware framework that decomposes counterfactual transitions by restricting the model to a local hypercube connecting baseline and counterfactual features. Each hypercube is discretized into a grid to construct an induced micro-player cooperative game in which elementary grid-step moves become players. Shapley and LES values on this TU-micro-game yield geometry-aware within-pot attributions that converge to the diagonal Aumann--Shapley / Integrated Gradients limit under grid refinement, and recover equal-split Shapley as the degenerate

m=1

special case. An exact grid-state closed form gives polynomial-time computation for fixed interaction order. On a synthetic benchmark with known ground truth, equal-split Shapley carries an irreducible bias while Aumann-SHAP converges to the correct decomposition. On German Credit, interaction geometry changes feature priority rankings in

12.3\%

of instances. On UCI Heart Disease, equal-split misattributes a cholesterol suppressor as a positive contributor, which is a sign error Aumann-SHAP corrects. On MNIST, game-theoretic attribution reaches target confidence with

3.5\times

fewer edits than magnitude-based ordering, with micro-game Shapley achieving the best efficiency across all budgets.

Abstract

Feature Attribution

Explore similar work

Aumann-SHAP: The Geometry of Counterfactual Interaction Explanations in Machine Learning

A Polynomial-Time Axiomatic Alternative to SHAP for Feature Attribution

Local Additive Feature Attribution: A Mathematical Taxonomy and Reporting Checklist