Title: Generalized Balancing Weights via Deep Neural Networks
Venue: NeurIPS 2023
Accepted: no
Keywords: causal inference

density ratio estimation
Abstract: Estimating causal effects from observational data is a central problem in many domains.
A general approach is to balance covariates with weights such that the distribution of the data mimics randomization.
We present generalized balancing weights, *Neural Balancing Weights* (NBW), to estimate the causal effects of an arbitrary mixture of discrete and continuous interventions.
The weights were obtained through direct estimation of the density ratio between the source and balanced distributions by optimizing the variational representation of $f$-divergence.
For this, we selected $\alpha$-divergence as it presents efficient optimization because it has an estimator whose sample complexity is independent of its ground truth value and unbiased mini-batch gradients;  moreover, it is advantageous for the vanishing-gradient problem.
In addition, we provide the following two methods for estimating the balancing weights: improving the generalization performance of the balancing weights and checking the balance of the distribution changed by the weights.
Finally, we discuss the sample size requirements for the weights as a general problem of a curse of dimensionality when balancing multidimensional data.
Our study provides a basic approach for estimating the balancing weights of multidimensional data using variational $f$-divergences.

Abstract

Estimating causal effects from observational data is a central problem in many domains. A general approach is to balance covariates with weights such that the distribution of the data mimics randomization. We present generalized balancing weights, Neural Balancing Weights (NBW), to estimate the causal effects of an arbitrary mixture of discrete and continuous interventions. The weights were obtained through direct estimation of the density ratio between the source and balanced distributions by optimizing the variational representation of f -divergence. For this, we selected α -divergence as it presents efficient optimization because it has an estimator whose sample complexity is independent of its ground truth value and unbiased mini-batch gradients; moreover, it is advantageous for the vanishing-gradient problem. In addition, we provide the following two methods for estimating the balancing weights: improving the generalization performance of the balancing weights and checking the balance of the distribution changed by the weights. Finally, we discuss the sample size requirements for the weights as a general problem of a curse of dimensionality when balancing multidimensional data. Our study provides a basic approach for estimating the balancing weights of multidimensional data using variational f -divergences.

1 Introduction

Estimating causal effects from observational data is a central problem in many application domains, including public health, social sciences, clinical pharmacology, and clinical decision-making. One standard approach is balancing covariates with weights that are the same as the density ratios between the source and balanced distributions, such that their distribution mimics randomization. Many methods have been developed to estimate the balancing weights, such as inverse propensity weighting (IPW) Rosenbaum and Rubin [24], augmented inverse propensity weighting (AIPW) [22], generalized propensity score (GPS) [10], covariate balancing propensity score (CBPS) [9], overlap weighting [13], and entropy balancing (EB) [8, 30]. However, these methods are limited to categorical or continuous interventions.

In this study, we propose generalized balancing weights to estimate the causal effects of an arbitrary 29 mixture of discrete and continuous interventions. To the best of our knowledge, no causal infer30 ence method focusing on the balancing weights exists for this problem. We approach this problem 31 by directly estimating the density ratio, more precisely, the Radon-Nikodým derivatives, between 32 the source and balanced distributions using a neural network algorithm by optimizing a variational 33 representation of a f -divergence. f -divergences, whose values are greater than or equal to zero and 34 considered zero if the two distributions are equal, are the statistics used to measure the closeness 35 of the two distributions. The optimal functions for the variational representations derived from f -36 divergences with the Legendre transform correspond to the density ratio between the distributions 37

.

[16]. An approach to estimate the density ratio by optimizing a variational representation of a f -38 divergence was developed in the domain adaptation region [29]. 39

However, optimizing the f -divergences, including estimating the density ratio, is challenging. This 40 is due to the following reasons. First, for KL-divergence, the dominant f -divergence, the require41 ments for sample size increase exponentially with the true amount of the divergence [14, 28]. Sec42 ond, a naive gradient estimate over mini-batch samples leads to a biased estimate of the full gradient 43 [4]. Third, gradients of neural networks often vanish when the estimated probability ratios are close 44 to zero [2]. 45

46

47

48

49

To avoid the first problem, we focus on α -divergence, which is a subgroup of f -divergence. α -divergence has an estimator whose sample complexity is independent of its ground truth value and unbiased mini-batch gradients. In addition, by selecting α from a particular interval, we avoid vanishing gradients of neural networks when the neural networks reach extreme local minima.

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

In addition, we provide two techniques for estimating the balancing weights. First, we propose a validation method using test data and an early stopping method to improve the generalization performance of balancing. The generalization performance of the weights worsens as the dimensions of the data increase, and the sample size requirements of the weights increase exponentially with the dimensions. Next, we present a method for measuring the performance of balancing weights by estimating the α -divergence information to check the balance of the distribution,

This study is divided into seven parts. First, we introduce the background of the study. Second, we review related studies. Third, we define the terminology and concepts for causal inferences. Fourth, we present our novel method for estimating balancing weights. Fifth, we provide techniques for estimating the weights. Sixth, we discuss the sample requirements for the weights. Finally, we conclude this paper. All the numerical experiments and proofs are described in the appendix.

2 Related Work

Balancing weight: Balancing weight: Many methods have been proposed to estimate the balancing weights. The following methods are proposed for binary intervention: IPW [24], AIPW [22], CBPS [9], and overlap weighting [13]. The following methods have been proposed for continuous intervention: GPS [10] and EB [8, 30]. Statistical divergences and density ratio estimation: Despite the abundance of classic studies [15, 29], we focused on studies that directly estimate density ratios or optimize statistical divergences using neural networks. In this review, these studies have beenclassified into four groups. First is the estimation of KL-divergence or mutual information [3, 18, 21]; the second is density ratio estimation [11]; the third is generative adversarial networks (GANs) [17, 31, 6, 32] (statistical divergences were used as discriminators for GANs); and the fourth is domain generation [27, 6, 35, 1]. In addition to these application studies, divergences were improved [5].

3 Terminologies and Definitions

Here, we briefly introduce the terminology and definitions used in this study.

Notations and Terminologies. Random variables are denoted by capital letters; for example, A . 75 Small letters are used for the values of random variables of the corresponding capital letters; a is 76 the value of the random variable A . Bold letters A or a represent a set of variables or random 77 variable values. In particular, V = { V 1 , . . . , V n } are used for the observed random variables and 78 U = { U 1 , . . . , U m } are used as unobserved random variables. For example, the domain of the 79 variable A is denoted by X A , and X A 1 × · · · × X A n is denoted by X A for A = A 1 × · · · × A n . 80 V ∪ U are assumed to be semi-Markovian models and G = G VU denotes the causal graph for 81 V ∪ U . Pa ( A ) G , Ch ( A ) G , An ( A ) G , and De ( A ) G represent parents, children, ancestors, and 82 descendants of the observed variables in G , respectively, for A ⊂ V . In this study, Pa ( A ) G , 83 Ch ( A ) G , An ( A ) G , and De ( A ) G do not include A . P and Q are used as the probability measures 84 on ( R d , F ) , where F denotes the σ -algebra of subsets of R d . E P [ · ] and E P [ ·|· ] denote expectation 85 and conditional expectation under the distribution P , respectively. For example, E P [ X ] = ∫ X X dP 86 and E P [ Y | X ] = ∫ X Y dP ( Y | X ) . ˆ E P [ · ] denotes the empirical expectation under P ; that is, the 87

sample mean of the finite observations drawn from P . P is called absolute continuous with respect 88 to Q , P ( A ) = 0 whenever Q ( A ) = 0 for any A ∈ F , which is represented as P /lessmuch Q . dP dQ denotes 89 the Radon-Nikodým derivative of P with respect to Q for P and Q with P /lessmuch Q . In this study, 90 we refer to density ratios as the Radon-Nikodým derivatives. µ denotes a probability measure 91 on R d with P /lessmuch µ and Q /lessmuch µ . X ( N ) = { X 1 , . . . , X N } denotes N i.i.d. random variables 92 from µ . X ( N ) P = { X 1 ∼ P , . . . , X N ∼ P } and X ( N ) Q = { X 1 ∼ Q , . . . , X N ∼ Q } denote variables defined as 93 P ( X i ∼ P ≤ x ) = µ ( X i ≤ x ) and Q ( X i ∼ Q ≤ x ) = µ ( X i ≤ x ) , ∀ x ∈ R d , for 1 ≤ i ≤ N . We 94 represent f /lessorsimilar g when lim sup n →∞ f ( n ) /g ( n ) < ∞ holds. The notation f /greaterorsimilar g is defined similarly. 95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

3.1 Definitions

In this study, we considered the causal effects of joint and multidimensional interventions. For clarity, we used different notations, @g do @h and @g do , @h for single-dimensional and multidimensional interventions, respectively. 1 For a single-dimensional intervention, a do symbol is used, which is the same as Pearl's do -calculation.

Definition 3.1 ( do -calculation, Pearl(2009)) . For the two given disjoint sets of X , Y ⊂ V , the causal effect on Y for intervention in X with values x , denoted by P ( Y | do ( X = x )) , is defined as the probability distribution, such that

where V ′ = V \ ( X ∪ Pa ( X ) G ∪ Y ) . The causal effect of X on Y under the conditions Z denoted by P ( Y = y | do ( X = x ) , Z = z ) is defined as the probability distribution, such that

Notably, from Definition 3.1, a do -calculation for a set of variables coincides with the simultaneous interventions for each variable:

Furthermore, we use the do symbol for multidimensional intervention. Intuitively, a do symbol represents the intervention of the variables that preserves the functional relationship within the variables.

where X = { X 1 , X 2 , . . . , X n } . Here, we refer to each intervention in (3) as a @g single-dimensional intervention @h .

Definition 3.2 ( do symbol) . do symbol defines the following probability distribution:

do symbols are useful, particularly when we consider interventions in a multivalued discrete variable expressed using one-hot encoding. In this case, we cannot express the causal effect effectively using do symbols. For example, let us consider the case of an intervention in the ternary variable X , X X = { x 1 , x 2 , x 3 } and let X be expressed by X ′ = ( X ′ 1 , X ′ 2 , X ′ 3 ) , such that X ′ i = 1 if X = x i otherwise X ′ i = 0 for i = 1 , 2 , 3 . Then, P ( ·| do ( X = x 3 )) is the same as P ( ·| do ( X ′ = (0 , 0 , 1))) , which differs from P ( ·| do ( X ′ = (0 , 0 , 1))) . We refer to this type of intervention as a @g multidimensional intervention @h .

Next, we provide definitions of the f -divergence and f -divergence information.

Definition 3.3 ( f -divergence) . The f -divergence D f between the two probability measures P and Q with Q /lessmuch P induced by a convex function f satisfying f (1) = 0 is defined by D f ( Q || P ) = E P [ f ( dQ/dP )] .

Many divergences are specific cases obtained by selecting a suitable generator function f . For 126 example, f ( u ) = u log u corresponds to the KL-divergence. In particular, we focus on α -divergence, 127 which is expressed as follows: 128

129

130

where α ∈ R \ { 0 , 1 } . From (5), Hellinger divergence is obtained as α = 1 / 2 , and χ 2 divergence by α = -1 .

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

From f -divergence, the f -divergence information is defined as the mutual information if we choose the KL-divergence as the f -divergence. Here, we present a definition of f -divergence information for multi-variables.

Definition 3.4 ( f -divergence information) . For disjoint variables X = { X 1 , X 2 , . . . , X n } ⊂ V , let P X be the joint probability measure for X . For each i = 1 , 2 , . . . , n , P X i = ∫ X X \ X i dP X is a measure of the marginal distribution of P X for X i . The f -divergence information for X 1 , X 2 , . . . , X n under P X and a convex function f satisfying f (1) = 0 is defined as the f -divergence between P X and P X 1 × P X 2 ×··· × P X n :

4 Problem Set Up

Before describing the details of the problem, we provide a notation for the probability distribution, which is the goal of balancing. Hereafter, P denotes the probability distribution of observational data. For the given disjoint sets X 1 , X 2 , . . . , X n , Y , Z ⊂ V , let ˜ P be a probability distribution, as follows:

where X = X 1 ∪ X 2 ∪ · · · ∪ X n . ˜ P is the probability distribution of the counterfactual data from simultaneous (multidimensional) interventions in X 1 , X 2 , . . . , X n under the condition Z .

Objective. The objective of this study is to obtain the balancing weights that transform P ( Y , X , Z ) into ˜ P ( Y , X , Z ) . More precisely, given the i.i.d. observational data { ( x i , z i ) | i = 1 , 2 , . . . , N } , we aim to estimate the weights BW ( X , Z ) , such that holds for any measurable function f on R d . If we obtain the weights, we estimate the conditional average causal effect (CACE) for P ( Y | do ( X 1 ) , do ( X 2 ) , . . . , do ( X n ) , Z ) , that is E ˜ P [ Y | X , Z ] , using state-of-the-art supervised machine learning algorithms, with the weights assigned as the individual weights for each sample.

Assumptions. We assumed the following to achieve our objective:

Assumption 1. The causal effect P ( Y | do ( X )) is identifiable, or equivalently, ˜ P from (7) can be identified. 2 3

Assumption 2. Let P = P ( X 1 , X 2 , . . . , X n , Z ) and let Q = P ( X 1 ) × P ( X 2 ) × · · · × P ( X n ) × P ( Z ) . Subsequently, we assume that Q /lessmuch P .

Assumption 2 is the same as the overlap assumption if we consider this a single-dimensional intervention. Here, we propose overlapped assumptions for joint and multidimensional interventions.

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

5 Estimation of Balancing Weights

In this section, we present the way to effectively estimate the probability density ratios by optimizing f -divergence.

Density Ratios as Balancing Weights. We first note that the density ratios, which are referred to as the Radon-Nikodým derivative in this paper, are equal to the balancing weight of the target. For a density ratio of P to ˜ P , that is d ˜ P dP , it holds that

for any measurable function f in R d . Then, (8) and (9) are equivalent. As an example of the aforementioned density ratio, let X be a binary variable with X X = { 1 , 0 } and let Z be covariates. Using propensity score e ( z ) = P ( X = 1 | Z = z ) , we observe that d ˜ P dP ( X = 1 , z ) = P ( X = 1) /e ( z ) and d ˜ P dP ( X = 0 , z ) = P ( X = 0) / (1 -e ( z )) . That is, d ˜ P dP is the stabilized inverse probability of the treatment weighting [23].

5.1 Our Approach

Our approach involves obtaining the density ratios as an optimal function for a variational representation of an f -divergence. This approach is based on the fact that the optimal function is connected to density ratios [15].

Variational representation. Using the Legendre transform of the convex conjugate of a twice differentiable convex function f , f ∗ ( ψ ) = sup r ∈ R { ψ · r -f ( r ) } , we obtain a variational representation of f -divergence:

where supremum is considered over all measurable functions with E Q [ f ′ ( φ )] < ∞ and E P [ f ∗ ( f ′ ( φ ))] < ∞ . The maximum value is achieved at φ = dQ/dP .

We obtained the optimal function for (10) by replacing φ in the equation with a neural network model φ θ and training it through back-propagation with a loss function, such that

Selecting α -divergence for Optimization. We select α -divergence for the following reasons. 182 First, the sample size requirements for α -divergence is independent of its ground truth value: second, 183 it has unbiased mini-batch gradients; third, it can avoid a vanishing gradient problem. 184

The variational representation of α -divergence is as follows (Lemma C.1 in Appendix C): 185

Sample size requirements for α -divergence. The α -divergence has an estimator with sample 186 complexity O (1) (Corollary 1 in Birrell et al., 2022, P19; Corollary C.10 in Appendix C). Con187 versely, the sample complexity of KL-divergence is O ( e KL ( Q || P ) ) [14, 28]: 188

where ̂ KL N ( Q || P ) is the KL-divergence estimator for sample size N using a variational represen189 tation of the divergence, and KL ( Q || P ) is the ground truth value. 190

Unbiasedness for mini-batch gradients. φ in (12) can be expressed in a Gibbs density form 191 (Proposition C.2 in Appendix C). Then, we observe that 192

193

194

where supremum is considered over all measurable functions T : R d → R with E P [ e ( α -1) · T ] < ∞ and E Q [ e α · T ] < ∞ .

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

From this equation, we obtain our loss function, which has unbiasedness for mini-batch gradients (Proposition C.8 in Appendix C), as follows :

Advantage in vanishing gradients problem. By setting α within (0 , 1) , we can avoid vanishing gradients of neural networks when they reach the extreme local minima. The vanishing-gradient problem for optimizing divergence is known in GANs [2]. Now, we consider the case where the probability ratio e T θ ( x ) in (15) is nearly zero or large for some point x , corresponding to cases in which the probabilities for P or Q at some points are much smaller than those for the other.

To show the relation between e T θ ( x ) and the learning of the neural networks, we obtain gradient of (15):

The behavior of ∇ θ L α ( θ ) when E Q [ e T θ ] → 0 or E Q [ e T θ ] →∞ , under some regular conditions for T θ and an assumption that P /lessmuch Q , can be summarized as follows: Let E [ · ] denote E P [ E Q [ · ]] , then

For α > 1 and α < 0 , cases exist where E [ ∇ θ L α ( θ )] → /vector 0 . This implies the possibility that the neural networks reach extreme local minima such that their estimations for density ratios are 0 or ∞ . However, this problem can be avoided by selecting α from interval (0 , 1) . We note that the selecting of α does not cause instability in numerical calculations for cases where E [ ∇ θ L α ( θ )] → /vector ∞-/vector ∞ . In Appendix D.1, we present numerical experimental results for different values of α .

6 Method

In this section, we first present the main theorem that summarizes the new balancing weight method proposed herein. Next, we present the balancing weight method.

6.1 Main Theorem

Here, we present the main theorem that summarizes the new balancing weight method proposed herein.

Theorem 6.1. Given disjoint sets of X = { X 1 , X 2 , . . . , X n } , Y , Z ⊂ V satisfying

Let P = P ( X 1 , X 2 , . . . , X n , Z ) and Q = P ( X 1 ) × P ( X 2 ) × · · · × P ( X n ) × P ( Z ) , and ˜ P = P ( Y | do ( X ) , Z ) × P ( X 1 ) × P ( X 2 ) ×···× P ( X n ) × P ( Z ) . Weassume that P satisfies Assumptions 1 and 2 in the aforementioned setting, and it holds that E P [ ( d Q /d P ) 1 -α ] < ∞ for some 0 < α < 1 , then, for the optimal function T ∗ , such that

Algorithm 1 Training a Neural Balancing Weight model

it holds that 228

229

Here, T α denotes the set of all non-constant functions T ( x ) : R d → R with E P [ e ( α -1) · T ( X ) ] < ∞ .

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

Here, we mention that the assumption (17) is necessary for the (19) to hold, which is derived from our Theorem C.15 in Appendix C.

6.2 Balancing Weight Method

We present the implementation of training a neural balancing weights (NBW) model in Algorithm 1. It is important to consider the stopping time K for neural network model T θ K in Algorithm 1, which is discussed in the next section. To obtain the sample mean under Q , that is, the estimator for E Q [ e α · T θ ] in (18), a shuffling operation can be used for the samples. Now, we define neural balancing weights (NBW). 4 5

Definition 6.2 (Neural Balancing Weights) . Let T θ K be a neural networks obtained from Algorithm 1. Then, the NBW of T θ K , expressed as BW ( X 1 , X 2 , . . . , X n , Z ; T θ K ) , are defined as

We estimate E ˜ P [ Y | X , Z ] , that is the CACE for P ( Y | do ( X 1 ) , do ( X 2 ) , . . . , do ( X n ) , Z ) , using BW ( X 1 , X 2 , . . . , X n , Z ; T θ K ) as the sample weights of the supervised algorithm:

where Z = ˆ E P [ e -T θ K ( X 1 , X 2 ,..., X n , Z ) ] .

̂ E ˜ P [ Y | X , Z ] = ̂ E P [ Y · BW θ K | X , Z ] . (21) Here, ̂ E P corresponds to the model of a supervised machine learning algorithm. As an example, we demonstrate a back-propagation algorithm using balancing weights for the mean squared error (MSE) loss in Algorithm 3 in Appendix E.

7 Techniques for NBW

We propose two techniques for estimating balancing weights: ( i ) improves generalization performance of the balancing weights. ( ii ) measures the performance of the balancing weights by estimating the α -divergence information.

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

7.1 Improving the Generalization Performance of the Balancing Weights

In this section, we first present an overfitting problem for balancing distributions. We then present two methods for improving the generalization performance of the weights: a validation method using test data and an early stopping method. Herein, let T θ t denote an NBW model at step t in Algorithm 1. Let ˆ X ( N ) Q ( t ) = e -T θ t · X ( N ) P , that is, the data balanced by the weights of e -T θ t . Subsequently, let ˆ Q ( N ) t and ˆ Q ( N ) denote the probability distributions of ˆ X ( N ) Q ( t ) and ˆ X ( N ) Q , respectively, which correspond to the estimated and true distributions for balancing.

An overfitting problem for balancing distributions. From Corollary C.12 in Appendix C, we observe ˆ X ( N ) Q ( t ) d - - → X ( N ) Q as t →∞ . Then, Theorem 1 in [33] shows that

where W 1 is the Wasserstein distance of order 1 and d is the lower Wasserstein dimension defined in [33]. (22) implies that, for balancing finite data, the destination of the balanced distribution is an empirical distribution, and the generalization performance of balancing worsens exponentially when the dimension of the data is larger. In view of optimizations of GANs, [34] referred to this phenomenon the @g momorization @h and proposed an early stopping method.

Validation method using test data. Wecan use a validation method using test data. Because ˆ Q ( N ) and ˆ P ( N ) are empirical probability distributions, we observe that d ˆ Q ( N ) /d ˆ P ( N ) ( x ) = dQ/dP ( x ) if x ∈ X ( N ) , otherwise d ˆ Q ( N ) /d ˆ P ( N ) ( x ) = 0 (Proposition C.17 in Appendix C). Then, the optimal function of (15) for both distributions, that is T ( N ) ∗ = -log ( d ˆ Q ( N ) /d ˆ P ( N ) ) , is infinite except for the observations, and the loss of the T ( N ) ∗ is infinite for data independent of the observations. This implies that the loss of T ( N ) t for the test data turns to increase from the middle of the training period, and we can determine the training step at which the generalization performance of the weights begins to worsen. In Section D.2 in Appendix D, we provide numerical experimental results to confirm the relationship between dimensions of data ( d ) and steps in training ( K ).

Early stopping method. In addition, we present an early stopping method for estimating the balancing weights as follows, which is inspired by the method developed in [34] (Corollary C.24 in Appendix C): for some δ > 0 , let

where C > 0 is constant. Then, we have W 1 ( Q, ˆ Q ( N ) K 0 ) /lessorsimilar N -1 / ( d + δ ) . Unfortunately, the curse of dimensionality remains in the proposed method. This will be discussed in the next section.

7.2 Measuring the Performance of the Balancing Weights

Let us assume that we obtain an NBW model T θ 0 and let BW θ 0 = BW ( X 1 , X 2 , . . . , X n , Z ; T θ 0 ) be the balancing weights of T θ 0 . If BW θ 0 successfully estimates dQ dP , then the α -divergence between Q and P 0 will be nearly zero. Conversely, if BW θ 0 fails to estimate dQ dP , the α -divergence between Q and P 0 is significantly different from zero. This implies that we can measure the performance of the balancing weights using the α -divergence information for P 0 .

Next, we present the definition of an α -divergence information estimator using neural networks.

Definition 7.1 (Neural α -divergence Information Estimator) . For disjoint variables X 1 , X 2 , . . . , X n ⊂ V , the neural α -divergence information estimator for P is defined as

To measure the performance of balancing the weights from the NBW model, we estimate the α -288 divergence information for balanced distribution from the weights. That is, we use the sample mean 289 under a balanced distribution, despite the sample mean under P for (24). For example, we assume 290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

Algorithm 2 Algorithm for checking the balance

Input: Train Data { ( x i 1 , . . . , x i n , z i ) } N i =1 , Test Data { ( ˜ x i 1 , . . . , ˜ x i n , ˜ z i ) } N i =1 , A Neural Balancing Weight Model T θ

Output: The estimated α -divergence information ̂ I α for the balanced distribution with the balancing weights from T θ

that we have certain weights BW ′ = { bw i : i = 1 , 2 , ..., N } , where bw i denotes the weight of sample i of N . The balanced distribution from the weights is

The α -divergence information for P ′ is estimated by replacing P with P ′ for (24) in the following manner: despite the sample mean ˆ E P [ e ( α -1) · T θ ] for these equations, we use the weighted sample mean, such that

Details on the implementation for measuring the performance of balancing weights from an NBW model are provided in Algorithm 2, which includes the validation method for the overfitting problem in Section 7.1.

8 Limitations: Sample Size Requirements.

In Section 7.1, we noted that our method has a curse of dimensionality. The sample size requirement of the proposed method is N > ( 1 ε ) d + δ for W 1 ( Q, ˆ Q ( N ) K 0 ) < ε (Corollary C.25 in Appendix C). However, the curse of dimensionality is an essential problem when balancing multivariate data owing to the following factors. Because the optimal balancing weights defined as (8) for (finite) observational data are the density ratios of the empirical distributions, the distribution of the data balanced by them is the empirical distribution. Subsequently, owing to the balancing of the weights, the curse of dimensionality of the empirical distribution occurs, which is the same as that described in Section 7.1. Therefore, to achieve high generalization performance, we need to obtain weights that differ from the ideal density ratio between the source and target of the empirical distribution. Further research is required to address this problem. In Appendix D.3, we present the numerical examination results in which the causal effects of joint and multidimensional interventions were estimated with different sample sizes.

9 Conclusion

We propose generalized balancing weights to estimate the causal effects of an arbitrary mixture of discrete and continuous interventions. Three methods for training the weights were provided: an optimization method to learn the weights, a method to improve the generalization performance of the balancing weights, and a method to measure the performance of the weights. We showed the sample size requirements for the weights and then discussed the curse of dimensionality that occurs as a general problem when balancing multidimensional data. Although the curse of dimensionality remains in our method, we believe that this study provides a basic approach for estimating the balancing weights of multidimensional data using variational f -divergence.

321

References

David Acuna, Guojun Zhang, Marc T Law, and Sanja Fidler. f-domain adversarial learn322 ing: Theory and algorithms. In International Conference on Machine Learning , pages 66-75. 323 PMLR, 2021. 324

Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adver325 sarial networks. arXiv preprint arXiv:1701.04862 , 2017. 326

Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, 327 Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International 328 conference on machine learning , pages 531-540. PMLR, 2018. 329

330

331

332

Marc G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance as a solution to biased wasserstein gradients. arXiv preprint arXiv:1705.10743 , 2017.

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

Jeremiah Birrell, Markos A Katsoulakis, and Yannis Pantazis. Optimizing variational representations of divergences and accelerating their statistical estimation. IEEE Transactions on Information Theory , 2022.

Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research , 17(1):2096-2030, 2016.

Gustavo L Gilardoni. On pinsker's and vajda's type inequalities for csiszár's f -divergences. IEEE Transactions on Information Theory , 56(11):5377-5386, 2010.

Jens Hainmueller. Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies. Political analysis , 20(1):25-46, 2012.

Kosuke Imai and Marc Ratkovic. Covariate balancing propensity score. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 76(1):243-263, 2014.

Kosuke Imai and David A Van Dyk. Causal inference with general treatment regimes: Generalizing the propensity score. Journal of the American Statistical Association , 99(467):854-866, 2004.

Masahiro Kato and Takeshi Teshima. Non-negative bregman divergence minimization for deep direct density ratio estimation. In International Conference on Machine Learning , pages 5320-5333. PMLR, 2021.

Simon Lacoste-Julien, Mark Schmidt, and Francis Bach. A simpler approach to obtaining an o (1/t) convergence rate for the projected stochastic subgradient method. arXiv preprint arXiv:1212.2002 , 2012.

Fan Li, Kari Lock Morgan, and Alan M Zaslavsky. Balancing covariates via propensity score weighting. Journal of the American Statistical Association , 113(521):390-400, 2018.

David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statistics , pages 875-884. PMLR, 2020.

XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by penalized convex risk minimization. Advances in neural information processing systems , 20, 2007.

XuanLong Nguyen, Martin J Wainwright, and Michael I Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory , 56(11):5847-5861, 2010.

Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural 365 samplers using variational divergence minimization. Advances in neural information process366 ing systems , 29, 2016. 367

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive 368 predictive coding. arXiv preprint arXiv:1807.03748 , 2018. 369

Judea Pearl. Causal diagrams for empirical research. Biometrika , 82(4):669-688, 1995. 370

Judea Pearl. Causality: Models, reasoning and inference, 2009. 371

372

Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational

373

374

bounds of mutual information. In

5180. PMLR, 2019.

James M Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of regression coefficients 375 when some regressors are not always observed. Journal of the American statistical Association , 376 89(427):846-866, 1994. 377

James M Robins, Miguel Angel Hernan, and Babette Brumback. Marginal structural models 378 and causal inference in epidemiology, 2000. 379

Paul R Rosenbaum and Donald B Rubin. Reducing bias in observational studies using sub380 classification on the propensity score. Journal of the American statistical Association , 79(387): 381 516-524, 1984. 382

Albert Nikolaevich Shiryaev. Probability, 1995. 383

Ilya Shpitser and Judea Pearl. Identification of conditional interventional distributions. arXiv 384 preprint arXiv:1206.6876 , 2012. 385

Si Si, Dacheng Tao, and Bo Geng. Bregman divergence-based regularization for transfer 386 subspace learning. IEEE Transactions on Knowledge and Data Engineering , 22(7):929-942, 387 2009. 388

Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual infor389 mation estimators. arXiv preprint arXiv:1910.06222 , 2019. 390

Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density-ratio matching under the 391 bregman divergence: a unified framework of density-ratio estimation. Annals of the Institute 392 of Statistical Mathematics , 64(5):1009-1044, 2012. 393

Stefan Tübbicke. Entropy balancing for continuous treatments. Journal of Econometric Meth394 ods , 11(1):71-89, 2022. 395

396

397

398

Masatoshi Uehara, Issei Sato, Masahiro Suzuki, Kotaro Nakayama, and Yutaka Matsuo. Generative adversarial nets from a density ratio estimation perspective. arXiv preprint arXiv:1610.02920 , 2016.

Ananya Uppal, Shashank Singh, and Barnabás Póczos. Nonparametric density estimation & 399 convergence rates for gans under besov ipm losses. Advances in neural information processing 400 systems , 32, 2019. 401

Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of 402 empirical measures in wasserstein distance. 2019. 403

Hongkang Yang and E Weinan. Generalization and memorization: The bias potential model. 404 In Mathematical and Scientific Machine Learning , pages 1013-1043. PMLR, 2022. 405

Yuchen Zhang, Tianle Liu, Mingsheng Long, and Michael Jordan. Bridging theory and al406 gorithm for domain adaptation. In International Conference on Machine Learning , pages 407 7404-7413. PMLR, 2019. 408

International Conference on Machine Learning

, pages 5171-

[21]

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method for estimating the balancing weights in causal effect estimation. The balancing weights are given by the density ratio of the counterfactual distribution and the observed distribution, and the density ratio can be estimated by solving the optimization of the variational expression of the f-divergence between the two distributions. The authors claim that using the $\\alpha$-divergence is useful because it addresses the vanishing-gradient problem when we use a neural network for the model. The paper presents some theoretical results.\n\nSTRENGTHS:\n- Expressing $\\phi$ as $\\exp(T)$ in Eq. (14) is interesting since the domain of the optimization variable becomes simpler.\n- The discussion about vanishing gradient is interesting, and the suggested way of setting $\\alpha$ may be useful to address this issue.\n\nWEAKNESSES:\n- There are several technical parts that might not be precise. See the Questions section below.\n- The writing could be improved. The paper rather looks like a collection of definitions of results. Connections between sections are not smooth, and the message of the paper as a whole is not clear.\n- The abstract says, \"we selected $\\alpha$-divergence as it presents efficient optimization because it has an estimator whose sample complexity is independent of its ground truth value and unbiased mini-batch gradients; moreover, it is advantageous for the vanishing-gradient problem.\" I am not sure what makes the $\\alpha$-divergence special in terms of these aspects. Also, the explanation about the vanishing-gradient issue in Section 5.1 is not convincing enough to make me believe that the authors' suggestion really addresses the issue.\n- The paper relies on references (including one without open access) for important definitions. The paper could be more self-contained.\n- Overall, the paper is mainly about estimating density ratios by the variational expression of the f-divergence, which is not novel.\n- There is no experiment in the main part of the paper.\n\nQUESTIONS:\n- Is Definition 3.1 really the same as the standard definition from Pearl (2009)? Which part of the book are the authors referring to?\n- I believe there is something wrong with Eq. (4) and Eq. (7). The right hand side depends on $P(\\boldsymbol{X}_1), ..., P(\\boldsymbol{X}_n)$ while the left hand side does not.\n- In lines 207-211, what is $\\overset{\\to}{\\infty}$?\n- \"We note that the selecting of α does not cause instability in numerical calculations for cases where $E[\\nabla_{\\theta} \\mathcal{L}_{\\alpha}(θ)] \\to \\infty - \\infty$.\" Is there any theoretical result to support this claim?\n- Theorem 6.1 looks strange because Eq. (19) concerns $\\tilde{P}$ although Eq. (18) has nothing concerning $\\tilde{P}$.\n- What does it mean by using test data for validation? Is this validation performed only for performance evaluation or also for hyper-parameter tuning?\n- As far as I know, people just perform as many epochs as they can afford while keeping the model with the best validation performance. So, fixing the early stopping time is not very useful in practice unless. Am I missing something?\n- The $\\alpha$-divergence \"has an estimator whose sample complexity is independent of its ground truth value\": is this true? Eq. (13) seems to depends on both $P$ and $Q$.\n- \"unbiased mini-batch gradients\": isn't this true for any $f$-divergence?\n\n-----\nEdit on Aug 2: There is another thing that I would like the authors to address if it is not too late:\nIs Eq. (10) correct? Is $f'$ the derivative of $f$ and are we optimizing the objective with respect to $\\phi$?\n\nI thought that the authors use the result from [15], but I cannot exactly match the equation with the one in the paper.\n\n[15] XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by penalized convex risk minimization. Advances in neural information processing systems, 20, 2007.\n\nLIMITATIONS:\nThe authors discuss limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an approach for estimating balancing weights. The ultimate target is to use these weights for the estimation of the casual effects of interventions. The search of balancing weights are formulated as density ratio estimation. To tackle this problem, the authors employ a variational representation of the $f$-divergence. . The density ratio is modeled by a neural network. Specifically, this paper advocates the use of $\\alpha$-divergence, and provide various ways to improve the practical performances of the proposed weights.\n\nSTRENGTHS:\n- This paper successfully makes use of the density ratio learning tools based on $f$ -divergence for causal inference applications.\n- This paper summarizes various related existing results on density ratio learning via variational form of $f$ -divergence.\n- This paper provides detailed discussion on various practical improvement of the proposed algorithm, which would benefit practitioners.\n\nWEAKNESSES:\n- This paper lacks theoretical results on the causal effect estimation (not the density ratio learning), which is regarded as a key target of this work.\n\t- It is unclear whether the proposed estimator obtain optimal convergence rates and (semi-parametric) efficiency, when degenerated to standard settings like average treatment estimation, and conditional average treatment estimation.\n- This work is not impressive in terms of novelty, as many significant components follow more directly from existing work.\n\t- The variational form of $f$ -divergence and the corresponding estimation of density ratio has been explored by references [15] and [16].\n\t- The main theoretical result (Theorem 6.1) seems to be a fairly standard dual result (e.g., Theorem 4.4 of https://arxiv.org/pdf/1003.5457.pdf)\n- The theoretical results are not stated clearly. Many results are not stated in rigorous statements with careful listing of assumptions. Some seem to be exclude certain important cases like mixture distributions, as claimed to be a major contribution of this proposed work. (see my question below.)\n\nQUESTIONS:\n- When the target causal effect is a scalar, does the proposed weighted estimator achieve $\\sqrt{N}$ rate? is the corresponding weighted estimator efficient?\n- For conditional treatment estimation, the target quantity is a function, and could be nonparametric in general. Would the weighted estimator converge? Does it converge with optimal rate?\n- How to initialize the algorithm?\n- How to choose $C$ in (23) in practice?\n- Line 277, can the author formalize their claim that $W_1(Q, \\hat{Q}_{K_0}^{(N)})$ (i.e., with clear conditions) and provide a proof? Since the estimator is parametrized by neural network, it leads to a non-convex optimization in general. The loss landscape in terms of the neural network parameters seem complicated, and the performance of the estimator resulting from early stopping seems to be related to initialization. (In the supplement (Section C.3), it seems that the authors analyze a different algorithm (instead of Algorithm 1) where the GD is not applied to the neural network parameters.)\n- While it is great that the authors provide a detailed discussion of various ways to improve the performance of the proposed estimator in Section 7.1, these strategies seems relatively standard.\n- The authors emphasize that their estimator can handle arbitrary mixture of discrete and continuous interventions, but I don't see much discussion/innovation related to this challenge throughout the paper. It seems to me the only trick that is advocated to deal with this challenge is to view the weights as Radon-Nikodym derivatives (which seems to be pretty well known) and then use NN to model. Am I missing anything?\n- Also, in the case of mixture interventions, can the authors confirm and discuss the applicability of the theoretical results in Section 7.1 and 8? It seems that the analysis of the early stopping does not apply to the mixture settings (Assumption E4).\n\nLIMITATIONS:\nSection 8 discusses limitations on sample size requirements. The authors claim that the sample size required for controlling the error of $\\hat{Q}_{K_0}^{(N)}$ has an exponential dependence on the dimension. Although I have questions regarding this theoretical results (as stated above), I would not be surprised about the curse of dimensionality. However, since the ultimate target is the causal effect, I am more interested in a direct error or sample complexity analysis of the causal effect. The lack of such analysis is a major limitation of this work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper derives a novel method of reweighting covariates for the purpose of causal inference. This is done through learning optimal change of measure weights modeled via neural networks and trained using alpha-divergence measures between the true joint distribution and a mutually independent one. Further tweaks are introduced to help enable this procedure to be more practical in its implementation. \n\nA brief disclaimer: I am not very familiar with causal inference (to be honest I am not quite sure how I received this paper to review); however, I am quite familiar with variational inference and related techniques. I will focus primarily on the latter during my review, but please keep all of this in mind while reading my comments.\n\nSTRENGTHS:\nAll of the work presented is very precise and detailed, with novel contributions that I can see with regards to the learning of the weights via the alpha-divergence objective. I found the result of the main algorithm to be particularly concise and intuitive for such an involved derivation (e.g., almost trivial to describe but definitely not to prove). As far as I could tell, all of the decisions made throughout the proposal were well justified and deliberate.\n\nWEAKNESSES:\nThe precision presented definitely came at the cost of readability in my opinion. The paper is very notation heavy and introduces many concepts rapidly and with great precision that lead me to having to re-read sections many times over to understand the message. It feels like the paper could benefit from summarizing some of the technical details in the main paper and reproduce the more exact version in the appendix.\n\nAdditionally, I understand the work faced space limitations; however, I believe it should definitely have the numerical experiments presented in the main paper. If pressing for pages, I could see including the experiments in the main paper and moving section 7.1 (either partially or completely) to the appendix.\n\nI found no weaknesses in the technical information itself from what I could understand.\n\nQUESTIONS:\nThe derivations showcased that optimal values of alpha range from 0 to 1, in your experience was there any intuition over what value in this range alpha should be? Or put differently, is there much difference between different values so long as they are all in this range? I am primarily interested both in terms of the resulting divergence value, but also in terms of the resulting downstream effect that such weights would have on applications.\n\nLIMITATIONS:\nThe authors very adequately discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents the way to estimate the balance weight, that represents the causal effects of arbitrary mixture of intervention, by using the neural network. Specifically, authors recognize this balance weight as the density ratio between the source and balanced distributions, and estimates this ratio by optimizing the variational representation of  $\\alpha$-divergence with $\\alpha \\in (0,1)$. In this procedure, authors justify why the $\\alpha$-divergence should be used in terms of the property of the estimator, such as the sample complexity and unbiasedness, and the vanishing gradient issue for training.\n\nSTRENGTHS:\n* This work seems to have a solid theoretical explanation for the balance weight estimation through neural network.\n\nWEAKNESSES:\nFirst of all, I confess that I am not expert in this field, and thus ask for understanding if my feedback does not make sense in this field.\n\n* The proposed has not been compared with other baselines for balance weights.\n> Authors claim that the proposed approach is a general method that can estimate balancing weight even when the datasets are generated from arbitrary mixture of discrete and continuous interventions. This means that the proposed method could be comparable with other baselines if either discrete interventions or continuous interventions exists on datasets. I believe that this comparison seems necessary because it can validate whether the proposed method can be regarded as the general method or not by showing that the proposed method is competitive with the existing methods.  However, since the proposed method was validated only for mixed variable interventions, its effectiveness on single intervention (either discrete or continuous) is not clear.\n\n \n*  The advantages of the proposed approach seems less clear as compared to the existing baseline.\n> In appendix, table 3 shows that the baseline approach (Entropy Balancing) outperforms the proposed method in most cases. Therefore, I am skeptical about why the proposed approach is meaningful. Does the proposed method have any other advantages over the baseline approach, that are not shared at current draft ?\n\nQUESTIONS:\nSee above Weaknesses.\n\nLIMITATIONS:\nSee above Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}