Flow-Disentangled Feature Importance
Authors: Xingshu Chen
Yifeng Guo
Jin-Hong Du
Venue: ICLR 2026 Poster (2026)
Keywords: Interpretability
Feature Importance
Statistical Inference
Correlation Distortion
Uncertainty Quantification
Primary area: interpretability and explainable AI

Abstract: Quantifying feature importance with valid statistical uncertainty is central to interpretable machine learning, yet classical model-agnostic methods often fail under feature correlation, producing unreliable attributions and compromising inference. Statistical approaches that address correlation through feature decorrelation have shown promise but remain restricted to $\ell_2$ loss, limiting their applicability across diverse machine learning tasks. We introduce Flow-Disentangled Feature Importance (FDFI), a model-agnostic framework that resolves these limitations by combining principled statistical inference with computational flexibility. FDFI leverages flow matching to learn flexible disentanglement maps that not only handle arbitrary feature distributions but also provide an interpretable pathway for understanding how importance is attributed through the data's correlation structure. The framework generalizes the decorrelation-based attribution to general differentiable loss functions, enabling statistically valid importance assessment for black-box predictors across regression and classification. We establish statistical inference theory, deriving semiparametric efficiency of FDFI estimators, which enables valid confidence intervals and hypothesis testing with Type I error control. Experiments demonstrate that FDFI achieves substantially higher statistical power than removal-based and conditional permutation approaches, while maintaining robust and interpretable attributions even under severe interdependence. These findings hold across synthetic benchmarks and a broad collection of real datasets spanning diverse domains.

Published as a conference paper at ICLR 2026
FLOV-DISENTANGLED FEATURE IMPORTANCE
Xingshu Chen
Yifeng Guo
School of Mathematics Sun Yat-sen University Guangzhou, China
Department of Computational Biology
St. Jude Children's Research Hospital TN, USA
Jin-Hong Du
Department of Statistics and Actuarial Science & Institute of Data Science The University of Hong Kong
Hong Kong SAR, China
ABSTRACT
Quantifying feature importance with valid statistical uncertainty is central to in- terpretable machine learning, yet classical model-agnostic methods often fail un- der feature correlation, producing unreliable attributions and compromising infer- ence. Statistical approaches that address correlation through feature decorrelation have shown promise but remain restricted E2 loss, limiting their applicability across diverse machine learning tasks. We introduce Flow- Hlow-Disentangled Feature Importance (FDFI), model-agnostic framework that resolves these limitations by combining principled statistical inference with computational flexibility. FDFI leverages flow matching to learn flexible disentanglement maps that not only han- dle arbitrary feature distributions but also provide an interpretable pathway for understanding how importance attributed through the data's correlation struc- ture. The framework generalizes the decorrelation-ben based attribution to general differentiable loss functions, enabling statistically valid importance assessment for black-box predictors across regression and classification. We establish statistical inference theory, deriving semiparametric efficiency of FDFI estimators, which enables valid confidence intervals and hypothesis testing with Type error con- trol. Experiments demonstrate that FDFI achieves substantially higher statistical power than removal-based and conditional permutation approaches, while main- taining robust and interpretable attributions even under severe interdependence. These findings hold across synthetic benchmarks and a broad collection of real datasets spanning diverse domains.
INTRODUCTION
Quantifying the importance of input features is fundamental to model interpretability and scientific discovery (Murdoch et al., 2019). However, standard aggossc methods falter when features are correlated. Removal-based approaches, such as 2018) and resample-based approaches such as Conditional Permutation Importance (CPI) (Strobl et al., 2008), can produce ambiguous attributions because they cannot cleanly isolate the unique predictive contribution of a single variable from that of its statistical dependents. This confounding effect of multicollinearity undermines the reliability of explanations derived from complex models (Williamson et al., 2021; Verdinelli & Wasserman, 2024a). Further, many attribution methods pro- vide only point estimates, lacking uncertainty quantification necessary for valid statistical inference, such as constructing confidence intervals or performing hypothesis testing (Chamma et al., 2023).
(LOCO) (Lei et al.,
To address the challenge of correlation, Disentangled Feature Importance (DFI) was recently pro- posed as a principled framework for attribution under dependence (Du et al., 2025). The core idea is first learn a transformation that maps the original correlated features into latent space where they
I Equal contribution
Corresponning author (j miinoon@hhkk.hi

---

Published as a conference paper at ICLR 2026
X
X3
FDFI
Decorrelation via flow matching
Latent importance assignment
T
Joa Z, l1 7/
X
Z
Black- box model
Feature importance attribution
LoCo/CPVShapley
correlation distortion
1 Xg
060 Nj
X, X,
X1 Xy X
Y
Figure 1: Overview of FDFI. black-box model takes correlated features X as input to predict Y. Conventional attribution methods (e.g.. LOCO, CPI, Shapley) underestimate the importance of correlated features. DFI addresses this issue for Gaussian features under squared-error loss. The proposed FDFI framework generalizes DFI by replacing the linear optimal transport with flexible flow matching and extending to general losses and black-box models. Operationally, FDFI addresses this by (I) decorrelating X into latent variables Z via flow matching, (2) assigning importance scores in the disentangled latent space, and (3) attributing them back to the original features with uncertainty quantification. The latent importance reveals the intrinsic predictive variability, and the disentangled map enables interpretable attributions under correlations.
become statistically independent (Genizi, 1993). By measuring importance in this disentangled rep- resentation and then mapping the scores back the original feature space, DFI effectively isolates the unique signal attributable to each feature. While powerful, the original DFI framework has two key limitations: (i) its reliance on an optimal transport (OT) map perform the disentanglement can be computationally intensive and less flexible for complex, high-dimensional distributions beyond Gaussianity, and (ii) its formulation is restricted to importance scores evaluated with the Eg loss.
In this work. we build upon this foundation to introduce Flow-Disentangled Feature Importance (FDFI), a significant generalization and enhancement of the DFI framework. An overview of the FDFI framework is provided in Figure 1. We replace the restrictive OT map with a more powerful and flexible transformation learned via flow matching, a tfate-of-the-are technique from generative modeling (Lipman et al., 2022). This allows our method to learn complex, nonlinear disentangle- ment maps between arbitrary feature distributions. Furthermore, we extend the DFI formulation to accommodate general differentiable loss functions, making it applicable to broader range of tasks, including regression and classification. FDFI thus provides a unified framework that not only deliv- robust feature importance under feature dependency but also enables valid statistical inference and uncertainty quantificatonn for these importance scores. Moreover, the generative nature of the flow-based map provides transparent mechanism for understanding how importance is attributed back through the data's correlation structure, providing reliable feature attribution.
1.1 SUMMARY OF CONTRIBUTIONS
general feature importance framework. We analyze the relationship between three founda- tional feature importance measures under general differentiable loss functions (Theorem 2.1). We establish their formal equivalence under the Eg loss (Lemma 2.2), result that underscores their shared fundamental limitation: a vulnerability to correlation distortion when features are depen- dent. To address this, we propose new framework, Flow- now-Disentangled Feature Importance (FDFI), which makes two key advances: (i) it replaces restrictive Gaussian transport map (Du et al., 2025; Genizi, 1993) with a more flexible and powerful transformation learned via flow matching (Lipman et al., 2022), enabling it to handle arbitrary feature distributions; and (ii) it extends the attribution framework beyond the f2 loss to any differentiable loss function (4), broadening its applicability to classification and other machine learning tasks.
Semiparametric statistical inference. We establish valid statistical inference to quantify the un- certainty of estimated importance based semiparametric efficiency theory. We derive the efficient influence functions and formally prove the asymptotic normality of our FDFI estimators for both the latent importance scores (Theorem 3.1) and the final original feature importance scores (Propo-
2

---

Published as a conference paper at ICLR 2026
sition 3.2). This theoretical foundation provides principled basis for valid statistical inference, enabling the construction of confidence intervals and hypothesis testing.
Extensive empirical validation. We conduct extensive experiments on both synthetic and real- world data, covering regression and classification tasks. Synthetic experiments systematically vary correlation strength, sample size, and data-gendn sata-generating processes (Section 4.1 and appendix E.1). Across nine real-world datasets spanning multiple domains and tasks (Section 4.3 and appendix E.2), FDFI produces robust and clinically interpretable importance profiles, consistently outperforming existing methods in identifying influential features under complex dependency structures.
1.2 RELATED WORK
Model-agnostic feature attributions. A central goal of explainable is to develop model-agnostic methods for quantifying feature importance. Permutation approaches (Breiman, 2001; Janitza al., 2018) remove information by randomly shuffling a variable, but were originally introduced as algo- rithmic heuristics without statistical guarantees. One refinement is Conditional Permutation Impor- tance (CPI) (Strobl et al., 2008; Hooker et al., 2021) that improves robustness through conditional resampling, and recent works (Chamma al., 2023; Reyero-Lobo et al., 2026) further explore the statistical inference problem. In parallel, Leave-One Oae---nnn----teeeee -ut (LOCO) defines importance as the change in predictive risk when variable is removed, which naturally supports statistical inference (Lei et al., 2018; Rinaldo et al., 2019; Verdinelli & Wasserman, 2024b). In contrast, Shapley-value methods, originating from cooperative game theory (Shapley, 1953) and popular- ized in machine learning through SHAP (Lundberg & Lee, 2017), provide axiomatic framework that ensures fairness and additivity; however, exact computation and statistical inference remain challenging. Subsequent work shows that SHAP values be expressed weighted averages of LOCO estimands under squared-error loss, thereby linking the two paradigms (Williamson & Feng, 2020). Despite these methodological advances, all approaches remain sensitive to collinearity, often underestimate the importance of correlated features (Verdinelli & Wasserman, 2024a).
Feature importance under correlated predictors. A growing body of work has sought address the correlation distortion that undermines classical feature importance measures (looss & Prieur, 2019; Williamson & Feng, 2020; Williamson al., 2021; Verdinelli & Wasserman, 2024b). Be- yond conditional resampling strategies, complementary line of work leverages knockoff construc- tions and conditional randomization tests to control the false discovery rate of variable selection (Candes et al., 2018; Gimenez et al., 2019; Mason & Fei, 2025); however, this differs from our objective of quantifying the marginal importance of features. Disentangled Feature Importance (DFI) (Du et al., 2025), as nonparametric extension of classic R2 de-oompoooiioonnn under linear models with correlated features (Genizi, 1993). It maps correlated predictors into independent latent space via OT, computes importance in this disentangled repre- sentation, and attributes the results back to the original variables. However, DFI less flexible for mapping between complex high-dimensional distributions and is also restricted to the square loss.
conceptually distinct direction
2
PRELIMINARIES
2.1 FOUNDATIONAL FEATURE IMPORTANCE MEASURES
Let EXXYCR Rd x be a random vector representing features and a target variable. For a loss function f: Y + R., the risk of the model is defined as R(f: X, Y) E(((Y, f(X)). For any E [1,...,d}, d], we denote by X-, = (XI,.... Xj-1, 11-1,X9p. Xi+th..., Xa) the feature vector except X, and (
X, X Xi,1............ Xa) the feature vector with jth feature replaced by X, ~ p(X, X independent of and X, conditional on X-j- For given loss function we study the behavior of prediction model f X + Y. In particular, we focus the analysis of the Bayes optimal predictors whose prediction values are defined as
f(x) = argmin E(((Y,)) X = x] and f-,(x-)) argmin E(E(Y,)) XX)==
WEY
There various types of importance measures. Below, restrict our analysis to three basic ones.
3

---

Published as a conference paper at ICLR 2026
Leave-One Pane-Covvriate-O0 Importance (LOCO) (Lei et al., 2018) that measures the increase in risk when the model is retrained without feature j:
R(f:
In particular, the Shapley value with prediction error the value function be expressed weighted average of LOCO estimators over all subsets (Verdinelli & Wasserman, 2024b). Conditional Permutation Importance (CPI) (Hooker et al., 2021; Reyero-Lobo et al., 2026) that measures the increase in error when feature is replaced by a random draw from its conditional distribution:
[R( XRIJX Y) - R(J:X,Y)
Cobol-Conditional Permutation Importance (SCPI) (Reyero-Lobo et al., 2026) that measures the portion of the model's prediction variance attributable feature, conditional other features:
SCCI OX2
Rigs; X-j, Y) - XY)
where g,(X-)) -If x is the condition mean of prediction through f given X_,-
For general function, f-j(X-)) () and gj (X=;) may not be identical. Their difference gives rise to the distinction between LOCO and CPI. We first establish an exact identity that decomposes the xx,1 into interpretable components, which requires A Assumption A1.
difference $1LCCO
Assumption A1 (Loss function). Let the loss function & RR x R + R. be function of the true label y and the prediction ÿ. We assume the function fly, -) is convex and differentiable with respect to its second argument for any fixed Furthermore, [(y, -) is M-smooth, i.e., 102 /0 S M.
This assumption imposes standard regularity conditions for analyzing the risk functional. Convexity and differentiability ensure that the Bayes predictors well-defined minimizers of the expected loss. The smoothness condition controls the second-order variability and holds for loss functions (e.g., the tg loss {yy. 7 (y ?) the binary cross-entropy loss defined the logit scale); though is not needed for the FDFI framework we introduce in the next section. Under this assumption, an identity involving LOCO and CPI is given in the following theorem.
Theorem 2.1 (Bound on LOCO and CPI Discrepancy). Under Assumption A1, the difference be- twen LOCO and CPI can be decomposed into two components:
QX
EMIE
Expproxs
where EMIE := 2 ( - Jg) is the model interaction effect and Eapprox := R(g,) R(f-s) is the approximation error, and Jn = - is the Jensen gap. The absolute difference bounded by:
L,00 1Øx,
CPI xx,1
< Eupprox + |EMIEl:
where individual components be further upper-bounded by
M
Eapprox I 2 11g; - f-5ll?
EMIEI
I 2 EV((((X) |X-)))
The approximation error Espprox quantifies the suboptimality from using the averaged model g instead of the actual optimal model f-j, while the Jensen gap Jg quantifies the risk difference due to the convexity of the loss function, based on Jensen's 's inequality. The term EMIE precisely captures the discrepancy that arises from using a general loss function instead the Ez loss. For the & loss, where M 2, one has J sCCI E[V
X-j)]. Substituting these into the definition
yields EMIE 0. Since Expprox is also zero for the Eg loss (as g, f-z). the bound is tight. These results are summarized in Lemma 2.2.
Lemma 2.2 (Equivalence of LOCO, CPI, and SCPI for lg loss). For the lg loss E 9) (y 7) the Bayes optimal predictor f satisfies gi = f-j. Further, the identity (1) in Theorem 2.1 equals zero, yielding an exact equivalence: gLOCO serl ox = =dx, EVV(((X) X-3)].
4

---

Published as a conference paper at ICLR 2026
Despite their equivalence under the lg loss, all three measures share a fundamental limitation: they suffer from correlation distortion when predictors are collinear. These methods cannot disentangle the predictive contribution of a feature from the shared signal of its statistical dependents, causing the importance scores be diluted or masked and yielding ambiguous attributions (Verdinelli & Wasserman, 2024a;b). For example, as shown by Du et al. (2025, Example 5), given linear model Y = + X2 + e where features X1 and X2 are near-perfectly correlated (X1 R X2), the above methods would assign near-zero importance to X1 (since removing it causes minimal performance drop, as X2 retains all predictive information) and, symmetrically, near-zero importance to X2. This contradicts the fact that X, and X2 are both critical.
2.2 DISENTANGLED FEATURE IMPORTANCE
To overcome this limitation, disentangled feature importance (DFT) (Du et al., 2025) was introduced to attribute importance scores while accounting for correlation under an Eg loss. More specifically, DFI seeks a Gaussian optimal transport map T IRd + RRd such that Z T(X) has independent coordinates, ideally matching simple reference distribution such as multivariate Gaussian Na (0, T)- Without loss of generality, assume Z, has zero mean and unit variance for all I,... .d. Once the disentangled representation is obtained, the latent DFI score DRI for Zz is defined as:
2 = = (V((X) Z-,) E (V((ZZ) Z-,)],
where (*xx) = E| X = x] and 77(= == (x) denote the regression functions in the latent and feature spaces, respectively. This measure aligns with both LOCO and CPI, except that it is defined in terms of the latent feature Z rather than the raw feature X. The DFI score & for each original feature X; by transferring importance from the disentangled features Z back X through the sensitivity of Zj with respect to X₁. Formally, the original DFI measure for the Eg loss is
0x oz;
2
Σ
E v(f(x) z-j)
(3)
where a8, denotes the partial derivative of Xi eTT with respect Zj. The inner term ((n(2) Z-j) ) is the first-order Sobol index of Z,. representing the "intrinsic" predictive signal uniquely attributable to that disentangled direction. Multiplying by (0X,/0Z,) gauges how strongly fluctuations in Z, are expressed through X,; integrating over the data distribution, averages these local sensitivities into a global importance score. Thus, ADF quantifies how much of the irreducible signal carried by all latent directions is channelled through Xi-
While DFI provides an alternative strategy for attribution under dependence, its formulation presents two key limitations that motivate our work: (i) By defining latent importance via conditional variance (2), the framework is intrinsically restricted to the & loss, precluding its application to classification tasks or models using general differentiable loss functions; and (ii) The reliance of DFI on optimal transport maps to learn the transformation T can be computationally intensive and less flexible for mapping the complex, digh-dimensional distributions encountered in modern machine learning.
FLOW-BASED DISENTANGLED FEATURE IMPORTANCE
3.1 MODEL- DFI WITH GENERAL LOSS FUNCTIONS
To extend DFI beyond the & loss, we generalize its two core components: the latent importance measure and the attribution rule. First, we require a latent importance measure Φz, that well- defined for general loss function f. For each disentangled feature Zj, we define the latent FDFT score as the expected increase in risk upon resampling its values conditional on other latent features:
2, : E[w(O;T)), where wO;T) 1 -
FDFI
Here, O (X, Y) is the observation, z is the latent vector with its j-th coordinate z76r resampled from p( (Z)), and T' is the nuisance transport map.
5

---

Published as a conference paper at ICLR 2026
Second, we generalize the aggregation rule in (3). The original formulation weights the intrinsic signal of Zj quantified by the conditional variance V(f(X) Z-ş) by a sensitivity term. The first-order approximation of the geometric influence of the latent
squared Jacobian is based variables on the original features. This sensitivity term is a property of the transport map 7 itself, independent of the specific loss function. For general loss, the analogous measure of intrinsic signal is the conditional expected increase in loss B[w(O;;) T) Z- By substituting this term for the conditional variance, we obtain a natural generalization:
0x, 27 0z,
PDFI
E
E
(O; 1-11
(4)
This definition is a principled extension, and recovers DFI (3) under &g loss; see Appendix A.3.
3.2 DISENTANGLED TRANSFORMATION WITH PROBABILISTIC FLOW
To overcome the limitation of (Gaussian) optimal transport, utilize flow matching to learns vector field that transports between a target distribution X into a simple source distribution Z (Lip- man et al., 2022). Suppose that Po and PI the densities for the source and target distributions, respectively. This amounts to constructing U, Ik - R such that u ~ Po. then Uz(u) Pr for some density satisfying Pt=o = Po- Pt=1 Pl. In particular, velocity field et is used construct the flow Uz( of the ordinary differential equation:
d
(.) = (u)) Uo(n)= = u.
dt
Flow matching (Lipman al., 2022; Liu al., 2022) was proposed to efficiently learn regression- based vector field along a predefined probability path. The central idea is to define an interpolation between Uo Po and U1 ~ PL. typically in the form Ur (1 (1-t)UG t)U + tU1, E [0, 1]. The velocity field for flow matching Lipman et al. (2022; 2024) is defined as
L(we Uo, U1), where C(w, Uo, U1) = / E w1) Us + Uol|² dl.
Ut. E arg min an w)Ela(pe)
The above optimization problem admits unique minimizer and can be formulated as vu( EIU, - Up U, = u| (Lipman et al., 2022; Liu et al., 2022; Hertrich et al., 2025), which further guarantees unique map U, by Theorem B.1. Hence, we can define FDFI (4) relative to this unique flow map T : U1 (not necessarily an optimal transport map) that transforms the original feature X to the latent feature Z, rather than discovering single ground-truth importance. Let U, be the flow map obtained by solving the above ODE with i, replaced with eü Then, the estimated transport map can be represented by T := U,. In particular, since the disentangled flow does not require labels, we can utilize independent, and potentially much larger, auxiliary unlabeled data to estimate T.
3.3
STATISTICAL ESTIMATION AND INFERENCE
After we obtain T from auxiliary samples, we construct an estimator for FDH using a set of ri i.i.d. samples {0:) For each data point O; = (X;, Yi), let 2. = =T(X)). The estimator is defined as:
FDD 02
[(Y. ( - 1
e(Y,
(5)
2M
where (2,6.)) M 1 is generated by resampling the j-th coordinate of 2, We analyze this as a cross- fit estimator for the parameter FFR defined in (2).
Before stating the main result, we briefly summarize the standing assumptions presented in Ap- pendix C. Assumption A2 imposes basic identifiability and regularity conditions ensuring that the latent representation and associated flow well-defined. Assumption A3 collects smoothness and integrability requirements on the velocity field e, that guarantee existence and stability of the flow T and its inverse. Assumption A4 (i)-(ii) encode the loss differentiability and standard complex- ity/rate conditions on the nuisance estimators, which together ensure the pathwise differentiability of (P)) and that the remainder of the cross-fit estimator (p

---

Published as a conference paper at ICLR 2026
Theorem 3.1 eficceeey of latent FDFI), Assume that Assumptions A2-A4 (i)-(iii) hold.
If the nuisance estimator satisfies Vo 1|0z - Telli, dt or(n-1/1), then the cross-fit estimator Bron given in Algorithm D.I is asymptotically linear. It satisfies the expansion:
2DDII 22
22 (2) (Pn, [P)) [P) + op(n -1/3),
where the efficient influence function (EIF) 42, (O;P)) is given by:
P2, (O:P) == w(;:T) 22, (P)
(6)
Consequently, vn(Ø2,
PD RDR) 0%
9 (O, ((pz;(O:P))) under the alternative Haj
R (P) +
+0. 0.
Theorem 3.1 establishes semiparametrie efficiency under the Neyman orthogonality condition that the estimand 22 is locally insensitive to first-order errors in the estimation of the nuisance transport map T. Consequently, the EIF (6) is identical to the EIF one would obtain if the actual map T were known priori. This permits us to use a flexible nonparametrically estimated T (or equiavlently the velocity field 0G through flow matching), which may converge at rate slower than n 11/2 and still achieve a /T-consistent, asymptotically normal, and efficient estimator for the latent importance FFDFI Without this property, the EIF would contain a complex correction term accounting for the O7
influence of estimating T, as shown in Appendix C.2.
To construct Wald-based confidence intervals for dFDRI we estimate the asymptotic covariance V(4z, (O:P)} by plugging in consistent estimators of T in (6) and evaluating the sample variance over independent observations. If d2. O, the observation contribution to the asymptotic variance will be zero, which leads additional complications that discuss further in Appendix D.3.2.
cross-fit fit estimator can also be used to estimate the importance 2 defined in (4) for raw features:
d'
(Y, (0 Ry((2)
Σ
Xi
n
2M
j=1
where (z) = (V is the square of estimated Jacobian of X, with respect to Zj. The statistical properties of this estimator are analyzed in the following proposition. In addition to the conditions of Theorem 3.1, Assumption A4 (iv) is mild strengthening of the complexity require- ment that ensures uniform convergence of the plug-in estimator of the flow estimator T.
Proposition 3.2 (Statistical inference of Φx,). Assume conditions in Theorem 3.1 and Assump- tion A4 (iv) hold. Then, the estimator FEDRI asymptotically normal under Hui : FDR (F) + 0:
vi(
2DH
OX FDFI (P)) 4 NOO,, (XPx;
where (O;PP) = (O:P)) is the EIF for (P) and Pjtt is the EIF for the component 631 = ()TT defined as Pa((;;) = (O)))) 6z)) + Co:(w,(O:T), IF nr, [)], where IF nt is the influence function of Hzt-
Proposition 3.2 provides inferential tools for the final FDFI scores. The EIF Øit for component at- tribution from Z, to Xi induces more complex structure than its latent counterpart in Theorem 3.1. consists of two parts: (i) a first-order approximation term, Wi (O; T) (X) - Φzt. which repre- sents the pointwise importance score centered by its mean, and (ii) second-order correction term, Cov(w, (O; T), H (O:-)) which arises because the squared Jacobian term H₃ is also estimated from data. The form of covariance reveals that the bias is driven by the covariance between point- wise importance scores and the wosnt-wise influence on the map's 's estimated geometry. Since H is estimated on the auxiliary sample of size m, its influence function IF Hy is of order Op(m-1/2
When m is large, the correction term is negligible, and approximate EIF can be used for practical inference (see Appendix D.3. Algorithm 1 outlines statistical inference procedure using FDFI.
7

---

Published as a conference paper at ICLR 2026
Algorithm FDFI (brief). See Algorithm D.I for full pseudocode and implementation details,
Require: Labeled data eest {X,, Y=) black-box predictor f and loss &; independent unla- beled covariates Dx {Xi flow training routine M and other hyperparameters.
Ensure: Latent importance scores (62, }3= and original feature scores {Øx, )(= with p-values.
1: Step 1: Disentangled representation. Train the flow on Dx to obtain transport map T M(Dx) such that = T(X) has approximately independent coordinates. For each labeled
point X;, set Z, == T())). J: vz) 2, and 37 (2,) == (Ji) for j, (1.....d) d). 2: Step 2: Point-wise latent perturbations. For each observation index i E (1...., n. and latent coordinate j {1,.... d), construct perturbed latent vectors { j== by resampling only the j-th coordinate of Z, from the reference pz,, keeping all other coordinates fixed. Define the point point-wise latent score Qij == 2N Ex=1 (Y,, ((T.f(TT (2( - f(x))
3: Step 3: Attribution original features. For each observation and original feature index E (1...., d), aggregate latent scores: 4 == - P.y (2.).
4: Step Latent importance and inference. For each latent coordinate j. compute Óz,
1 - Say in (5) and one-sided p-values based on estimated EIF components.
5: Step 5: Original feature importance and inference. For each original feature index I, compute
xx := - V in (7) and one-sided p p-values based on its approximate EIF.
6: return Importance scores and p-values.
EXPERIMENTS
4.1 THE IMPACT OF DISENTANGLEMENT
We evaluate the performance of LOCO, CPI, DFI, FDFI (SCPI), and FDFI (CPI) on identifying important raw features Xi's 's across varying feature correlation strengths and sample sizes through simulations. DFI uses the Gaussian optimal transport map. while the two variants of FDFI estimators given in (14) and (15) use flow matching. The simulated data is generated from a nonlinear response model y = arctan(Xe + X1) [[x,00} + sin(X X4) 1 <xx<<0) +e with ~ No Σ(p)) and e ~ (O, 1). The covariance matrix (p) := 1,8", E 500050 consists of 5 blocks of the equicorrelated matrix Σp E R 10 with correlation coefficient p. i.e., (Ep)is = if i = j and p otherwise. The = 50 features are partitioned into three non-overlapping subsets {1...., .d} = L Based on the response model, C1 {0, .......4 4) contains the active features that directly generate y. The set C2 = (5.... (5.....9) .9} contains features from the first correlation block that are correlated nulls, i.e, dependent on C1 but with no direct predictive effect. The set C3 (10, ..,,....44 49) contains the remaining independent null features from the other four blocks.
We evaluate all methods based four key metrics: (1) AUC (Area Under the ROC Curve) on C Cs. (2) Power on C, (3) Type error on C3. and (4) Power on C₁ U C2. AUC is computed by treating the estimated feature importance scores as prediction values and the ground-truth feature labels (informative 1. null = 0) as binary outcomes. For statistical inference, we test Hoj feature X, is not important versus H1;- feature Xj is important for each feature j. Based on a p-value Ps. statistical power is defined as the probability of rejecting the null when H1, is true, and type-I error the probability of incorrectly rejecting Hoj when is true, i.e., declaring null feature significant:
Power = P(P, S a Hi is true),
Type I Error = P(P3 S a Hoj is true).
In all experiments, we present results using random forests (Breiman, 2001) as the regres- sor/classifier for clarity, while additional comparisons with alternative predictors (Lasso and neural networks) and with more complex dependency structures are provided in Appendices E.1.
As shown in Figure 2, all methods control the Type-I error at the nominal 5% level on the indepen- dent null features (Ca) as expected. However, FDFI and DFI consistently outperform LOCO and CPI. attaining higher AUC and statistical power. Their performance is notably robust to increas- ing correlation, whereas LOCO and CPI's performance degrades significantly. Though both FDFI variants (SCPI and CPI) are theoretically equivalent under the & loss (Lemma 2.2), the CPI variant demonstrates superior sample power in low-sample or low-correlation regimes; thus select
8

---

Published as a conference paper at ICLR 2026
LOCO
CPI
FDFI(SCPI)
FDFI(CPI)
DFI
0.05
AUC Score
Power(C1)
Type Error
Power(C1uC2)
1.0
0.10
1.0
0.8
0.8
0.8
0.5
0.5
0.4
0.03
0.2
0.00
0.0
200
400
600
800
1000
200
400
600
800
1000
200
400
600
800
1000
200
400
600
800
1000
Sample
Sample
Sample
Sample
1.0
1.0
1.0
0.9
0.8
0.05
0.6
0.8
0.6
0.03
0.4
0.6
0.6
0.6
0.6
Correlation p
Correlation P
Correlation p
Correlation p
Figure 2: Simulation results of Section 4.1. Top: varying sample size with fixed = 0.8. Bottom: varying correlation with fixed 1000. We report AUC, power, and type-I error for each method over 100 runs. Shaded regions denote 95% bootstrap confidence intervals, and the dashed line indicates the nominal type-I error level (a 0.05).
FDFI (CPI) the representative method for subsequent experiments. These results demonstrate the effectiveness of disentanglement for assessing feature importance.
4.2 VALIDATION ON RÑA- GNQ-NEQUNNCING DATASETS
To further validate FDFT's performance with high-dimensional and complex correlated features, we evaluate it on two RNA-seq datasets: (i) the TCGA-PANCAN--HSSe bulk RNA-seq dataset 20531) for classifying five tumor types (BRCA, KIRC, COAD, LUAD, PRAD) (Weinstein et al., 2013); and (ii) a human single-cell RNA-seq dataset (n = 632, d = 23257) dis- tinguishing neoplastic cells from tumor core versus periphery (Darmanis et al., 2017). In practice, common ad-hoc approach for mitigating multicollinearity is to apply CPI or LOCO after hier- archical clustering. We select highly variable genes and compare the average prediction accuracy across datasets for important features selected by FDFI, DFI, and the ad-hoc approach, using two- fold splits and reporting the mean accuracy for each dataset; see Appendix E.2.5 for further details. As shown in Figure E10, FDFI consistently outperforms both DFI and the ad hoc method, high- lighting its superiority in identifying more predictive and biologically representative gene sets; see Appendix E.2.5 for specific biomarkers identified by FDFI and their clinical relevance.
(n 801,
4.3 CASE STUDY ON CARDIOTOCOGRAPHH DATASET
In clinical applications, feature dependency is ubiquitous, and decisions are high-stakes, making reliable interpretability essential for clinical decision support. To evaluate FDFI's practical utility, we focus here on the Cardiotocography (CTG) dataset (n = 2126, d = 21) for a case study, which uses fetal heart rate (FHR) and uterine contraction (UC) features to classify fetal state (normal y O, non-normal 1) (Campos& & Bernardes, 2000). We evaluate feature importance with the binary cross-entropy loss: I(y, Ø) = y log(ū) (1-y) log(1 ÿ). Additional large-scall real data studies, covering diverse domains and high nigh-dimensional settings, are provided in Appendix E.2.
On the CTG dataset, nonlinear feature correlations cause LOCO and CPI to identify only few es- sential features, whereas FDFI demonstrates substantially higher statistical power (Figure E7). Fig- ure 3(a) visualizes the FDFT attribution pipeline: latent importance scores (left bar plot) are mapped via the squared Jacobian heatmap to produce the final FDFI scores. The heatmap (ax,/02) )² veals the first-order influences among features, capturing a strong block-diagonal relationship among the FHR histogram features LB, Mean, Mode, and Median, which shows how the predictive impor- tance from underlying latent features is distributed across correlated features.
As illustrated in Figure 3(b), the ad-hoc approach clusters features using Spearman correlation and selects a medoid to represent each cluster. However, choosing the top-k features with FDFI or DFI achieves higher predictive accuracy than relying on cluster representatives from the ad-hoc method

---

Published as a conference paper at ICLR 2026
(E-0.88
0.55
D30
DOE
2.04
0.00
0.00
002
0.2
2.02
003
FDR-- (5=0.6)) 63)
0.00
3.75
6072
1.00
125
800
1.00
2.0
0:31
9.06
622
0.21
0.50
TDT1
LB
Mean
Mode
Median
MLTV
Temdennc
AC
Max
Noerros
Mmax
Wilt
MSTV
DS
Varfance
ALVV
ASTV
(b)
(c)
Hierorchical Clusteiing
Spearman
053
Moan Moanacconn acuracy
023
LB
Ac
Wours
ATV
Figure 3: Data analysis of the CTG dataset. (a) FDFI estimation. Barplots show the estimated latent importance scores o2 (left) and original importance scores FDDRI (top). The heatmap shows
the squared Jacobian weight /0 (b) Hierarchical clustering results based on Spearman correlation. (c) Prediction accuracy with selected important features for FDFI, DFI, and ad hoc method that applies CPI on the cluster-representatie features.
(Figure 3())). Moreover, FDFI consistently outperforms DFI, underscoring the benefit of its flexible flow-based mapping DFI's more restrictive Gaussian assumption.
CONCLUSION
We introduce model-agnostic framework, FDFI, that uses flow matching assist feature impor- tance attribution under general differentiable loss functions. We establish the semiparametric effi ciency of our FDFI estimators and provide valid statistical inference. Empirically, FDFI resolves the correlation distortion problem (Verdinelli & Wasserman, 2024b), successfully recovering estab- lished and correlated diagnostic features that classical methods overlook.
Beyond the empirical evaluations above, FDFI also enables several practical downstream workflows that benefit from statistically principled importance scores under dependence. First, by ranking features in way that is robust to correlation, FDFI can be used for feature pruning and model compression, preserve predictive performance while reducing dimensionality, inference cost, and model complexity (Han et al., 2015; Nelson et al., 2022; Ranek et al., 2024). Second, in data col- lection and experimental design, FDFI highlights variables or modalities with the highest marginal utility to prioritize operations under resource constraints (Wang et al., 2024). Third, by mitigating spurious attributions that arise from correlated covariates, FDFI improves model debugging and the detection of implausible signals, unlike traditional importance measures that assign non-zero relevance to purely null variables (e.g., Chen et al., 2022). summary, FDFT reliably guides what to measure, keep, discard in high-stakes resource-limited decision-making pipelines.
Lastly, FDFI inherently local sensitivity measure in a learned latent space and is most appropriate when the flow and predictor f --1 are locally smooth and well-approximated by first-order expansions. In highly non-smooth or purely combinatorial regimes (c.g.. parity/XOR- -tpe rules with near piecewise-constal predictors), infinitesimal latent perturbations can fail to probe the truly influential directions. Developing extensions of fow-disentangled importance reliable under such non-smooth, combinatorial interaction structures is interesting direction for future work.
10

---

Published as a conference paper at ICLR 2026
Acknoledgement. The corresponding author gratefully acknowledges support from the Institute for Data Science (IDS) at the University of Hong Kong through the HKU100 award. Part of the computations was performed using research computing facilities offered by Information Technology Services, the University of Hong Kong.
Ethics statement. The authors acknowledge and adhere to the ICLR Code of Ethics. This research utilizes publicty available, anonymized biomedical datasets for empirical validation, which stan- dard in the machine learning literature, and their use in this context, to the best of our knowledge, presents no new risks patient privacy. We do not foresee any direct negative societal impacts stemming from our methodological contribution. Furthermore, in line with ICLR policy, our use of Large Language Models as writing assistants has been disclosed in the Appendix.
Reproducibility statement. We have taken several steps to ensure the reproducibility of work. The main text provides clear descriptions of the proposed method, experi- mental setup, and evaluation protocols. Complete proofs of the theoretical results and detailed explanations of assumptions are included in the Appendix. For experiments, we deseribe datasets and hyperparameter settings the Appendix. The code for re- producing results be found https: /ttp://gttttt/tt/ cwm//orao/ia seforAndahher/ cert
WWW.-DTTSENNTNNNLLA-EEEEEEEEEE IMPORTANCE/ and the FDFI package can be found at https://gf com/
REFERENCES
Leo Breiman. Random forests. Machine learning, 45(1):5-32, 2001.
D. Campos and J. Bernardes. Cardiotocography. UCI Machine Learning Repository, 2000.
Emmanuel Candes, Yingying Fan, Lucas Janson, and Jinchi Lv. Panning for gold: mmdel- d'knockoffss for high dimensional controlled variable selection. Journal of the Royal Statistical Society Series B: Statistical Methodology, 80(3)5551-577, 2018.
Ahmad Chamma, Denis A Engemann, and Bertrand Thirion. Statistically valid variable impor- tance assessment through conditional permutations. Advances in Neural Information Processing Systems, 36:67622-67685, 2023.
Hugh Chen, Scott M Lundberg, and Su-In Lee. Explaining series of models by propagating shapley values. Nature communications, 13(1):4512, 2022.
Spyros Darmanis, Steven A Sloan, Derek Croote, Marco Mignardi, Sophia Chernikova, Peyman Samghababi, Ye Zhang, Norma Neff, Mark Kowarsky, Christine Caneda, et al. Single-cell ma- seq analysis of infiltrating neoplastic cells the migrating front of human glioblastoma. Cell reports, 13355 1410, 2017.
Jin-Hong Du, Kathryn Roeder, and Larry Wasserman. Disentangled feature importance. arXiv preprint wwXiv:2507.0026// 2025.
Abraham Genizi. Decomposition of R² in multiple regression with correlated regressors. Statistica Sinica, pp. 407-420, 1993.
Jaime Roquero Gimenez, Amirata Ghorbani, and James Zou. Knockoffs for the mass: new feature importance statistics with false discovery guarantees. In The 22nd international conference on artificial intelligence and statistics, pp. 2125-2133, PMLR, 2019.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Johannes Hertrich, Antonin Chambolle, and Julie Delon. On the relation between rectified flows and optimal transport. arXiv preprint arXiv:2505. 19712, 2025.

---

Published as a conference paper at ICLR 2026
Giles Hooker, Lucas Mentch, and Siyu Zhou. Unrestricted permutation forces extrapolation: vari- able importance requires at least one more model, or there is no free variable importance. Statistics and Computing, 31(6):82, 2021.
Bertrand looss and Clémentine Prieur. Shapley effects for sensitivity analysis with correlated inputs: comparisons with sobol indices, numerical estimation and applications. International Journal for Uncertainty Quantification, 9(5), 2019.
Silke Janitza, Ender Celik, and Anne-Laure Boulesteix. A computationally fast variable importance test for random forests for high-dimensional data. Advances in Data Analysis and Classification, 12(4):885 915,,915,222 2018.
Jing Lei, Max G'Sell, Alessandro Rinaldo, Ryan J Tibshirani, and Larry Wasserman. Distribution- free predictive inference for regression. Journal of the American Statistical Association, 113 (543):109-11111, 2018.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:24/2.06266 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.033: 200000, 2022.
Scott Lundberg and Su-In Lee. A unified approach to interpreting model predictions. Advances in neural information processing systems, 30, 2017.
Evan Mason and Zhe Fei. Novel knockoff generation and importance measures with heterogeneous data via conditional residuals and local gradients. arXiv preprini arXiv:2508. 14882, 2025.
W James Murdoch, Chandan Singh, Karl Kumbier, Reza Abbasi-Asl, and Bin Yu. Definitions, meth- ods, and applications in interpretable machine learning. Proceedings of the National Academy of Sciences, 116644)220711-22000 2019.
Michael E Nelson, Simone G Riva, and Ana Cvejic. Smash: scalable, general marker gene identi- fication framework for single-cell ma-sequencing. BMC bioinformatics, 23(1):328, 2022.
Jolene S Ranck, Wayne Stallaert, Justin Milner, Margaret Redick, Samuel C Wolff, Adriana Beltran, Natalic Stanley, and Jeremy Purvis. Delve: feature selection for preserving biological trajectories in single-cell data. Nature Communications, 15(1):2765, 2024.
Angel Reyero-Lobo, Pierre Neuvial, and Bertrand Thirion. Conditional feature importance revisited: Double robustness, efficiency and inference, 2026. URL https: http:///axvvvv.////eeeeeee.aa..a org/ abs//2/ 2501. 17520.
Alessandro Rinaldo, Larry Wasserman, and Max G'Sell. Bootstrapping and sample splitting for high-dimensional, assumption-lean inference. The Annals of Statistics, 47(6))3438-344 -469,, 2019.
L. S. Shapley. A value for n-person games. In Contributions to the Theory of Games (AM-28), Volume I1, pp. 307-318. Princeton University Press, Princeton, NJ, 1953. ISBN 9780691079355
Carolin Strobl, Anne-Laure Boulesteix, Thomas Kneib, Thomas Augustin, and Achim Zeileis. Con- ditional variable importance for random forests. BMC bioinformatics, 9:1- 2008
Isabella Verdinelli and Larry Wasserman. Decorrelated variable importance. Journal of Machine Learning Research, 15(7): 2024a.
Isabella Verdinelli and Larry Wasserman. Feature importance: A closer look at shapley values and loco. Statistical Science, 39(4)623-636, 2024b.
Yu Wang, Nan Liang, and Ge Gao. Quantifying the regulatory potential of genetic variants via a hybrid sequence-oriented model with sven. Nature Communications, 15(1):10917, 2024.
12

---

Published as a conference paper at ICLR 2026
John N Weinstein, Eric A Collisson, Gordon B Mills, Kenna R Shaw, Brad A Ozenberger, Kyle Ellrott, Ilya Shmulevich, Chris Sander, and Joshua M Stuart. The cancer genome atlas pan-cancer analysis project. Nature genetics, 45(10): 11000:1:111111111 111-111110011 2013.
Brian Williamson and Jean Feng. Efficient nonparametric statistical inference population feature importance using shapley values. In International conference on machine learming, pp. 10282- 10291, PMLR, 2020.
Brian D Williamson, Peter B Gilbert, Marco Carone, and Noah Simon. Nonparametric variable importance assessment using machine learning techniques. Biometrics, 77()):--2,, 2021.
13

---

Published as a conference paper at ICLR 2026
Appendix
This serves appendix to the main paper. Below, provide an outline for the appendix along with a summary of the notation used the main paper and the appendix.
Organization.
The content of the appendix organized as follows.
Appendix
Content
Appendix A.
Proof of Theorem 2.1.
Appendix A
Appendix A.2
Proof of Lemma 2.2.
Appendix A.3
Equivalence of FDFI and DFI under &2 loss.
Appendix B.1
Uniqueness flow solution (Theorem B.1). Estimation error of transport map (Lemma B.2).
Appendix B
Appendix B.2
Appendix C.1
Assumptions.
Appendix C
Appendix C.2
Proof of Theorem C.I and Theorem 3.1.
Appendix C.3
Proof of Proposition 3.2.
Appendix D.I
The full FDFI algorithm and computational devices. Technical details flow matching FDFI. Technical details of statistical inference for FDFL Computational time comparison.
Appendix D.2
Appendix D
Appendix D.3
Appendix D.4
Appendix E.
Extra simulation studies.
Appendix
Appendix F.2
Extra real data studies.
Notation. An overview of some general notation used in the main paper and the appendix is as follows.
For vector X E Rd. Xs E Rls denotes a sub-vector of X indexed by S C [d]. and X-3 == In Rd, the j-th standard basis vector is denoted by ej and the zero vector is denoted by 0g or simply 0 if the dimension is clear from the context. The cardinality of S is denoted by |S). The indicator function is denoted by [.)--
(11....j1..j11........
For a tuple of random vectors O = (X, Y), ), the expectation and probability over the joint distribution P denoted by IE()) and P(-), respectively. For (potentially random) measurable functions ʃ. we denote expectations with respect to the data-generating distribution of O alone by P.((O) = fdlP, while [O)] marginalizes out all randomness from both O and any nuisance functions dependent on. The empirical expectation over samples is denoted by (O. En= The L2 norm of function is denoted by [l
For statistical estimand c, we write ((P) to emphasise its dependence on the underlying distri- bution P. The population and empirical variances (covariances) are denoted by V and V, (C and C.). The d-dimensional multivariate normal distribution with mean μ and covariance Σ is denoted by Na(,, Σ). For matrices A1, Ag...., Ag, the notation diag(A,, A2.... Ag) represents block diagonal matrix that combines all of the matrices.
We use "o" and "O" to denote the little-o and big-O notations; "up" and "O>" " are their probabilistic counterparts. For sequences } and {bn), we write an S b, if a, = Obb.): and an x bn a, Obb,) and b, = Oa,)) Convergence in distribution is denoted by " "
The Use of Large Language Models. In preparing this manuscript, Large Language Models (LLMs) were used strictly as auxiliary tools for: (i) Language editing and polishing: improving clarity, grammar, and academic style of author-written text without changing technical content. (ii) Literature search assistance: surfacing potentially relevant references and recent work; all cited ma- terials were independently verified and read by the authors. No parts of the scientific contribution (problem formulation, methodology, experiments, results, or analysis) were generated by LLMs. The authors take full responsibility for all contents of the paper, including any text refined with LLM support. LLMs are not eligible for authorship.
14

---

Published as a conference paper at ICLR 2026
A COMPARISON BETWEEN LOCO- AND CPI-BASED DFI
A.1 PROOF OF THEOREM 2.1
Proof of Theorem 2./. The proof proceeds in three parts. First, we prove the exact identity. Second, we apply the triangle inequality. Third, we derive the bounds for the error terms Fapprox and Jg-
For simplicity, write X(f:X,Y) and R(f-j:X-jj+ Y) R(f) and R(f-j), respectively. Part 1: Proof the Identity. We start with the right- hand side the identity and substitute the definitions of @CCPI Eapprox, and Jg-
1 Jg -Eapprox
SCPI X,
2 Φ
2
(R(g)) R(()) (B(C(Y, R(g;)) (R(gs) R(S-s))
= (3 R(())) - R(g,) +
1 (((O)) + +((f-) R(f) Ele(y, f(x(0))
2
R(f-)) R(J) FlC(Y,
(R(J-)) ((()))))((())) 2 (R(S; x(), R(J:X))
CPI OX, Xx,'
CPL This confirms the identity Øx, Øx, = 2 [ - Jo) Eapprox.
Part 2: Application of the Triangle Inequality. Taking the absolute value of the identity and applying the triangle inequality yields:
10 0x,
10xx, Jaa/22 Expprox
1 [ScPI
Jul +1 Eapprox]
2
EMIE| Empproxx
The last step follows from the fact that Expprox 0, because the function f-y is the Bayes optimal predictor, meaning it is the minimizer of the risk functional.
Part 3: Bounding the Error Terms. We now derive the bounds for Eupprox and EMIE-
(i) Bounding Eapprox = R(g,) - R(f-)): The risk R(h) = E[E(Y, h(x-))) is a convex functional of the predictor h. For convex and Mmsoooth functional, the difference in value between any point and the minimizer is bounded. Because R(h) is M-smooth with respect to the L2 norm on functions, we have:
M 1g - s-sll
R(g;) 10-R((J)))) 2
- (
2
This provides the bound for the approximation error.
(ii) Bounding J, = E[((Y, E(((Y, - R(gg): We can write Jy using iterated expectation:
Jn By,x B(e(Y, X-5] {() (((()) x-)D]
(8)
The inner term is the Jensen gap for the convex function E(y, :) and the random variable f(x()) (where the randomness comes from X, ~ P(- X-))). For an M-smooth convex function ø, the Jensen gap is bounded by E[Ø(Z)] - (E < Mf WZZ). Applying this to the inner expectation (conditional on Y, X-x)):
E f(X())) X-,] -Xx,)--(YYEEE E[ (1) x x-11) S M
v(((X() (x-))).
Taking the expectation of both sides with respect to Y, X-j gives the final bound:
Jg < E
v(((x)
x-1) =
E (V(((X) ((((X)Xx)))) x-,))-
2
2
15

---

Published as a conference paper at ICLR 2026
(iii) Bounding |EMIEl: By the triangle inequality on the definition of EMIE:
1 ( اوl -
EMIEI
SCHI x;
SCPI + Jg).
2
2
From step (ii), the bound for Jg is: J9 S ME ( Xx,]]] To bound (xx
[SSPI we use a second-
order Taylor expansion of ()) around f(X). By Assumption A1, for some & on the line seg- ment between Øj (x- and ((X), we have:
+ C( ((x f(x)) _ - - ((x))²
f(Y. ((Y,g, =
Taking the expectation and rearranging gives
dSCPL ox;
R(g,) R(f))
B(((Y, e(Y,
F(((V,((()((sxxx))) -X + [ ((X))
Since f is the Bayes optimal predictor. ((X(),((X)) X] = 0. The first term vanishes by iterated expectation: (X, | (X_,) -
= O. This leaves:
(E) - n3 & "1((((()) | - f) &
1
M Bl(g, - f)²l.
scCpi =
2
Recognizing that [X(ox(X-;)) -(x))))==VV(1( EVV((()) x-1)), get
|0,,1| by (((((((
X-))]-
Combining the bounds for | and Jg:
1 |EMIEE= S 2 (10x, +g) - E (((((X) x-x)] + 4 E[V(((X) x-;)]
M
M
M
E (V(((X) X-,)).
Substituting the final bounds for Eapprux and |MIE into the result from Part completes the proof.
A.2 PROOF OF LEMMA 2.2
Proof of Lemma 2.2. For the &g loss, we analyze each term of the identity &1oco 2x = (
J2)/2 Eapprox- Since [x=(X_)) BI X-y) and gj (X-x) EE X] X_y) ELy
X X_ we have J-, = gy- This implies the approximation error Fapprox R(g,) - R((-s) = and by definition,
LOCO 6X,
ΦX,
(9)
For (8) under the & loss, from the derivation in Du et al. (2025a, Lemma 2.2), we also have that
xx,' Jn ((V((()) X-3))-
(10)
Substituting these into the identity gives eyoco Øx,
J
JJ 0 = 0, i.e.,
1,00 ΦX,
CPI = ØX2-
(11)
Combining (9) and (11) completes the proof.
A.3 EQUIVALENCE OF FDFI AND DFI UNDER l LOSS
Recall that the DFI and FDFI estimands for the original features are defined in (3) and (4), respec- tively.
16

---

Published as a conference paper at ICLR 2026
For l2 loss, the influence function component in (4) becomes
w(O;)) = [( (01,11011220))-00000001111
= -
Taking expectation conditional on Z-j, we further have
(O:(O;T) Z-y] 2 E - -x-)) Z-jl
(Y
-1/(011220))--11111110
- 2-5]
- -1117(()22)) (Z))- 2-1 V(F(T (2)) -- 2 (2) 2-1),
(Z)(
which reduces to the component in (3). Therefore, the DFI and FDFI estimands coincide under & loss.
17

---

Published as a conference paper at ICLR 2026
B PROPERTIES OF PROBABILISTIC FLOW
B.1 UNIQUENESS OF SOLUTIONS
Assumption A2, Let (Uo. U1.U1) E 3dd admit a joint density pe(uo, u1, for & E [0, 11. Define
pe(u) == | pi(uo. u1, u) dugduu,
m²(u) = / (ur - uuo)p un, u) dugdui,
and
ve(u)= = mi(u)
pr(u)
Assume that
(i) Continuity: my are continuous on [0, 10,1]xB0 XRx, are C in u, and Wa V VuPt are continuous on [0, 1] x Rd:
(ä) Uniform lower bound and bounded score on compact supports: For every compact set K C Rd, there exist 3K > 0 and MK < 3o such that
inf (L,))E(0.1IxXK pu(u)= AK:
sup (
\o s Mx;
(iti) Bounded Vumy on compact supports: For every compact set K, there exists JK < oo with
sup
[Vmm{(u)l S Jk:
(1,
(iv) Conditional second moment with linear growth (uniform in t): There exist bounded continuous c0 [0, 1] [0, oo) such that
E[ - Uoll U. = u] S ct( + (t)llll for all (t, u).
: sup; v cc((t) oo and supt v (t((t) <0.
The conditions in Assumption A2 standard in the theory of ordinary differential equations, As- sumption A2(i)-(iii ensure that the velocity field (u(u) is continuous and uniformly Lipschitz in u. This is precisely the regularity required by the istence and uniqueness of the ODE solution (Hartman, 2002; Amann, 201 1). Assumption A2(iv) fur- ther imposes a linear growth bound, ruling out finite-time blow-up and thus ensuring that the solution extends globally to [0, 1] (Hartman, 2002; Amann, 2011). In the Gaussian mixture setting, Assump- tion A2(i)-íii) hold since the density and score are smooth and bounded on compacts (Bishop & Nasrabadi, 2006; Hyvärinen, 2005). Moreover, the conditional expectation EE - Un U, u] is convex combination of affine functions u, so that vy (u) remains smooth and m,(u) = p: (u())e (u) has gradient V, (u) continuous and bounded on compacts (Anderson, 2003; Bishop & Nasrabadi, 2006). Assumption A2(iv) verified in the final part of the proof.
theorem, which guarantees local ex-
Theorem B.1 (Local and global uniqueness of Ut). Under Assumption A2 (i)-(iii), the velocity field ve (u) is continuous on [0, 1] x Rd and, for every compuct K C Rd, is uniformly (in t) Lipschitz in Hence, for every initial R.d, there exists a such that the ODE
d U.( = (v,(())). Uo(u) = u,
dt
admits locally unique solution on (0,
lf, in addition, Assumption A2 (iv) holds, then the solution extends to the whole interval [0, 1]. and globally unique. Moreover, this condition automatically satisfied when the source distribution Un is Gaussian and the target distribution Un follows a Gaussian mixture distribution.
In the standard theory of ordinary differential equations (ODEs), local uniqueness of the solution follows immediately if one directly assumes that the velocity field vx(u) is continuous in & and
18

---

Published as a conference paper at ICLR 2026
locally Lipschitz in This is precisely the setting of the classical these conditions, the ODE admits a unique local solution. A further linear growth condition ensures global uniqueness.
theorem. Under
In practice, when U is parameterized by a neural network, Lipschitz continuity can be enforced by architectural choices: for instance, using I-Lipschitz activations such ReLU tanh together with spectral norm constraints on each weight matrix. In this case the Lipschitz constant of the network upper bounded by the product of the spectral norms of the layers (Virmaux Scaman, 2018; Bartlett et al., 2017; Gouk et al., 2021). Compared to directly assuming Lipschitz continuity of or. our Assumption A2 adopts distributional perspective: it imposes regularity and boundedness on P: and my, which in turn imply that U/ is Lipschitz continuous. The advantage of this formulation is that it does not rely on a particular parameterization of o and can be verified for broad distribution families such as Gaussian mixtures.
Proof of Theorem B.I. The proof proceeds in six parts. The first five parts establish the proof of Theorem B.1, and the last part shows that the Gaussian and Gaussian mixture case automatically satisfies Assumption A2 (iv).
Part 1: Continuity of oz By the quotient rule (componentwise),
Vuvi(u) Vumi(u) (z(u) & V (())
pe(u)
Assumption A2 (i) guarantees that mu(u) and (e(u) are continuous on [0. 1] Rd they are differ- entiable in and their derivatives V.m,(u) and V (u) are also continuous. Since (r > O, the reciprocal 1 I/p.(u) (u) is continuous, and thus (,),(u)/ Pt (u) is continuous V. log (u) = V..p.(u) (u)/px()) is continuous by the same reasoning. Consequently, both terms on the right-hand side above are continuous, which shows that (,vv(u) is continuous. Together with (u) m,(u)/po(u) being continuous, conclude that both i, and V. continuous [0, 1] x Rd.
well. Moreover,
Part 2: Uniform Lipschitzness of Ut. Fix compact K with radius RK suPack |//|]. From Assumption A2 (ii) and (iii),
Vumi(u) pr(u)
JK 3x
sup I
From Cauchy-Schwand inequality and Assumption A2 (iv).
[[e(u))] S voo(t) + (t) ||||| S A + B [|u]]
so (j. [[r(())] A+ BRK. Using ||a & bl| I ||a|| |bb|| and Assumption A2 (ii),
JK 3K
sup (
[V <
+ Mr (A BRx) LK < Co.
Thus, [lo.(u) (.(ll < Lx lu- wll for all and u, u E K.
Part 3: Local existence and uniqueness of the ODE solution. By Part 1 and Part 2, (u(u) is continuous and uniformly Lipschitz in on compacts; Picard-Lindelöf theorem yields unique solution on T),
Part 4: Proof of priori bound and extension. Using ||v.(U.))l < A BI
|| S A B|lla|| + ||U||| < (||=l) + Bt 8 < [[|u| + A)e (0,
dt
Hence, there no finite-time blow-up, the solution extends to [0, 1].
Part 5: Global uniqueness of the ODE solution. Let Ug and V, be two solutions with the same initial condition Uo Vg. By the linear-growth bound [x][]]]
A + B||z|| and Grönwall's inequality, there exists R < oo such that [|Uel < Rt and [[all < R for all t E [0, 1]. Hence
both trajectories remain the common compact ball K {2 ||ж|| s R). on which the field uniformly (in t) Lipschitz in w:
(w))(w) vx(v)) S Lx ||u oll
1I E [0, 1]. Vu, E K.
19

---

Published as a conference paper at ICLR 2026
Writing the integral form of the ODE and subtracting,
Ut Vi 1 (vo(Ue) - (v)(V)) (c) ds,
we obtain, with Z = U2= - Vt=
1Zell - / [Io((U))- (V.)ll 1 L'' LK ||
By Grönwall's inequality in integral form, letting ((t) == |Ze| yields q(t) I ((0) exp(Lxt) = 0, so (t)) and hence U, = V,. Therefore the solution on [0, is unique.
Part 6: Proof of the Gaussian and Gaussian mixture case. We first consider the case where U1 follows a Gaussian distribution, and then extend the argument to show that the result also holds when U1 follows Gaussian mixture distribution. Note that Gaussian distribution can be regarded as a special case of Gaussian mixture distribution.
Gaussian. Consider Y
[Uo; U1] E I2dd such that Y NMu,, E) with mean and covariance:
(D (1))
Euo E10
Eo1
= Σ
> 0.
Σ11
Define U+ t) +tU = C.Y.D== U1- BY, where C.= [[1 - t)da, tla E Rdx2d and B I--,,, Ia] E Rdx2d, By the Gaussian conditioning formula,
Y U, =u N(++ + EC, Ac(u C(p), E-ECTACEE
with A, = (C,ECT) y Consequently,
D =u ~ N(mllt, S(O),
where
m(t,u) alt) + M(tJu, S(t) = B(E A.GE)
with ((1) =Bp- BTCT A,CrN and M(I) BEC, A,. Thus,
DI U. = u] = ||m[t, u) w)ll Ar 5()).
Since [|m(t, < 1 + 2 ( || /ul we obtain
2 U, w| s (211a(011 (t)())))++))))+--))) + 21)
||]|
=:0o(1)
=te((t)
Here c2()) are continuous and bounded on [0, 1].
Gaussian mixture. Consider U1 ~ K and Z (1,.... K} with P(Z k) Tk. Conditioned on Z k. (Un, U1) is jointly Gaussian, and the bound from the previous Gaussian case yields
Et= nu ( Σ1k), where Ta > 0, E Th =1 ΣΙ > 0,
D|| U. u, Z k] Ctx(t) (t
Denote the posterior weight by wx(t.. u) = P(Z = k U. = u). Hence
k
2|| U. = uj =
Σ
wk (t, u) El| D1 U, = u, Z = k)
k=1
<
max Iskck
con(t)
+
max CIA (t) 11/112. ISKSK
This completes the proof.
20

---

Published as a conference paper at ICLR 2026
B.2 ESTIMATION ERROR
In this subsection, we characterize the estimation error of the resulting transport map in terms of the estimation accuracy of the probabilistic flow. We require the smoothness condition of the estimated velocity field associated with the learned flow.
Assumption A3 (Smoothness of velocity field). Assume that 0e (u) is continuously differentiable in (t, u) and Lipschitz in uniformily on (t, u) [0, 1] x R with Lipschitz constant K.
Under the uniqueness assumption and the above smoothness condition, we can show that the estima- tion error of the transport map can be controlled by the uniform (over the time horizon) estimation error of the velocity field. This suggests that if one can estimate the velocity field well enough, then T can also be well approximated in L2 norm.
Lemma B.2 (Estimation error of transport map). Under Assumptions A2 and A3, it holds that
T - S e222k [." 110e - -.el12, dt.
Proof of Lemma B.2. By Assumptions A2 and A3, Albergo & Vanden-Fijnden (2023, Proposition 3) gives that
? IT - -T11 = 1 -
po(u) du
gzd
s e2 [./.. / (va(Ui(w)) (U(((()) po(u) du
e f 1|0. Fell
Fell2ggdl. dl.
21

---

Published as a conference paper at ICLR 2026
C SEMIPARAMETRIC THEORY
C.I. ASSUMPTIONS
To formalize the asymptotic analysis of our estimator, rely on the von Mises expansion, which is a form of Taylor series for functionals. This requires a precise notion of functional derivative. We introduce two related concepts: the Gateaux derivative, which directional, and the stronger Fréchet derivative, which ensures a uniform linear approximation.
The Gateaux derivative of the functional with respect to the function T in the direction of another function h is defined as:
@TT + ch) ((T)
V [T) lim
describes how the functional behaves along specific linear path. Gateaux differentiability is also referred to as pathwise differentiability. where the path specifies the direction in the function space. To control the remainder term of the Taylor expansion for any perturbation, not just along straight lines, a stronger condition of Fréchet differentiability is needed to guarantee that the functional can be well-approximated by linear map in a neighborhood of T. Formally, let V and W be normed vector spaces. let U V be an open set, and let f U + W. The function is Fréchet differentiable at E U if there exists bounded linear operator A V - W such that:
[l + h) (x) Ahllw
lim [| Iv
||bllv
such a bounded linear operator A exists, unique and called the Fréchet derivative of f at x. denoted by Df()).
Within this framework, a pivotal concept Neyman orthogonallity (Chernozhukov al., 2018). This condition is defined using the Gateaux derivative: the functional é is Neyman-orthogona with respect to the nuisance parameter if its Gateaux derivative the zero functional, meaning evaluates to zero for any valid perturbation function h:
Vo(T) -0
Vo(T) 0, wh.
When this condition holds, the first-order bias from the estimation of the nuisance function van- ishes, which simplifies the asymptotic analysis and leads to estimators with only higher-order bias.
For the remaining section, we will write FDDII and Φx. as Φz, and Φx,, respectively, for notation simplicty. We require assumptions about the data generating process and nuisance function estima- tion.
Assumption A4 (Regularity conditions). For any P, assume the following holds:
(i) (Smoothness) The map T + w(; is twice Fréchet differentiable with respect T. The loss function E(y, 7) is differentiable with respect to its second argument, and the function T continuously differentiable.
(ii (Donsker Class) The class of functions (w(O:T") IT - -Tlle, < 8} for some > i
P. Donsker class.
(ii) (Neyman orthogonality) Gateaux derivative of Φz, with respect to T satisfies that Væz, (T)
(iv) (Nuisance estimator influence function) The nuisance function estimator T, trained on an in-
dependent auxiliary dataset with size m, admits an efficient influence function representation. For any point ity estimation error can be lineurized as:
-(x)-T(x) = paa [I==)[= ( O= -T112,,
where paux is the empirical average over the auxiliary data, and
=
Assumption A4(i smoothness condition, and Assumption A4ii)-(iii) are analogous to those nec- essary for the results in Chernozhukov et al. (2018); Williamson & Feng (2020).
Assumption A4(ii) ensures that empirical process terms of the form (P,, well- bell-behaved. It can be weakened to proper bounded moment conditions with cross- cross-fitting, in double machine learning literature (Chermozhukov al., 2018). However, we stick with the former assumption to simplify the exposition.
22

---

Published as a conference paper at ICLR 2026
Neyman orthogonality condition (Assumption A4(ii)) is the conceptual equivalent of the core quirement in Williamson & Feng (2020), as both require the parameter interest to be locally quadratic (i.e., first-order insensitive) to estimation errors in the nuisance function. The primary distinction arises from the nature of the nuisance parameter itself: our framework must manage the estimation of the complex generative transport map T, whereas related works typically analyze the estimation of the regression function f.
The Neyman orthogonality condition Assumption A4(ii)) holds in important cases. For instance, for the f2 loss, if the predictor ((x is the true conditional expectation, f()) [ = æ]. the Gateaux derivative of the parameter Øz, with respect to is zero. This occurs because the derivative term involves B| - f(X) X]. which is zero by definition. In this scenario, our estimator first-order insensitive to errors in estimating T. The practical implication is that even though T is estimated nonparametrically and may converge slowly, the final estimator áz, remains vn- consistent and asymptotically normal under the alternative without requiring a correction term for the estimation of T. When the Neyman orthogonality condition does not hold, the first-order error from estimating the nuisance map T no longer vanishes. To achieve vnn-consistent estimator for Z we must explicitly correct for this bias using a one-step estimator can be used to correct for first-order bias, which is possible the nuisance estimator T itself admits a linear expansion, in Assumption A4(iv).
C.2 SEMIPARAMETRIC EFFICIENCY
Under Assumption A4 (i), (ii) and (iv), we first derive general EIF in Theorem C.I without the Neyman orthogonality assumption. We will then specialize this result to prove Theorem 3.1.
Theorem C.1 (General EIF for Latent FDFI), Let the data be randomly partitioned into a main sample of size n and an auxiliary sample of size m. Let T be an estimator of the nuisance function T computed on auxiliary samples. Under Assumption A4 (i), (ii) and (iv), the cross-fit estimator ( Pr (P([((;;T) T) for the parameter Φ2, (P) [O: (O:T) satisfies the following asymptotic
expansion:
(P) 12 (P) (O:P) + (Pm (PP + Rn,m.
where Pn and F. are the empirical measures for the main and auxiliary samples, respectively. The efficient influence function Φz, (O;P) is given by:
92, OO; = Pmain (O; IP) + Pcorr (O;
where the components of the efficient influence function are:
imaann (O;I) = wOO - - (IP)
(2((2))))))--((2()2))222)) 22, (P), Pcorr (O:P) (O(T)
= 2 --
The remainder term Rn,m is of order Op(n-1/2E +En), where &n = - is the L2 error of the nuisance estimator computed on the auxiliary sample. Here, X = x0) = (Z()), (+)()))) is the influence function of evaluated at observation O, and
Dx = Ey ((((() T'
X=
with El (u, y) being the derivative of l with respect to its first argument. The expectations Ex|:] and E Exool are taken over the marginal distributions of X and x(),, respectively.
The EIF consists of the "naive" influence function, (O;T) oz (P), which would be correct if T were known, plus correction term. This correction term, (T(T)(IFT(O;.)))) )), directly accounts for
23

---

Published as a conference paper at ICLR 2026
the first-order impact of using estimate T instead of the true T. When the Neyman orthogonality condition holds (Vo(T) = 0), this correction term vanishes, and we recover the simplified EIF from Theorem 3.1. Therefore, this theorem provides a complete characterization of the estimator's asymptotic behavior, nesting the orthogonal case as a special instance.
While the EIF itself is independent of sample sizes, the asymptotic variance of the cross- fit estimator depends on n and m. For a cross-fit estimator, its variance is approximated by a sum of variances of the EIF components, scaled by the respective sample sizes:
v(žz,) ~ ((Pmiu)) +
1 =((ooor))
when the remainder term Rm,m is negligible.
C.2.1 PROOF OF THEOREM C.I
Proof of Theorem C.I. Note that (5) is sample average of M copies of Zj. Below, we present the proof for the simplified case M = 1, and the conclusion holds for general M using the linearity and independence across Z(ak) The proof proceeds by deriving the von Mises expansion (Du et al., 2025b) for the estimator Øz, (); We decompose the estimation error as follows:
oz, - ~z =P -B(((O))
= (I,, - - + (P,, -))))((())))
Term 1
Term II
Torm III
Term is the desired empirical average of the influence function's 's main part, centered to have zero mean. By the central limit theorem, it converges to a normal distribution when V(v) >
Term is the first-order bias term due to plugging in the estimated nuisance function î. By
Assumption A4 (i), the Taylor expansion of (TT) == [O; around the true T' gives:
Term II = o(f) o(T) = - -
where [ is the Gateaux derivative of &i in the direction 87 = ff - given by:
d
Vo(T)[67] = de
B[w(O:T+ cōT]]
d
= E
- (6)
de
l==22
d le= + cdr) (2) = GI(T 1z))
Using the chain rule and the identity
-( {())>> we obtain:
de
1 (((()), (-Tr(A) (F((X))) ( (- (-m(xon)
Vo(T)]OT] = 2
where X = -T) and x = (Z()). Using the law of iterated expectations and the
definition of Dx this simplifies to:
[o(T)[or] = [
By Assumption A4 (iv), 7- T can be represented via its influence function: 7(x) T(x) = panx [1F)((; + O+[[]] TIE,)) Substituting ô7 T - T and its IF representation:
VO(T) T VT(() a [F; + (+ -TII2.)
1 TF - Op(E
P Poore (0:) (zh)),
P
24

---

Published as a conference paper at ICLR 2026
where Peoo (O) - (F Note that Elycorr (O)] - VV (0:-)]l 0.
Thus, Pmn Fml- PPoor]] = (Pm - P and Term II reduces (Fm - Pcorr- + (E
Term III is empirical process term. Under Assumption A4 (īi) and (iv), this term is of smaller order. Specifically, (IPn (T)) g()) is stochastically equicontinuous, leading to Term III being O3(n-1/211)) Til = Op(n (pnnn1//0) 11/2
Combining terms, have:
Φz, Q2, = ,I,,, - (w((:)) oz, (P)} (Pm - (Pcooc( + Rm,
= (I - ()) oz, (P)) + (Fm - ()V RR Rn,ms
where the remainder Rn.m - form of the efficient influence function oz, (O;P).
+ Em. This completes the proof and establishes the
C.2.2 PROOF THEOREM 3.1
Proof of Theorem 3.1. The derivation of the EIF follows from Theorem C.I by noting that the cor- rection term in the definition of efficient influence function vanishes when the Neyman orthogonality condition holds (Vo(T)) = 0). This shows that pz, (O; P)) == wO; - Qz, (P) is the EIF:
(P) - (2, (P) = (F,, P) + Oz(n 1/22 + ) (P)
where Em = ||T - - -1
Recall that Ü, is the flow map obtained by solving the above ODE with replaced with 3N and the transport map can be represented by T = U1.
From Lemma B.2 and the rate condition, we have
& = -T -T12, S eel22K / -el1,,, ot(n
Therefore, we conclude that
ēz, (P) oz, (PP) (P,, (0)()))()))) op op 1/23)
and consequently, the asymptotic normality follows.
C.3 PROOF OF PROPOSITION 3.2
Proof of Proposition 3.2. The proof proceeds by first deriving the efficient influence function (EIF) for single component (IP) of the total importance score éx,, (F), and then aggregating the results. By the linearity of the influence function operator, the EIF for dx,, (P) = ΣS (PP is simply the sum of the EIFs for each component, i.e., Px, (O;P) -E PP)...
Part I: Deriving the EIF for single component () Let (g) (T) = (lo,
where we make the dependence of the sensitivity term (x)()) = (0x,/02,)² on the transport map T explicit. The cross-fit fit estimator Φ, Pnlw, (O:T)H((X;T)) We perform von Mises expansion of the estimation error (P) (P))
(IP) - (P) = P,, - (IP,, +,, -))
+ (P,, - -
Tem
Next, analyze each term separately.
25

---

Published as a conference paper at ICLR 2026
Term I is the standard empirical process term that converges to normal distribution after vn scaling by the Central Limit Theorem.
Term II is the first-order bias from due to the nuisance estimator T. We linearize this term using the functional derivative of @T( Let öT = Tf - T. By Fréchet differentiability in Assumption A4 (i) and the functional product rule:
Term II = TT((T) OT(TT
(1
When the predictor f is Bayes optimal, the parameter Φz, (T) = E(w,(O;T) is Neyyman-rrthooona with respect to T. This implies that its Gateaux derivative VrØz, (T) = B[VTw;
is zero. Under regularity conditions allowing the interchange of derivative and expectation, the T)] = O.
weighted expectation also vanishes, i.e., ([(V (
Therefore, the bias term simplifies to its second component. Using the influence function represen- tation for T from Assumption A4 (iv), 8y T - T = Paux [IFr-O00x +)] + Op(E)),, we have:
Term E (O.T(
O+(())
1
Raux
E
(0; (V,())))((V+),xxx)))1 ))) Op(E))
Faux
(IPaux
P) (Eor (X,
)D] Or()),
where the final step uses E|IFT] = O. This term is the empirical average of the correction term
evaluated on the auxiliary data.
Term III is a higher-order empirical process term. Under Donsker condition (Assumption A4 (ii)), this term is of order (n-n1/2)
Part 2: Asymptotic normality. Combining the terms, and noting that the main sample (for P.,) and the auxiliary sample (for panx) are independent, the total estimation error can be written as an empirical average over single sample:
(T((P) (P) = (P, - PP) (w, (O:T)H)()))
+Eo (O' T) [IVT
O+E))
By subtracting the mean (P) from the first part, we identify the EIF for dyt as:
P)((PP) = (x: (P)) + Eo (O;; (V, T) IF (o; (I;
Naive Term
Correction Term
By defining IF Hyy V [[ the influence function of the estimator for the function HGt. and using the fact that influence functions have zero mean, the term becomes a covariance
CO:(w,(O:T), IF H, (O:-)). Thus, the total EIF for éx, (IP) is xxx (O;P) = Σ 2j=1 (O; P). The
vector of estimators xx thus has the asymptotic linear expansion:
(T(xx(P) (P))
Vn (O;:))+o((1)
where (O,;P)) = (O; Na (T:)): By the multivariate Central Limit Theorem, it follows that V (PP) Φx (P) 4 (o Es), where ΣS = P)) with entries (Ex)ik Cov(px //: P): This completes the proof.
26

---

Published as a conference paper at ICLR 2026
D COMPUTATIONAL DETAILS
D.1 ALGORITHM AND COMPUTATIONAL DEVICES
Algorithm. The full algorithm of FDFI given in Algorithm D.1.
Algorithm D.1 Flow-disentangled feature importance
Require: Labeled data Dest = {O; = (X,, Y)) a black-box model f. a loss function é. In- dependent auxiliary unlabeled data Dx [X;] and flow model training procedure MM() that returns a map T. The null adjustment constant e and the Monte Carlo sample size M.
Ensure: FDFI scores fZz, 1 fox, Vi=1 and p-values {pz, 1==1- (Px.
Obtain transport map through flow matching f M(Dx). b Disentangled map estimation
2: Initialize point-wise score storage: v +- empty d matrix, f2 t- empty n X d matrix. P Compute scores on the labeled data 4: Compute latent vector 2, - f(x,) and generative Jacobian J; +- V
3: for = 1 to n do
z=
5: Compute squared sensitivities (2.( - (Je)īz for all j. E {l....,d d}.
6: for j to do
> Point-wise latent scores be M copies of 2 where the j-th coord is resampled from pz,- (T- (2( ((,,
7: Let (2.
JM Jm=1
8:
Sig + 2M
Σ
9:
end for
10: for l = 1 to d do
Attribution original features
U/ j=n Lzz (2)(
d
end for 13: end for
for
== to do
> Compute latent feature importance Oz,
15:
Ōz, + 1 - Paj.
n
D Cross-fit estimator (5)
16:
Py + S4y - -2, for 1...n
EIF components
17:
sez, + (V, + c)/n.
> Adjust standard error for inference
18:
pz, t @ (0zz /sez,)
> Compute p-values
19: end for
20: for
1= to d do
> Compute original feature importance Φxx Estimator (7)
21:
Txi + 1 Σい Vua
22:
Pg t Va 0x. for = 1... n.
> Approximate EIF components
23:
SCX, + (H + cc/nn
> Adjust standard error for inference
24:
px, <- 1 - 6(xxx//cxxi)
Compute p-values
end for
26: return Estimated importance scores and their uncertainty.
Variants FDFI estimators. There are two different ways to construct an estimator for latent FDFI (2). One can adopt the idea of CPI (Strobl et al., 2008), SCPI (Reyero-Lobo et al., 2026), and LOCO (Lei et al., 2018) to estimate this quantity. More specifically, we can define
M
Tem (P) 22
((Y f(T (26.))))))))))))))))) ( -
(12)
2M
1=1
kml
07
Σ
Y, M
Σ [((T)))()))))
e(V (YY, -
(13)
where (12) coincides with the latent FDFT estimator (5) presented in the main text.
27

---

Published as a conference paper at ICLR 2026
The final FDFI estimator of the original feature can be constructed as follows:
M
Tom Ox,
(P)
(
(2(4)) Bm(Zs)
(14)
2M
)=1
1=1
k=1
d
M
Ox,
SCP (P)
=
Yi,
1
[ [(())))))))
-exx, Hz)( (15)
π
M
1=1
where H2)(Z) = (Z)l3, the square of estimated Jacobian of X; with respect to Z,.
One can, in principle, construct LOCO-type estimator; however, this requires refitting submodels for foT"- which is computationally expensive in general. Except in the special case of Ez loss, the LOCO-type estimator coincides with the SCPI-type estimator if f is the Bayes optimal predictor. For this reason, we didn't 't explore this variant in the current paper.
Computational devices. All experiments were conducted on dedicated computing platforms. All experiments, except for Appendix D.4, were executed on server equipped with an AMD EPYC 7542 32-Core Processor CPU and NVIDIA RTX 3090 GPUs. Appendix D.4 was carried out on a personal computer with an Intel Core i5-14600KF CPU and an NVIDIA RTX 5070 Ti GPU.
D.2 FLOW MATCHING MODEL
Flow matching is a concise and powerful framework for generative modeling that has advanced the state of the art across various domains and applications. Following Lipman et al. (2022; 2024), adopt the Conditional Flow Matching (CFM) model as the backbone of our approach. Similar to the setting in Lipman al. (2024), we train the model for 5000 steps per 1000 samples. Because performance is sensitive to architectural and optimization choices-includdig dimensions, and batch size-we systematically explore multiple hyperparameter combinations to identify a configuration that offers a favorable cerformance-efficieee trade-off for training.
network depth, hidden
In statistics, the Maximum Mean Discrepancy (MMD) is widely used metric for quantifying the difference between two probability distributions. A smaller MMD value indicates that the two dis- tributions are more similar, whereas a larger value reflects greater divergence. For the target data, we split the dataset into training and testing subsets with a 1:1 ratio. Using the training set, we trained the model to generate samples from a standard normal distribution, thereby obtaining the generated target distribution. We then computed the MMD between the real and generated distributions on both the training and testing sets. These results were used as the criterion for selecting appropriate configurations.
To select a suitable hyperparameter configuration, we proceeded as follows. As an example, in Section 4.1 we fixed the training sample size at 1000 and trained the model for 5000 steps using the Adam optimizer with a learning rate of 1 x 10-3 In each trial, a distinct random seed was used to generate a new simulated dataset; for that dataset, we evaluated multiple hyperparameter configurations. This procedure was repeated for 50 seeds in total, and for every configuration-seed pair, we recorded the resulting MMD2 on both the training and test sets.
suitable set of hyperparameters should yield relatively small MMD? values on both the training and test sets. Moreover, the gap between training and test MMD? values should remain modest; otherwise, the model risks underfitting or overfitting. As shown in Table DI, the configuration with hidden dimension 128, batch size 256, and a two-layer network satisñies the above criteria: it attains relatively small MMD" on both training and test sets with only modest gap between them. By contrast, smaller hidden dimensions and single-layer network yield large MMD² on both splits, indicating underfitting. Increasing the hidden dimension and depth further reduces the training MMD² but enlarges the train-test gap, signaling an overfitting risk. Although the setting with hidden dimension 256, batch size 128, and two layers performs comparably, we ultimately select the didden-iimension-11 balance between model performance and computational efficiency.
batch-size-256, two-layer configuration as it offers a more favorable
After identifying an appropriate hyperparameter configuration based on the preceding selection pro- cedure, examined the effect of proportionally increasing the training sample size and steps, scal-
28

---

Published as a conference paper at ICLR 2026
Table D1: Train and test MMD? (mean + std) across hidden dimensions, batch sizes, and network depths.
Hidden dim
Batch size
Train (1 layer)
layer)
Truin (2 layers)
(2 layers)
Train (3 layers)
(3 layers)
64
32
0.0100+000009
0.0109+0.0056
0.0099=0.0039
0.00%--0.002
0.0076++00044
0.0084++0. 0042
0.0065:000040
0.0073=00.044
0.0072=0.0038
0.008000.0039
0.0055=0.0037
0.0069+000005
64
128
0.0071+0.0044.
0.0076+0.0049
0.0053+0.0028
0.0058-0.00.3
0.0038+0.0020
0.0046++0,0023
256
0.0052::000028
0.0064=0.0031
0.0047=0.0025
0.0099=0.0027
0.0041=0.0029
0.0046++.00225
64
384
0.0048++0.022
0.0054+0.0027
0.0041=0.0031
0.0057-0.0032
0.0038+0.0022
0.0046++0.0020
128
32
0.0062-+0.0025
0.0068+00.0024
0.0057-00.004
0.0067-00.045
0.0061 +00003-
0.0070+0.00330
128
0.0052100.0032
0.006210.0034
000022
0.0059_0.000 0.0041
0.0043.1..025
0.005810.0025
128
128
0.0050+-0.0024
0.0061 +-0 0.003
0.0036-0.0023
0.0046- 0.0026
0.0027+40.022
0.0035+40.0020
128
256
0.0044100.002
0.005210.0019
0,0024 0.0022
0.00300 0.0022
10.006
0.003510.0015
128
384
0.0041+00.000
0.0052+-0.0020
0.0030-0.0022
1.040-0.0022
0.0023-40.0014
0.0033++0.0011
256
0.0036:0.0017
0.0035=0.0018 0.0041=000018
0.001100.002
0.0032+0.0019
0.0039+0.00028
256
64
0.00261000016
0.00341:0.011
0.0024:0.00199
0.0022-000026
0.0035:10.0019 0.002510.0021
256
128
0.0018=000012
0.0028:0.0015 0016
0.0023=0.0016
0.0028=0.0015
0.0020=0.0012
0.0051+0.0015
256
2566
0.00171000014
0.0027:0.0019
0.0014:00.012
0.0025
0.001110.0006
0.002611000008
256
0.00241000018
0:0011
0.0013:0.0011
0.001
0.001210.0011
0.00291000001
384
32
0.0016=0.0010
0.0024=0.0019
0.0015=0.0009
0.0028=0.0012
0.0027=0.0021
0,0040=0.0024
384
0.0016+0.0011
0.0026+0.0017
0.0013+0.0011
0.0025-0.0013
0.0022+0.0010
0.0029+000013
384
128
0.0015=000013
0.0030=0.000
0.0010=0.000
0.001000.0110
0.0016=000011
0.0023+0.0011
384
256
0.001612000133
0.0025:20.0014
0.0012=000009
0.0023=0.0012
0.0020=1.00099
0.005620.0015
384
384
0.0018=0.0008
0.0027=0.0011
0.0015=0.0011
0.0027=0.0009
0.0023=000014
0,0099+0.0018
ing from (1000, 5000) to (2000, 10000) and beyond. For comparability, the procedure was repeated 50 times, each with a fixed random seed; within each run, datasets of different sizes were generated using the same seed to ensure consistency. The results in Figure D1 show that when the sample size reaches 3000 or above, the MMD values on both training and test sets stabilize. While larger sample sizes can still improve generation quality, we adopt 3000 as the sample size, balancing performance and computational cost.
MMD vs. Sample size
MMD Test MMD = std
0.005
0.004
MMD
0.003
0.002
1000
2000
3000
4000
5000
Sample
Figure DI: Comparison of Train and Test MMD across varying sample sizes.
D.3 STATITSICAL INFERENCE
D.3.1 APPROXIMATION FOR EIF OF Φx,
Recall that the component EIF from Proposition 3.2 is given by
(O; = (x)( + (O) =: (O:P) + C3.
The first-order approximate term 3 (O; P) can be estimated using a cross-fit estimator; the correc- tion term C, however, requires evaluation of the influence function of Hzt- Analytically deriving IF, (and thus IF )1,)) for a flow-based model is generally intractable. In practice, the correction term must be estimated using resampling methods.
29

---

Published as a conference paper at ICLR 2026
Algorithm D.2 Jackknife estimation of the EIF correction term (Cji)
Require: Dataset D {Oi)== where O (X,, Yi). Flow model training procedure M(-) that returns a map T. Point-wise latent score function Wy (O;
Ensure: Estimated correction term C2: (O; T), IF H,,
1: Initialize score vector W +- []-
> Step Compute point-wise scores using f.
2: for i to n do
3:
W; + W (O;; 7)
4: Append W; to W.
5: end for
6: Let (-)(()) be the estimator function derived from the full map T.
7: Initialize influence vector I <- []-
8: for n do
(X1, Xi-1, X Xn). 10: Retrain model: 7(-i) + M(D
9: Let Dx
b Step 2: Compute leave-one-our estimate.
11: Derive the leave-one-out estimator function
(.) () from 7(-().
12:
1 (n 1) (X,( (x.)
Step 3: Estimate point-wise influence.
Append 4
14: end for
Cj + CW,,
Step 4: Compute the sample covariance.
16: return Cjt
Gaussian
Gaussian Mixture
Value
10
1.000
2,06
2335
230L
0
2
5
8
4
7
Correction
Feature
Figure D2: Comparison between first-order approximate term (O; and the correction term Cjp. The correction term C5 of smaller order than the first-order approximate term 45) (O)
The influence function IF H measures the sensitivity of the estimator ПG to the data. To measure this, must re-estimate H₃ on perturbed versions of the data. direct, though computationally intensive, method to estimate the covariance term is the jackknife, as outlined in Algorithm D.2.
We fix the dimension d 10 and study the response model y 5X, where No 1). The covariates X E IR: are sampled under two regimes: (i) Gaussian. X ~ N(O, Eo.), and (ii) Gaussian mixture. X ~ 0.2 Eo.s) + 0.8.V(O, Eo.2), where Σ, E 100110 with [_p],]
if i = and otherwise, is the ARI covariance matrix with parameter p.
In this part, we choose the sample size as 1000 and compute the approximate term (); P) and the correction term Cjt. respectively. Crucially, this experiment is conducted without sample splitting; the same data is used estimate the transport map T and to subsequently compute the im- portance scores. As shown Figure D2, the correction term is of smaller order than the approximate term. For this reason, we can approximate the EIF of fxx with the individual first-order approx- imate EIFs Pa. This provides strong empirical evidence that while sample splitting is the most theoretically rigorous approach, practitioners may often able to use the simpler approximate EIF for inference without splitting the data, which offers significant computational advantage.
30

---

Published as a conference paper at ICLR 2026
D.3.2 INFERENCE NEAR THE NULL
As stated in Verdinelli & Wasserman (2024a), when dealing with the quadratic functional, under the null hypothesis, the influence function for the parameter vanishes, the coverage rate becomes n -1 and the limiting distribution no longer a Gaussian distribution. This hinders variable selection with direct statistical guarantees, which is an essential component of reliable scientific discovery (Reyero- Lobo et al. (2026)). In Verdinelli Wasserman (2024a), they propose to replace the standard error sen with / se + c/n to ensure control of the type-I error, where e is constant, and many constants could be used. In this paper, we choose с = मh min VVY), V(Y), V(Y)², V(Y)", where V(Y) is the variance of the output and is the dimension all features.
Due to the quadratic nature of the statistic ē, under the null hypothesis, the variance vanishes as 1. It can be found that under the null hypothesis, using Cantelli's v(0)/ Za se + c/ni)
v(δ) = O( with
inequality, we have Poo Za se + guarantees that the expanded confidence interval controls the Type error.
+ This
D.4 COMPUTATION TIME COMPARISON
For the computational efficiency study, we followed the experimental protocol of Experiment 2. However, computing Shapley value for all 50 features is prohibitively tome-consuming therefore, for Shapley value (Shapley, 1953). we restrict the feature dimension to d = 10 while keeping all other settings identical. Under this setting, we measured the runtime CPI, LOCO, nLOCO, dLOCO, DFI, FDFI-Z, FDFI, and Shapley value across range of sample sizes characterize scalability and comparative efficiency as data volume grows. Each method was run 10 times for every sample size configuration.
Time (veconds))
1.P
10!
109
CPI
LOCO
ALOCO
dLOCO
DFI
FDFL
FDR
SHAP
Figure D3: Computation time in logarithmic scale of different methods. The bars indicate standard deviations across 10 random seeds.
The computation-tinet results show that CPI, FDFI, and DFI are substantially more efficient than LOCO-based and Shapley value methods. Across all sample sizes, both approaches consistently require only a fraction of the runtime compared with LOCO-based and Shapley value estimators, while maintaining stable variance reflected by the narrow error bars. This highlights their scala- bility and practical advantage for large-sample applications. lower runtime than correlation-aware methods such dLOCO, further underscoring its efficiency. Although DFI allows for an arbitrary optimal transport map in theory, its implementation the orig- inal paper only uses Gaussian transport map, which admits closed-form solution and is indeed efficient, as shown in the above figure. However, if one needs to transform features to another latent distribution, general OT solver must be used, and it is usually computationally expensive when the feature dimension is large.
particular, FDFI achieves markedly
31

---

Published as a conference paper at ICLR 2026
E ADDITIONAL EMPIRICAL RESULTS
E.1 EXTRA SIMULATION STUDIES
E.1.1 BENCHMARKING WITH DIFFERENT PREDICTORS
We conduct comparative evaluation of LOCO, CPI. DFI, FDFI (SCPI), and FDFI (CPI) employ- ing three prediction models: random forests, neural networks, and Lasso in the same setting of Section 4.1, fixing the correlation at p 0.4 and the sample size at = 1000. Each configura- tion is repeated 100 times with different random seeds. We choose = 0.4 because the perfor- mance methods such as LOCO and CPI tends to deteriorate as p increases; therefore, among E {0.4, 0.6, 0.83. the smallest correlation coefficient was selected to allow more objective com- parison of the performance differences across different predictors.
random forest
1.0
AUC Score
Power(C1)
0.10
Type Error
1.0
Power(C1uC2)
0.9
0.05
0.8
0.7
0.00
0.0
0.0 LOCO
LOCO
CoI FDFILSCHI!
FDF(CC)
DF
LOCO
EDFI(SCPD
cPi
FOFNCPD
DFi
LOCO
FDF(SCH))
CPI
FDPICPI)
DR
CPI FDRRSCHD FoTUCc))
DFi
neural network
1.0
AUC Score
Power(C1)
0.10
Type
Error
1.0
Power(C1uC2)
0.8
0.6
0.00
LOCO
FD (OFICCP
CPI
DFI
LOCO
FDF FDFNISCP
CPI
EDPUCPP
DFi
LOCO
CPI FDFI(SCP))
FDFICPU
DR
LOCO
CPI FDFM(scR PDPICPI)
OFi
AUC Score
Power(C1)
Type Error
Power(C1uC2)
1.0
0.10
1.0
Lasso
0.8
0.05
0.6
0.00
LOCO
CPI PORLCC))
DFI
LOCO
CPI FOFISCPP)
PDFUCPD
DFI
LOCO
FDF(SCCP
CPI
FOHICPN
DF
LOCO
CPI FDFN(CCR) (DFICPI)
DF1
LOCO
CPI
FDFI(SPPI)
FDFI(CPI)
DFI
a= 0.05
Figure E4: Performance comparison across three different predictors under the settings of Sec- tion 4.1. Top: results obtained with the random forest predictor. Middle: results obtained with the neural network predictor. Bottom: results obtained with the Lasso predictor. Points indicate mean values, with error bars representing 95% bootstrap confidence intervals over 100 runs.
According to Figure E4, when using the random forest predictor, different methods demonstrate consistently better performance in terms of AUC, Power, and Type I error control. Therefore, adopt random forests as the predictor in our experiments.
1.1 DISENTANGLED TRANSFORM MAP
The term (ax1/OZ,) in Equation (4) quantifies how strongly fluctuations in Z, are transmitted through X, Collecting these sensitivities across all pairs (1,j) yields matrix that characterizes how latent variations propagate into the observed space. This sensitivity matrix reveals the block structure of X, where coherent feature groups manifest as localized blocks. To illustrate this phenomenon, visualize the matrix heatmap. Specifically, we consider the nonlinear response model
y arctan(X,, (Xo) sin(X2) +
with ~ (0, E). and ~~ (0, 1). The covariance matrix Σ is block-diagonal with two blocks of equal size (5 x 5). The first block (features X is subdivided into uncorrelated groups,
32

---

Published as a conference paper at ICLR 2026
X1 and X2, X3, X4, while the second block (features X5-X9) forms single equicorrelated
cluster, as in Section 4.1. The within- block correlation fixed at = 0.8.
6.93
0.7
0.72
0.5
Restires
60
60
930
010
0.9
0.2
9
Feat unes
Figure ES: Sensitivity matrix under block structure (p = 0.8). The color bar denotes pairwise correlation strength, with darker red indicating higher correlation.
As shown in Figure E5, the heatmap exhibits clear block patterns consistent with this design. These structures provide direct empirical evidence of grouped dependencies, supporting the claim that the disentangled representation successfully recovers the latent dependency structure of X.
E.1.3 BENCHMARKING COMPLEX FEATURE DISTRIBUTIONS
To evaluate in more challenging settings that approximate real-world data, we move from sim- ple Gaussians to Gaussian mixtures, which possess universal approximation capability to capture complex patterns commonly observed real data (Goodfellow et al., 2016). We retain the same response model from Section 4.1 but generate covariates from two-component Gaussian mixture: X ~ 0.2.N E(0.8) + 0.8 (O. Both covariance matrices, [(0.8) and (0.2)). follow the block structure introduced in Section 4.1, with correlation parameters p 0.8 and p 0.2, respectively. In this non-Gaussian setting, we benchmark 5 feature important measures, includ- ing LOCO, nLOCO (normalized LOCO; Verdinelli & Wasserman, 2024b), dLOCO (decorrelated LOCO; Verdinelli & Wasserman, 2024a), DFI, and FDFI, which provide both point estimates and uncertainty quantification. We exclude Shapley-value-
methods from this comparison due to their prohibitive computational cost, as detailed in our runtime analysis (Appendix D.4).
AUC Score
Power(C1)
Type Error
Power(C1uC2)
1.0
1.0
0.10
1.0
0.5
0.5
0.00
0.0
Lo LoCooo Roco
Dfl PDFI
LOCO pLOCO Loco of FDFI
LOCO
LLoC &occo
DFI FDFI
LOCO
FLoC g Df FDF!
noCoo
dLoco
0.05
Figure F6: Benchmarking results of Gaussian mixture setting with a sample size of n = 1000. Points indicate mean values, with error bars representing 95% bootstrap confidence intervals over 100 runs.
The results in Figure E6 demonstrate that FDFI consistently achieves the strongest performance, attaining the highest AUC scores and statistical power while robustly controlling the Type I error. DFI also improves over the LOCO-based methods and maintains well-controleed Type error. However, DFT's reliance on the Gaussian optimal transport map, which is less flexible than flow matching. limits its discriminative ability this complex non-Gaussian data, resulting in inferior AUC and power compared to FDFL. By contrast, the LOCO, nLOCO, and dLOCO baselines remain substantially weaker in both AUC and power. Overall, these results highlight that FDFT's flow-based
33

---

Published as a conference paper at ICLR 2026
disentanglement provides superior flexibility and statistical power complex distributions.
Moreover, also adapt more sophisticated design for the covariance matrices and construct X to exhibit heavy-tailed distribution characteristics. We first generate latent Gaussian covariates Z (Z,,... ........ 11000) E IR 1000x50 from two-component Gaussian mixture
Z: 0.2 E(0.8)) 0., E(0.2)),
i= 1.... 1000,
where Σ(p) Is @ E, and (Ep)s -pi->/// To induce heavy tails, then transform Z into X via Gaussian-t scale mixture: for each i, draw We xi0 independently and set
10
X;
Zz.
W
so that the rows of X follow a heavy-tailed distribution of multivariate t-type with 10 degrees of freedom and covariance proportional to E(p). We retain the same response model in Section 4.1 and perform 100 independent experiments, each with a randomly chosen seed. We then compare the performance of LOCO, nLOCO, dLOCO, DFI, and FDFI in terms of AUC score, Type error, and Power, as reported in Table E2. It can be observed that even under a sophisticated covariance de- sign and heavy-tailed distributions, FDFI still achieves the best performance among all methods by leveraging flow-base model. This further highlights the advantage FDFI over DFI in handling complex dependence structures and non-Gaussian data.
Table E2: Comparison of LOCO, nLOCO, dLOCO, DFI, and FDFI in terms of AUC score, Type I error, and Power.
Method
AUC score
Type I error
Power (C1)
Power (C1 C2)
LOCO
0.9536
0.0063
0.5180
0.2640
nLOCO
0.9549
0.0095
0.5300
0.2700
dLOCO
0.9893
0.0357
0.9200
0.5080
DFI
1.0000
0.0037
0.9280
0.5730
FDFI
1,0000
0.0307
1.0000
0.8660
E. LATENTE INDEPENDENCE DIAGNOSTICS
To inspect the effect of the latent representation Z on FDFI, we compute the pairwise normalized HSIC (nHSIC) with Gaussian kernel:
HSTC(Z,, 7x)
nHSIC(Z,, Zx)
HSSCCZZ, HSTCCZZ Zk)
and the distance correlation (dCor):
dCov 2Z,, Zx.
dCor²
/Var dVar (
between two coordinates Zj and Z₁ for the learned latent representation under different settings. Here, dCov and dVar the covariance and variance of the Euclidean distance matrix.
To illustrate how one can diagnose the latent independence and its effects on statistical inference, we retain the experimental setup of Section 4.1 and keep all other settings fixed. By reducing the width of the neural network and the number of training steps (hidden dimension 128 - 16, training steps 5000 - 5). we deliberately deteriorate the training performance so that the correlations among the components of Z increase, and then investigate the resulting impact. For both the "good" and "bad" training regimes, we perform 50 independent runs with different random seeds. In each run, we compute the two metrics between the coordinates of Z and summarize the dependence structure by averaging all off-diagonal entries (i.c., excluding the self-dependence terms with j = k) of the resulting pairwise nHSIC and distance-correlation matrices. Finally, we average these run-wise summaries across the 50 runs and relate them to the average empirical power and the standard deviation of the feature importance of the first covariate Xp. The results are reported in Table E3. As expected, bad training fails to decorrelate Z s coordinates (large nHSIC and dCor), and hence deteriorates the statistical inference.
34

---

Published as a conference paper at ICLR 2026
Table E3: Effect of average dependence in (pairwise nHSIC and distance correlation) on test performance and feature importance stability.
Regime
n
dCor(Z)
Power(Ci U C2)
sd (o(())
Good training
0.0032
0.0543 0.1872
1.0000
0.0147
Bad training 0.0711
0.4480
0.0229
E. SENSITIVITY TO NUISANCE QUALITY
To analyze the sensitivity of FDFI to different flow training, retain the model construction in Section 4.1 and inspect the performance with different parameter values. According to Figure 2, FDFI becomes highly stable when the sample size is sufficiently large. Therefore, better investi- gate how these hyperparameters affect performance, we fix the sample size at 300 in the following experiments.
In Section 4.1, the response is generated as
y arctan(Xo X1) [(x,00) -in(Xsnnn)))) 1l (x,<00 +c.
Therefore, for Xo and X1, their feature importances should be identical in theory. Motivated by this, we introduce an additional metric to assess the effect of hyperparameters on the training outcome, namely the geometric-relative discrepancy between Φu and Φ1, defined by
100 $1
geoo (ou, Φ1)
V V
Ideally, this quantity should as small as possible.
We consider different configurations of each hyperparameter (and keep all other settings identical Section 4.1): (i) The size of auxiliary set. We vary the auxiliary set size m, while keeping the training steps proportional
to ensure sufficient training. (ii) Latent resample size. We vary the number latent resamples M from to 50. (iii) Model depth/width. We evaluate 5 configurations of the hidden dimension (width) and the number of layers (depth) of the neural network.
The statistical power, geometric-relative discrepancy, and Type error in all configurations are pre- sented in Table E4. We observe that increasing the auxiliary size m, the latent resample size M, or the network width and depth consistently improves power and reduces geometric discrepancy while maintaining well-controlled Type error. Notably, as long the hyperparameters lie within a reasonable range, the overall performance remains highly stable, showing that FDFI maintains stable performance without being overly sensitive the specific choice of hyperparameters.
Table E4: Sensitivity analysis of FDFI to flow training hyperparameters Power, geometric-
relative discrepancy, and Type error.
Hyperparameter
Power( C1 U2)
Agoo(do,00 01)
Type error
300
0,8560
0.2275
0.009
Size of auxiliary
500
0.8880
0.2161
0.0180
1000
0.9080
0.1821
0.0165
0.7700
0,2522
0.0115
10
0.8820
0,2143
0.0185
Latent resample size M
50
0.9080
0.1821
0.0165
(16, 1)
0.4340
0.5637
0,0215
(64, 1)
0.8860
0,2402
0.0125
(128, 1)
0.8980
0.2251
0.0205
(width, depth)
(128, 2)
0.9080
0.1821
0.0165
(128, 3)
0.9040
0.2023
0.0135
E.2 EXTRA REAL DATA STUDIES
We conducted data analysis on nine different real-worke datasets, which vary in sample sizes, feature dimensions, task types, variable types, and domains. We summarize their information in Table E5.
35

---

Published as a conference paper at ICLR 2026
Table E5: Summary of all real datasets studied in the paper. In the 'Task' column, 'Cls' denotes classification and 'Reg' denotes regression.
Dataset
(n, d)
Task
Variable type
Domain
Cardiotocography
(2126, 21)
Cls
Continuous Continuous
Medical
Pima Indians Diabetes
(768, 8)
Cls
Medical
MicroMass
(571, 1300)
Cls
Discrete Continuous & Discrete
Biological
Codon usage
(13028, 69)
Cls
Biological
Default of Credit Card Clients
(30000, 23)
Cls
Continuous & Discrete
Commercial
Superconductivity Data
(21263, 81)
Reg
Continuous
Industrial
Video Transcoding
(68784, 19)
Reg
Continuous & Discrete
Industrial
TCCGAPAANCAN-HiSee
(801, 20531)
Cls
Discrete
Biological
human single-cell RNA-seq
(632, 23257)
Cls
Discrete
Biological
E.2. CARDIOTOCOGRAPHY DATASET
The Cardiotocography (CTG) dataset (n 2126, 21) (Campos Bernardes, 2000) consists of 2,126 computer-processed processed cardiotocograms with 21 diagnostic features derived from fetal heart rate (FHR) and uterine contraction (UC) signals. These features capture baseline FHR, counts of accelerations and decelerations, short- and long-term variability, and FHR-histogram statistics. We use this dataset to classify fetuses as normal or abnormal and to compare feature importance esti- mates produced by different methods. The Shapley value method was not applied here because its computation time exceeded one day, even with only 21 features.
As shown in Figure E7, LOCO and CPI identify only a small subset of features as significant, in sharp contrast to FDFI and DFI. This divergence stems from the block-diagonal structure of the feature correlation matrix (Section 4.3). which is manifestation of correlation distortion under multicollinearity. In such settings, LOCO and CPI attenuate or mask true effects, whereas FDFI and DFI are less affected (Verdinelli & Wasserman, 2024b). In the CTG dataset, I.B anchors the tracing around the expected 110-160 bpm baseline for reassuring status, while central central-tendency descrip- tors (Mean, Median, Mode) provide concordant evidence of the baseline neighborhood (Ayres-de Campos et al., 2015; Jia et al., 2023). Reactivity is reflected by AC (accelerations), often accom- panied by FM (fetal movements), both indicating intact autonomic control and low short-term risk. Excessive UC (uterine contractions), i.e., tachysystole, can compromise uteroplacental perfusion and precipitate decelerations (Ayres-de Campos al., 2015). Beat-to-beet and longer-scale vari- ability are captured by MSTV/ASTV (short-term variability mean and abnormal-tine proportion) and MLTVIALTV (long-term variability mean and abnormal-time proportion). Adequate variability (higher MSTV, MLTV and lower ASTV, ALTV) is reassuring, whereas depressed variability is non- reassuring (Ayres-de Campos et al., 2015; Jia et al., 2023; Stampalija et al., 2023). Decelcration phenotypes-DD -DL (mild/early). DS (severe/variable), and DP (prolonged) span spectrum from benign positional patterns to forms associated with cord compression and hypoxia/acidemia, with DP (> 2 min) and repetitive deep DS most strongly linked to metabolic risk (Parer & Ikeda, 2007). Histogram-based morphology provides supportive, context-dependent information: overall spread via Width and Variance (very low values echo variability loss; excessively wide or erratic distri- butions indicate nonreassurance), extremes via Min/Max (sustained brady- or tachycardic periods), modal structure via Nmax/Nzeros (multimodality and sparsity cues), and slow drift via Tendency. These descriptors are standard in computer-aided CTG (e.g.. SisPorto) and are highlighted in con- temporary reviews (Ayres-de Campos et al., 2000; Romano et al., 2016; Zhao et al., 2018).
FDFI consistently isolates the physiologically meaningful signal and preserves it under strong cor- relation and Bonferroni control, yielding a stable, guideline-comcordant importance profile. This pattern demonstrates clear advantages in interpretability and robustness, retaining clinically relevant effects despite multicollinearity and multiple testing, thereby providing a more reliable attribution map than methods that prone correlation distortion.
Case-wise wise attributions: Recall that the importance of feature X₁ is estimated by Algorithm D.I as:
Va Σ 1=1
(2))
Tx
1 n
36

---

Published as a conference paper at ICLR 2026
Table E6: Description of 21 features in the Cardiotocography dataset (n 2126, d= 21). Where is the sample size and d is the feature dimension. 'FHR' means 'Fetal Heart Rate".
Broad category
Feature name
Meaning
LB
FHR baseline (beats per minute) Number of accelerations per second
AC
Baseline & counts
FM
Number of fetal movements per second
UC
Number of uterine contractions per sccond
DL
Number of light decelerations per second Number of severe decelerations per second Number prolonged decelerations per second
Decelerations
DS
DP
ASTV
Percentage of time with abnormal short-term variability Mean value short-term variability
Short-term variability
MSTV
ALTV
Percentage of time with abnormal long-term variability
Long-term variability
MLTV
Mean value long-term variability
Width
Width of the FHR histogram
Min
Minimum value the FHR histogram
Max
Maximum value of FHR histogram
Nmax
Number of histogram peaks
Nzeros
Number of histogram zeros
Histogram (FHR)
Mode
Mode of the histogram
Mean
Mean the histogram
Median
Median of the histogram
Variance
Variance of the histogram
Tendency
Tendency of the histogram
0.150
Feature Importanee
0.100
0.075
0.050
0.025
0.000
ALTV
ASTV
Varianca
Mode
Median
Mean
DL.
AC
MITV
MSTV
FM
Max
Whith
Bnddancy
UC
Nmax
Nzeraa
1..0
11
FDFI
Figure E7: Bar plot of feature importances of the random forest classifier the Cardiotocography (CTG) dataset. The symbols and ** on the bars denote statistical significance at a level of = 0.05 and a 0.05/21, 22, respectively. The meaning of all features is provided in Table E6.
In Figure 3(a), the heatmap exhibits strong block-diagonal structure among LB, Mean, Mode, and Median. For feature LB,, we define the block-restricted contribution for LB as
black W
Σ Oig HL JEBLB
BLB == {LB, Mean, Mode, Median].
LLB
Averaging these restricted sample-wise contributions over all i yields block-based feature impor- tance block 1 Σ, V/olock For comparison, ΦXLE Σ, VII denotes the overall feature XLB n
importance LB computed using features. In real data analysis, the overall importance ΦXLB = 0.0286, the block-based importance biock 0.0228, and their ratio is 0.7972, meaning that the four-feature block accounts for nearly 80% the importance of LB.
E.2.2 PIAM DIABETES DATASET
The Pima Indians Diabetes dataset (n 768, d 8) (Smith et al., 1988) contains medical records from female patients of Pima Indian heritage. Each record includes eight clinical attributes such the number of pregnancies, plasma glucose concentration. blood pressure, skinfold thickness, insulin level, body mass index (BMI), diabetes pedigree function, and age. The task to predict the onset of Type diabetes, providing widely used benchmark for classification in clinical settings.
37

---

Published as a conference paper at ICLR 2026
Table E7: Description of features in the Pima Indians Diabetes dataset (n 768, 8) Where is the sample size and d is the feature dimension.
Broad category
Feature name
Meaning
Obstetric history
Pregnancies
Number of times pregnant
Glycemia
Glucose
Plasma glucose concentration 2 hours in OGTT (mg/dL)
Blood pressure
BloodPresssure
Diastolic blood pressure (mm Hg)
SkinThickness
Triceps skinfold thickness (mm)
Adiposity
BMI
Body mass index (kg/m")
Insulinemia
Insulin
2-hour serum insulin (mU/L)
Family history
DPF
Diahetes pedigree function (family history-hased risk score)
Demographics
Age
Age (years)
Importance
LOCD
DFI
0.06
CPI
0.04
Shapley
0.00
FDFI
Glucose
BMI
Age
Pregnancies
DPF
Insulin
Skin Thickness
Bloo BloodPressure
Figure ES: Bar plot of feature importances of the random forest classifier on the Piam diabetes dataset. The symbols and ** the bars denote statistical significance at level of 0.05 and 0.05/8, respectively. The meaning of all features is provided in Table E7.
On the Piam diabetes dataset, the results in Figure E8 align with clinical knowledge while under- scoring FDFI's robustness under correlation and multiple testing (Bonferroni a 0.00625). (i) Diagnostic marker: Glucose is consistently identified as significant by all methods except Shap- ley value and remains so after Bonferroni adjustment, reflecting its established diagnostic role in diabetes (ElSayed et al., 2023). (ii) Risk factors: BMI shows a nearly identical pattern, being re- tained after correction all methods except Shapley value, consistent with evidence that adiposity strongly elevates diabetes risk, though its impact diminishes once baseline glycemia is considered (Jayedi et al., 2022). FDFI additionally detects Age, Pregnancies,
DPF after correction and background
assigns non-negligible weight to BloodPressure, aligning with their interpretation screening factors rather than diagnostic markers (Davidson et al., 2021; Valdez et al., 2007; Emdin al., 2015). (iii) Physiology-alignet lights Insulin at the nominal level and assigns non-negligible weight to SkinThickness,, consistent with prior evidence linking parity, OGTT-based insulin, and skinfold adiposity with type 2 diabetes risk (Guo et al., 2017; Nicholson et al., 2006; Hanley et al., 2003; Ruiz-Alejos et al., 2020).
signals: Beyond these established predictors, FDFI also high-
E.2.3 MICROMASS DATASET
The MicroMass dataset (n 571,d 1300) (Mah & Veyrieras, 2014) contains mass spectrometry measurements from 571 bacterial samples. Each sample is encoded as a 5000-dimen spectrum of mass-to-charge intensity values. which serve as high-dimensional fingerprints of bacterial compo- sition. While originally designed for species-level identification, the dataset also provides a natural binary partition between Graampositin and Gram-negative bacteria. our work, we leverage this property frame the problem binary classification task, aiming to discriminate Gram-positiv from Gram-negaa reea-iveeatve organisms. This setting provides biologically meaningful benchmark for evaluating the effectiveness of the proposed method in high-dimensional, structurally complex data.
For the MicroMass dataset, due to the large number of features, we report the counts of features identified as significant by the four methods (LOCO, CPI, DFI, and FDFI) under two significance thresholds: the nominal 0.05 level and the Bonferroni-aduusted 0.05/ /1300 level, as summarized in Figure (9). Results show that DFI and FDFI identify considerably larger number of significant
38

---

Published as a conference paper at ICLR 2026
(a)
(b)
Number Significant Features
1400
0.975
FDFI
DF)
hoc
1201
1200
0..50
10821082
1066
1000
Mean accuracy
0.900
0.675
600
478
0.850
400
0.625
200
0.800
0 LOCO
0
0.775
CPI
DFI
FDFI
50
30
Topk variables
1300 features
0.05 level
0.05/1300 level
Figure E9: Data analysis of the MicroMass dataset. (a) Number of significant features identified by LOCO, CPI, DFI, and FDFI under different significance levels on the MicroMass dataset. (b) Prediction accuracy with selected important features for FDFI, DFI, and an ad hoc method that applies CPI on the clusser-rerresentatie features.
features, whereas CPI yields only a few significant findings and LOCO fails to detect any, reflecting their susceptibility to correlation distortion under multicollinearity.
Consistent with the preceding analysis in Appendix E.2, also compare FDFI, DFI, and an ad- hoc procedure that first clusters features on the Spearman correlation matrix, selects one medoid per cluster as a representative, and then applies CPI on this reduced feature set. As illustrated in Figure E9(b), on the MicroMass dataset, FDFI exhibits substantial superiority over both DFI and ad hoc method, empirically confirming the effectiveness of the proposed approach under high- dimensional, structured-feature settings. Notably, by operating directly on the full feature space, FDFI avoids the information loss inherent to cluster selection and remains robust to correlation distortion, yielding consistently higher accuracy across the range of k.
E.2.4 RESULTS ON FOUR ADDITIONAL LARGE- DATASETS
In this section, we conduct experiments on four real-world datasets of moderately large sample sizes. These datasets cover both classification and regression tasks. The larger sample sizes allow for more reliable results, making it possible to assess the robustness and generalizability of the models across different task types. We compare average prediction performance across datasets for features selected by FDFI, DFI, and an ad hoc CPI-based clustering approach, using prediction accuracy for classification and root mean squared error (RMSE) for regression, with results averaged over two-fold splits for each dataset. For the classification task, we conduct experiments on datasets Codon usage (n = 13028, = 69) and Default of Credit Card Clients (DCCC) (n = 30000, d 23), and for the regression task, conduct experiments 21263, d = 81) and Video Transcoding (n 68784, = 19). The results reported in Table E8 clearly demonstrate the superiority of our method. FDFI consistently outperforms both DFI and the ad hoc CPI-based clustering approach across these datasets for both classification and regression tasks. This superiority is evident in its ability to more effectively select the top-k features, which leads to better model performance.
datasets Superconductivity Data (n
E.2.5 EXTRA RESULTS ON TWo RNA-SEQ DATASETS
To further demonstrate that our method can achieve superior performance over existing approaches in genuinely -igehdimensional regimes with intricate inter-feature correlations, we additionally con- duct experiments on two challenging RNA-seq datasets. (i) The PAGA---AN bulk RNA-seq dataset (n 801, d = 20531) (Weinstein et al., 2013), where the goal is to classify samples into five tumor types: breast invasive carcinoma (BRCA), kidney renal clear cell carcinoma (KIRC), colon adenocarcinoma (COAD), lung adenocarcinoma (LUAD), and prostate adenocarci- noma (PRAD). This dataset is representative expression structures and severe p imbalance. (ii) The human single-cell RNA-seq dataset (n = 632, d = 23257) (Darmanis et al., 2017). which is used to distinguish ncoplastic cells orig- inating from the tumor core versus those from the periphery. This dataset exemplifies ultra-high-
large-scale transcriptomic studies with strong co-
39

---

Published as a conference paper at ICLR 2026
Table E8: Classification performance (accuracy) of three methods with different top-k feature sets on 4 datasets with large sample sizes.
Dataset
Codon usage
Default of Credit Card Clients
Type
Method
Top-8
Top-4
Top-2
Top-8
Top-4
Top-2
FDFI
0.7642
0.6448
0,5401
0.8182
0.8197
0.8197
Classification
DFI
0.7588
0.6391
0,5142
0.8150
0.8137
0.8173
ad-hoc
0.7132
0.5839
0.4479
0.8173
0.8118
0.8173
Dataset
Superconductivty Data
Video Transcoding
Type
Method
Top-8
Top-4
Top-2
Top-8
Top-4
Top-2
FDFI 11.0347
12.7988
15.0577
6.8195
7.9727
9.0607
Regression
DFI 11.4675
13.4683
16.9657
6.8588
8.2657
9.3723
ad-hoc
12.0083
13.1241
17.8346
6.8991
8.2052
9.1645
TOGA-PANCAN
Human sriM//ssg
Wein sccaa
0.94
0.03
0.7
Number varisbles DO
Number variabies (K)
FDF
+ DF nd hoc
Figure E10: Prediction accuracy vs. top-k selected variables across two RNA-seq datasets.
dimensional single-cell measurements with heterogeneous cell populations and complex gene-gene dependencies, providing stringent testbed for evaluating feature-importance methods. Following the preprocessing procedures in (Yan et al., 2025), we selected 1,500 highly variable genes (HVGs) bulk RNA-seq dataset and 2,000 HVGs for the human single-cell RNA-seq dataset, and performed all downstream analyses pared the average prediction accuracy across datasets for the important features selected by FDFI, DFI, and an ad hoc CPI-based clustering approach, using two-fold splits and reporting the mean accuracy for each dataset.
for the TCGA-PANCAN-
these HVG subsets. We then com-
In particular, for the human single-cell RNA-seq dataset, found that the Top-20 genes selected by our method include ALDOC, HES6, and CPE, all of which have been previously implicated in glioma biology. ALDOC, glycolytic enzyme highly enriched in neural tissue, has been im- plicated in glioma metabolic reprogramming, where altered ALDOC expression associated with enhanced glycolytic activity, tumor progression, and increased migratory potential of glioma cells (Chang et al., 2024). HES6, a basic helix-loop-helix transcription factor, selectively overex- pressed in glioma and functions an important regulator driving tumor-cell proliferation, migration, and lineage plasticity (Haapa-Paama) Paananen et al., 2012). CPE which encodes neuroendocrine peptide- processing enzyme, has more recently been recognized as an oncogenic factor in high bigh-grade gliomas, where its elevated expression promotes tumor-cell survival, invasion, and stress adaptation through enhanced metabolic resilience and extracellular-man matrix remodeling (Hareendran al., 2022).
These concordances between known glioma biology and the features automatically selected by our method provide additional support for the clinical relevance of our approach. In contrast, for the CAGAA-PACAN-HHSSS dataset, only gene indices (rather than gene symbols) are provided, so we report prediction accuracy only and do not attempt to interpret the selected genes terms of their clinical relevance.
40

---

Published as a conference paper at ICLR 2026
REFERENCES
Michael S Albergo and Eric Sunden-Eijnden. Building normalizing flows with stochastic inter- polants. In 11th International Conference on Learning Representations, 2023.
Herbert Amann. Ordinary differential equations: an introduction to nonlinear analysis, volume 13. Walter de gruyter, 2011.
W. Anderson. An Introduction to Multivariate Statistical Analysis. Wiley, 3rd edition, 2003.
Diogo Ayres-de Campos, Joao Bernardes, Antonio Garrido, Joaquim Marques-de Sa, and Luis Pereira-Leite, Sisporto 2.0: a program for automated analysis of cardiotocograms. Journal of Maternal- Medicine, 915)311--3 2000.
Diogo Ayres-de Campos, Catherine Y Spong, Edwin Chandraharan, et al. Figo consensus guide- lines on intrapartum fetal monitoring: Cardiotocography, International Journal of Gynecology & Obstetrics, (1):31233 2015.
Peter L Bartlett, Dylan J Foster, and Matus Telgarsky. Spectrally-normalized margin bounds for neural networks. Advances in neural information processing systems, 30, 2017.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, vol- ume 4. Springer, 2006.
D. Campos and Bernardes. Cardiotocography. UCI Machine Learning Repository, 2000.
Yu-Chan Chang, Ming-Hsien Chan, Chien-Hsiu Li, Chi-Long Chen, Wen-Chiuan Tsai, and Michael Hsiao. Ppar-r agonists reactivate the aldoc-nr2fl axis to enhance sensitivity to temozolomide and suppress glioblastoma progression. Cell Communication and Signaling. 22(1):266, 2024.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural pa- rameters, 2018.
Spyros Darmanis, Steven A Sloan, Derek Croote, Marco Mignardi, Sophia Chernikova, Peyman Samghababi, Ye Zhang, Norma Neff, Mark Kowarsky, Christine Caneda, et al. Single-cell ma- seq analysis of infiltrating neoplastic cells the migrating front of human glioblastoma. Cell reports, 21(5): 2017.
Karina W Davidson. Michael Barry, Carol M Mangione, Michael Cabana, Aaron B Caughey, Esa M Davis, Katrina Donahue, Chyke A Doubeni, Alex Krist, Martha Kubik, et al. Sereening for prediabetes and type 2 diabetes: Us preventive services task force recommendation statement. Jama, (2366))73,7743, 2021.
Jin-Hong Du, Kathryn Roeder, and Larry Wasserman. Disentangled feature importance. arXiv preprint arXiv:2507.00260, 2025a.
Jin-Hong Du, Zhenghao Zeng, Edward H Kennedy, Larry Wasserman, and Kathryn Roeder, Causal inference for genomic data with multiple heterogeneous outcomes. Journal of the American Statistical Association, pp. 1-24, 2025b.
Nuha A ElSayed, Grazia Aleppo. Vanita R Aroda, Raveendhara Bannuru, Florence M Brown, Dennis Bruemmer, Billy S Collins, Jason L Gaglia, Marisa E Hilliard, Diana Isaacs, et al. 2. classification and diagnosis diabetes: standards of care in diabetes-200 Diabetes care, 46 (Supplement -.S19-S40, 2023.
Connor A Emdin, Simon G Anderson, Mark Woodward, and Kazem Rahimi. Usual blood pressure and risk of new-onset diabetes: evidence from 4.1 million adults and meta-analysis of prospec- tive studies. Journal of the American College of Cardiology, 66(14):1552-1562 2015.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1(2). MIT press Cambridge, 2016.
41

---

Published as a conference paper at ICLR 2026
Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael J Cree. Regularisation of neural net- works by enforcing lipschitz continuity. Machine Learning, 110(2):333444 2021.
Peng Guo, Quan Zhou, Lei Ren, Yu Chen, and Yue Hui. Higher parity is associated with increased risk of type 2 diabetes mellitus in women: linear dose-response meta-analysis of cohort studies. Journal of diabetes and its complications, 3 31(1):58-66, 2017.
Saija Haapa-PPanaana S Kiviluoto, M Waltari, Marjut Puputti, John Patrick Mpindi, Pekka Koho- nen, Olli Tynninen, H Haapasalo, Heikki Joensuu, Merja Perälä, et al. Hes6 gene is selectively overexpressed in glioma and represents an important transcriptional regulator of glioma prolifer- ation. Oncogene, 11110:1199911101111 2012.
Anthony JG Hanley, Ken Williams, Clicerio Gonzalez, Ralph Agostino Jr, Lynne E Wa- genknecht, Michael P Stern, and Steven M Haffner. Prediction of type diabetes using simple measures of insulin resistance: combined results from the san antonio heart study, the mexico city diabetes study, and the insulin resistance atherosclerosis study. Diabetes, 52(2):463-469, 2003.
Sangeetha Hareendran, Xuyu Yang. Vinay Kumar Sharma, and Y Peng Loh. Carboxypeptidase and its splice variants: key regulators of growth and metastasis in multiple cancer types. Cancer letters, 548:215882, 2022.
Philip Hartman. Ordinary differential equations. SIAM, 2002.
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6:695-709, 2005.
Ahmad Jayedi, Sepideh Soltani, Sheida Zeraat-talab Motlagh, Alireza Emadi, Hosein Shahinfar, Hanieh Moosavi, and Sakineh Shab-Bidar. Anthropometric and adiposity indicators and risk of type 2 diabetes: systematic review and dose-response meta-analysis of cohort studies, Bmj, 376, 2022.
Yan-Ju Jia, Tullio Ghi, Susana Pereira, Anna Gracia Perez-Bonfils, and Edwin Chandraharan. Patho- physiological interpretation of fetal heart rate tracings in clinical practice. American journal of obstetrics and gynecology. 228(6):622 644, 2023.
Jing Lei, Max G'Sell, Alessandro Rinaldo, Ryan J Tibshirani, and Larry Wasserman. Distribution- free predictive inference for regression. Journal of the American Statistical Association, 113 (523):1094-1111, 2018.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXhv:2210.02747, 2022.
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.06264 2024.
Pierre Mah and Jean-Baptiste Veyrieras. Micromass. UCI Machine Learning Repository, 2014.
Wanda K Nicholson, Keiko Asao, Frederick Brancati, Josef Coresh, James S Pankow, and Neil R Powe. Parity and risk of type 2 diabetes: the atherosclerosis risk in communities study. Diabetes care, 29(11):2349 -354,, 2006.
Julian T Parer and Tomoaki Ikeda. A framework for standardized management of intrapartum fetal heart rate patterns. American journal of obstetrics and gynecology, 197(1):26-e1, 2007.
Angel Reyero-Lobo, Pierre Neuvial, and Bertrand Thirion. Conditional feature importance revisited: Double robustness, efficiency and inference, 2026.
17520.
Maria Romano, Luigi luppariello, Alfonso Maria Ponsiglione, Giovanni Improta, Paolo Bifulco, and Mario Cesarelli. Frequency and time domain analysis of foetal heart rate variability with traditional indexes: a critical survey. Computational and mathematical methods medicine, (016(1):958543) 2016.
42

---

Published as a conference paper at ICLR 2026
Andrea Ruiz-Alejos, Rodrigo M Carrillo-Larco, J Jaime Miranda, Robert H Gilman, Liam Smeeth, and Antonio Sernabé-rrizz Skinfold thickness and the incidence of type diabetes mellitus and hypertension: analysis of the peru migrant study. Public health nutrition, 23(1):63-71, 2020.
L. Shapley, A value for n-person games. In Contributions to the Theory of Games (AM-28), Volume II, pp. 307-318, Princeton University Press, Princeton, NJ, 1953. ISBN 9780691079356.
Jack W Smith, James E Everhart, William C Dickson, William C Knowler, and Robert Scott Jo- hannes. Using the adap learning algorithm forecast the onset diabetes mellitus. Proceed- ings of the annual symposium computer application in medical care, pp. 261, 1988.
T Stampalija, A Bhide, AEP Heazell, A Sharp, C Lees, al. Computerized cardiotocography and dawes-redman criteria: how should we interpret criteria not met? Ultrasound in Obstetrics & Gynecology, 61(6):6616666 2023.
Carolin Strobl, Anne-Laure Boulesteix, Thomas Kneib, Thomas Augustin, and Achim Zeileis. Con- ditional variable importance for random forests. BMC bioinformatics, 9:1-11 2008.
Rodolfo Valdez, Paula W Yoon, Tiebin Liu, and Muin J Khoury. Family history and prevalence of diabetes in the us population: the 6-year results from the national health and nutrition examination survey (1999-2004). Diabetes care, 30(10):5117-2222 2007.
Isabella Verdinelli and Larry Wasserman. Decorrelated variable importance. Journal of Machine Learning Research, 25(7):1-27, 2024a.
Isabella Verdinelli and Larry Wasserman. Feature importance: A closer look at shapley values and loco. Statistical Science, 39(4):623-633 636, 2024b.
Aladin Virmaux and Kevin Scaman. Lipschitz regularity of deep neural networks: analysis and efficient estimation. Advances in Neural Information Processing Systems, 31, 2018.
John N Weinstein, Erie Collisson, Gordon B Mills, Kenna Shaw, Brad A Ozenberger, Kyle Ellrott, Ilya Shmulevich, Chris Sander, and Joshua M Stuart. The cancer genome atlas pan-cancer analysis project. Nature genetics, 15(10):1113-123) 2013.
Brian Williamson and Jean Feng. Efficient nonparametric statistical inference on population feature importance using shapley values. In International conference on machine learning, pp. 10282- 10291. PMLR, 2020.
Guanao Yan, Shuo Harper Hua, and Jingyi Jessica Li. Categorization of 34 computational methods to detect spatially variable genes from spatially resolved transcriptomics data. Nature Communi- cations, 16(1):1141, 2025.
Zhidong Zhao, Yang Zhang, and Yanjun Deng. A comprehensive feature analysis of the fetal heart rate signal for the intelligent assessment of fetal state. Journal of clinical medicine, 7(8)2223, 2018.
43


Reviews:
[{"review_id": "iO8RaThDe6", "reviewer": "Reviewer_rfpb", "summary": "The paper propose a new method to reliably quantify feature importance using flow-disentangled. The method follows Disentangled Feature Importance (DFI) and resolve two critical problems of the previously proposed principled framework for attribution under dependence: 1) relying on optimal transport (OT) and 2) restriction to important score with L2 loss. The paper develop statistical inference theory and show empirical results.", "questions": "Is there any downstream task that could be used or improved by the novel feature attribution method?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper has a strong methodological innovation\n  - The previous method Disentangled Feature Importance (DFI) is limited by the optimal transport (OT)  and L-2 loss with restricted applications.\n- The evaluation is comprehensive. The experiment settings include both synthetic data and real data settings. The method outperformed other baselines significantly. \n- The paper contains solid theory analysis.", "weaknesses": "- The experiment setting is over-simplified with some simple correlation and the real experiment setting is too small with less than 3000 data points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MTQmGEAr8R", "reviewer": "Reviewer_8E6r", "summary": "The paper introduces Flow-Disentangled Feature Importance (FDFI), a new model-agnostic framework for estimating feature importance and its statistical uncertainty in ML models. Unlike standard model-agnostic attribution methods (e.g., permutation or removal-based approaches) that become unreliable when features are correlated (correlation distortion), FDFI provides statistically valid and interpretable attributions even under severe multicollinearity.\n\n\nPrevious decorrelation-based methods (e.g., Disentangled Feature Importance, DFI) were limited to l2 loss, but FDFI generalizes this to any differentiable loss function, allowing its use in both regression and classification settings. FDFI employs flow matching, a recent generative modeling technique (Lipman et al.), to learn a disentanglement map that transforms correlated features X into statistically independent latent variables Z. This approach replaces the restrictive Gaussian otpimal transport assumption used in prior work. It enables flexible modeling of arbitrary feature distributions while preserving interpretability of the transformation.\n\nThe authors derive a semiparametric inference theory showing that FDFI estimators are asymptotically efficient and allow valid confidence intervals and hypothesis testing with controlled Type I error. The FDFI estimator remains sqrt(n)-consistent even when the flow map T is estimated nonparametrically at slower rates. Furthermore, FDFI quantifies how each feature’s contribution to model loss propagates through the correlation structure of the data. It provides both latent-level and original-feature-level importance scores, linking disentangled representations back to raw variables.\n\nEvaluations on synthetic and biomedical datasets compare FDFI with LOCO (Leave-One-Covariate-Out), CPI (Conditional Permutation Importance), and DFI. FDFI achieves higher statistical power in detecting relevant features, controlled type I error, and robust performance under high feature correlation. FDFI also recovers correlated yet predictive features that other methods fail to identify.", "questions": "Please see weaknesses section.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main contribution of the paper is well-motivated, clearly situating its novel approach within an interesting problem setting. The authors provide a compelling analysis of existing methods, outlining its strengths and weaknesses, clearly identifying specific gaps and limitations that effectively frame the need for their proposed framework. This strong introduction is complemented by a solid theoretical foundation. The authors build their first theoretical results on reasonable assumption (A1) to establish strong initial guarantees, such as the tight bounds presented in Theorem 2.1, ensuring the soundness of their method.\n\nThe primary theoretical contribution is Theorem 3.1, establishing the semiparametric efficiency of the latent FDFI estimator, $\\hat{\\phi}^{\\text{FDFI}}_{Z_j}$. This is a solid result. It proves that the estimator achieves the optimal $\\sqrt{n}$-convergence rate. A significant practical implication of this efficiency is that this $\\sqrt{n}$-consistency holds even when using a flexible, nonparametrically-estimated nuisance map, $\\hat{T}$, which itself may converge at a slower rate.\n\nThis property is important indeed as it allows for the use of complex machine learning models to estimate the nuisance components without sacrificing the statistical validity of the final importance score. Finally, the authors complete their framework by providing the necessary inferential tools in Proposition 3.2. This proposition grounds the theory in practice, allowing for the construction of valid confidence intervals and hypothesis tests for the FDFI scores, making the entire method statistically sound and practically applicable.\n\nBefore proceeding to the weaknesses I would like to point out that, although the theoretical results seem sound, I have not checked their validity (especially regarding the proofs).", "weaknesses": "The paper extends the DFI framework and identifies two of its key limitations: its reliance on the optimal transport map for disentanglement, which can be computationally expensive and less flexible for complex, high-dimensional distributions. However, this claim seems somewhat not well tackled by FDFI. From the appendix, it becomes evident that the proposed FDFI methods are actually slower in terms of computation time, and the experiments are applied only to relatively simple settings—Gaussian and mixture-of-Gaussian data, along with a real-world dataset of max 21 dimensions. When reading this section, I was hoping that FDFI would be faster than DFI or at laset that the authors would provide stronger evidence demonstrating its advantages on complex, high-dimensional data. Neither of these expectations is convincingly addressed.\n\nFigure 1 focuses primarily on highlighting the benefits of FDFI compared to LOCO, CPI, and SHAP, yet it would be more informative if it also emphasized improvements over DFI. This addition would help clarify and disentangle which of the observed gains are due to the new FDFI framework itself and which are inherited from DFI.\n\nIn line 233, the authors claim that DFI is limited because it cannot be applied to classification tasks or models using general differentiable loss functions. However, in Assumption 1 they assume that the loss is M-smooth, which already implies differentiability. It is unclear whether this assumption applies only to the first two theoretical results or to the broader framework. If the former case is the true one, it would be useful for the authors to clarify this connection and more explicitly relate Section 2.1 to the FDFI formulation.\n\nIn line 260, the authors mention that Equation (4) recovers the DFI formulation under the l2 loss. It would strengthen the paper if they explicitly demonstrated this equivalence or at least referred readers to an appendix section where the derivation is shown.\n\nAnother issue is that Assumptions A2–A4, which seem to be key to the theoretical results in Section 3, are introduced in Theorem 3.1 without any prior mention/discussion or intuitive explanation. Unless I have missed this (in which case I would appreciate if the authors can point me to this), the absence of even a brief description of what these assumptions entail and how restrictive they are leaves the reader with little understanding/comprehension of the conditions under which the results hold. The same issue appears in Proposition 3.2, which refers to Assumption A4(iv) without elaboration. A more detailed or intuitive analysis/explanations of these assumptions in the main text would greatly improve the clarity of the manuscript.\n\nWhen reviewing Algorithm D.1 in the appendix, it appears quite complex and, when compared to DFI, even slower in execution. Given its importance, a concise yet clear overview of how the inference procedure operates under FDFI should be included in the main text rather than entirely postponed to the appendix. This would make the paper more accessible and self-contained.\n\nRegarding the experiments, the authors conduct three sets—on Gaussian data, mixtures of Gaussians, and a real-world dataset. While these are informative, they do not convincingly support the paper’s claims about handling complex, high-dimensional data distributions. The authors should either perform experiments on higher-dimensional datasets or tone down their claims about performance in such settings. Alternatively, they could provide a justification explaining why their chosen real-world dataset adequately reflects the complexity they aim to address.\n\nIn addition, the time analysis presented in Figure D.3 raises further questions. Out of curiosity, could the authors explain the sudden increase in DFI’s runtime between sample sizes 400 and 600? Also, the slope of FDFI’s runtime growth with respect to sample size appears steeper than that of DFI.. These observations suggest scalability issues, and it would be appropriate for the authors to acknowledge this limitation explicitly in the main text or conduct further experiments.\n\nOverall, while the paper makes a strong theoretical contribution and proposes a well-founded framework, the implementation appears surprisingly complex and may benefit from simplification or at least justification of its current form. Some crucial content that supports understanding—such as explanations of key assumptions and the inference algorithm—should be moved from the appendix into the main text. On the other hand, less important material, like parts of the Gaussian mixture experiments (which I do not believe add much), could be moved to the appendix. Doing so would improve  the paper’s structure, its readability and transparency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ySOvnxYYxN", "reviewer": "Reviewer_akD3", "summary": "This paper introduces Flow-Disentangled Feature Importance (FDFI), a model-agnostic interpretability framework designed to provide statistically valid feature importance estimation even in the presence of strong feature correlations. The key idea is to use flow matching—a flexible generative modeling technique—to learn a disentanglement map that transforms correlated features into an approximately independent latent space. Once features are decorrelated in this latent space, feature importance is computed and then projected back to the original space, along with uncertainty quantification based on semiparametric efficiency theory.\nFDFI generalizes earlier Disentangled Feature Importance (DFI) work, extending it from the restrictive $\\ell_2\\text{-loss}$ setting to arbitrary differentiable loss functions, thereby supporting both regression and classification tasks. The paper provides theoretical results proving the asymptotic normality and semiparametric efficiency of the proposed estimators, along with empirical results on synthetic and biomedical datasets (e.g., Cardiotocography, Diabetes, and MicroMass). Empirically, FDFI achieves higher AUC, power, and more robust inference than competing model-agnostic methods such as LOCO, CPI, and DFI, especially under multicollinearity.", "questions": "1. Latent independence diagnostics: Can you report quantitative independence scores (e.g., pairwise HSIC or distance correlation) for the learned $Z$ in the main experiments, and relate them to attribution stability and power?\n2. Sensitivity to flow training: In Appendix D.2 you select hyperparameters with MMD; could you add (a) an ablation on auxiliary set size $m$, (b) $M$ (latent resamples) vs. variance/power, and (c) model depth/width vs. AUC/power?\n3. Interpretability of Jacobians: The CTG heatmap (Fig. 4a) is compelling. Can you provide case‑wise attributions (e.g., conditional on ranges of clinical variables) to illustrate how Jacobian structure plus latent scores produce the final importance?\n4.When FDFI may fail: Are there data‑generating processes (e.g., strong higher‑order interactions that can’t be linearized locally) where the squared‑Jacobian weighting is systematically biased, even if the latent scores are well estimated?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Sound theoretical framing: The paper provides a clear unification of LOCO, CPI, and SCPI under general differentiable loss functions, and introduces a principled FDFI definition that subsumes DFI as the ℓ₂ special case. The semiparametric efficiency and asymptotic normality proofs are technically rigorous and well-documented.\n2. Flexible, scalable transport: The flow-matching approach yields expressive and stable mappings with a well-defined training objective.\n3. Compelling empirical results: FDFI consistently achieves higher AUC and statistical power while maintaining Type-I error control across correlated and non-Gaussian settings (Figures 2–3). The CTG, Pima, and MicroMass studies further confirm interpretability and clinical relevance, including recovery of known physiological patterns and robustness in high-dimensional regimes (d ≈ 1300).", "weaknesses": "1. Because the disentanglement map $T$ is learned via flow matching rather than optimal transport, different flow paths could, in principle, lead to distinct ‘independent’ latent representations $Z$. While Appendix B establishes uniqueness of the ODE solution for a chosen interpolation path, semantic identifiability of the latent axes remains uncertain. This limits how literally one can interpret the Jacobian heatmaps as explanatory structures. Incorporating quantitative independence diagnostics (e.g., distance correlation or energy tests on $Z$) would strengthen interpretability claims.\n2. Although the asymptotic theory accommodates $n^{-\\frac{1}{4}}$ convergence rates, finite-sample behavior may be sensitive to inaccuracies in the estimated flow $T$. The paper could be strengthened by an ablation over auxiliary-sample size $m$ and architectural choices (depth, width) to examine effects on MMD independence, feature scores, and variance estimates. Appendix D.2 provides training details, but robustness to nuisance quality remains empirically under-explored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)