Counterfactual Density Estimation using Kernel Stein Discrepancies
Authors: Diego Martinez-Taboada
Edward Kennedy
Venue: ICLR 2024 poster (2024)
Keywords: counterfactual density estimation
kernel Stein discrepancy
causal inference
kernel methods
Primary area: causal reasoning
TL;DR: Given a class of densities known up to normalizing constants, we propose to model counterfactual distributions by minimizing kernel Stein discrepancies in a doubly robust manner.

Abstract: Causal effects are usually studied in terms of the means of counterfactual distributions, which may be insufficient in many scenarios. Given a class of densities known up to normalizing constants, we propose to model counterfactual distributions by minimizing kernel Stein discrepancies in a doubly robust manner. This enables the estimation of counterfactuals over large classes of distributions while exploiting the desired double robustness. We present a theoretical analysis of the proposed estimator, providing sufficient conditions for consistency and asymptotic normality, as well as an examination of its empirical performance.

Published as a conference paper at ICLR 2024
COUNTERFACTUAL DENSITY ESTIMATION USING KERNEL STEIN DISCREPANCIES
Diego Martinez-Taboada Department of Statistics & Data Science Carnegie Mellon University Pittsurgh, PA 15213, USA die diegomar& andrew. cmu. edu
Edward H. Kennedy
Department of Statistics Data Science Carnegie Mellon University Pittsburg,, PA 15213, USA edward@stat cmu. edu
ABSTRACT
Causal effects are usually studied in terms of the means of counterfactual distri- butions, which may be insufficient in many scenarios. Given a class of densities known up to normalizing constants, we propose to model counterfactual distribu- tions by minimizing kernel Stein discrepancies in doubly robust manner. This enables the estimation of counterfactuals over large classes of distributions while exploiting the desired double robustness. We present a theoretical analysis of the proposed estimator, providing sufficient conditions for consistency and asymptotic normality, as well as an examination of its empirical performance.
INTRODUCTION
Causal targets examine the outcomes that might have occurred if specific treatment had been administered to a group of individuals. Generally, only the expected value of these outcomes is analyzed, as seen in the widely used average treatment effect. Nonetheless, focusing solely on means proves insufficient in many scenarios, Important attributes of the counterfactuals, such as their variance or skewness, are often disregarded. Modelling the entire distribution gives a complete picture of the counterfactual mechanisms, which opens the door to richer analysis. For example, a multimodal structure in the counterfactual may indicate the presence of distinct subgroups with varying responses to the treatment (Kennedy et al. (02T). This profound understanding of the counterfactual may ultimately exploited the design of new treatments.
In order to model a distribution, one may consider estimating either its cumulative distribution func- tion (CDF) or its probability density function (PDF). While the statistical analysis of the former is generally more straightforward, the latter tends to be more attractive for practitioners given its appealing interpretability. Although density estimation may be conducted non non-parametrically, mod- elling probability density functions based parametric class of distributions has gamered signif- icant attention. These models allow easy integration of prior data knowledge and offer interpretable parameters describing distribution characteristics.
However, a number of interesting parametric density functions are only known up to their normal- izing constant. Energy-based models (LeCun al. 2006) establish probability density functions qe (v()) x exp(-Ea (v)). where Eg is parametrized potential fulfilling f exp(-Ex(x) Sometimes, the energy is expressed uet of experts or restricted Boltzmann machines (Ackley et al. 1985 Hinton 2002 2010). In con- trast, energy-based models may link inputs directly to outputs (Mnih & Hinton 2005|]Hintt 2006). Gibbs distributions and Markov random fields (Geman & Geman 1984 Clifford 1990), as well as exponential random graph models (Robins et al. ( al. 2013) are also examples of this form of parameterization. We highlight the generality of energy-based models, which allow for modelling distributions with outstanding flexibility.
< co.
combination of hidden and visible variables, namely prod-
al.,
Generally, the primary difficulty in manipulating energy-based models stems from the need to pre- cisely estimate the normalizing constant. Nonetheless, this challenge may be circumvented by using the so-called kernel Stein discrepancies, which only require the computation of the score function V. log qa (z). Given class distributions Q = {4o loce and samples from a base distribution Q, one may model the latter by the distribution go,, that minimizes its kernel Stein discrepancy with

---

Published as a conference paper at ICLR 2024
respect to the empirical distribution Qn. However, these minimum kernel Stein discrepancy esti- mators have not been explored in the challenging counterfactual setting, where outcomes are not always observed, and the treatment assignment procedure ought to be taken into account.
In this work, we propose to model counterfactuals by minimizing kernel Stein discrepancies in doubly robust manner. While the presented estimator retains the desired properties of double robustness, it enables flexible modelling of the counterfactual via density functions with normalizing constants that need not be specified. Our contributions are two-fold. First, we present novel estimator for modelling counterfactual distributions given parametric class of distributions, along with its theoretical analysis. We provide sufficient conditions for both consistency and asymptotic normality. Second, we illustrate the empirical performance of the estimator in a variety of scenarios.
RELATED WORK
Counterfactual distribution estimation has been mainly addressed based on CDF approximation
(Abadiel 2002 Chernozhukov & Hansen 2005) Chernozhukov et al. 2013 Diaz, 2017), where
the predominant approach reduces to counterfactual mean estimation. contrast, counterfactual PDF estimation generally relies on kernel smoothing (Robins & Rotnitzky 2001 Kim et al.) 2018 8)
or projecting the empirical distribution onto a finitedimeenionnl model using f --dverrgences or L norms (Westling& & Carone| 2020; Kennedy al. 2021 Melnychuk et al. 2023). We highlight that none of these alternatives can in principle handle families of densities with unknown normalizing constants. For instance, using the general f f-divergences
L norms in the projection stage of the estimation (Kennedy et al.] (2021) requires access to the exact evaluation of the densities, which in- cludes the normalizing constants. This is precisely what motivated the authors of this contribution to explore kernel Stein discrepancies in the counterfactual setting.
Kernel Stein discrepancies (KSD). which build on the general Stein's method (Stein 1972 Gorham & Mackey| (2015), were first introduced for conducting goodness-of ooddness-of-fit tests and sample quality anal-
ysis (Liu al. 2016]] [Chialkkoski al. 2016 Gorham & Mackey| (017): 7) they may be understood as kernelized version of score- mateh-na divergence (Hyvärinen & Dayan 2003). Minimum kernel Stein discrepancy (MKSD) estimators were subsequently proposed (Barp et al. 2019 Matsubara eet al. 2022), which project the empirical distribution onto finite-dimensional model using the KSD. We highlight that MKSD estimators had not been proposed in the counterfactual settings prior to this work. Lam & Zhang (2023) suggested doubly-robust procedure to estimate expectations via Monte Carlo simulation using KSD, but their motivation and analysis significantly diverge from our own contributions: while MKSD estimators minimise the KSD over class of distributions, (quasi) Monte Carlo methods exploit KSD by transporting the sampling distribution (Oates et al., [017]] Fisher al. 2022T] Korba al.) 2021). We refer the reader to Anastasiou al. (2023) for review on Stein's 's methods, and to Oates et al. (2022) for an overview of MKSD estimators.
Kernel methods have been gaining interest in causal inference for assessing whole counterfactual distributions. In order to test for (conditional) distributional treatment effects, Muandet et al. (202T) and Park et al. (202I) made use of kernel mean embeddings via inverse propensity weighting and plug-in estimators, respectively. This line of work was later generalized to doubly robust estimators by Fawkes a (2022) and |Martinez- Startinez-Taboada et al. 2023). Beyond distributional representation, kernel regressors have have found extensive use in counterfactual tasks (Singh et al., 2019 2020
2021
Zhu et al. 2022).
BACKGROUND
Let AY) P. where X E X. Y C Rd, and A E {0,1) represent the covariates, outcome, and binary treatment respectively, We frame the problem in terms of the potential outcome framework (Rubin) 2005 Imbens& & Rubin 2015). assuming A1) Y = AY + (1 A)Y (where yl ~ Q! and n Qa are the potential outcomes or counterfactuals) A2) y y A X, A3) 1|X) < - almost surely for some > 0.
< (X) = P PAIX
Conditions A1-A3 are ubiquitous in causal inference, but other identifiability assumptions are also possible. Condition 1 holds true when potential outcomes are exclusively determined by an indi-
2

---

Published as a conference paper at ICLR 2024
vidual's own treatment (i.e., interference), and Condition A2 applies when there unmea- sured confounders. Condition A3 means treatment is not allocated deterministically.
Under these three assumptions, it is known that the distribution of either counterfactual may be expressed in terms of observational data. The ultimate goal of this contribution to model either distribution Y using a parametric class of distributions Q {go Yeee which only need to be specified up to normalizing constants. Without loss of generality, we restrict our analysis to the potential outcome of the treatment Y Q². For conducting such task, we draw upon minimum kernel Stein discrepancy estimators, which build on the concepts of reproducing kernel Hilbert space and Stein's 's method.
Reproducing kernel Hilbert spaces (RKHS): Consider non-empty set ) and a Hilbert space H. of functions f y + R equipped with the inner product )). The Hilbert space H. is called an RKHS if there exists function & y x ) + IR (referred to as reproducing kernel) satisfying (i) k(,y) E H for all y E V. (ii) (f, k(: y) = f(y) for all y E y and f E H. We denote by H the product RKHS containing elements h == (h1.-..,ha) ha) with h; E H and (h, h) Ei= (hi, hT)n-
Kernel Stein discrepancies (KSD); Assume Qo has density go on y C Rd, and let H be an RKHS with reproducing kernel & such that Vk( exists for all y V. Let so(y) V, log (v()) and define Eo(: v) = (so(v))(.,)) + Vk( v)] E Hd, so that
ho(y. 5) (Eo(- Eo(- W)
= (so(v), so (o)) 0) + (so(0), V Vyk(y. U))
(so(y), Vy 0+ + (Vyk(..)). Vyh((, ())
reproducing kernel. The KSD defined KSD(Q, [Q) = (Ey, uolyiolll the (Y. 1/2 Under certain regularity conditions, KSD(Qe = 0 if (and only if) Q = Q Chwialkowski al. 2016); other properties such weak convergence dominance may also be established (Gorham & Mackey 2017 Further, if EQy / hY, Y) < co, then KSD(Qu [[) |Byy~ (5 Y)] 11
In non-causal settings, given YI,..., Yn ~ Q and the closed form evaluation of he presented in equation I the V-statistic
1
(O)
1 Eo(. 10
he ho,
may be considered for estimating KSD2 (SD*(Q+)I()) (Q+ Similarly, removing the diagonal elements of this V-statistic gives way to unbiased U-statistic, which counts with similar properties.
Minimum kernel Stein discrepancy (MKSD) estimators: Given Q = {ya}eee known up to
normalizing constants, the scores 89 can be computed. MKSD estimators model Q by Qo... where
On E arg min gm (0; Yi,..., 0g0
and 9n is the V-statistic detined in equation 2 or its unbiased version. The MKSD estimator 0, is consistent and asymptotically normal under regularity conditions (Oates et al. 2022).
4 MAIN RESULTS
We consider the problem of modeling the counterfactual distribution Y1 Q² given a paramet- ric class of distributions Q = {aelece with potentially unknown normalizing constants. We work under the potential outcomes framework, assuming that we have access observations (X;, A,, 1 E X' x {0, 1} x y sampled as (X, A, Y) P, such that Conditions A1-A3 hold. Throughout, we denote (X, Y) and
x (0. x y for ease of presentation.
Like the previously introduced MKSD estimators, our approach includes modeling Q' by choosing a such that
0, E arg min In (0: Zi,... Zn), leB
(3)
3

---

Published as a conference paper at ICLR 2024
where 9n is a proxy for KSD² (')) In contrast such MKSD estimators, defining In. the V- statistic introduced in equation would lead to inconsistent estimations, given that the distribution of Y may very likely differ from that of y (this is, indeed, the very essence of counterfactual settings).
In order to define an appropriate In. we first note that, by the law iterated expectation,
= Ey [Eo y²]] Ez-p [o(Z)],
where
π(x) (5o(:33) (z) + +0(2), π (2x) = [AA x], (x(x)) = E [so( Y)|4= 1, x
x(z) = a
(4)
An analogous result holds for y or any other discrete treatment level. Note that the yo does not affect neither nor 3g. In fact, the problem could have been presented missing outcome data problem, where the data from Y is missing. The authors have posed the problem in counterfactual outcome terms simply for motivational purposes.
The embedding & induces a reproducing kernel =) (0(=), (()) ze, which will be used in subsequent theoretical analyses. We highlight that c the efficient influence function for v. This implies that the resulting estimators have desired statistical properties such as double robustness, as shown later. That is, V = Ps will be consistent if either á or ßg is consistent, and its convergence rate corrssonds to the product of the learning rates of â and Bg.
Nonetheless, it may not be feasible to estimate Be individually for each @ of interest. In turn, let us assume for now that have access to estimators á and B, where the latter approximates B(x)) E [k(, () = 1, = 2r). and is of the form
B(x) = Σ (x)
(5)
We highlight that many prominent algorithms for conducting Hd-valued regression, such as con- ditional mean embeddings (Song al. 2009; Grünewälder al. (012) and distributional random forests (Cevid et al. 2022 När et al. 2023), are of the form exhibited in equation 5 In order to
avoid refitting the estimator 30 for each value of 6, propose define Bo from 3 follows:
Bo(x) [8o(8)] = Σ (d( Y.).
(6)
Intuitively, we could expect that if estimator 9 is consistent, then 3p is also consistent if the mapping k(,v) > [so(v) + Vk(.v) is regular enough. As a result, we propose the statistic
(0; Y Yn)
So
(o))) Zz), n
T2
H
where 3o is constructed from as established in equation 6 Although this statistic resembles the one presented in Martiiee--Taboada al. (2023) and Fawkes al. (2022), several important differences arise between the contributions. First and foremost, they studied a testing problem, so the causal target was different. Second, their work did not extend to embeddings that depend on a parameter 0. Lastly, they focused on the distribution of their statistic under the null in order calibrate a two-sample test; in contrast, our approach minimizes a related statistic for directly modelling the counterfactual. The main theoretical contribution of Martinez- Taboada et al. (2023) the extension of cross U-statistics to the causal setting; in stark contrast, our contribution deals with the theoretical properties of the minimizer of 'debiased' V-statistie. Both the motivation and theoretical challenges our study diverge from those in the carlier works; we view our research as complementary and orthogonal to these prior contributions.
While we have assumed far that estimators â and Bg are given, defining optimal strategy for training such estimators is key when seeking to maximize the performance. A simple approach, such using half of the data estimate and So. and the other half on the empirical averages of the
4

---

Published as a conference paper at ICLR 2024
Algorithm 1 DR-MKSD
1: input Data D = (X; A, Yo) and Q toee such that that so is known for all in O.
2: output A distribution go, that approximates the potential outcome Y
3: Choose kernel I and estimators 3, #.
4: Split data in two sets D1 = (Xi, A,, 10 and D2 = (X,, A,, /2n///++11
5: Train #(1,, (00) on D2 and , on D1. Define ((r) from 3ir) as in equation
(00
6: Define de as 0p on 1-r for E (1, 2}, and In. as in equation
7: Take On E arg mintge In (0; YI.................... Yn).
8: return 9a..
statistic In. would lead to an increase in the variance of the latter. So, we use cross-fitting (Robins) et al. 2008 Zheng & van der Laan, 2010 Chernozhukov et al. 2018). This is, split the data in half, use the two folds to train different estimators separately, and then evaluate each estimator on the fold that was not used to train it.
Based on these theoretical considerations, we propose a novel estimator called DR-MKSD (Doubly Robust Minimum Kernel Stein Discrepancy) outlined in Algorithm [ Two key observations regard- ing Algorithm | are noteworthy. Although we have presented 9n as an abstract inner product in H., equation 7 has closed-form expression long as hg can be evaluated. Further details can be found in Appendix A Additionally, the estimator 0, is defined through a minimization problem, the complexity of which depends on Eg. general, In need not be convex with respect to 8, and thus, estimating O,, may involve typical non-convex optimization challenges. However, this is inherent to our approach, aiming for flexible modeling the counterfactual distribution. The potential Eg itself could be a neural network, making non-convexity unavoidable aspect of the problem.
We note that evaluating 9n has a time complexity of O(n²). Nonetheless, in stark contrast to the methods presented in Kim et al. (2018); Kennedy et al. (2021). the proposed DR-MKSD procedure only requires the nuisance estimators to be fitted once. This enables DR-MKSD to leverage com- putationally expensive estimators, such deep neural networks, providing significant advantage with respect to these previous contributions.
While it may difficult to find global minimizer of 9n. turn to study the properties of the estimator as if that optimization task posed no problem, with the scope that this analysis sheds light on the expected behaviour of the procedure if the minimization problem yields a good enough estimate of 0. We thus provide sufficient conditions for consistency and inference properties of the optimal On, We defer the proofs to Appendix C well as an exhaustive description of the notation used.
Theorem (Consistency) Assume that θ C IRP open, convex, and bounded. Further, let Condi- tions AA-A3 hold, as well as
A5) / sup hz 2z) < oc, 0:0
A4)
sup 17 ()dP()d)(() < co, 8EB
A6) /ʃ
xc, A7) / [0011 (z)==
sup oEe
sup 0 (z, =)(lp=dP(=)) < oc.
0ee
(i) (( E 1-,, = (ii) supace 1 - = or and (iil) [|((() suelll 113gr))
3ell = (1) [1, then
IID(Q&,11Qt) 3, min KSD(Q= dEB
Conditions A4-A7 build on the supremum of an abstract kernel hig but we note that the properties stem from those of k and sa. If sa (v),, dgs k(, y). and V V,k(,,)) y) are bounded, then so are hū and Öehg. and Conditions A4-A7 are fulfilled. In particular, if sy and are smooth and y is compact, such assumptions are attained. We highlight the weakness of Condition (i), which only requires that our estimates are bounded away from 0 and 1, and Condition (ii), which does not even require the estimates of Φy to be consistent, and is usually implied by Condition (iii).
Condition (iii) implicitly conveys the so-called double robustness. That is, as long the estimator of is consistent the estimators of Bo are uniformly consistent, then so will be the procedure. We
5

---

Published as a conference paper at ICLR 2024
highlight that not estimating 3p independently for each 8, but are rather constructing all of them based on an estimate of 3. This opens the door to having immediate uniform consistency as long the estimator of / is consistent (depending the underlying structure between ( and 3g) in a similar spirit to the work on uniform consistency rates presented in Hardle et al. (1988).
The consistency properties of doubly robust estimator important characteristic is the convergence rate, which we show next corresponds to the product of the convergence rates of the two estimators upon which it is constructed.
highly desirable. However, most
Theorem 2 (Asymptotic normality) Ler O C R' be open, convex. and bounded, and assume On 4 0.. where KSD ||Q') = min KSD(Qe
Suppose Conditions Al-A5 hold, as well as A8) the maps {8 - 0g غ E ] are differentiable on θ A9) the maps {6 + 0ghā(2, 8):
2, E 2) are uniformly continuous at f,, A10) I[] supage |0ghī 8) 3)
< 00, All) IJ (10gh 2). (() l0-0. < 00, A12) ʃ [uupcee 1 2)) 2) < oc, A13) 00, 00, A14) Jf 0gh5(2, ()dP(z)dP(3) 10=0,> O.
If supece [0gh,(2, d)
(f') P(((r) (c, d)- = (ii) supage {1 Φoll + 1100 ) 0o0ol or(1), and
(ii')) [| -| sup { {1.36² - Boll + 110 - 00.3ol| } = or(
1
vn
8
for {1, then
Vn(O, - 0*) S (0
where Σ = Co [3) l0=0.-
As proposed in Oates et al. 2022). a sandwich estimator I for the asymptotic variance
4r -T- can be established (Freedman 2006), where En and I, are obtained by substituting 0. by On and replacing the expectations by empirical averages. If the estimator 4I - - is consistent at any rate, then //" (0. 0,) & 10, in view of Slutsky's 's theorem.
Conditions A8-A13 can be satisfied under diverse regularity assumptions, analogously Conditions A4-A7. Further, we note that assumptions akin to Condition A14 are commonly encountered in the context of asymptotic normality. While Condition (i') is carried over the consistency analysis, Condition (ii') now requires the estimators of do and doog to be uniformly consistent. However, we highlight the weakness of this assumption. For instance, it is attained as long as म is consistent and 3g uniformly bounded. Lastly, we again put the spotlight on Condition (iii'): if the product of the rates is or-nn-122) asymptotic normality of the estimate 0n can be established.
The double robustness of In has two profound implications in the DR-MKSD procedure. First, gn (0) converges to KSDD(QulIQ"" faster than if solely relying on nuisance estimators a or Bo. which translates to a better estimate 0g- Further, it opens the door to conducting inference the product of the rates is or (n-1/3) which may be achieved by a rich family of estimators.
We underscore that On L, 0, need not be a consequence of Theorem 1 Theorem establishes consistency on KSD(Qs,, I|Q*), not O,, itself. Consistent estimators for KSD(Qe. " but not for On may arise, for example, if the minimizer is not unique. In such a case, one cannot theoretically derive the consistency of the estimate (at least, while remaining agnostic about the optimization solver). For instance, if there are two minimizers and the optimization solver alternates between the two of them. KSD(Qe, IIQ') is consistent but 6₂ does not even converge. Nonetheless, we emphasize that the theorems provide sufficient, but not necessary, conditions for consistency and asymptotic normality.
Furthermore, a potential lack of consistency does not necessarily translate to poor performance of the proposed estimator. Two distributions characterized by very different parameters 0 may be close in the space of distributions defined by the KSD metric (analogously to two neural networks with very different weights having similar empirical performance). Estimating the distribution char- acterized via the parameter, not the parameter itself, is the ultimate goal of this work. Hence, the proposed estimator can yield good models for the counterfactual distribution even with inconsistent O..

---

Published as a conference paper at ICLR 2024
Lastly, we highlight that the DR-MKSD procedure can be easily redefined by estimating So, for each 8g belonging to a finite grid g8g] The uniform consistency assumptions would conse- quently translate on usual consistency for each 0g of the set { geg) and the minimization problem would reduce to take the arg min of a finite set. The statistical and computational trade-off is clear: estimating Bo, for each 0g may lead to stronger theoretical guarantees and improved performance, but the computational cost of estimating Se, increases by factor of |G].
EXPERIMENTS
We provide a number of experiments with (semi)synthetic data. Given that this is the first estimator to handle unnormalized densities in the counterfactual setting, we found no natural benchmark to compare the empirical performance of the proposed estimator with. Hence, we focus on illustrating the theoretical properties and exploring the empirical performance of the estimator in variety of settings. Throughout, we take k(z,y) (c² + |-3 || y y to be the inverse multi-quadratic
(IMQ) kernel, based on the discussion Gorham Mackey (2017), with 8 0.5, 0.1 and 1. We estimate the minimizer of gr. (6) by gradient descent. We defer an exhaustive description of all simulations and further experiments Appendix
Consistency of the DR-MKSD estimator: To elucidate the double robustness of DR-MKSD, we define its inverse probability weighting and plug-in versions, given by ΦIr,,0(2) ((x) So y) and (=) = (x( respectively. We draw (X,, A., Y:) such that Y NN 1). Further, we let X ~~ NYY, 11) and we sample A using logistic model that depends on X. Note that this sampling procedure consistent with Conditions A1-A3. We consider the set of normal distributions with unit variance Q (W(0, 1)R T
the mean squared error of the procedures across 100 bootstrap samples, different sample sizes n and various choices of estimators. The procedure shows consistency long as either the IPW or PI versions consistent, even if the other is not. We defer to Appendix B an analysis of the confidence intervals for On given by Theorem 2 where illustrate that they have the desired empirical coverage.
Asymptotic normality of the DR-MKSD estimator: Following the examples in Liu et al. (2019) and Matsubara al. (2022), consider the intractable model with potential function Bo(y) = (V))) where E RR² V E RR,,, n7()) == (-0.5, 0.6, 0.2, 0, 0, 0)", and (v) (E v². 31322. Eiss Vivi, tanh(u))*. The normalizing constant is not tractable except for 0 = O, where we recover the density of a Gaussian distribution No. We draw (X;, Ai, Y, Y)),0n so that Y No and X .N(Y1, 4s). Lastly, A follows logistic model that depends on X X. Figure 2 displays the empirical estimates of 0 = (01, 02), which shows approximately normal.
Counterfactual Restricted Boltzmann Machines (RBM): We let (X,, Y2) such that y is drawn from -vo-dimensional RBM with one hidden variable h such that go hh) x exp(h + (0, y ) n2a 2|1y' (12) Additionally. X N(YY 0.25Iz) and follows logistic model that depends
(A)
(B)
(C)
0.07
DR-MKSD
DR-MKSD
DR-MKSD
0.06
IPW-MKSD
0.10
PW-MKSD
0.04
PWWMMKDD
PL-KKDD
P--KSDD
PI-KSDD
0.05
0.08
0.04
0.06
0.03
0.02
0.04
0.02
0.01
0.02
0.00
0.00
200
600
200
Figure 1: Mean squared error (MSE) of 6, approximated with 100 bootstrap samples. π and are estimated via (A) boosting and conditional mean embeddings (CME), (B) logistic regression and I-NN, (C) logistic regression and CME. The doubly robust approach proves consistent in all cases.
7

---

Published as a conference paper at ICLR 2024
(A)
(8)
(C)
0.15
0.4
DR-MKSD
Standard normal pdf
0.10
Ordered Values
0.05
0.00
0.2
-0.05
-0.10
-0.15
0.1 0.0 66
-0.1
0.0
-2
0
2
6,,1
Theoretical quantiles
Figure 2: Illustration 100 simulations of the DR-MKSD estimator 0, = (0,,1 2n,2), constructed
on estimates of and 3 fitted random forests and CME for 500, with (A) scatter plot of empirical distribution for 0n, (B) histogram of the standardized values of 8,,1, (C) Normal QQ plot of the standardized values of 0 n, We highlight the two-dimensional Gaussian behaviour of 0,,
on X. Although estimating the exact 0 is hopeless (h is unobservable), it may still be possible to uncover the underlying structure governing the behavior of Y which depends on the ratio between shows that minimizing 9n. does indeed recover the direction of 6
81 and 02. Figure 3
Experimental data: Suppose one has access to a noisy version of Y denoted as X, and the process of denoising this data to recover Y is costly. Due to budget constraints, one can only afford to denoise approximately half the data. The objective is to estimate the distribution of y E Rdd which can be described by model qa (Y') ε cxp((T((Y)) 0)). Here, T(Y!) E R2 represents a lower-dimensional representation Y To achieve this, one randomly selects observations for denoising with a probability of 1/2 and employs DR-MKSD to model y
For illustration we consider ten distinct counterfactuals, denoted as y E R220 for i E (1...., 10)- These counterfactuals are defined y1,, f(W'), where each W represents observation from the MNIST dataset corresponding to digit i, and f is a pretrained neural network. Additionally, we utilize another pretrained neural network 7' _ 1810, In Figure we display random MÑIST dataset observations in the top row, alongside those with the lowest unnormalized density of fWW) in the bottom row, estimated by DR-MKSD independently for each digit. Notably, the latter images appear more distorted, which aligns with the expectation that DR-MKSD accurately models Y
(A)
(B)
(C)
5
5
4-
4
3-
3
2-
2
-1
1
0
0
1
1
1
2
2
33
4
01
0
Figure 3: Values of the statistic 9n, (f) for the RBM with (A) 0 = 0 = 0, (B) 0
0 = 1.
-02 = Minimizing gn leads to recover the direction
8

---

Published as a conference paper at ICLR 2024
7
2
3
4
5
8
9
2
3
4
5
<
7
8
9
Figure Observations of the MNIST dataset chosen randomly (top row), and with the lowest estimated density of f(W) (bottom row). We highlight the potential outlier detection in digits 0, 2, and 6.
DISCUSSION
We have presented an estimator for modeling counterfactual distributions given flexible set of distributions, which only need to be known up to normalizing constants. The procedure builds on minimizing the kernel Stein discrepancy between such set and the counterfactual, while simul- taneously accounting for sampling bias in a doubly robust manner. We have provided sufficient conditions for the consistency and asymptotic normality of the estimator, and we have illustrated its performance in various scenarios, showing the empirical validity of the procedure.
There several avenues for future research. Employing energy-based models for estimating the counterfactual enables the generation of synthetic observations using a rich representation of the potential outcome. Exploring the empirical performance of sampling methods that do not require the normalizing constants, such as Hamiltonian Monte Carlo or the Metropolis-Hasting algorithm, holds particular promise in domains where collecting additional real-world data is challenging.
Furthermore, extensions could involve incorporating instrumental variables and conditional effects. Of particular interest would be the expansion our framework accommodate time-vate time-nayying treat- ments. Lastly, we highlight that, while we have framed the problem within a causal inference con- text, analogous scenarios arise in off off-policy evaluation (Dudík et al. [001]] Thomas & Brunskill 2016). Extending our contributions to that domain could yield intriguing insights.
REPRODUCIBILITY STATEMENT
Reproducible code for all experiments provided in the supplementary materials.
ACKNOWLEDGMENTS
DMT gratefully acknowledges that the project that gave rise to these results received the sup- port of fellowship from 'la Caixa' Foundation (ID 100010434). The fellowship code is LCF/BQ/EU22/11190077,. EK was supported by NSF CAREER Award 2047444.
REFERENCES
Alberto Abadie. Bootstrap tests for distributional treatment effects in instrumental variable models. Journal of the American statistical Association, 97(457):284-29) 2002.
David H Ackley, Geoffrey E Hinton, and Terrence J Sejnowski. A learning algorithm for boltzmann machines. Cognitive science, 91):1177-169 1985.
Andreas Anastasiou, Alessandro Barp, François- Briol, Bruno Ebner, Robert Gaunt, Fate- meh Ghaderinezhad, Jackson Gorham, Arthur Gretton, Christophe Ley, Qiang Liu, et al. Stein's method meets computational statistics: A review of some recent developments. Statistical Sci- ence, 1801)9120-139 2023.
Alessandro Barp, Francois- Xavier Briol, Andrew Duncan. Mark Girolami, and Lester Mackey. Min- imum stein discrepancy estimators. Advances Neural Information Processing Systems, 32, 2019.

---

Published as a conference paper at ICLR 2024
Domagoj Čevid, Loris Michel, Jeffrey Näf, Peter Bühlmann, and Nicolai Meinshausen. Distribu- tional random forests: Heterogeneity adjustment and multivariate distributional regression. The Journal of Machine Learning Research, 23(1):14987-1500 2022.
Victor Chernozhukov and Christian Hansen. An iv model quantile treatment effects. Economet- rica, 73(1):245-261, 2005.
Victor Chernozhukov, Iván Fernández- Val, and Blaise Melly. Inference on counterfactual distribu- tions. Econometrica, 81(6):2205-235 2268, 2013.
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. The Econometrics Journal, 21(1)CC1-C6, 2018.
Kacper Chwialkowski, Heiko Strathmann, and Arthur Gretton. A kernel test of goodness of fit. In International Conference Machine Learning, pp. 2606-2615. PMLR, 2016.
Peter Clifford. Markov random fields in statistics. Disorder in physical systems: A volume in honour of John M. Hammersley, pp. 19-32, 1990.
Iván Díaz. Efficient estimation of quantiles in missing data models. Journal of Statistical Planning and Inference, 190:39 51, 2017.
Miroslav Dudik, John Langford, and Lihong Li. Doubly robust policy evaluation and learning. In Proceedings of the 28th International Conference on International Conference on Machine Learning. Omnipress, 2011.
Jake Fawkes, Robert Hu, Robin Evans, and Dino Sejdinovic. Doubly robust kernel statis- tics for testing distributional treatment effects even under one sided overlap. arXiv preprint arXiv:2212.04922, 2022.
Matthew Fisher, Tui Nolan, Matthew Graham, Dennis Prangle, and Chris Oates. Measure transport with kernel stein discrepancy. In International Conference on Artificial Intelligence and Statistics, pp. 1054-1062. PMLR, 2021.
David A Freedman. On the so-called "huber sandwich estimator" and "robust standard errors". The American Statistician, 60(4):299-300 2006.
Stuart Geman and Donald Geman Stochastic relaxation, gibbs distributions, and the bayesian restoration of images. IEEE Transactions on Pattern Analysis and Machine Intelligence, PAMI-6 (6):721-741, 1984. doi: 10. 199.///TAM..9.9.777777533
Jackson Gorham and Lester Mackey, Mcasuring sample quality with stein's method. Advances in neural information processing systems, 28, 2015.
Jackson Gorham and Lester Mackey. Measuring sample quality with kernels. In International Conference on Machine Learning, pp. 1292-1301. PMLR, 2017.
Steffen Grünewälder, Guy Lever, Luca Baldassarre, Sam Patterson, Arthur Gretton, and Massimi-
lano Pontil. Conditional mean embeddings regressors. In Proceedings the 29th International Coference on International Conference on Machine Learning, pp. 1803-1810, 2012.
W Hardle, Paul Janssen, and Robert Sertling. Strong uniform consistency rates for estimators of conditional functionals. The Annals of Statistics, pp. 1428-1449, 1988
Geoffrey Hinton. A practical guide to training restricted boltzmann machines. Momentum, 9(1): 926, 2010.
Geoffrey Hinton, Simon Osindero, Max Welling, and Yee-Whye Tch. Unsupervised discovery of nonlinear structure using contrastive backpropagation. Cognitive science. 30(4):725-731, 2006.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 17(8)1177--1800, 2002.
10

---

Published as a conference paper at ICLR 2024
Aapo Hyvärinen and Peter Dayan. Estimation of non-noomalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4). 2005.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sci- ences. Cambridge University Press, 2015.
Edward H Kennedy, Sivaraman Balakrishnan, and Larry Wasserman. Semiparametric counterfcctll density estimation. arXiv preprint arXiv:2102. 12034, 2021.
Kwangho Kim, Jisu Kim, and Edward H Kennedy. Causal effects based on distributional distances.
arXiv preprint arXiv: 08630033,,, 2018.
Anna Korba, Pierre-Cyril Aubin-Frankowkki. Szymon Majewski, and Pierre Ablin. Kernel stein discrepancy descent. In International Conference on Machine Learning, pp. 5719-5730. PMLR, 2021.
Henry Lam and Haofeng Zhang. Doubly robust stein-kernelized monte carlo estimator: Simulta- neous biri-var reduction and supercanonical convergence. Journal of Machine Learning Research, 24(85):1-58, 2023.
Yann LeCun, Sumit Chopra, Raia Hadsell, Mare' Aurelio Ranzato, and Fu-Jie Huang. Predicting structured data, chapter tutorial on mnegy--based learning, 2006.
Qiang Liu, Jason Lee, and Michael Jordan. A kemnelized stein discrepancy for goodness-of-fit fit tests. In International Conference on Machine Learning, pp. 276-284. PMIR, 2016.
Song Liu, Takafumi Kanamori, Wittawat Jitkrittum,, and Chen. Fisher efficient inference of intractable models. Advances in Neural Information Processing Systems, 32, 2019.
Dean Lusher, Johan Koskinen, and Garry Robins. Exponential random graph models for social networks: Theory, methods, and applications. Cambridge University Press, 2013.
Diego Martinez- Taboada, Aaditya Ramdas, and Edward H Kennedy, An efficient doubly-robust test for the kernel treatment effect. arXiv preprint arXiv:2304. 13237, 2023.
Takuo Matsubara, Jeremias Knoblauch, François-Xavier Briol, and Chris J Oates. Robust gen- eralised bayesian inference for intractable likelihoods. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(3):977-1022, 2022.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Normalizing flows for interventional density estimation. In International Conference on Machine Learning, pp. 24361-24397. PMLR, 2023.
Andriy Mnih and Geoffrey Hinton. Learning nonlinear constraints with contrastive backpropagation. In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pp. 1307 IEEE, 2005.
Krikamol Muandet, Motonobu Kanagawa, Sorawit Saengkyongam, and Sanparith Marukatat. Coun- terfactual mean embeddings. The Journal of Machine Learning Research, 22(1):732--739, 2021.
Jeffrey Năf, Corinne Emmenegger, Peter Bühhlmnn, and Nicolai Meinshausen. Confidence and uncertainty assessment for distributional random forests. arXiv preprint arXin:2330.05561122222
Chris Oates et al. Minimum kernel discrepancy estimators. arXiv preprint 1rXiv:2210.163:: 1635,, 2022.
Chris J Oates, Mark Girolami, and Nicolas Chopin. Control functionals for monte carlo integration. Journal of the Royal Statistical Society Series B: Statistical Methodology, 79(3):695-718, 2017.
Junhyung Park, Uri Shalit, Bernhard Schölkopf, and Krikamol Muandet. Conditional distributional treatment effect with kernel conditional mean embeddings and U-statistic regression. In Interna- tional Conference on Machine Learning, pp. 8401-8412, PMLR, 2021.
Gany Robins, Pip Pattison, Yuval Kalish, and Dean Lusher. An introduction to exponential random graph (p*) models for social networks. Social networks, 19 2007.
11

---

Published as a conference paper at ICLR 2024
James Robins and Andrea Rotnitzky. Inference for semiparametrie models: Some questions and an answer comments, Statistica Sinica, 111:90093366 10 2001.
James Robins, Lingling Li, Eric Tchetgen, and Aad van der Vaart. Higher order influence func- tions and minimax estimation of nonlinear functionals. In Institute of Mathematical Statistics Collections, pp. 335-421. Institute of Mathematical Statistics, 2008.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-333 2005.
Rahul Singh, Mancesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. Ad- vances Neural Information Processing Systems, 32, 2019.
Rahul Singh, Liyuan Xu, and Arthur Gretton. Kernel methods for causal functions: Dose, heteroge- neous, and incremental response curves, arXiv preprint UrXiv:2010.048355 2020.
Rahul Singh, Liyuan Xu, and Arthur Gretton. Kernel methods for multistage causal inference: Mediation analysis and dynamic treatment effects. arXiv preprint arXiv:21 2021.
Le Song, Jonathan Huang, Alex Smola, and Kenji Fukumizu. Hilbert space embeddings of condi- tional distributions with applications dynamical systems. In Proceedings of the 26th Anmual International Conference on Machine Learning, pp. 961-968, 2009.
Charles Stein. A bound for the error in the normal approximation to the distribution of sum of dependent random variables. In Proceedings of the Sixth Berkeley Symposium on Mathematical Statistics and Probability, Volume 2: Probability Theory, volume 6, pp. 583-603. University of California Press, 1972.
Philip Thomas and Emma Brunskill. Data-efficient off-policy policy evaluation for reinforcement learning. In International Conference Machine Learning, pp. 2139-2148. PMLR, 2016.
Ted Westling and Marco Carone. A unified study of nonparametric inference for monotone func- tions. Annals of statistics, 4802):1001, 2020.
Wenjing Zheng and Mark J van der Laan. Asymptotic theory for cross-validated targeted maximum likelihood estimation. U.C. Berkeley Division of Biostatisties Working Paper Series, 2010.
Yuchen Zhu, Limor Gultchin, Arthur Gretton, Matt Kusner, and Ricardo Silva. Causal inference with treatment measurement error: nonparametric instrumental variable approach. In Uncer- tuinty in Artificial Intelligence, pp. 2414-2424. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2024
A CLOSED FORM OF THE STATISTIC
Under estimators of the form Bo(x) = >i=1 (i(())((( Yi), we derive that both inner products
(Bo()), H¹
(1 x)), Σ (1)(())(()))))
(x())) (8) (Sol: Eol, (b)),) R
im
Yz),
( (Bolm). Eo(, (2)) Ha
d ()( (o),
(x) (Eo Yo), Eo(, Eo), Y) 20
Σ (h))
count with closed forms as long as he can be computed. Consequently
1
1 Σ 6o(2,)
9n (0: ..)). =
60
(Z:),
1=1
34a
72 A
{so(h -BB(())) + (X,),,
+(+))
Ai
{so(. M3) - Ba(x3)) + (oo(x))) H1
*(X,)
Ai
1 :((1- *(X,)
Ai
} (Bo(x)).
Eo(,Y))
*(X,)
Aj
+ {1- *(X,)
} /olxx))
*(X,)
A:A,
(So(, (o. X3) )
+
*(x))((X,)
32
2
4
{1
{1- *(X,)
}
Aj
( (3(())).B0((3)) )+
#(X,)
A
}
Ag ) (elx2)..6(xx))
#(x,)
A,A
(x)) ho((i, +
11 1
*(,,)+({,))
2
A
{1- *(X,))
Aj
} Niv (X,)ho) Y2) +
+
n2
#X,,)
[1-
A
}
{1- *(X,)
Aj
} N (X,
+
2
#(X.)
is closed form as long as ho, ât and w can be evaluated.
13

---

Published as a conference paper at ICLR 2024
B EXPERIMENTS
B.1 CONSISTENCY OF THE DR-MKSD ESTIMATOR
We consider the set of normal distributions with unit variance Q {N(0, 1)JBER. We draw (X, A,, Y=)== for n E {200,
800} such that Y1 [N 1). Further, X ~ , and A follows Bernoulli distribution with log odds X.
For obtaining Figure
we estimate π with
Default LogisticRegression from the scikit-learn package with C le5 and max iter 1000,
Default AdusovxiClsssiter from the scikit-learn package,
and B with
Conditional Mean Embeddings (CME) with the radial basis function kernel and regulariza- tion parameter le- 3,
Vector-valued One Nearest Neighbor (1-NN).
Parameter 0 estimated gradient descent for number of 1000 steps. total number of 100 experiments with different random seeds were run in order to yield Figure T
We further analyze the empirical coverage of the estimated 0.95-confidence intervals 0g + n>
where A, 1.96 Er, and E,, and r, are defined described in Section Figure 5 exhibits the confidence intervals for n 200 and n = 300. We highlight that the respective empirical coverage for those sample sizes 0.94 and 0.96 respectively, hence presenting desired empirical coverage.
(A)
100
Randoom
(B)
-02
-0.0
Random seed
Figure 5: Estimated 0.95-confidence intervals 0,, = An, where A, = 1.96 * V 4r I /n for (A) 200, and (B) 300. The empirical coverages of the confidence intervals are (A) 0.94,
and (B) 0.96.
14

---

Published as a conference paper at ICLR 2024
B.2 ASYMPTOTIC NORMALITY OF THE DR-MKSS ESTIMATOR
We consider the intractable model qo (o()) (xp(-Fo(w) with potential function -Fa (n(0), (y))zz, where
(0)) (-0.5, 1-0.0,066,00200000000 OER?,
5
Jly Vi. V192, Σ Div tamh(y)) VERS.
The normalizing constant not tractable except for 0 = O, where we recover the density Gaussian distribution (0, E), with
0.6
0.2
0.2
0.2
0.6
0
Σ-1
0.2
0
0.2
0
0
1
0
0.2
We draw (X,, A,, 1 so that Y ~ N(0, and X ~ NNY! 1s). Lastly, A follows a Bernoulli distribution with log odds Σ (X} - 1).
For obtaining Figure we estimate x with the default RandomForest FrreddmeorexxIlassfffer from the seikit-learn package, and 3 with Conditional Mean Embeddings (CMH) with the radial basis function kernel and regularization parameter le - 3. Parameter was estimated by gradient descent for a number of 1000 steps. A total number of 100 experiments with different random seeds were run.
B.3 COUNTERFACTUAL RESTRICTED BOLTZMANN MACHINES (RBM)
We let (X,, A,, Y,)500 such that y drawn from two-dimensional RBM with one hidden variable such that go(y² h) exp(h + (0, y )²))== 2 [12) In order to draw from such RBMs, we make use of Gibbs sampling and burn the first 1000 sam- ples. Additionally. X
for (0, 0). 0 (1, 1), and (1, -1).
(02552) and follows Bernoulli distribution with log odds
DE - 0.5).
We estimate with the default LogisticRegression from the scikit-learn package with C le5 and max iter 1000, and 3 with Conditional Mean Embeddings (CME) with the radial basis func- tion kernel and regularization parameter 1e 3. Figure exhibits the values of In
grid Figure 6 exhibits the values of 9n for further pairs (0,, 02), illustrating the behaviour of gn with different magnitudes of 0, and 02.
{(11,0) 01, 02 E (-5, 449, -4.8 ..,8....3 5}}. 1.8....5)),
(A)
(B)
(C)
5
-5
4
4
3
3
15.0
2
2
12.5
0
7.5
4
25
4 3
5
4
8
Figure 6: Values of the statistic 9n (0) for the RBM with (A) 0 -0.1,02 5, (B) 0 -1,02
0.1, (C) A = 1, fg = 0.1.
15

---

Published as a conference paper at ICLR 2024
B.4 EXPERIMENTAL DATA
We start by training a dense neural network with layers of size [784, 100, 20, 10] on the MNIST train dataset. For this, minimize the log entropy loss 80% of such train data and parameters that minimize the log entropy loss for the remaining validation data (remaining 20% of the MNIST train dataset).
the
We then evaluate this trained neural network on the MNIST test dataset. For each digit, we take the 20-dimensional layer to be y The covariates X defined Y + G, where NN(0, /20). Treatment A follows a Bernoulli distribution with probability 0.5.
We consider the intractable model qo (Y() x exp (op((T(Y') (. 0)), where the lower dimensional repre- sentation T(Y') is taken as the evaluation of the 10-dimensional layer of the neural network. Note that 7 therefore differentiable, so we can apply the DR-MKSL estimator.
Out of the whole MNIST test dataset, we take the first 500 observations to train the DR-MKSD estimator (i.e., find 8,,)) with (X, A, Y2) We define
0.5 and estimate with Conditional Mean Embeddings (CME) with the radial basis function kernel and regularization parameter le 3.3. 3.
In order to obtain Figure we take the minimizers of the estimated unnormalized density of y (the evaluation of the 20-dimensional layer) for the remaining observations of the MNIST test dataset. The top row images are taken randomly for the same subset of the MNIST test dataset.
16

---

Published as a conference paper at ICLR 2024
C PROOFS
We now present the proofs of the theorems stated in the main body of the paper. For this, we start by introducing the notation that will be used throughout. The proofs of Theorem I [and Theorem subsequently follow.
C.I NOTATION
Given Hilbert space K. and K-valued function ((z) E K., we denote its norm in the Hilbert space by [|o(z)] ||k. Furthermore, || I [|x())] (z) denotes the squared L2(Q) norm of the K- valued function norm. We highlight that the expectation is only taken with respect to the randomness of Z, while ò is considered to be fixed. Note that in the case K. = I, || denotes the usual L2 (Q) norm real-valued function.
We let P. P. (o(2)) = E, denote the sample average and P(ő) = P[)(2)} = (2) d=()) the expected value of (Bochner integral) with respect to Z, treating & as fixed. If function & takes two arguments (this the case when k is a kernel), then we denote P2 1 n² Σi, k(Y, Yj). Further, Tr denotes the trace of matrix.
Lastly, we make use of standard big-oh and little-oh notation, where X, = Op(rn) implies that the ratio Xn/r,, is bounded in probability, and X,, = (p(rn) indicates that Xn converges in probability to O. Throughout, we make use of 'calculus' with this stochastic order notation, such as (1) + (1() = (1()) and (1)00)))) -1)00pp1pppp))) (1()).
C.2 PROOF OF THEOREM 1
We prove the theorem in three steps:
1. First, show that supage PP Proollxr 2 0.
2. We then prove that supace 19, (0) KSD(Q= 3
3. We finally arrive KSD(Qe, IQ')) Fs minsee KSD(Q=lIQ))
We now need to demonstrate the validity of each of the steps.
Details of step 1. For simplicity, let us assume that n is even and hence n/ /2 = n//2. We note that
Pn - Pn =
(0o(Z)) - oo(2,)}
n/2
Σ
(oo(Z)) oo(2.)) +
Σ (ōo(Z.) -o(())))(Z2))
1=1
t=n/2
a/2
m
(o())) - -60
Σ
(do(Z)) - so (2i)
n/2
2
n/2
imn/2
n/2
1
2
(1)
(Za)
2o (zi)
2
Σ
((2)
(Z,) ((2)(())))
oo
99
90
2
11
i=1
i=n/2
())
(II)
(1)
Let us first work with term (I), and note that it may be rewritten as Pa Pa/200 P Pa/200. Additionally,
Pa/200 process term and T2( Q(2(1)) Φb)) is the bias term. Given that 2(1) was trained independently
201) Pu/200 Ti (0) (2))), where (1) (Pn/2 (Ø
(1) Φ0) is the empirical
(1)
17

---

Published as a conference paper at ICLR 2024
from {Z, /2, ..... 2n following Martinez-Taboadk et al. (2023 Lemma C.7) and the arguments in Martinez- Taboada et al. (2023 Theorem C.9), we deduce that
((1) 6oll [T2 S -1| 111.3 Bg 1|-
[Ti())) Op
V n/2
Given Conditions (i), we deduce that
sup ||Z' (0)ll + |122(0) pppe I sup eg
[[Ti
+ sup [12(
8E8
supace
10g
Φoll
1 || + 111 1/
Op
sup aE0
vn/2
=r((1) (1(
= ol(())
Hence, infer that supage || 1P supsce |I(I1)l op()). We thus conclude
2(1) Pn/ Pn/200llze 0p(1). Analogously, we deduce that
sup Pr PRØelAx = sup |1 (I) + (1
Age
2
sup
2 (1) + sup
|| (IT)
AE0
AE0
0P( + OP
= [1()
Details of step 2. Denoting (0) Paδo Proo, have that
1
1
(0; Y,
Vn) =
dal (7:)
dol (Zi)
i=1
1=1
Pnoo Pno0)) 7d
= (Pado + (f), Pnoo + Xm (f))
(Pado; Pnoo)n + 2 (Pn00, (n())) (xn(0), Xa (0)) z())
We now work with these three terms separately. First, note that assumption A5 implies that, for all 0 in θ,
f (g(2,2)dP(2) < suph)(z,)))P(3) < <0,
Me0
Hence, by Jensen's inequality,
/ (ngters)dP(E) < / nā(e)))dP(3) < v0 e.
Given that we have also assumed Conditions A6-A7, we can apply Oates et al. 2002 establish that
Lemma 11)
sup (Pago, Paoein - = 6EB
Second, we have that | = (h < (? hih. hence || 1| < P2 (supece hz). Based on the law of large numbers for V-statistics and Conditions A4-A5, we deduce that p2 (supace 1i) 2, ff supace 12(2, ())P()))P()) < co, so
sup || Pr 2 (r(1).
(8)
#20
18

---

Published as a conference paper at ICLR 2024
Further, remind that have proven supaee llxx(0)llx= 4 0 on Step 1. Consequently,
(0)
sup lee
(Proo, Xx(()) sea
S
sup
9c8
(sup Aee
sup || 8E8
1lxx
(p(1)op(1)
op(1),
where (i) is obtained by Cauchy-Schwart inequality. Lastly, we highlight that
sup
llx, (0) = o = sup (0)] = op()).
A
It suffices to note that supoce (0.(0) KSD(QO, |Q\)) may be rewritten
sup (Pasho, Pnoo a (o Xn (n) (Xn((), Xa))) KSD(Qo)1Q*),
6EB
which is upper bounded by
sup 8
(Pado; Paoo) KSD(Qol(=*)+ 2
(Pa0o: Xn (n) sup (x,(0), (0))
which is (1)),, concluding
sup | (0) =SD(Q lE
Details of step 3. In order conclude, we extend the argument presented in Oates al. (2022) Lemma 7) to convergence in probability. Take H" arg min KSD(Qa Given that supace (0) KSD(QalQ*)) = (1),), for any and >> 0, there exists n such that 1g.. KSD(Q, (()) < with probability at least 6 for all n 2 n'. Consequently,
KSD(Q I|Q") - KSD(Qe., (1 + 9. (0,) S (0' S KSD(QZ)1Q") +
for n. 2 n² with probability 1 - 5. This is, (0.) - - < for 2 n² with probability
- 8, concluding that gn gn.(0,) (0..) 3, KSD(Q=-11Q")
C.3 PROOF OF THEOREM 2
We prove the theorem in three steps:
1. First, prove supace Plndo P. Proellze =0 op (tx) and
supace 1h (náo Pnoso) ll op (t)))
2. Second, we show that < (Pnog 3 Proe, ( D An/2, with || = op (1)
3. We then conclude that vo((0. 0-) 4 N(O, -1----
We now need to demonstrate the validity of each of the steps.
Details of step 1. With analogous argument used in Proof
C.2 (step 1), yield
n/2
P3 Pr.0 =
2
(1) P0
(Zi) - (Z))
+
1
2
(
90
(2,) - 80 (z))
2
1=1
i=n/
(1)
(II)
19

---

Published as a conference paper at ICLR 2024
where
supace 1100 Toll 2
sup (l = Op 8EG
sup 8EB
|| # 111.30 -88
v Vn/2
c
supnce 1(00) dell
|| - #12) 11 30 188113
sup [I1) = =p 8E0
+ sup 8
Vn/2
Assumptions (ii') and (iii") imply
sup (1)) or 111 +op 8E
(vn22)
(12))
Vn
sup
[ITT Hd =
(nn22 (
+ or
1 vn
so
1
sup AE
PP Paooll op vn/22
+
VT
vn
Similarly, based on Assumptions (ii') and (iii',, we obtain supace
Pros
Pactoo
llengg
OP
Details of step 2. Given that On is the minimizer of differentiable function Ja, have that (0.) 0. Denoting (0) P. Pa Po. have that
0 (4mm)
00 g(m))
①
(Prio,, + Xn (n((n); Pudu, + Xn (0(On)
00
①
(Prdo. xm(0n)) Pro.. + xx(01))
00
Hd
θ
+
Paoe, + Xx(()).
(Pa00 + xx((m)))
00
30 (
+
(xm(0n)) Xx (0n))
20
Hd
Hd
+ Patoo. 00 (xa(0a)))
(Pr ; Xx (0n) +
00
Hd
Let us upper bound the last three terms of the latter addition. First, we note that
00 (20 (x.(000)) xxn(xxn(())
00 (xx (x)
sup dEe
zd
Hd 1e
sup 860
(xn (0)}
(1<80990 |1xx (O))ee)))))))))
20
sup
(x))
sup fllxa (0) aee
00
e
1
1
op
0.p
n
0p
(10)
20

---

Published as a conference paper at ICLR 2024
where (i) is obtained by Cauchy-Schin Sechaa- inequality. Second,
① Keado.: 20 (xn(Ow)))
8 Pnov, (xn(01)
S sup PE
R
00
(1)
||Prooll 00 (xx(0) (0)) 11ozegp
sup 858
30 (xn( lleweyy } AEE
< sup AE
{l| 1h00llzz0
|1
sup
0p(1)op
(11)
where (i) is obtained based on Cauchy-Schwand inequality, and (ii) is derived as in Equation equa- tion g given Conditions A4-A5.
Third, we have that
θ (Prio Ilireg I sup
||
{Pa
do
li
08
AE0
00
sup AE
Pr
n
00
sup AE
Pr
00p
82
sup AE
p2
002 k
h1
P
82
sup p2 OEG
h
003
82
< p2
sup
002 h
0Ee
k=1
82
p2
sup Tr
ag2 h
AEE
Note that Tr [i is dominated by the L norm (sum of the absolute values of the entries) the 0² Hessian 20a hi- Further, the I norm is equivalent to the T.2 norm || llexx (all norms equivalent
in finite dimensional Banach spaces). Hence, based on the law of large numbers for V-statistics and Conditions A12-A13, we deduce that
82
B. //
y2 h hø 2) d) co.
p2 sup lEe
002 hi
BoxD }
sup ae
302
Consequently,
2
sup
||
(Pado.) 1ir = Or(1),
Ae 00
21

---

Published as a conference paper at ICLR 2024
and hence
|1
0 (Prdo) xm xx((1)
20 (Prdon) xalOn))
sup AE
00
Ha
20a Re
[Pr oo) 11(7e000 (x))
8Ee
00
0
sup
(Pr00)
sup {llxn (0)l MEB
00
#ce
Op(1)op
1
or
Vn
(12)
where (i) is obtained based Cauchy-Schwand inequality.
Combining Equation equation with the upper bounds from Equations equation TO equation and equation [2 derive that
11
2
(Pn (Proon) J Puodo.
A...
30
with ||nn |R+ - 0p
+ Up
(ta)) + Up (t) = 0p (#1z
Details of step 3. Further, denoting 9. (0) == (Proo, Passo) Ha E IR, we arrive to
( 00 (Prdo..: Prov.
2
Ha 00 (P.0,,, Pa Paoun)
a0 In (0n).
Based on the mean value theorem for convex open sets, there exists O, bn + (1 Ln) with tn [0, 1] for which
8
8²
00 9n (0.)
In (0.) (0.. -0)) (o.) An.
Oge
8² If 50 U (0m) is non-singular, then rewriting this expression obtain
Vn(On - 0.) 0.)= - [ I (G.)) aga
[VnA.]]
82 Öga I (0..)
Vn ug (0.)
Under slightly weaker assumptions than stated in Theorem
2 Oates al. 2022
Theorem 12)
showed the non-singularity of B0
8²
97 (0,) by convergence to the matrix I, as well as
82
In (0m)
11
V (00 $ N0
Given that | op
and
P. I non-singular, we deduce that 387 9/h (0m)
vt
82
003 9h (Bm)
[Vnd.]
=op(1),
RP
hence concluding
VT(( - -0.) -N0. -N
22


Reviews:
[{"review_id": "Py0xJzrrPZ", "reviewer": "Reviewer_d5Rs", "summary": "Under the assumption of conditional ignorability of a dichotomous treatment, this paper studies counterfactual densities. The densities are estimated within parametric classes and are only estimated up to a constant. The procedure involves doubly robust estimation, cross fitting, and kernel Stein discrepancies. The authors prove asymptotic normality for the finite density parameter with product rate conditions on nuisances.", "questions": "1. Many doubly robust methods exist for counterfactual distributions. When is it a prohibitively difficult problem to transform a counterfactual distribution to a counterfactual density? This seems central to the paper’s motivation.\n\n2. What are use cases in which we care to recover the parameter of the counterfactual density (up to a constant in some sense) rather than the counterfactual density? This seems central to the paper’s motivation.\n\n3. How do the results provided for the parameter translate to results for the counterfactual density?\n\n3. How is the asymptotic variance estimated?\n\n4. Do the corresponding confidence intervals have the desired coverage in Monte Carlo experiments?\n\n5. Is not Hd the sum kernel rather than the product kernel? I believe the sum kernel is not characteristic, and wonder if that is an issue.\n\nI will improve my score if these items are addressed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality: It appears that the estimator is a new combination of known tools, namely doubly robust estimation, cross fitting, and kernel Stein discrepancies, for a different task than previous work that combined these tools (Lam and Zhang 2023). The closest pieces of work are Fawkes et al. (2022) and Martinez-Taboada et al. (2023), which use kernel mean embeddings instead of kernel Stein discrepancies. All of these papers are well cited.\n \nQuality: The results are generally high quality, though some aspects seemed incomplete; see my comments below.\n\nClarity: The paper is written very well.\n\nSignificance: I would like the authors to provide further discussion in this regard; see my comments and questions below.", "weaknesses": "1. In my evaluation, the analysis is a direct extension of Martinez-Taboada et al. (2023), replacing features of Y with the xi object evaluated at Y.\n\n2. I would like to see more discussion of how good the optimization of the parameter must be for these results to be applicable.\n\n3. The paragraph beginning with “we underscore that…” was confusing.", "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": "ct2G0tKfgC", "reviewer": "Reviewer_bMPy", "summary": "Inference causal effect is an important question among many domains. The paper introduces an estimator for modeling counterfactual distributions given a flexible set of distributions that only need to be known up to their normalizing constants. The procedure is based on minimizing the kernel Stein discrepancy between this flexible set and the counterfactual distribution. Simultaneously, it accounts for sampling bias in a doubly robust manner. Besides,  the paper provides theoretical foundation for the consistency and asymptotic normality of the estimator, ensuring its reliability and statistical soundness.", "questions": "See above \"weaknesses\" section.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-   MKSD estimators have primarily been employed for conducting goodness-of-fit tests and sample quality analysis. Nevertheless, in the counterfactual context of this study, the MKSD estimator had not been previously proposed.\n-   Conversely, under certain assumptions, the distribution of either counterfactual can be expressed in terms of observational data. This opens the possibility of using MKSD as the primary tool to address the counterfactual distribution estimation problem.\n\n- The paper is well-organized and maintains a clear, logical flow throughout. The structure makes it easy for readers to follow the research from start to finish.\n    \n- The paper  incorporates an abundance of related work. This comprehensive overview of prior research provides valuable context for the study and underscores the authors' deep understanding of the field.\n    \n- The abstract and introduction effectively set the stage for the paper's major content. They provide a consistent and coherent overview of the research objectives, making it clear to the reader what to expect in the subsequent sections.", "weaknesses": "Numerous other studies have explored semiparametric estimators within the debiased machine learning framework for counterfactual density estimation, such as:\n\n[1] Mou, Wenlong, Martin J. Wainwright, and Peter L. Bartlett. \"Off-policy estimation of linear functionals: Non-asymptotic theory for semi-parametric efficiency.\" _arXiv preprint arXiv:2209.13075_ (2022).\n\n[2] Mou, Wenlong, et al. \"Kernel-based off-policy estimation without overlap: Instance optimality beyond semiparametric efficiency.\" _arXiv preprint arXiv:2301.06240_ (2023).\n\n[3] Kennedy, Edward H., Sivaraman Balakrishnan, and Larry Wasserman. \"Semiparametric counterfactual density estimation.\" _arXiv preprint arXiv:2102.12034_ (2021).\n\n- How do you compare the MKSD estimator with other estimators?\n- What are the advantages of using MKSD estimators in comparison to other methods?", "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": "48VgGVjR5Y", "reviewer": "Reviewer_kcoz", "summary": "This paper presents a method for estimating the counterfactual density using Kernel Stein's method. The estimated method exhibits doubly robustness w.r.t. nuisance paraeter estimation.", "questions": "1. Is $w_i(x)$ is known? If so, is it dependent on $A$? I am asking this because I think the quantity in Equation 5 should be dependent on $A$ too. \n2. Need more explanation on Equation 6 from Equation 5. Specifically, in Equation 6, what’s the meaning of $[\\hat{\\beta}_{\\theta}(\\hat{\\beta})]$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper is technically strong. It presents a complicated theory on semiparametric inference in a comprehensive manner, which is well-written.\n2. Related works provide a comprehensive summary of the relevant literature.", "weaknesses": "1. This paper could be improved by adding a discussion on the comparison between the KSD-based method and the projecting-based method (e.g., Kennedy et al., 2021). Readers may be interested in understanding the reasons or practical guidelines for choosing the KSD-based method over other methods.\n2. I am concerned about the sample complexity of the KSD-based method, as it appears to have a time complexity of $O(n^2)$. Could you please discuss how the time complexity of this method compares to other competitive methods (e.g., Kennedy et al., 2021 or Kim et al., 2018)?\n3. In the experiment, it would be interesting to compare it with other competitive methods.\n4. It would be great if there were some fixed working examples by specifying $Q_{\\theta}$ for better comprehensibility.", "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": "tvuoOAiuli", "reviewer": "Reviewer_nhik", "summary": "This paper proposes a method for estimating counterfactual densities by minimizing a Kernel Stein Discrepancy. It describes an estimator and algorithm, and conducts statistical analysis proving the consistency and asymptotic normality under certain conditions. Experiments on synthetic data are done illustrating the theoretical results and an experiment on MNIST shows that the results are intuitively reasonable and realistic.", "questions": "1. If the minimizer is unique, then is the estimate consistent?\n2. Can the distribution of $Y^0$ affect the experimental results? It is not mentioned in Section 5.\n3. In figure 3, what happens if $\\theta_1$ and $\\theta_2$ are of different signs and magnitudes?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "# Originality and significance #\nAs far as I know, this is this first work applying Stein kernel discrepancies to counterfactual estimation. It has moderate novelty, transferring previous work on Kernel Stein discrepancies to the causal estimation framework. Causal inference has become a popular research topic, with applications in healthcare, reinforcement learning, and others.\n\n# Quality #\nThe proposed estimator is intuitive, simple, and has good theoretical properties without overly strong assumptions. Experiments on synthetic data confirm the theory.\n\n# Clarity #\nThe estimator, algorithm, and theoretical results are presented clearly and thoroughly.", "weaknesses": "# Quality #\nThe number of baselines that the algorithm is compared to is fairly limited, they only include kernel-based estimators. It would strengthen the work to include more categories of counterfactuals, and discuss the pros and cons of all, such as the computation burden.\n\n# Clarity #\nThe experimental set-up is a bit unclear (see my question below).\n\n# Significance #\nThe experimental baselines are fairly limited, making it difficult to gauge the performance of the algorithm compared to SOTA.", "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)