DIGNet: Learning Decomposed Patterns in Representation Balancing for Treatment Effect Estimation
Authors: Yiyan HUANG
WANG Siyi
Cheuk Hang LEUNG
Qi WU
Dongdong WANG
Zhixiang Huang
Venue: Accepted by TMLR ()

Abstract: Estimating treatment effects from observational data is often subject to a covariate shift problem incurred by selection bias. Recent research has sought to mitigate this problem by leveraging representation balancing methods that aim to extract balancing patterns from observational data and utilize them for outcome prediction. The underlying theoretical rationale is that minimizing the unobserved counterfactual error can be achieved through two principles: (I) reducing the risk associated with predicting factual outcomes and (II) mitigating the distributional discrepancy between the treated and controlled samples. However, an inherent trade-off between the two principles can lead to a potential loss of information useful for factual outcome predictions and, consequently, deteriorating treatment effect estimations. In this paper, we propose a novel representation balancing model, DIGNet, for treatment effect estimation. DIGNet incorporates two key components, PDIG and PPBR, which effectively mitigate the trade-off problem by improving one aforementioned principle without sacrificing the other. Specifically, PDIG captures more effective balancing patterns (Principle II) without affecting factual outcome predictions (Principle I), while PPBR enhances factual outcome prediction (Principle I) without affecting the learning of balancing patterns (Principle II). The ablation studies verify the effectiveness of PDIG and PPBR in improving treatment effect estimation, and experimental results on benchmark datasets demonstrate the superior performance of our DIGNet model compared to baseline models.

Published in Transactions on Machine Learning Research (06/2024)
DIGNet: Learning Decomposed Patterns in Representation Balancing for Treatment Effect Estimation
Yiyan Huang
vighasang appllr..duu..
Department of Applied Mathematics, The Hong Kong Polytechnic University
Siyi Wang
siyi, www.ggnngggciinn..eeuu....
School of Data Science, City University of Hong Kong
Cheuk Hang Leung
chleung87@@cityn.eeu.k
School of Data Science, City University of Hong Kong
Qi Wu
qi. ceru@iitty..dda.h
School of Data Science, City University of Hong Kong
Dongdong Wang
wangdongdong00@jj.ccmmm com
JD Digits
Zhixiang Huang
SuangzhirianggBjd.coo com
JD Digits
Reviewed on OpenReview: https: ht openrevee.. net/ pormmo Z20FInfWIm
Abstract
Estimating treatment effects from observational data is often subject to a covariate shift problem incurred by selection bias. Recent rescarch has sought to mitigate this problem by leveraging representation balancing methods that aim to extract balancing patterns from observational data and utilize them for outcome prediction. The underlying theoretical ra- tionale is that minimizing the unobserved counterfactual error can be achieved through two principles: (I) reducing the risk associated with predicting factual outcomes and (II) miti- gating the distributional discrepancy between the treated and controlled samples. However, an inherent trade-off between the two principles can lead to a potential loss of information useful for factual outcome predictions and, consequently, deteriorating treatment effect es- timations. In this paper, we propose a novel representation balancing model, DIGNet, for treatment effect estimation. DIGNet incorporates two key components, PDIG and PPBR, which effectively mitigate the trade-off problem by improving one aforementioned principle without sacrificing the other. Specifically, PDIG captures more effective balancing patterns (Principle IT) without affecting factual outcome predictions (Principle I), while PPBR en- hances factual outeome prediction (Principle I) without affecting the lcarning of balancing patterns (Principle II). The ablation studies verify the effectiveness of PDIG and PPBR in improving treatment effect estimation, and experimental results on benchmark datasets demonstrate the superior performance of our DIGNet model compared bascline models.
1 Introduction
In the context of the ubiquity of personalized decision-making causal inference has sparked a surge of research exploring causal machine learning in many disciplines, including economics and statisties (Wager &
"The co-first authors.
T'e corresponding author.

---

Published in Transactions on Machine Learning Research (06/2024)
Athey, 2018; Athey & Wager, 2019; Farrell, 2015; Chernozhukov et al., 2018; Huang et al., 2021), healthcare (Qian et al., 2021; Bica et al., 2021a;b), and commercial applications (Guo et al., 2020a;b; Chu et al., 2021). The core of causal inference is to estimate treatment effects, which is closely related to the factual outcomes (observed outcomes) and counterfactual outcomes. The concept of the counterfactual outcome is closely linked to a fundamental hypothetical question: What would the outcome be if an alternative treatment were reccived? Answering this question challenging because counterfactual outcomes are unobservable in reality, making it impossible to directly access ground-truth treatment effects from observational data. Consequently, an increasing amount of recent research has focused on developing innovative machine learning models that aim to enhance the estimation of counterfactual outcomes to obtain more accurate treatment effect estimates.
One of the challenges in estimating counterfactual outcomes lies in the covoriate shift problem. In observa- tional data, the population can be typically divided into two groups: (i) individuals who received treatment (T = 1), referred to as treated samples or treatment samples, and (ii) individuals who did not reccive treat- ment (T- = 0), referred to as controlled samples or control samples. The covariate shift problem indicates the difference between the distribution of covariate in the treated group and that in the controlled group, mcaning P(X[T- = 1) + P = 0). This phenomenon is a result of the non-random treatment assignment mechanism, where the decision to receive treatment (e.g., heart medicine) is often determined by the covari- ate (e.g., age). For example, people receiving heart medicine treatment tend to be much older compared to those who do not receive such treatment, because the doctor's 's decision- regarding whether to undergo heart medicine treatment highly depends on the patients' age. Such non-random treatment assignment is known as the selection bias phenomenon in the causal inference literature.
Although the covariate shift arises from the association between covariate and treatment, this issue can significantly exacerbate the difficulty in inferring counterfactual outcomes, as traditional machine learning models can be invalid in estimating potential outcomes when a covariate shift is present (Yao et al., 2018; Hassanpour & Greiner, 2019a). Specifically, to infer the potential outcome y for treated (T = 1) samples, the conventional approach is to first train a model (X) using controlled (T = 0) samples, and then utilize (X) to predict y for treated (T 1) samples. This approach, known as the T-learner in the causal inference literature (Curth & Van Der Schaar, 2023; Mahajan et al., 2024), becomes problematic because the training data (control samples) used for model training do not have the same distribution the test data (treated samples), i.e., P(XIT = 1) + P(XIT= =0. 0). This violates the assumption in machine learning that training data and test data should be independent and identically distributed.
To alleviate the covariate shift problem, recent advancements in representation balancing research have ex- plored the representation learning model, such as CounterFactual Regression Network (CFRNet) (Shalit et al., 2017), to estimate individual treatment effects (ITEs). These representation balancing models aim to extract balancing patterns from observational data and utilize these patterns to predict outcomes. The cor- responding objective function is typically concerned with minimizing the empirical risk of factual outcomes while concurrently minimizing the distributional distance between the trcatment and control groups in the representation space (Shalit et al., 2017; Johansson et al., 2022a). The underlying theoretical logic behind these studies is that minimizing counterfactual error can be achieved by two principles in the representa- tion space: (Principle 1) manimiziig the risk associated with factual outcome prediction, and (Principle 1I) reducing the distributional discreponcy between the treated and controlled samples.
While the representation balancing framework provides a powerful tool to address the covariate shift issue, models based on classical structures such as CFRNet (Figure I(a)) still encounter a trade-off between the aforementioned two principles: Enforcing models to focus solely on balancing can undermine the predictive power of the outcome function (Zhang et al., 2020; Assaad et al., 2021; Huang et al., 2023). A detailed discussion of this trade-off problem can be found in Appendix A.1. This inherent trade-off motivates us to explore a pivotal question: considering the inherent trade-off between the two principles, is il possible to explore scheme that enhances one principle without sacrificing the other? More specifically, can we explore improving treatment effect. estimation through the following two paths: (Path T) learning more effective balancing patterns without sacrificing factual outcome prediction and (Path II) enhancing factual outcome prediction without sacrificing the learning of balancing patterns? In the following, we present the proposed solutions and the rationale behind the underlying intuitions.
2

---

Published in Transactions on Machine Learning Research (06/2024)
In classic representation balancing models, the process of learning balancing patterns can lead to the loss of silcome-predicttvg information. It. is therefore natural to consider module that can preserve the pre- balancing information before the representation balancing step. This increases the model's complexity to maintain the useful prodictive knowlodge while still benefiting from the covariate balancing properties of the representation balancing framework. Furthermore, multi-lask learning, distinet representations are learned for different tasks, with each task involving its own objective function. An important step in multi- task learning is integrating the information from these separately learned representations into a unified representation. Therefore, following the multi-task learning paradigm (Li et al., 2018; Baltrušaitis et al., 2018; Crawshaw, 2020; Yan et al., 2021; Xu et al., 2023b), we propose concatenating the representations learned with Wasserstein distance and H-divergence to form a joint representation. This joint representation can effectively capture the lask-specific balancing information provided by each distance metric without adversely affecting the outcome modeling task. A detailed discussion of these intuitions can be found in Section 5.
Based on the above motivations, we introduce a novel representation balancing model, DIGNet (Section 5.2.2), neural network that incorporates Decomposed pallerns with Individual propensily confusion and Group distance minimization. Decomposed patterns denote distinct components disentangled from specific reprosentations in DIGNet (Section 5.2). The individual propensity confusion aspect of DIGNet aims to learn representations that obscure the propensity of each individual being treated or controlled (Section 5.1.2), grounded in our derived H-divergence guided counterfactual and ITE error bounds (Section 4.2). The group distance minimization aspect focuses on learning representations that minimize the distance between treated and controlled groups (Section 5.1.1), supported by previous work on Wasserstein distance guided counterfactual and ITE error bounds (Shalit et al., 2017) (Section 4.1). Figure 1 visually depicts these introduced concepts and their relationships.
Contributions. Our main contributions are summarized follows:
We derive theoretical upper bounds for counterfactual error and ITE error based on H-divergence (Soction 4.2). In particular, this thcoretical foundation highlights the important role of propensity score for representation balancing models, connecting the representation balancing with the concept of individual propensity confusion.
2. We suggest learning decomposed patterns in representation balancing models (Section 5.2.1) to miti- gate the trade-off problem rooted in classic causal reprosentation balancing models. First, we propose a PDIG method (Figure (bb), which aims to learn Patterns Decomposed with Individual propen- sity confusion and Group distance minimization to improve treatment effect estimation through Path 1. Second, we propose a PPBR method (Figure I(c)), which aims learn Patterns of Pre-balancing and Balancing Representations to improve treatment effect estimation through Path II.
3. Building upon PDIG and PPBR, we propose a novel representation balancing model, DIGNet (Figure I(d)), for treatment effect estimation. In Section 6, ablation studies verify the efficacy of PDIG and PPBR in improving ITE estimation through Path I and Path 11, respectively. Furthermore, experimental results on benchmark datasets demonstrate that DIGNet surpasses the performance of baseline models in terms of treatment effect estimation.
2 Related Work
The presence of a covariate shift problem stimulates the line of representation balancing works (Johansson et al., 2016; Shalit et al, 2017; Johansson et al., 2022a). These works aim to balance the distributions of reprosentations between treated and controlled groups and simultancously try to maintain representations predictive of factual outcomes. This idea is closely connected with domain adaptation. In particular, the ITE error bound based on Wasserstein distance is similar to the generalization bound in Ben-David et al. (2010); Long et al. (2014); Shen et al. (2018). The thcoretical foundation and the classic CFRNet structure proposed in Shalit al. (2017) have inspired many subsequent studies on representation balancing methods for treatment effect estimation, including Yao et al. (2018); Shi et al. (2019); Zhang et al. (2020); Hassanpour
3

---

Published in Transactions on Machine Learning Research (06/2024)
Balanging
Balancing
1(9g. 1)
A(
M
hf g0 0)
D
h( 0:
D
M(000)
(O.0)
Balancing
Balancing
Balancing
Balancing
(a) Classic model
(b) PDIG (yellow)
(c) PPBR (yellow)
(d) DIGNet
Figure 1: (a): The classic model (e.g., GNet in Section 5.1.1 and INet. in Section 5.1.2) maps the original data D into representations Dg to achieve representation balancing. The balanced representations are referred to as balancing patterns. These balancing patterns are also used for outcome prediction. (b): The PDIG (Section 5.2.1) is illustrated as the yellow part, where balancing patterns are decomposed into two distinct components, & and oT- &e serves for group distance minimization (Section 5.1.1) and $, serves for individual propensity confusion (Section 5.1.2). The balancing patterns &g and 1 are concatenated for predicting outcomes. (c): The PPBR (Section 5.2.1) is represented by the yellow section, where & is used for feature extraction and & is used for representation balancing. Here representations are decomposed into pre-balancing patterns &g and balancing patterns &G. &g and &G are concatenated for predicting outcomes. (d): The proposed model DIGNet (Section 5.2.2) integrates both PDIG and PPBR. Specifically, DIGNet decomposes balancing patterns into two distinct components, &g and 2y. The outcome predictors are further formed by concatenating &e, 2 and pre- nrebbiann patterns &g.
& Greiner (2019a); Assaad et al. (2021); Huang et al. (2022a). This paper derives a new ITE error bound based on H-divergence (Ben-David et al., 2006; 2010; Ganin et al., 2016). In addition to the connection to domain adaptation, causal representation learning is also linked to the field of fair representation learning, which aims to ensure that machine learning algorithms make fair decisions by learning fair representations. The main goal of these studies is to enforce a classification model to be less sensitive to certain sensitive variables when the representations & Storkey, 2015; Beutel et al., 2017; Madras et al., 2018; Zhang et al., 2018; Adel et al., 2019; Feng et al., 2019; Zhao al., 2019a; Zhao Gordon, 2022). Notably, the original idea of adversarial learned fair reprosentations in Edwards & Storkey (2015) is also motived by the domain adaptation work (Ben-David et al., 2006; 2010; Ganin et al., 2016), sharing similar motivation to our utilization of INet, which relies on H-divergence guided error bounds for ITE estimation. Moreover, Wasserstein distance has also been employed for learning fair representations in Jiang et al. (2020).
different groups are sufficiently similar (Zemel et al., 2013; Edwards
Another recent line of causal representation learning literature investigates efficient neural network structures for treatment effect estimation. Kuang et al. (2017); Hassanpour & Greiner (2019b) extract the original co- variates into speetii factors, outcome-specific factors, and confounding factors; X-learner (Künzel et al., 2019) and R-learner (Nie & Wager, 2021) are developed beyond the classic S-learner and T-lcarner; Curth & van der Schaar (2021) leverage structures for end-to-end learners to counteract the inductive bias towards treatment effect estimation, which is motivated by Makar et al. (2020). There are some other deep neural network models that have been employed in treatment effect estimation Louizos et al. (2017); Yao et al. (2018); Yoon et al. (2018); Shi et al. (2019); Du et al. (2021). To ensure comparability and consistency, we rigorously follow the same framework as these cansal inference works. The causal graph these studies satisfies the standard setup T + X + Y and T + Y. Additionally, it is also wort noting that there are many other causal inference works exploring treatment effect estimation under more complex causal graphs. For instance, studies such Kallus et al. (2019); Jesson et al. (2021); Miao et al. (2023) specifically tackle the treatment effect estimation when unobserved confounders U present. In this case, the causal graph sctup exlends to T + X + Y, T - Y, T + U + Y. A recent work (Cao el al., 2023) further expands this static causal graph to a dynamic setting. Moreover, some studies such as Angrist et al. (1996); Burgess et al. (2017); Wu et al. (2022); Yuan et al. (2023) estimate treatment effects with instrumental variables I involved. In this case, there are various causal graph setups such as T + X + Y, I + T + Y, and T <- 7 Y. More complex causal graph settings (Nogueira et al., 2021; Vowels al., 2022; Zanga et al.,
4

---

Published in Transactions on Machine Learning Research (06/2024)
2022) have been studied with the development of Directed Graphical Models (Pearl, 2009), which represents another significant research direction known as causal discovery.
Our method is highly motivated by the trade-off problem between outcome prediction and representation balancing. In the causal representation learning literature, a similar trade-off phenomenon has been noticed by Zhang et al. (2020); Assaad et al. (2021); Huang et al. (2022a), where the researchers argue that highly- balanced representations can have adverse effects on outcome modeling. However, the explanations for this phenomenon and its connections with other related literature are not extensively provided in their work. We highlight that the trade-off between outcome prediction and representation balancing is also connected with tradc-offs observed in other rescarch domains. In representation balancing models, representation bal- ancing helps improve the model's ability to generalize to counterfactual estimates. However, representation balancing can potentially sacrifice information necessary for predicting factual outcomes. In supervised ma- chine learning, penalizing model complexity during model training helps the model to learn simpler patterns, thereby promoting generalization ability (reducing its variance) to unseon data. However, a bias-variance trade-off occurs because less flexible models tend to exhibit, higher bias in training data (Geman et al., 1992; Domingos, 2000; Valentini & Dietterich, 2004; Yang et al., 2020). In the literature of domain adap- tation (Shen et al., 2018; Zhao et al., 2019b), transfer learning (Long et al., 2015; 2017; Ma et al., 2023), out-of-distribution detection (Kumar et al., 2021; 2022), and fair representation learning (Zliobaite, 2015; Hardt et al., 2016), enforcing a model to capture proxy features that are domain-invariant helps the model to generalize well to unseen target (also known as out-of-distribution) data. However, a trade-off between classification accuracy and aamaii-invaaiiane (or fairness in fair representation learning literature) occurs because the pursuit of domain-invariant features may lead to a loss of classification accuracy on the source (also known as in-distribution) data (Zhao et al., 2019a; Zhao & Gordon, 2022; Zhao ct al., 2022).
3 Preliminaries
Notations. Suppose there are N i..d. random variables D (XX, Ti, YOJ with observed realizations
{(x,, f,. yi) m where there are N1 treated units and No controlled units. For cach unit i, X; E X C Rd denotes d-dimensional covariates and T. E {0, 1} denotes the binary treatment, with e(x,) == p(T; = X; x;) defined as the propensity score (Rosenbaum & Rubin, 1983). Potential outcome framework (Rubin, 2005) defines the potential outcomes y1, Y E y C R for treatment = and T = 0, respectively. We let the observed outcome (factual outcome) be Y = T- T- + outcome) be Y T y + (1 T) y1. For (1-T))-Y" TY,Y,, and the unobserved outcome (counterfactual
(0, let 7(x)) [ X x] be function of Y w.r.t. X, then our goal is to estimate the individual treatment effect (TTE) 7(x) = E - -
(x) - (x) 1, and the average Treatment effect (ATE) TATE = E 1 introduced concepts PPBR and PDIG are illustrated in Figure 1, and the necessary representation functions gg, 2g and 1 as well as different model structures, are illustrated in Figure 2. Throughout the paper, we refer to patterns as meaningful representations. For instance, decomposed patterns are distinct components disentangled from some specific representations.
= fxx (x))(xxxxx The
3.1 Problem setup
In causal representation balancing works, we denote representation space by R C Rd, and $ : X - R is assumed to be a diiee-iiffeeentiable, one-to-one and invertible function with its inverse V : R + X' such that
(((())) = 1. The densities of the treated and controlled covariates are denoted by p I= -==1(() (x) == p(x T = 1) and p == = (x"= =- p(x T = 0), respectively. Correspondingly, the densities of the treated (r
and controlled covariates in the representation space are denoted by pT- p/ (r) : (r T 1) and p p 0 =p (r) (r T = 0), respectively.
'The term Y r x = x is commonly known as the Conditional Average Treatment Effect (CATE). In order to
maintain consistency with the notion used in the existing causal representation balancing literature, egg, Shalit et al. (2017), refer to this term as ITE throughout this paper. Note that the original definition of ITE for the i-th individull is commonly expressed the difference between their potential outcomes, represented as Y, YP
Theoretically, the invertibility is necessary for deriving the upper hounds of ITE error, specifically for equation 39 and equation 47. However, the invertibility can be hard verify in practice (Johansson et al., 2022b).
5

---

Published in Transactions on Machine Learning Research (06/2024)
Our study is based on the potential outcome framework (Rubin, 2005). Assumption states standard and necessary assumptions to ensure treatment effects are identifiable. Before proceeding with theoretical analysis, we also present some necessary terms and definitions in Definition 1.
Assumption 1 (Consistency, Overlap, and Unconfoundedness). Consistency: If the treatment is t, then the observed oulcome equals Y Overlop: The propensily score is bounded away from O to 1, i.e, 0 < e(x) < 1. Unconfoundedness: Y 11 T X, Vt E {0.
Definition 1. Let h R {0,1} + y be hypothesis defined over the representation space R such that h()()),t) estimates y", and L y x y + R. be loss function (e.g., the squared loss L(y. v') (y - y'y² or the absolute loss L(y. v/) = by v.. If we define the expected loss for (x, t) as (x, = fy L(y", ((((x),t)))(C (app(y(())dd (x)y', we then have factual and counterfactual errors, as well as them on the treated and controlled:
ECF(h,4) = JEx (x,(xxt)p(x 1 t)dxdt,
EF(,, = xx,1011 (x t)p(x, t)dxdt,
1 ln,o(x, (oo(xx1)p"==11 1xp"
S l,,0(x, 0)p""=<x))=)
(h,
(h,
T= () / (x, (x)xxxxxxxx
f (n,
(,, b)
If we let f(x,t) be h(o(x),)), t), where f X x {0.1} + y is a prediction function for outcome, then the estimated ITE over I defined as (x(x) == f(x, 1) f(x, 0). We can measure the error in ITE estimation with the metric, Precision in the expected Estimation of Heterogeneous Effect (PEHE):
(PEH J (()))))
(1)
Here, (fEnE(f) can also be denoted by <PEHE(h, +) if we let f(x,t) be h(((x),
4 Theoretical Results
In this section, we first prove epene is bounded by eF and ECF in Lemma 1. Next, we revisit the upper bound for Wasserstein distance guided representation balancing method in Section 4.1. Furthermore, we state the new theoretical results concerning H-divergence guided representation balancing method in Section 4.2.
Lemma 1. Let functions h and & be defined in Definition 1. Recall that r'(x) = E E[ X= x). Define o = minfo? (p(x,t)), t)).
and Ag = max(Ay t)), Ay (p(x, +))} vt E (0,1), Ixxx0x1xx1(1* - (x))) (x)()xxxx)) t)dy dxdt and Ag (p(x,t)) Sx ly
where of ()(,
(x)))p(()') (x)p(x t)dy' dxdt Vt E 10.
Let loss function be the squared loss. Then we have:
EPEHE (h, 4) S 2(ecr(h, ②) + (er(,, -2
(2)
Let loss function be the absolute loss. Then we have:
EPERE( 4) S ecr(h, ②) + er(h, ②) +AA
(3)
Lemma 1 reveals that the ITE error EPEHE is closely connected with the factual error eF and counterfactual CCF, as well as a constant of (or Ay) that is unrelated with functions h and T. Here, of is the smaller value of the variance in Y w.r.t. the distribution p(x,)) () and the variance in y1 w.r.t. p(x, - (), and A, is the larger value the absolute deviation in Y w.r.t the distribution p(x, t) and the absolute deviation in y w.r.t. the distribution p(x, 1 t). The proof of Lemma 1 is deferred to Scction A.2. Note that equation (2) corresponds to the result presented in Shalit et al. (2017), while equation (3) is our new result, which supplements the when T denotes the absolute loss.

---

Published in Transactions on Machine Learning Research (06/2024)
4.1 Wasserstein Distance Guided Error Bounds
Previous causal learning models commonly adopt the Wasserstein distance guided approach to seek rep- resentation balancing, In this subsection, we first give the definition of Wasserstein distance (Cuturi & Doucot, 2014) by introducing the Integral Probability Metric (IPM) (Sriperumbudur et al., 2012) defined in Definition 2. Then we state the theorem regarding the upper bounds for counterfactual error ECF and ITE error <PEHE using Wasserstein distance Theorem 1.
- R. For a pair of distributions pt, Definition Let G be function family consisting of functions g : S P2 over S, the Integral Probability Metric is defined as
IPM P2) = sup I (())(1(()) p2(s))
gEg
II G is the lamily of 1-Lipschitz functions, we can obtain the so-called I-Wasserstein distance, denoted by Wass(pi,p) p2). Next, we present the bounds for counterfactual error ecr and ITE error OPEHE using Wasserstein distance in Theorem
Theorem 1. Let
X R be an invertible representation with & being its inverse. Defone of = minfoy ()) ofx 1))} and Ay = max max (pxx(xxxx ()), Ay (p(x, ())} AI {0,1} where (p(x,)) = (xx - (x))) t)dy' dxdt and (p(x,)) = Sxx les (x))p((y' (x)xxxx)pxxx))xxx)x)(( x))((xt)dyyde dxdtt Vt 10. Let (r), pd (r) be as defined before, h : R {0, 13 y, u: Pr(T =1) and G be the family of I-Lipschitz functions. Assume there cxists constant B& 2 0, such that for 1 E {0,1}, the function (r, == I. - oV,,(V( EG Given loss function L,, we have
(cr 1) <(1 - u) o (h,
(h, b) + tuuc oF (,, 1) + Bg W Wass(pp T=1 T=0).
P4 ,p4
(4)
Let loss function be the squared loss. Then we have:
EPEHE(h, D) 2 ( 4) + o (h, Ba Wass(pp T=1 ,p'
-00=- 20
(5)
Let loss function L be the absolute loss. Then have:
EPEnE(h, 4) S (h, 1) + - 1) + BB W Wass(p' T=1 -b==0) + 2Ag
(6)
Theorem reveals that the ITE error is closely tied to the factual error eF and the Wasserstein distance between treated and controlled groups in the representation space. This theorem provides a theoretical foundation for representation balancing models based on group distance minimization (Section 5.1.1). The proof of Theorem is deferred to Section A.3. Note that equation (5) corresponds to the result presented in Shalit et al. (2017), while equation (6) our new result, which supplements the case when L denotes the absolute loss.
4.2 H-divergence Guided Error Bounds
In most representation balancing literature, the models mainly rely on Wasserstein distance guided error bounds as discussed in Section 4.1. In this subsection, we will focus on establishing H-divergence guided error bounds for counterfactual and ITE estimations in representation balancing approach. We first give the definition of H-divergence (Ben-David et al., 2006) in Definition 3. Then we state the theorem regarding the upper bounds for counterfactual error ECF and ITE error ÉPEHE using H-divergence in Theorem 2.
Definition 3. Given a pair of distributions PL, P2 over S, and a hypothesis binary function class H. the
H- -idivvereence between p1 and p2 is defined as
dx((p)p2) = 2 sup (P))) 1] Prppln((s 1]1 MEH
Theorem 2. Let X R be an invertible representation with v being its inverse. Define oy = min (p(x, 1)),07 oz (p(x1 f))} and Ay = maxx(A (p(x, 0)), Ay (x(x, 1)) Vt. (0,1),

---

Published in Transactions on Machine Learning Research (06/2024)
where (p(x,))) Sxx1o.))xxll' - x')))) xx)((x, Jddydxxdt and Ag (p(x,t)) fxx /y ('))((( |x)p(x, d)dy' Vt E {0,1}. Let (r), (r) be as defined before, h : R x {0,1} + y, u == Pr(T = 1) and H be the family of binary functions. Assume that there cxists constand K 2 0 such that Jy Liy. S K Vy E V. Given a loss function L, we have
( 1) 5(1 - u) .F T- (1,, (), (0,)))
u- e (h, 4) + dd T=1 PP4 0)
(8)
Let loss function be the squared loss. Then we have:
K T=1
CPEHE(,, 3) S ( + (h, dx(p
PP -p4-00 -20
(9)
Let loss function L be the absolute loss. Then have:
K εPEHE(h, @) S T + (h, +)+
+
2 du(pb
d Pp ) 2Ay
(10)
Theorem reveals that the ITE error is closcly connected with the factual error op and the H-divergence between treated and controlled samples in the represenlation space. This new theoretical result provides theoretical foundation for representation balancing models based individual propensity confusion (Section 5.1.2). The proof of Theorem 2 is deferred to Section A.4.
5
Method
In the preceding section, we have stated the theoretical foundations for representation balancing methods, which are the Wasserstein distance guided error bounds (results in Shalit et al. (2017)) and H-divergence guided error bounds (Our results). Moving on to Section 5.1, we will begin by introducing representation balancing methods without decomposed patterns. Specifically, Scction 5.1.1 revisits a Wasscrstein distance based representation balancing network GNet, and Section 5.1.2 demonnstrate how Theorem 2 can be con- nected with individual propensity confusion, helping us to build a H-divergence based representation balanc- ing network INct. Subsequently, in Section 5.2, will introduce how to design representation balancing method within the scheme of decomposed patterns, based the PDIG and PPBR methods (Section 5.2.1). The final proposed model DIGNet. is presented in Section 5.2.2.
5.1 Representation Balancing without Decomposed Patterns
In representation balancing models, given the input data tuples (x,t,y) (x,,tt,, covariates x are extracted by some representation function Φ(-), and representations ((x) are then fed into the outcome functions h (+):= h(., 1) and (-)- == h(-:0) that estimate the potential outcome y' and yh respectively. Finally, the factual outcome can be predicted by 1²(()) = ())()) + (1 t)))(-), and the corresponding outcome loss
t,, w)) the original
N ((x)))
C, 1x
(11)
The loss function by approximates the factual error er appeared Theorems and 2. Minimizing by also corresponds to the Principle mentioned in the Introduction.
5.1.1
GNet: Group Distance Minimization Guided Network
The group distance minimization focuses on learning representations that minimize the distance between the treated and controlled groups, and the corresponding thooretical foundation supported by Wasscrstein distance guided counterfactual and ITE error bounds (Theorem 1). Previous causal inference methods (c.g., Shalit et al. (2017); Yao et al. (2018); Zhang et al. (2020); Huang et al. (2022a)) commonly adopt Wasserstein
8

---

Published in Transactions on Machine Learning Research (06/2024)
distance to achieve group distance minimization. Specifically, these methods aim to minimize the empirical approximation of La(x, t; b) W Wass( (ass(((( (() (o(x)))tt;=)) to learn balancing patterns. If we denote (-(-) by the feature extractor that extracts the original covariates x, then the objective function designed on Theorem is
min $g
h²) + (x t;
(12)
Since the objective is to learn balancing patterns by minimizing the distributional distance between treated and controlled groups, i.e, group distance minimization, we refer to model with the objective in equation (12) as GNet. For the reader's convenience, we illustrate the structure of GNet in Figure 2(a). Note that CFRNet (Shalit et al., 2017) is also the category of GNet.
5.1.2 INet: Individual Propensity Confusion Guided Network
In the field of causal inference, the propensity score plays central role because it characterizes the probability that one receives treatment (Rosenbaum & Rubin, 1983). For example, the propensity score has been widely employed in prior literature for matching (Caliendo & Kopeinig, 2008) or weighting (Austin & Stuart, 2015) purposes. In this paper, we emphasize that the propensity score also plays an important role in representation balancing, where serves as a natural indicator of the adequacy of leanred balancing patterns. Specifically, we propose the concept of individual propensity confusion, which aims to learn representations that are difficult. to utilize for characterizing the propensity of each individual being treated or controlled. The underlying theoretical foundation is upon the H-divergence guided ITE error bounds derived in Theorem 2. Specifically, equations 9 and 10 in Theorem 2 highlight the significance of minimizing the generalization bound associated with factual outcome error and the H-divergence between treated and controlled representations in reducing ITE errors. Subsequently, we will present the details of achieving representation balancing by reducing the H-divergence between treated and controlled samples in the representation space.
Let (a) be indicator function that gives if a is true, and H. be the family of binary functions defined in Theorem 2. To achieve representation balancing, the objective function designed on Theorem 2 should aim to minimize the empirical H-divergence dr (p -p such that
1|t, = 0| 0)
du(p
T= 2 T-1 Ps
min nEH
Σ 1|4
Σ (x)()(x)))==
(13)
N
N
(x))==
The "min" part in equation (13) indicates that the optimal classifier 1f" E minimizes the classification error between the estimated treatment v)" (*x()) and the observed treatment 4, i.e., discriminating whether (xx;) is control (T = 0) or treatment (T= 1). Equation (13) suggests that du (pb pg =T=0) will be large if 1f" can easily distinguish whether ((x:) is treated or controlled. In contrast, dx(2g T= .p'=0) will be small if it is hard for 17" to determine whether 2(x;) is treated or controlled. Therefore, the prerequisite of a small H-divergence to find map & such that any classifier 17 E H. will get confused about the probability of (xx,) being treated or controlled. To achieve this goal, similar the strategy of empirical approximation of H-divergence (Ganin et al., 2016), we define a discriminator æ(r): R [0, 1] that estimates the propensity score of r, which can be regarded a surrogate for (/r). The classification error for the 7th individual can be empirically approximated by the cross-entropy loss between T(((x;)) and t;:
Le(ts, ((tx)) = - [t, log (xx))) + (1 t.) log -1((())))
(14)
As a consequence, we aim to find an optimal discriminator n" for equation (13) such that xt maximizes the probability that treatment is correctly classified:
1 N
max MEH
max WEH
Σ c.tt, c,(,,,
(15)
N
Given the feature extractor $ (.), the objective of INet can be formulated as a min-max game:
min max ty, h²) (g.
(16)
$ kl
9

---

Published in Transactions on Machine Learning Research (06/2024)
In the maximization, the discriminator is trained to maximize the probability that treatment is correctly classified. This forces <x;)(xx)) closer to the true propensity score e(x;). In the minimization, the feature extractor gg is trained to fool the discriminator л. This confuses т such that (x.) cannot correctly specify the true propensity score e(x;). Eventually, the representations are balanced as the adversarial process makes it difficult for и to determine the propensity of each individual being treated or controlled. We refer to this process as individual propensity confusion. Such an adversarial learning technique has been widely used in domain adaptation (e.g,, Ganin al. (2016); Tzeng al. (2017)) and fair representation learning and fair representations.
(e.g.: Edwards & Storkey (2015); Madras et al. (2018)) to learn For the reader's convenience, we illustrate the structure of INet in Figure 2(b).
5.2 Representation Balancing with Decomposed Patterns
5.2.1 The Proposed PDIG and PPBR Methods
PDIG. Although Theorems 1 and 2 provide solid theoretical foundation for GNet (model proposed by Shalit et al. (2017) and INet (model proposed by us); both of these model types still encounter the inherent trade-off between representation balancing and outcome modeling. To this end, we expect to capture more effective balancing patterns by learning Patterns Decomposed with Individual propensity confusion and Group distance minimization (PDIG). More specifically, the covariates are extracted by the feature extractor & (g(-), and then (g(x) are fed into two distinct balancing networks Pc(-) and Φ7(-) for group distance minimization and individual propensily confusion, respectively. In summary, PDIG decomposes the balancing patterns into two distinct parts, group distance minimization and individual propensity confusion, which are respectively achieved by the following loss functions:
min Lo(x, t; 2 2g)
(17)
26
min max (x DE,T)
(18)
Here, denotes the composition of two functions, indicating that ((.) in Lo(x, t; b) and Lr(x, t; d x)) are replaced by Pa((e(()) and ()){x(-)), respectively.
PPBR. Motivated by the discussion in Section 1, we design a framework that is capable of capturing Patterns of Pre-balancing and Balancing Representations (PPBR) to mitigate potential over-balancing issue mentioned in the Introduction, aiming preserve information that is useful for outcome predictions. In the PPBR method, the balancing patterns (c((g(x)) and ( are first learned over &G and #1, while 0g is remained fixed as pre-balancing patterns. Furthermore, we concatenate the balancing patterns & (x)() and 4 (x)))xx) with the pre-balancing representations & (x) as attributes for outcome prediction. As result, the proxy features used for outcome predictions are e (x(()()))))))))))))))))(((( (x) (Q where indicates the concatenation by column. For example, if a [1, 2] and b [3, 4), then a $ b [1,2, 3, 4]. Consequently, representation balancing accomplished over &G and Φ,, rather than &g. Even if there may be a loss of information relevant to outcome prediction in 4G and Ø,, the pre-balancing patterns ΦE can still effectively preserve such information. Finally, the objective function with regard to outcome modeling under PPER method becomes
min 01.01.00MM
Ly(x, yx; D (Φ, +E) @ (oo 1)),
(19)
5.2.2 The Proposed DIGNet
Combining with PDIG and PPBR, we propose a new neural Network model that incorporates Decomposed patterns with Individual propensity confusion and Group distance mininization, which we call DIGNet.
10

---

Published in Transactions on Machine Learning Research (06/2024)
Colx, t; 4o o o
b( I)
A(@.1)
A(B.1)
by k)
1
D 00
k(
4P
K[0,0
D-
((
Colx, t: PP)) (a) GNet (CFR-Wass)
Ly t: xp)x)
Ca(x, t: 9g)
1x (a.R)
Cylx, t #a.5)
(b) INet
(c) DGNet
(d) DINet
(e) DIGNet
Figure 2: Illustaations of the network architecture of the five models studied in Section 6.
The objective of DIGNet is separated into four stages:
min an t: to o dE),
(20)
tG
max (x; o 0 1g,
(21)
min 4,
(xL Fg.
(22)
min bc. Cy(x,t,y: D (0) 48) 2 (#6 #g),
(23)
Within each iteration, DIGNet minimizes the group distance through equation 20, and plays an adversarial game to achieve propensity confusion through equation 21 and equation 22. In equation 23, DIGNet updates both the pre-balancing patterns &E and balancing patterns Φ,, la. along with the outcome function h to
minimize the outcome prediction loss. For the reader's convenience, we illustrate the structure of DIGNet in Figure 2(e).
5.3 Insights of Representation Balancing with Decomposed Patterns
Our proposed DIGNet model builds upon the PDIG and PPBR methods. The PPBR method is relatively straightforward, as it forms more flexible predictor ( @ (Ø, o +E)) (or (o @ ( +E))) compared to the solely predictor &g. Therefore, incorporating both pre-balancing and balancing patterns is helpful in enhancing the model's complexity and its ability to capture more useful information for outcome prediction. However, there still remains further exploration to better understand why the PDIG method is effective. The DIGNet model aims to learn balancing patterns based on both Wasserstein distance and H-divergence. At first glance, one might assume that incorporating both distances could be redundant, as one distance seems naturally to imply the other. In this section, we gain some insights of these two divergence metrics. First, we provide systematic discussion addition, utilize toy example to illustrate their distinct abilitics in capturing distributional disparity. Further, use this example to aid readers in better understanding the trade-off problem encountered in representation balancing models (Figure 5). Finally, we establish a connection between our method and the Elastic Net method and Multi-task learning approach, which offers valuable insights and explanations regarding the intuition behind involving both metrics as regularizations.
the properties of Wasserstein distance and H-diveegence. Tn
5.3.1 Properties of Wasserstein Distance and H- P--Divergence
Wasserstein distance and H-Divergence possess distinct theoretical properties. The effectiveness of the Wasserstein distance in measuring distributional differences for classification tasks in domain adaptation has been demonstrated in Shen et al. (2018). Furthermore, Shalit al. (2017) highlights the potential of Wasserstein distance in representation balancing models for ITE estimation, which significantly outperforms traditional ITE estimation methods. Wasserstein distance is also widely adopted in other research domains, such as fair representation learning (Jiang et al., 2020), as discussed in Section Its prevalence stems from strong capability to capture better diversities compared to H-Divergence (Shui et al., 2020). Studies have proven that under certain conditions, it is possible to bound H-Divergence using Wasserstein distance (Villani et al., 2009; Shui et al., 2020), which provides a reasonable explanation for the overall superiority of
11

---

Published in Transactions on Machine Learning Research (06/2024)
3.
3.0
3.0
25-
2.5
2.5
P100
P.6g
1.5
1.0
1.0
1.0
0.5
0.5
0.25
0.50
0.75
1.00
0.23
0.75
1.00
Figure 3: Distributtins of (1(),, (2(2), and p2(r) in the example of Section 5.3.1.
Treated
Controlled
Treated
Controlled
35
Estimated
3.-
Estimated
Estimated
15-
Estimased
30
30-
16-
11-
12-
13-
15-
00
24
(1
(6x)
(a) Model fitting using (X,Y). The probability density functions of Xtreat and Xeontrol are Pı and Ps: respec- tively.
(b) Model fitting using (o(X), Y). The probability den-
sity functions of grreat (X) and o ccoutrol (X) are Pi and p2, respectively.
Figure 4: Model fitting using (X, and (.() based on the example in Section 5.3.1.
the Wasserstein distance in learning domain-invariant features (Zhiri et al., 2022). However, is important to note that this bound does not hold in general (Chae & Walker, 2020), suggesting that smaller H-divergence does not necessarily imply a smaller Wasserstein distance. To better illustrate the difference between these two measures, we provide a concrete example below.
Toy example. Consider the following three probability density functions (1(),, P2(x)), and ps(x) defined over [0, 1]:
2.5, < 0.25
0.5, <0<x<<0.>
0.5, if < 2 0.25
0.5, if25 < < 0.5 Pe(x)
2.5, 0f x 0.5
0.5, if <
(1(x)
0.5, 5.775 Pa(x)
0.5, 0.5 0.75
2.5, <
0.5,
0.5, 1.0.75
0.5, 0.75
The above three distributions are depicted in Figure 3. Further, we set the classifier H-divergence 7²)) 1 2 p}, and set the order in Wasserstein distance p 1. By utilizing the definitions of H-
divergence and 1- L-Wasserstein distance, one can make direct comparison between the discrepancy in (pr, P2) and the discrepancy in Pa):
dx(pi, P2) = du((i, Ps): W p2)))p),p2) < Wass(p1. (p1.
(24)
Equation 24 confirms that Wasscrstein distance is able to capture more diverse distributional disparitics compared H-divergence. However, in the subsequent content, we will demonstrate that such an advantage might be limitation in causal representation learning due to the trade-off problem.
12

---

Published in Transactions on Machine Learning Research (06/2024)
Understanding the trade-off. The above example serves as simple evidence that supports the conclusion that Wasserstein distance can capture better diversities between distributions compared to H-divergence (Shui et al., 2020). However, as discussed in Section 1, it is important to note that achieving a more balanced distribution does not necessarily ensure favorable generalization to counterfactuals. This is because the pursuit of balanced representations may inadvertently lead to a loss of information useful for factual outcome estimates. We will now use the above example to gain further understanding on this matter.
Consider simple data-generating process where X the covariate in the treated group, follows the distribution (z),, and Xcontrol, the covariate in the controlled group, follows the distribution pz(x). Let
the potential outcomes are generated by y (X) + and y = Tx(X + co, where 61 (V 0.1) and 6o ~ (0,,0, Let the true potential outcome functions T² and (z) be as follows:
- (x) = (x² + 1)1(0 x < 0.53 + (4r (055)
(25)
In addition, consider representation function (xx) such that
(x)) = ((t)==1(000) 0.25) + (2 0.25) < x < 0.5) + (z-0..5)1 < 0.75) + x1 2 S 1].
(26)
We can find achieves representation balancing under Wasserstein distance measure, but does not under H-divergence measure. In original data, wrtratt follows pl and conttol follows P3- After mapping to ((x),
(treat(x)) follows p1 and (control( follows p2. Consequently, based on the results in equation 24, we have
dx trent ..gontrol) = dx(pX treat .PX :pxutroo); .P
(27)
W Wass(pl (plres Pe control) < W Wass(pxxe ,pX poontrol).
We now investigate the fitting performance of models using y) and ($(x). y) to check whether there is a loss of outcome-related information during representation balancing. In Figure 4a and Figure 4b, we present scatter plots of samples from (x. y) and (f(x). y) respectively, depicted as gray points. Following the approach of Kennedy (2023), we employ smoothing spline functions to fit these samples, and the estimated functions are illustrated in blue.
In Figure 4a, we observe that both Ti and To are well fitted using (I,w), with their estimates being very close to each other. This is consistent with the setup of T₁ = To. In contrast, Figure 4b reveals that the fittings of Ti and To are inadequate using (#(z), v), resulting in substantially different estimates. The result of different estimates violates the setup of Ti = To- In this case, a model based on Wasserstein distance would retain & due to its achievement of representation balancing. Unfortunately, & suffers from a loss of valuable information that is crucial for outcome prediction. In contrast, model based on H-divergence would not keep & since it does not contribute to roducing the domain distance compared to the original data. Fortunately, the original data preserve the information necessary for outcome modeling. Therefore, this example not only emphasizes the significance of incorporating both metrics but also highlights the importance of considering both pre-balancing patterns and balancing patterns.
5.3.2 Connection with other machine learning methods
In the previous sections, we have discussed the trade-off between factual outcome prediction and represen- tation balancing in classic representation learning models. As part of our proposed improvements, DIGNet involves learning two distinct representations using Wasserstein distance and H-divergence separately and concatenates the learned representations for outcome modcling. In this section, we will explore more detailed connections between our design and other machine learning methods.
Connection with Elastic Net: balancing on two discrepancies. Our DIGNet model involves two discrepancy metrics: Wasserstein distance and H-divergence. We will now provide additional explanations on its connection with the Elastic Net method. In supervised lcarning, regularization term often incor- porated during model training to mitigate the hlas-vrrianee trade-off. In the case of linear regression, Lasso (Tibshirani, 1996) and Ridge (Hoerl& & Kennard, 1970) are proposed to improve the Ordinary Least Squares
13

---

Published in Transactions on Machine Learning Research (06/2024)
(OLS) method, with Lasso involving 1 regularization while Ridge involving l2 regularization:
N
d -w)² +a
Lasso:
min BER
lly x,31 +a11811 min BER
(x13
Σ
|8
N
N
N
d
Ridge:
min BER N
by x + a11311
min BER
Σ 1m1
[x m1)² + Σ 33.
N
The different properties between 1, rogularization and l regularization lead to distinct advantages and disadvantages between Lasso method and Ridge method. Given their differences, a method of Elastic Net (Zou & Hastie, 2005) is proposed by combing both I regularization and l2 regularization:
N Σ x y1)² + 01 BER
1
Elastic Net:
min BER
by x 2 a || + 02 11111
min
18g + 02
83.
The Elastic Net method integrates the strengths of two distinct approaches: the 4 regularization term enforces sparsity, while the la regularization maintains the grouping effect (Zhou, 2013; Narisctty, 2020). The Elastic Net has also mot ivateeded some research studies to adopt the idea of combining I and l2 regularizations in of deep neural networks (DNNs) (Kang et al., 2017; Chen et al., 2018; Hu et al., 2023; Xu et al., 2023a). Notably, a recent study (Xu et al., 2023a) presents an excess risk bound for Elastic Net Regularized DNNs. This finding provides supporting evidence that incorporating both I and l2 regularizations in a DNN model is reasonable. The insights gained from (Xu et al., 2023a) shed light on the thelretical explanation of our method, and even pave the way for exploring the integration of different divergence metrics in other research arcas, such domain adaptation, transfer learning, and fair representation learning.
Connection with multi-task learning: balancing on two representations. Our DIGNet model per- forms representation balancing on two distinct reprosentations using Wasscrstein distance and H-divergence separately, and the learned representations are then concatenated for outcome modeling. We will now pro- vide additional explanations regarding its connection with the multi-task learning method. In multi-task learning, distinct representations are learned for different tasks, with cach task involving its own objective function. An important step in multi-lask learning is integrating the information from these separately learned representations into unified representation. One common approach is to concatenate the task- specific representations to form joint representation, which effectively preserves the information from each task for outcome modeling (e.g..Li et al. (2018); Baltrušaitis et al. (2018); Crawshaw (2020); Yan et al. (2021); Xu et al. (2023b)). For example, in an E-commerce application Liu et al. (2023), diverse types of user footprints are encoded using different representations with diverse objectives. The learned representa- tions are then concatenated to make the final larget prediction. Similarly, in another application Wu et al. (2018), user and product attentions are separately learned on two distinct representations, which are later concatenated for the final outcome prediction. In tasks such as image and text classification Hao et al. (2023), concatenating multiple representations has shown effective improvements in the classification performance, which is brought by the complementary information of each representation. Furthermore, recent study on multi-view learning (Li et al., 2024) has also demonstrated that concatenating both the non-attention and attention representations of cach view can prevent information loss in the final classification task.
Summary of strengths and limitations. Our method combines Wasserstein distance and H-divergence for representation balancing to capture different types of balancing patterns compared to classic represen- tation balancing models. This shares a similar intuition with the Elastic Net, which combines I and l2 regularizations to learn features with different properties. Notably, the two regularizations in Elastic Net are learned on single parameter space with one objective, this provides more interpretability but might introduce a new trade-off. Different. from Elastic Net, our DIGNet model concalenates the two distinct representations that are learned from two different tasks: Wasserstein distance guided and H-divergence guided representation balancing. This aligns with the principle of multi-task learning. The concatenation fusion technique is extensively employed in numerous multi-task learning studies (Baltrušaitis al., 2018), it effectively preserves and integrates information from different. tasks, leading to improved performance
14

---

Published in Transactions on Machine Learning Research (06/2024)
7= 0.25
7= 0.5
7= 0.75
7= 1,5
3
0-
U-
0-
-2
25
25
25
-5
25
-25
25
25
-25
25
Figure 5: T-SNE visualizations of the covariates as y varies. Red represents the treatment group and blue represents the control group. A larger y indicates greater imbalance between the two groups.
in the final prediction objective (Hao et al., 2023; Li et al., 2024). However, it is crucial to acknowledge that this straightforward concatenation approach can present challenges when interpreting the specific role of each representation and can also increase model complexity (Jia et al., 2020).
6 Experiments
In non-randomized observational data, the ground truth regarding treatment effects remains inaccessible due to the absence of counterfactual information. Therefore, use simulated data and semi-synthetic benchmark data to test the performance of our methods and other baseline models. In this section, we primarily investigate the three following questions:
Q1. Is PDIG helpful ITE estimation through Path the Introduction, i.e., learning more effective balancing patterns without affecting factual outcome prediction?
Q2. Is PPBR helpful in ITE estimation through Path II in the Introduction, i.e., improving factual outcome prediction without affecting learning balancing patterns?
Q3. Can the proposed DIGNet model outperform other baseline models on benchmark dataset?
Ablation models. To investigate QI and Q2, we conducted ablation studies and designed two ablation models, DGNet and DINet, where DGNet (or DINet) can be considered as DIGNet without PDIG, and GNet (or INet) can be considered as DGNet (or DINet) without PPBR. The structures of DGNet and DINet are shown in Figure 2(c) and Figure 2(d), and the objectives of DGNet and DINet are deferred to Section A.6.
6.1 Experimental Settings
Simulation data. Previous causal inference works assess the model effectiveness by varying the distribu- tion imbalance of covariates in treated and controlled groups at different levels (Yao et al., 2018; Yoon al., 2018; Du et al., 2021). As suggested by Assaad al. (2021), we draw 1000 observational data points from the following data generating strategy:
X; 0O 1p1,1 (1-p))pp])
T; X; ~ Bernoulli exp(-
Y," 3 X; +
Y,1 = 3|x; X; + & N(0,1).
Here, 1p denotes the p-dimensional all-ones vector and I, denotes the identity matrix of size We fix 11.3, [1.3...., 1..,......]. and vary 7 E (0.25, 0.5, 0.75, 1, 1,5,2,3) 1,,
= 10, 0.3, x² = 2 = |0.3...., different levels of selection bias. As seen in Figure 5, selection bias becomes more severe with y increasing. For each 7, we repeat the above data generating process to generate 30 different datasets, with each dataset split by the ratio of 56%/24% 2246624%//2 24% 200 as moaninn/validaaionteeen sets.
to yield
Semi-synthetic data. The IHDP dataset, introduced by Hill (2011), originates from the Infant Health and Development Program (IHDP). This program conducted randomized controlled experiment in 1985
15

---

Published in Transactions on Machine Learning Research (06/2024)
V EPENE
CP
CCF
dæ
Wa8s
1.96
DINet
DGNet
DGNett
DGNet
INcl
GNet DiGNeet
INet
INct
30-
INet
194
DIGNet
0.08
GNet
GNet
GNet
DIGNO
DIGNet
DIGNe
1.9
0.05
3.0
1.90
0.02
0. 0.75
0.25
test set for different metrics Y varies in {0.25, 0.5, 0.75, 1,5, 2,3}. Each graph shows the average of 30 runs with standard errors shaded. Lower lines indicate lower values of the metric.
Figure 6:
Plots of model performances
to investigate whether there positive causal effect of frequent high-quality child care and home visits (treatment) on cognitive scores (outcome). The collected data comprise 25-dimensional pre-treatment co- variates, including measurements on the infants (e.g., birth weight, gender, head circumference), as well as measurements on the mothers during pregnancy (e.g., age, marital status, education, smoking and drinking habits). In order to create an observational dataset that involves selection bias, Hill excluded a subpopu- lation (children with nonwhite mothers) from the treated group. Consequently, the IHDP dataset exhibits covariate shift, resulting in imbalancod treated and controlled groups. The final IHDP datasct consists of 747 samples, comprising 139 treated samples and 608 controlled samples. The potential outcomes were generated using setting A in the NPCI package Dorie (2021). We use the same 1000 datasets as used in sets.
Shalit et al. (2017), with each dataset split by the ratio of 63%/27%/ 27%%/0% as
Models and metrics. In simulation experiments, we perform comprehensive comparisons between INet, GNet, DINet, DGNet, and DIGNet in terms of the mean and standard error for the following metrics: VEPEHE, VECF, and VEF with L defined in Definition 1 being the squared loss, well the empirical approximations of Wass(pp=1 - and dx< -1"=") (denoted by Wass and dr, respectively). Note that as shown in Figure 2, Wass is over DE for GNet while over 4G for DGNet and DIGNet; dx is over & for INet while over Φ for DINet and DIGNet. To analyze the source of gain and ensure fair comparison in simulation studies, we fix hyperparameters across models. This way is consistent with Curth & van der Schaar (2021). We apply an early stopping rule to all models as Shalit al. (2017) do. In IHDP experiment, we use VEPEHE, as well as an additional metric EATE |fATE- TATE! to evaluate performances of various causal models (see them in Table 6). More descriptions of the implementation details, as well as the analysis of training time, training stability, and hyperparameter sensitivity, are deferred to Section A.5.
Device. All the experiments are run on Dell 7920 with one 16-core Intel Xeon Gold 6250 3.90GHz CPU and three NVIDIA Quadro RTX 6000 GPUs.
6.2 Results and Analysis
6.2.1 Preliminary Experimental Results
In this part, we first make a general comparison between different models with the degree of covariate im- balance increasing, and the relevant results are shown in Figure 6. There are four main observations: (1) DIGNet attains the lowest /CPEH across all datasets, while GNet have inferior performances than other models; (2) DINet and DGNet outperform INet and GNet regarding VECF and VEPEht; (3) INet, DINet, and DGNet have comparable performance to DIGNet in terms of factual outcome estimations (VEF), but cannot compete with DIGNet in terms of counterfactual estimations (VECF) ITE estimations ( (<PEE)): (4) DIGNet achieves smaller du (or Wass) than DINet and INet (or DGNet and GNet), especially when the covariate shift problem severe (c.g., when 7 > 1). In conclusion, the above study has produced several noteworthy findings. Firstly, finding (1) reveals that our proposed DIGNet model consistently performs well in ITE estimation. Secondly, indicated by finding (2), implementing the PPBR approach can enhance the
16

---

Published in Transactions on Machine Learning Research (06/2024)
Table 1: Ablation study for PDIG: Mean + std of each Table 2: Ablation study for PPBR: Mean + std
metric averaged across 30 runs on test set when y = 3. each metric averaged across 30 runs on test set when Lower value is better.
7 = Lower value is better.
VCF
VCF
dy
Wass
VCF
VCF
di
Wass
DIGNet
1.07 + 0.01
2.89 + 0.07
1.94 + 0.00
0.06 0.00
DGNet
1.07 + 0.01
3.08 + 0.07
0.10 0.00
DINet 1.07 + 0.01
2.95 + 0.07
1.94 + 0.00
GNet
1.12 - 0.03
3.55 0.14
0.10 = 0.00
DGNet
1.07 0.01
3,08 + 0.07
0.10 :: 0.00
DINet
1.07 = 0.01
2.95 + 0.07
1.96 0.00
INet
1.08 + 0.01
3.47 + 0.12
1.96 0.00
predictive accuracy of factual and counterfactual outcomes. Lastly, findings (3) and (4) highlight the role of PDIG structure in enhancing the simultancous reinforeement and complementarity of group distance mini- mization and individual propensity confusion, resulting in more balanced representations. Our subsequent analysis will step beyond these preliminary conclusions gain a deeper understanding of the effectiveness of the proposed methods.
6.2.2 Further Ablation Studies
So far our preliminary observations have show that the relationship between the ITE errors of each model is: DINNet<DINee<TNNt and DIGNet<DGNet-GGNG to the improvement of ITE estimations, we choose the case with high selection bias (7 = 3) to analyze the source of gain for PDIG and PPBR. We report model performances (mean = std) of each specific metric averaged across 30 runs on test set in Table and Table We also report model performances (mean + std) averagod over 30 training and test sets in Table 3. Below we discuss the source of gain in detail.
To further explore how PDIG and PPBR contribute
A Allaioo study for PDIG. The PDIG structure is manifest to be effective in capturing more effective balancing patterns, without affecting factual outcome predictions. As depicted in Figure 6. DIGNet exhibits more balanced reprosentations, irrespective of whether the discrepancy is measured by du or Wass, while DIGNet, DINet, and DGNet demonstrate comparable estimates of factual outcomes (Ver). Two additional pieces of specific evidence can be observed from Table 1: (1) Despite the absence of PDIG in DLNet and DGNet when compared to DIGNet, these three models exhibit very similar performance regarding /Cr, with the performance being 1.07 + 0.01. This indicates that PDIG does not impact the factual estimation. (2) DIGNet achieves smaller dx with a |1.94/1.96 - 1 = 1.0% reduction (or Wass with a 10.06/0.10 - 1| = 40% reduction) compared with DINet (or DGNet). This indicates that PDIG enables the model to learn more effective balancing patterus. The above two points indicate that PDIG can capture more effective balancing patterns, without affecting factual outcome predictions. This advantage translates into superior counterfactual estimation, with DIGNet reduceing VECF by 12.89/2.95 1| = 2.0% and 12.89/3.08 1| 6.2% compared to DINet and DGNet, respectively. Correspondingly, DIGNet also shows superiority in treatment effect estimation (Verehe and <ATE) compared to DINet. (or DGNet), demonstrated in Table 3.
Ablation study for PPBR. The PPBR approach contributes to enhancing factual outcome predictions, without affecting learning balancing patterns. From Table 2, we gain two important insights: (1) The difference in learned balancing patterns, measured by du (or Wass), between DINet and INet (or DGNet and GNet), is negligible. This implies that PPBR does not affect learning balancing patterns. (2) Compared with INet, DINet achieves smaller VEF, with |1.07/1.08- 1| 0.9% error reduction. Similarly, compared with GNet, DGNet achieves smaller /EF, with |1.07/1.12 - 1| 4.5% error reduction. These two observations reveal that PPBR can improve factual outcome predictions, without affecting learning balancing patterns. Benefiting from the advantage of PPBR, the improvement is particularly pronounced in counterfactual estimation. Comparing DINet with INet, the reduction in VECF amounts to [2.95/3.47 15.0%. Similarly, comparing DGNet with GNet, the reduction is /3.08/3.55 1 13.2%. Correspondingly, DINet (or DGNet) attains smaller treatment effect errors (VEPEHE and CATE) compared with INet (or GNet), as shown in Table 3.
17

---

Published in Transactions on Machine Learning Research (06/2024)
Table 3: Training- & test- set VEPEHE EATE when Table 4: Training- & test- set VEPEHE & EATE on IHDP. Mean + standard error of 100 runs.
? 3. Mean = standard error of 30 runs.
Training set
Test set
Training set
Test. set
Vepgbe
CATE
VEPEER
CATE
V'PEEE
CaTE
VOPERE
CA7'
GNet
3.30+0.15
2.58=0.14
3.30+0.16
2.59:0.14
GNet
0.71+0.15
0.1240.01
0.77=0.18
0.15+0.02
INct 3.24.10.11
2.4610.00
3.2210012
2.47.10.10
INet
0.6610.09
0.13110.01
0.7210.11
0.15.10.02
DCNet
2.86:10066
2.13:0.03
2.8310.07
2.15:10.04
DGNet
0,5310.07
0.1110.01
0.6010.09
0,1310001
DINet
2.70+0.06
2.12++0.04
2.69++008
2.13+00.5
DINet
0.57+0.12
0.13+0.01
0.60+0.11
0.14+0.01
DIGNet
2.66:10.07
2.04=0.05
2.6310.07
2.03:0.04
DIGNet
0.4210.02
0.11#0.01
0.45:0.04
0.1210.01
Table 5: Significance analysis regarding the achieved improvements by comparing GNet and DGNet, INet and DINet, DGNet and DIGNet, DINet and DIGNet. The p-value S 0.05 indicates difference is statistically significant.
Training set
Test set
EPEEE
AAEE
VEEEHE
CATE
t-value
p-value
t-value
p-value
t-value
p-value
t-value
p-value
GNet vs.
DGNet
2.7435
0.0081
2.9844
0.0042
2.7073
0.0089
2.9269
0.0049
INel vs. DINet 4.0812
0.0001
3.5222
0.0008
3.5665
0.0007
3.0824
0.0031
DGNet vs. DIGNet
2.0240
0.0476
1.8888
0.0639
2.0650
0.0434
2.0935
0.0407
DINet vs. DIGNet
0.4513
0.6535
1.3525
0.1815
0.6079
0.5456
1.5473
0.1272
Significance analysis for the improvements. To assess the significance the improvements observed in the above ablation studies, we conducted an additional significance analysis by recording the values of VEPEHE and CATE for 30 runs of each of the models (GNet, INet, DGNet, DINet, and DIGNet), Subsequently, we performed t-test for GNet vs. DGNet, INet vs. DINet, DGNet DIGNet, and DINet vs. DIGNet, to investigate the statistical significance of their differences. The relevant results are reported in Table 5. The results reveal statistically significant difference between GNet and DGNet, INet and DINet, as well as DGNet and DIGNet. Note that the difference between DINet and DIGNet is not statistically significant, despite DIGNet exhibiting smallor trcatment effect ostimation errors on average compared to DINet.
6.2.3 Comparisons on IHDP benchmark.
In this part, we perform experiments on the THDP benchmark dataset to compare the performances of different models. The corresponding results are reported in Table 4 and 6.
First, report the ablation results on 1-100 IHDP datasets in Table aiming to examine the consistent effectiveness of PDIG and PPBR. Specifically, Table 4 shows that DINet and DGNet are superior to INet and GNet but inferior to DIGNet concerning treatment effect estimation, suggesting that both PDIG and PPBR are advantageous for treatment effect estimation. For example, on the test set, DINet reduces VEPEНE by 0.60/0.72 - 1| = 16.7% for INet, and DIGNet reduces VEPEHE by [0.45/0.60 1| = 25% for DINet. This is consistent with the findings before: PDIG and PPBR are beneficial to treatment effect estimation.
Furthermore, undergo comparisons between DIGNet and other causal models on 1-1000 IHDP datasets and report the results in Table 6. The results highlight the superior performance of the proposed DIGNet across all the models. Specifically, in comparison to the second-bost method in test-sample performance, DIGNet achieves a substantial improvement, with error reducted by |0.45/0.57 - 1 21% in terms of -
VEPEHE and |0.12/0.13 - || = 7.7% in terms of CATE. Moreover, it is worth noting that DIGNet consistently achieves the lowest errors across various datasets and metrics, revealing its robust performance. We also conduct an additional experiments on another benchmark dataset Twins. The details and results are deferred to Section A.5
18

---

Published in Transactions on Machine Learning Research (06/2024)
7 Conclusion
This paper ostablishes a thcoret- ical foundation by deriving coun- terfactual and ITE error bounds based on H-divergence. This the- oretical foundation builds connec- tion between representation balanc- ing and individual propensity con- fusion. Furthermore, based on in- dividual propensity confusion and group distance minimization, we suggest learning decomposed pat- terns for representation balancing models using the PDIG and PPBR methods. Further, building upon PDIG and PPBR, we propose a novel model DIGNet, for treatment effect estimation. Comprchensive experiments verify that PDIG and PPBR follow different pathways to improve counterfactual and ITE estimation. In particular, PDIG enables the model to capture more effec- tive balancing patterns without affecting factual outcome prediction, while PPBR contributes to improving factual outcome predictions with influencing learning balancing patterns. We hope these findings can con- stitute an important step to inspire more research concerning the generalization of representation balancing models for counterfactual and ITE estimation.
Table 6: Training- & test- VEPEHE & CATE on IHDP. Mean standard error of 1000 runs.
Training set
Test set
VYPERE
CATE
VEPEНE
CATE
OLS/LR; (Johansson et. al., 2016)
5.8 +
.73 + .04
5.8 + 3
.94 + .06
OLS/LRg (Johansson et al., 2016)
244 1
.14 01
2.5 11
.31 + 02
k-NN (Crump et al., 2008)
2.1 .1
.14 011
4.1 .2
.79 + .05
BART (Chipman et al. 2010)
2.1 1
.23 001
2.3 .1
.34 + .02
CF (Wager & Athey 2018)
2
.18 01
3.8 .2
.40 + .03
CRVAE (Lomizoe et al., 2017)
2.7
34 + 101
2.6
.46 + .02
SITE (Yao et al., 2018)
.69 4 0
222 + 01
.75 +
.24 .01
GANITE (Yoou et al., 2018)
1
.4
2.4
.49
BLR (Johansson et, al., 2016)
5.8 1884
.72 04
58
.8
BNN (Johansson et al., 2016)
2.2
37 03
211
.42
TARNet (Shalit al., 2017)
.2 01
28
CFR- CFR-Wass (GNet) (Shalit et al., 2017)
.73
.12
.11.0
.15 00
Dragonnct (Shi al., 2019)
133 A
.14 01
13455
20
MBRL (Huung al., 2022a)
.52
112
57
.13
DIGNet (Ours)
.11
.46
Limitations and futuure works. Our paper verifies the effectiveness of PDIG and PPBR in improving ITE estimation, it is also important step beyond empirical insights into future theoretical studies aimed at addressing the trade-off challenge mentioned in the introduction, e.g,, exploring the possibility of deriving tighter theoretical error bounds based on learning decomposed patterns, and involving the or- thogonal machine learning (Chernozhukov et al., 2018; Oprescu et al., 2019; Nie & Wager, 2021; Huang et al., 2022b) into the representation learning model to improve model's robustness to the misspecification. Furthermore, remains challenging to analytically determine the best divergence metric for representation balancing methods. A promising avenue for future theoretical investigations would involve developing new distributional divergences or exploring unified theory that enables models to select appropriate divergence metrics based the distinct data. Empirical studies can focus discouraging the redundancy of the concatenation fusion of each decomposed pattern and improving the efficacy of the multi-task learning ob- jectives. While we have followed the same approach as previous studies by evaluating model performance using simulated and semi-synthetic data, is crucial for future research to explore the creation appro- priate benchmark datasets (Athey & Wager, 2019; Curth et al., 2021) for assessing the performance of ITE estimation methods in real-world scenarios.
Acknowledgement
We are thankful for the constructive and helpful comments provided by the reviewers and action editor during the reviewing process of TMLR, which has contributed lot to the improvement of our work.
Qi WU acknowledges the support from The CityU-JD Digits Joint Laboratory Financial Technology and Engineering and The Hong Kong Research Grants Council [General Research Fund 11219420 11219420/9043008 /99430044 and 11200219/9042900 The work described this paper was partially supported by the InnoHK initiative, the Government of the HKSAR, and the Laboratory for AT-Powered Financial Technologies.
19

---

Published in Transactions on Machine Learning Research (06/2024)
References
Tameem Adel, Isabel Valera, Zoubin Ghahramani, and Adrian Weller. One-network adversarial fairness. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 2412 2420, 2019.
Joshua Angrist, Guido W Imbens, and Donald B Rubin. Identification of causal effects using instrumental variables. Journal of the American statistical Association, 91(434):444 455, 1996.
Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin. Counterfactual representation learning with balancing weights. In International Conference on Artificial Intelligence and Statistics, pp. 1972 1980. PMLR, 2021.
Susan Athoy and Stefan Wager. Estimating treatment effects with causal forests: An application.
Observational Studies, 5(2):37-55 -1, 2019.
Peter C Austin and Elizabeth A Stuart. Moving towards best practice when using inverse probability of treatment weighting (iptw) using the propensity score to estimate causal treatment effects in observational studies. Statistics in medicine, 34(28):3661 3679, 2015.
Tadas Baltrušaitis, Chaitanya Ahuja, and Louis-Philippe Morency. Multimodal machine learning: A survey and taxonomy. IEEE transactions on pattern analysis and machine intelligence, 41(2):423-44 -432,,444,,- 2018.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. Advances in neural information processing systems, 19, 2006.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1):151 175, 2010.
Alex Beutel, Jilin Chen, Zhe Zhao, and Ed H Chi. Data decisions and thcoretical implications when adver- sarially learning fair representations. arXiv preprint 70700:1707.00075, 2017.
Toana Bica, Ahmed M Alaa, Craig Lambert, and Mihaela Van Der Schaar. From real-world patient data to individualized treatment effects using machine learning: current and future methods to address underlying challenges. Clinical Pharmacology & Therapeutics, 109(1):87 100, 2021a.
Ioana Bica, Danicl Jarrett, Alihan Hüyük, and Mihacla van der Schaar. Learning "what-if" explanations for sequential meaksion-making. In International Conference on Learning Representations, 2021b. URL htttps://pperrevie/ net/foruur/tddhhteeelli
Stephen Burgess, Dylan S Small, and Simon G Thompson. A review of instrumental variable estimators for mendelian randomization. Statistical methods in medical research, 26(5):2333 2355, 2017.
Marco Caliendo and Sabine Kopeinig. Some practical guidance for the implementation of propensity score matching. Journal of economic surveys, 22(1):31-72, 2008.
Defu Cao, James Enouen, Yujing Wang, Xiangchen Song, Chuizheng Meng, Hao Niu, and Yan Liu. Esti- mating treatment effects from irregular time series observations with hidden confounders. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 6897 6905, 2023.
Minwoo Chac and Stephen G Walker. Wasserstein upper bounds of the total variation for smooth densitics. Statisties & Probability Letters, 163:108771, 2020.
Pin-Yu Chen, Yash Sharma, Huan Zhang, Jinfeng Yi, and Cho-Jui Hsieh. Ead: elastic-net attacks to deep neural networks via adversarial examples. In Proceedings of the AAAT conference on artificial intelligence. volume 32, 2018.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Gouble/debiasse machine learning for treatment and structural parameters. The Fconometrics Journal, (1(1)CC1-C6 2018.
20

---

Published in Transactions on Machine Learning Research (06/2024)
Hugh A Chipman, Edward I George, and Robert E McCulloch. Bart: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1):266 220)2 2010.
Zhixuan Chu, Stephen L. Rathbun, and Sheng Li. Graph infomax adversarial learning for treatment effect
estimation with networked observational data. In KDD, pp. 176-184, 2021. URL https://doll http////oo.ooo1....
1145/3447548 3467302.
Michael Crawshaw. Multi-task learning with deep neural networks: A survey,
arXiv preprint
arXiv:2009.09796 2020.
Richard K Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametric tests for treatment effect heterogeneity. The Review of Economics and Statistics, 90(3):389-444 2008.
Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogencous treatment effect estimation. Advances in Neural Information Processing Systems, 34:15883 34:15883-15894 2021.
Alicia Curth and Mihaela Van Der Schaar. In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation. In Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research,
pp. 6623 0623-6642. PMLR, 23-29 Jul 2023.
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating cate? critical look ml benchmarking practices in treatment effect estimation. In Thirty-fith conference neural information processing systems datasets and benchmarks track (round 2), 2021.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pp. 685 693. PMLR, 2014.
Pedro Domingos. A unified barsanariance decomposition. In Proceedings of 17th international conference on machine learning, pp. 231-238. Morgan Kaufmann Stanford, 2000.
Vincent Dorie. Nonparametric methods for causal inference.
wttptt/tww.../ooriieeeeeee 2021.
Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy. Adversarial balancing- based representation learning for causal effect inference with observational data. Data Mining and Knowledge Discovery, 35(4):17133 173888111778,, 2021.
Harrison Edwards and Amos Storkey. Censoring representations with an adversary. arXiv preprint arXiv:1511005897, 2015.
Max H Farrell. Robust inference average treatment. effects with possibly more covariates than observations. Journal of Econometrics, 189(1):1-2) 1-1,,,2,55 2015.
Rui Feng, Yang Yang, Yuehan Lyu, Chenhao Tan, Yizhou Sun, and Chunping Wang. Learning fair repre- scntations via an adversarial framework. arXiv preprint arXiv:1904. 13341, 2019.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarill training of ncural networks. The journal of machine learning research, 17(1):2096-200 2030, 2016.
Stuart Geman, Elie Bienenstock, and René Doursat. Neural networks and the clas/varance dilemma. Neural computation, 4(1):1-58, 18,155,,111 1992.
Ruocheng Guo, Jundong Li, Yichuan Li, К Selçuk Candan, Adrienne Raglin, and Huan Liu. Ignite: A minimax game toward learning individual treatment effects from networked observational data. In IJCAL. pp. 4534 4540, 2020a.
Ruocheng Guo, Jundong Li, and Huan Liu. Learning individual causal effects from networked observational data. In Proceedings of the 13th International Conference on Web Search and Data Mining, pp. 232 2222200 2020b.
21

---

Published in Transactions on Machine Learning Research (06/2024)
Yaru Hao, Xiao Yuan Jing, Runhang Chen, and Wei Liu. Learning enhanced specific representations for multi-view feature learning. Hhowledge-Baset Systems, 272:110590 2023.
Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportunity in supervised learning. Advances in ncural information processing systems, 29, 2016.
Negar Hassanpour and Russell Greiner. Counterfactual regression with importance sampling weights. IJCAI, pp. 5880-5887, 2019a.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference Learning Representations, 2019b.
Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217 240, 2011.
Arthur E Hoerl and Robert W Kennard. Ridge regression: Biased estimation for nonorthogonal problems. Technometrics, 12(1):55-61) 67, 1970.
Cong Hu, Yuanbo Li, Zhenhua Feng, and Xiaojun Wu. Attention-guided evolutionary attack with elastic-net regularization on face recognition. Pattern rocognition, 143: 143:109760, 2023.
Yiyan Huang, Cheuk Hang Leung, Xing Yan, Qi Wu, Nanbo Peng, Dongdong Wang, and Zhixiang Huang, The causal learning of retail delinquency. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 204-22 212-212,- 2021.
Yiyan Huang, Cheuk Hang Loung, Shumin Ma, Qi Wu, Dongdong Wang, and Zhixiang Huang. Moderately- balanced representation learning for treatment effects with orthogonality information. In Pacific Rim International Conference on Artificial Intelligence, pp. 3-16. Springer, 2022a.
Yiyan Huang, Cheuk Hang Leung, Qi Wu, Xing Yan, Shumin Ma, Zhiri Yuan, Dongdong Wang, and Zhixiang Huang. Robust causal learning for the estimation of average treatment effects. In 2022 International Joint Conference on Neural Networks (IJCNN 2022). IEEE, 2022b.
Yiyan Huang, Cheuk Hang Leung, Shumin Ma, Zhiri Yuan, Qi Wu, Siyi Wang, Dongdong Wang, and Zhixiang Huang. Towards balanced representation learning for credit policy evaluation. In International Conference on Artificial Intelligence and Statistics, pp. 3677 3692. PMLR, 2023.
Andrew Jesson, Sören Mindermann, Yarin Gal, and Uri Shalit. Quantifying ignorance in individual-level causal-effect estimates under hidden confounding. In International Conference Machine Learning, pp. PMLR, 2021.
Xiaodong Jia, Xiao Jing, Xiaoke Zhu, Songcan Chen, Bo Du, Ziyun Cai, Zhenyu He, and Dong Yue. Semi-supervised multi-view deep discriminant representation learning, IEEE transactions on pattern analysis and machine intelligence, 43(7):2496-25000 2020.
Ray Jiang, Aldo Pacchiano, Tom Stepleton, Heinrich Jiang, and Silvia Chiappa. Wasserstein fair classifica- tion. In Uncertainty in artificial intelligence, pp. 862 872. PMLR, 2020.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference machine learning, pp. 3020 3029. PMLR, 2016.
Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and repre- sentation learning for estimation of potential outcomes and causal effects. Journal of Machine Learning Research, 13166), 2022a. URL /ttp://jmlr. httpaaar.orgrrrreeeeeeee
Fredrik Johansson,, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and represen- tation learning for estimation of potential outcomes and causal effects. Journal of Machine Learning Research, 23(166):1-500 2022b.
22

---

Published in Transactions on Machine Learning Research (06/2024)
Nathan Kallus, Xiaojie Mao, and Angela Zhou. Interval estimation of individual-level causal effects under unobserved confounding. In The 22nd internat ional conference on artificial intelligence and statistics, pp. 2281 2290. PMLR, 2019.
Guoliang Kang, Jun Li, and Dacheng Tao. Shakcout: A new approach to regularized doep ncural network training. IEEE transactions on pattern analysis and machine intelligence, 40(5):1245 1258, 2017.
Edward H Kennedy. Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2):3008 -0488, 2023.
Kun Kuang, Peng Cui, Bo Li, Meng Jiang, Shiqiang Yang, and Fei Wang. Treatment effect estimation with data-driven variable decomposition. In Procoedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang, Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations, 2021.
Ananya Kumar, Tengyu Ma, Percy Liang, and Aditi Raghunathan. Calibrated ensembles can mitigate accuracy tradeoffs under distribution shift. In Uncertainty in Artificial Intelligence, pp. 1041-1051. PMLR, 2022.
Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Procoedings the national academy of sciences, 116(10):4156- 4165, 2019.
Jinxing Li, Chuhao Zhou, Xiaoqiang Ji, Mu Li, Guangming Lu, Yong Xu, and David Zhang. Multi-view
instance attention fusion network for classification. Information Fusion, 101:101974, 2024.
Yingming Li, Ming Yang, and Zhongfei Zhang. A survey of multi-view representation learning. IEEE transactions on knowledge and data engineering, (1(1)):1863 1883, 2018.
Qi Liu, Zhilong Zhou, Gangwei Jiang, Tiezheng Ge, and Defu Lian. Deep task-specific bollom representation network for multi-task recommendation. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pp. 1637 1646, 2023.
Mingshong Long, Jianmin Wang, Guiguang Ding, Jiaguang Sun, and Philip S Yu. Transfer joint matching for unsupervised domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recoguition, pp. 1410 1417, 2014.
Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In International conference on machine learning, pp. 97 105. PMLR, 2015.
Mingsheng Long, Han Zhn, Jianmin Wang, and Michael I Jordan. Deep transfer learning with joint adapta- tion networks. In International conference on machine learning, pp. 2208 2217. PMTR, 2017.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30, 2017.
Shumin Ma, Zhiri Yuan, Qi Wa, Yiyan Huang, Xixu Hu, Cheuk Hang Leung, Dongdong Wang, and Zhix- iang Huang. Deep into the domain shift: Transfer learning through dependence regularization. TEEE Transactions on Neural Networks and Learning Systems, 2023.
David Madras, Elliot Creager, Toniann Pitassi, and Richard Zemel. Learning adversarially fair and trans- ferable representations. International Conference on Machine Learning, pp. 3384 3393. PMLR, 2018.
Divyat Mahajan, Ioannis Mitliagkas, Brady Neal, and Vasilis Syrgkanis. Empirical analysis of model selection for heterogenous causal effect estimation. International Conference on Learning Representations, 2024.
23

---

Published in Transactions on Machine Learning Research (06/2024)
Maggic Makar, Fredrik Johansson, John Guttag, and David Sontag. Estimation of bounds potential outcomes for decision making. In International Conference Machine Learning, pp. 0661-6677 PMT.R, 2020.
Wang Miao, Wenjie Hu, Elizabeth Ogburn, and Xiao-Hua Zhou. Identifying effects of multiple treatments in the presence of unmeasured confounding. Journal of the American Statistical Association, 118(543): 1953 1967, 2023.
Naveen Naidu Narisetty. Bayesian model selection for high-dimensional data. In Handbook of statistios. volume 43, pp. 207 248. Elsevier, 2020.
Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108 (2):299-31)) 319, 2021.
Ana Rita Nogueira, João Gama, and Carlos Abreu Ferreira. Causal discovery in machine learning: Theories and applications. Journal of Dynamics & Games, 8(3), 2021.
Miruna Oprescu, Vasilis Syrgkanis, and Zhiwei Steven Wu. Orthogonal random forest for causal inference. International Conference on Machine Learning, pp. 4932 4941. PMLR, 2019.
Judea Pearl. Causality. Cambridge university press, 2009.
Zhaozhi Qian, Yao Zhang, Ioana Bica, Augela Wood, and Mihaela van der Schaar. Synctwin: Treatment effect estimation with longitudinal outcomes. Advances Neural Information Processing Systems, 34: 3178 3190, 2021.
Paul R. Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322 331, 2005.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076 3085. PMLR, 2017.
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018,
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems, 32, 2019.
Changjian Shui, Fan Zhou, Christian Gagné, and Boyu Wang. Deep active learning: Unified and principled method for query and training. In International Conference on Artificial Intelligence and Statistics, pp. 1308-1318, PMLR, 2020.
Bharath K Sriperunbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On the empirical estimation of integral probability metrics. Electronic Journal of Statistics, 6:1550-1599, 2012.
Robert Tibshirani. Regression shrinkage and selection via. the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267 288, 1996.
Eric Tzeng, Judy Hoffmman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proccedings of the IEEE conference on computer vision and pattern recognition, pp. 7167-7176, 2017.
Giorgio Valentini and Thomas G Dictterich. Bias-variance analysis of support vector machines for the development of svm- based ensemble methods. Journal of Machine Learning Research, ((Jul):725-775, 2004.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
24

---

Published in Transactions on Machine Learning Research (06/2024)
Matthew J Vowels, Necati Cihan Camgoz, and Richard Bowden. D'ya like dags? survey structure learning and causal discovery. ACM Computing Surveys, 55(4):1-36, 2022.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 113((23):12288 122-1122)22--122, 2018.
Anpeng WU, Kun Kuang, Bo Li, and Fei Wu. Instrumental variable regression with confounder balancing. In International Conference on Machine Learning, pp. 24056 24075. PMLR, 2022.
Zhen Wu, Xin-Yu Yu Dai, Cunyan Yin, Shujian Huang, and Jiajun Chen. Improving review representations with user attention and product attention for sentiment classification. In Proceodings of the AAAI conference on artificial intelligence, volume 32, 2018
Lihu Xu, Fang Yao, Qiuran Yao, and Huiming Zhang. Non-asymptotic guarantecs for robust statistical learning under infinite variance assumption. Journal of Machine Learning Research, 24(92): 2023a.
Peng Xu, Xiatian Zhu, and David A Clifion. Multimodal learning with transformers: survey. IEEE
Transactions on Pattern Analysis and Machine Intelligence, 2023b.
Xiaoqiang Yan, Shizhe Hu, Yiqiao Mao, Yangdong Ye, and Hui Yu. Deep multi-view learning methods: A
review. Neurocomputing, 448:106-129, 2021.
Zitong Yang, Yaodong Yu, Chong You, Jacob Steinhardt, and Yi Ma. Rethinking bias-variance trade-off for gencralization of neural networks. In International Conference on Machine Learning, pp. 10767-10777. PMLR, 2020.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. Advances in Neural Information Processing Systems, 31, 2018.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018.
Junkun Yuan, Xu Ma, Ruoxuan Xiong, Mingming Gong, Xiangyu Liu, Fei Wi,, Lanfen Lin, and Kun Kuang, Instrumental variable-driven domain gencralization with unobserved confounders. ACM Transactions on Knowledge Discovery from Data, 17(8):1 1-1), 2023.
Alessio Zanga, Elif Ozkirimli, and Fabio Stella. A survey on causal discovery: Theory and practice. International Journal of Approximate Reasoning, 151:101-129, 2022.
Rich Zemel, Yu Wu, Kevin Swersky, Toni Pitassi, and Cynthia Dwork. Learning fair representations. In International conference machine learning, pp. 325-333. PMLR, 2013.
Brian Hu Zhang, Blake Lemoine, and Margaret Mitchell. Mitigating unwanted biases with adversarial learning. In Proccedings of the 2018 AAAI/ACM Conference on AI, Ethics, and Society, pp. 335-340, 2018.
Yao Zhang, Alexis Bellot, and Mihacla Schaar. Learning overlapping reprosontations for the estimation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1014-1014. PMLR, 2020.
Han Zhao and Geoffrey J Gordon. Inherent tradeoffs in learning fair representations. Journal of Machine Learning Research, 23(5))11-26, 2022.
Han Zhao, Amanda Coston, Tameem Adel, and Geoffrey J Gordon. Conditional learning of fair representa- tions. In International Conference on Learning Representations, 2019a.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant represen- tations for domain adaptation. In International conference on machine learning, pp. 7523-7532. PMLR, 2019b.
25

---

Published in Transactions on Machine Learning Research (06/2024)
Han Zhao, Chen Dan, Bryon Aragam, Tommi S Jaakkola, Geoffrey J Gordon, and Pradeep Ravikumar. Fun- damental limits and tradeoffs in invariant representation learning. Journal of Machine Learning Research, 23(340):1 49, 2022.
YUAN Zhiri, HU Xixu, WU Qi, MA Shumin, Cheuk Hang Leung, Xin Shen, and Yiyan Huang. A unified domain adaptation framework with distinctive divergence analysis. 2022.
Ding-Xuan Zhou. On grouping effect of clastic net. Statistics & Probability Letters, 83(9):2106-2112, 2013.
Indre Zliobaite. On the relation between accuracy and fairness in binary classification. arXiv preprint 150.00330 2015.
Hui Zou and Trevor Hastie. Regularization and variable selection via the elastic net. Journal of the Royal Statistical Society Series B: Statistical Methodology. 67(2):301-322 2005.
A
Appendix
A.1 Discussion of the Trade-off Problem
We will now discuss two cases to gain deeper understanding of this trado-off phenomenon. (1) The case without representation balancing: In this case, the outcome functions are fitted by Y f² and y = (X using treated and controlled samples, respectively. 4 ('(xtreat) and (X(Xoontrol can be good estimates of factual outcomes based on the well-preserved pre- rre-balancing information (group informa- tion). However, the estimated counterfactual outcomes 50 (Xttreet) and + (Xemntrol) can be problematic due to the presence of the covariate shift problem P(XIT = 1) # P(X = =), where the dist.ribution of training data P(X, Y|T = t) differs from that of the test data P(X, Y|T 1-- t) for {0, 2. (2) The case with representation balancing: In this case, the outcome functions are fitted by Y1 (x(((e))) and y = (2((xoootrr)) using treated and controlled samples, respectively. Using ((X) to fit factual outcomes can improve the accuracy of the counterfactual estimates (=(((( and h () because rep- resentation balancing enforces the distributions of treated and controlled samples to be as close as possible in the representation space. As a result, reprosentation balancing effectively tackles the covariate shift issuc, resulting in training data and test data following the same distribution 3. However, executing representation balancing can inevitably lead to a loss of outcome-predictive information in p(X). This occurs naturally as becomes insensitive the treatment variable, thereby sacrificing pre-balaneing information (group information) that contributes to factual outcome predictions. To illustrate the negative impact of losing pre-balancing information in balanced representations, we present motivating example below.
Motivating example. Suppose there is a vaccine available to prevent a certain disease. We define X as the covariate, T = as the treatment (receiving the vaccine), T= 0 as the control (not receiving the vaccine), and Y as the outcome (the level of specific antibodies). Assume that the outcome is determined by Y T TXxp(XX) + (1 (1-7) 0 = T exx((X), which means that if an individual receives the treatment, the level of antibodies will be y exp(æ); otherwise, it will be y = O. In observational data, the treatment is assigned based on the covariate of each individual. The left graph of Figure 7 illustrates the distributions of X in the treated and control groups. We observe that individuals with positive x values are more likely to receive the vaccine, resulting in a higher level of antibodies. Given some sample i, a well trained model first determines whether the sample is more likely to be in the treatment or control group based on its covariate X, = x. If it determines the sample to be in the treatment group, the model then predicts y = exp(x); otherwise, it prodicts y = For example, if x 1, the model would classify the sample as more likely to be in the treatment group and predict y e. Therefore, in this case, the pre-balancing covariate remain informative for predicting the outcome. Now, let's consider representation function & that achieves
By unconfoundedlness, we have X,7 X,T POYiX, ()P(X|T = t), ie., P(X, Y t) 1 -), 'y one-to-one and invertible propertics of & and unconfounddnness, have ((Y P(Y))))) f). Given an ideal representation halancing P(*(X)
P(Y|X, 1-11--4 (). Due to the covariate shift P(X t) # P(X
t), we have P(YIX, t)P( =))xXI =) +
P( 1- 1), have -) P(Y) - 1-1,, i.e., P(O(X), P/
26

---

Published in Transactions on Machine Learning Research (06/2024)
Before bulancing
Aller balancing
Traated Costrolled
Trestel
Contolled
Improper balancing
((0(s))
((x)
Figure 7: Motivating example (similar to the example demonstrated in Huang et al. (2023)) for illustrating the trade-off between outcome prediction and representation balancing.
improper representation balancing between the treated and control samples. The right graph of Figure 7 shows the distributions of (XX) in the treated and control groups. In this case, it becomes challenging for model to accurately predict Y using (XX because the model may become confused about whether sample is more likely to receive the treatment or the control. Consequently, improperly balanced representations can lead to a loss of cntcome-prediction information.
The above discussions and motivating example illustrate the inherent tradc-off problem between outcome prediction (Principle I) and representation balancing (Principle II), which arises due the fact that repre- sentation balancing models alleviate covariate shift at the expense of factual outcome prediction.
A.2 Proof of Lemma 1
Proof of taking the squared loss, i.c., L(y,, 92) (y1 -y2)²
Proof. We denote (pEne()) = ()) = er(h, Φ), ecr(J) ecr(h, 3) for s(x,1) 1(((x),().
cr(s)
= 3xx(0,1xx (f(x,t) - p()" x))(xx t)dy' dxdt
- (f(x,) + (xx) p(y 1/) pp(x lydy' dxdi.
* [ (7'(x) '')2p(' p(u xx) t) dxdt
+2 [ 1XX1O1IXO (f(x,t) (x()) y') (x)/xx)pxxxx, 1)dy' dxdi
(28)
(f(x,t) - (xx)))>x(x t)dxdt + ozz (p(,)t)
(29)
1xx1011
Equation (29) is by the definition of oz (p(x, ()) in Lemma and equation (28) equaling zero since r'(x) Sy y pl'p(y (x)de. A similar result can be obtained for ECF:
ecr(S) = - - (f(x,1) - ((x)))x(xx 1 1)dxdl + 0g - -1)).
27

---

Published in Transactions on Machine Learning Research (06/2024)
(PEHE(F)
=/x N
- (+*(x) -xxx))
22 / (U((x, 1) - (x))) ()(x), 0) ((x))))((x()))
(30)
/ (f(x. 1) (xxx)) 1)dx + / (f(x, - pxxx T 00jxx
22
x
+ 2 / (f(x,1) 1) - r- (x))) 0)dx + 2 (f(x.0) 0) - -xxx)) T= 1)dx
(31)
-2 I (f(x,)) 8) - -' t)dxxt + 2 / (0.1)
(f(x,t) +' (x)) (xx))x(xx,x))))) p(xx
1)dxdt
(xx(0.1)
= - oy - ()(, +))c)) - (
(32)
Inequality (30) is by (x + y)² s 2(x² + y²): equation (31) is by p(x) = p(x,T T 0) + p(x,T 1). By
(equation 32) and the definition of of in Lemma we have
(PENE(1, 5 2lecr(h, 3) + trr Q) - -20
Proof of L taking the absolute loss, i.e., Lyy, y2) 11 32/
Proof. We denote (PEHE()) = ()) = cr(h, $), ccr(f) ecr(h, +) for f(x,t) = h(b(x),
cf(f)
= -xxxOoLXX 1s(x,1) - v'lp(''\ x))(xx Ddy'd dxdl
f SExtOLyXO 1s(x,t) (x)))((( xx xxx, thy dxdt
- (7((x) v v'ln x)p(xx dddt dxdl
(33)
1xx10.11 lf(x,t) ('xx))p(xx t)dxdt Aye (p(x, t))
(34)
Inequality (33) is by |a + vl 2 |x| - lel. equation (34) is by the definition Ay (p(x, t)) in Lemma 1. A similar result can be obtained for cCpt
(cr(S) 2 / Ixx1on1 1((x,) - (xxx)xx(x, - 1)dxdl - Ave (p(x, -1)).
28

---

Published in Transactions on Machine Learning Research (06/2024)
{PEHε(F)
= 1K f(x - (x' - (xx))x((xxxx
x
$ J (15(x,1) 1) - (x(x) 11(x,) 0) - -x(x))D(xxdx
(35)
L [f 1) - -x'(xxx(xxx 7' +) + / 1f(x 1) - (x) O0dxx
(36)
/ [f(x. 0) (x))xxx, = 0)dx + / 1f(x, 0) (x))lx(x, T = 1)dx x
(37)
x
[xx(0.1) 15(x,)) () - -- r'(x)l((x, ()dxdl +
[s(x, t) - (x))lx(xx 1 1)dxdl
xxx00.1)
Sef(f) Ay (p(x. t)) + ECF(F) Ay (p(x,1-t))
(38)
Inequality (35) by b |x| + |y]. Equation (36) and equation (37) are by p(x) - p(x, T +xx,
1). By equation (38) and the definition of Ay in Lemma 1, have
(HEHE 2) S {f() + A + +EC((f) + (pxx, t)) Sccr(h, 8) cr(h, @) 2A,-
A.3 Proof of Theorem 1
Proof of equation (4):
Proof.
(CF(h, - ((1 u)
+ ()) 4)]
=((1 u) . (h, t) + -. +)) - (11 u) FF
T=1 (h, p) + CF
T=0 (h, ( $)1
=(1 u) 1F (h, 4) - - (1, 4)] + u | (h, 4) Ex (h, 2)]
=(1 u)
x ( 1) (xx
-= (x)= (x))
(x, T=
(x) -rT=o
-(1 u) /
p/ -= (r)( + (V) (r)) p5 (r))dr
(39)
R
/ 1 ( 1) (r)
-B (1 - u)
p4 T=1 (r))dr
Bo
+ BB u
(n, ())
T= (r) pb - dd
R BB
/
<Bo (1 - u) sup
uusppp
(1
10 (r) - - T-1 (r))dr
ges
R
+ Bş u sup / ((x)
T-1 (r) p4
T=0 (r) dr
(40)
9E9
R
=Bo Wass(pp P4 T
(41)
Equation (39) is by the change of formula, (r) is the absolute of the determinant of the Jacobian of (I)). Equation (41) is by Definition 2.
pr= = (r))))())) where
Proof of equation (5):
29

---

Published in Transactions on Machine Learning Research (06/2024)
Proof.
CPEHE(h, $)
<2(cor(h, 4) + ch ②) -20 <2 (hi 1) + cp (h, 4) + Bo Wass(ph T=1 :PP )) 20
(42)
(43)
Inequality (42) is by equation (2) in Lemma 1. Inequality (43) is by equation in Theorem 1.
Proof of equation (6):
Proof.
<PEHE(1, $)
Secr (h, 4) + c +)+ 2A, <ef ②) + = (h, 1) + Bg W Wass(pp T=1 .p -b==0) + 2A,
(44)
(45)
Incquality (44) is by equation (3) in Lemma 1. Incquality (45) is by equation in Theorem 1.
A.4 Proof of Theorem 2
We first introduce Lemma 2 that is useful for proving Theorem 2.
Lemma 2. Let G that is defined in Definition be the family of binary functions. Then we obtain supyen Iss (p)((s)))) p2(s))ds} (PL.
Proof. Let I(.) denotes indicator function.
dx (P1, P2)
=2 sup wEH
(pi(s) (s))ds
(()==1
22 sup wEW
I[(n(s) ()(r((s) (2))
22 sup nEH
((s)(pp P2(s))ds
(46)
The last equation is because an indicator function also a binary function.
Proof of equation (8):
30

---

Published in Transactions on Machine Learning Research (06/2024)
Proof.
ECF(h, 2) - 1((1 - _. (h, 2) + u u. (h, Φ)
=(1 - u) (0 (r(). (p$ (r) p/ ---1 (r)) + S ln (V1)
p4 (r))drr
(47)
<(1 u)
b (V(() 11's (r) P4 (r))dr
+
o (V(r), (r) - p& - (r)
(r)
(48)
P6
5(1 K)
(Po
T=0 (r) p4 1= (r))dr + a K
(po (r)
(r) -/- (r))drr
(r)=
(49)
=(1 K)K
I(p' 1=0
> P4
T=
(P
T=0 (r) p4
T=1 T-1((r)) (r) dr
+ uu- K
I(p0 T=1 >
(r) )
T=U
(ps
T=1
p T=0 (r) (r))dr
R
/
5(1 K sup
(T( (r) - - -- (r))
wEA
+ u. K sup nEAL
(p( T=
T=1 (r) - pp
T-0 - (r) dr
(50)
R
/ ((P)
<K sup sup WEH
T- (r) - T-0 T=0 (r)
K
dn (r T=1 2
P4 0)
(51)
Equation (47) is derived in the same way as equation (39). Equation (48) is by 6, 2 for all r and t. Incquality (49) is by the definition of K Theorem 2. Incquality (50) because an indicator function is also a binary function. Equation (51) is by Lemma 2.
Proof of equation (9):
Proof-
SPENE(h, $)
<2(Ecr(h, 1) + e 1) -20z)
(52)
K
(h. 1) + 2 <2(c (h, + (
dx(pb
T=1 -p-=0 20
(53)
Inequality (52) is by equation Lemma 1. Inequality (53) by equation 8 Theorem
Proof of equation (10):
Proof.
<PEnE(h, x)
Secr(h, 1) + (n +) + 2A,
(54)
K
sef= 1) + + - Φ) + 2
dn (p T=1 - + 2Ag
(55)
Inequality (54) is by equation in Lemma 1. Inequality (55) is by equation 8 in Theorem 2.
31

---

Published in Transactions on Machine Learning Research (06/2024)
A.5 Additional Experimental details
Additional results on Twins Benchmark. To investigate the applicability of our model DIGNet benchmark datascts beyond the commonly used IHDP benchmark, conducted additional comparisons with several baseline models, including linear, tree, matching, and representation learning methods, on the Twins benchmark, as presented in Table 7.
The Twins datasct comprises records of twin births in the USA between 1989 and 1991. After preprocessing, each unit contains 30 covariates relevant to parents, pregnancy, and birth. The treatment D = indicates the heavier twin, while D indicates the lighter twin. The binary outcome variable Y represents 1-year mortality, For more comprehensive details on this dataset and the limitation of IHDP, refer to Curth et al. (2021).
Notably, for CATE, the simple linear or matching estimator performs best across different methods. On the other hand, when assessing ITE performance using the AUC of potential outcomes, representation learning models all demonstrate strong performance, with AUC values exceeding 0.800 on both training and test sets. The observation might stem from the fact that representation balancing models are based on ITE error bounds, rather than ATE error bounds, thereby optimizing for AUC instcad of CATE- Morcover, among all the models, our DIGNet achieves the second-best AUC results. The best results are achieved by MBRL, which involves the orthogonality information (similar to doubly robust estimators) in representation balancing. This, in turn, inspires us to explore ATE error bounds, or consider involving doubly robust methods in future research.
Table 7: Training- & test- set AUC & CATE on Twins. Mean + standard error of 100 runs.
Training set
Test set
AUC
CATE
AUC
CATE
OLS/LR, Johansson et al. (2016)
.660 + .005
.004 .003
.500 + .028
.007 + .006
OLS/LR2 Johansson et al. (2016)
.660 + .004
.044
.500 .016
.007 + .006
k-NN Crump al. (2008)
.609 .010
+000
.492 .012
.005 + .004
BART Chipman et al. (2010)
.56 .014
.121
.500 5.11
.127 + .024
CEVAE Louizos et al. (2017)
.45 .003
.022
.841 .004
.032 = .003
SITE Yao et al. (2018)
.62 - .002
.016
.853 .006
.020 = .002
BLR Johansson et al. (2016)
.611 + .009
.006
.10 .018
.033 - .009
BNN Johansson et al. (2016)
.690 + .008
..06
.676 .008
.020 + .007
TARNet Shalit et al. (2017)
.849 .002
.011
.840 .006
.015 = .002
CFR-Wase (GNet) Shalit et al. (2017)
.850 .002
.011 .011
.82 .005
.028 + .003
MBRL (Huang et al., 2022a)
.89 .000
.003 = .000
.74 .001
.007 = .00q
DIGNet (Ours)
.874 + .001
.004 + .001
.871 + .001
.008 + .001
Implementation details. In simulation studies, ensure fair comparison by fixing all the hyperpa- rameters in all datasets across different models. The relevant details are stated in Table 8. In IHDP studies,
Table 8: Hyperparameters of different models in simulation studies.
4g
40
A-
bntchsloe
Itenation
loarninn rate
learning rate for
Goet: [100, 100, 100, 180)
(100, 100)
(100,100) 100)
Inet
[100,100,,10, 100, 100, 100)
100, 100)
(100, 100)
(100,100)
0.1
100
300
le
le-t
DGNet
1,00,000,,100 11000000000 100)
100)
(100, 100)
(100,100)
1e-d
DINe
(100, 100, 100,1000 106)
(100, 100)
(100, 100, 100)
(100, 100)
(100, 100)
0.1
100
300
1e-3
1e
DIGNet
100, 100..00)) 100)
(100, 100)
(100, 100)
(100,
100)
1e-3
16-4
to compare with the baseline model CFR-Was (GNet), we remain the hyperparameters of INet, DGNet, DINet and the early stopping rule the same as those used in CFR- Wass Shalit et al. (2017). Since DIGNet
32

---

Published in Transactions on Machine Learning Research (06/2024)
is more complex than other four models, we adjust the hyperparameters of 4g. PG. Φ1, a1; and 0g for DIGNet as Shalit et al. (2017) do. The relevant details are stated in Table 9.
Table 9: Hyperparameters of different models in IHDP experiments.
batzhelioe
Itceation:
leening
Jearning for
Groet
(100, 100, 100, 100)
(100, 100, 1000
(1000 100, 100)
1a
(100,
(200,
1100,
1x-8
(3000 110), 1MM)
(100,
(100,
[100,
1e
DINes
(300, 100. 100)
(200, 200)
(150, 100)
(100, 100)
DIGNE
(100,
100), 190, 100)
(100, 100, 100) (100, 100, 100) (200), 2000 200)
(100) 100. 100)
(100, 100, 100)
Analysis of training time and training stability. We record the time it took for different models to run through 100 IHDP datasets in Table 10, and each model trained within 600 epochs. Following Shalit et al. (2017), all models adopt the early stopping rule. We also record the average early stopping epoch on 100 runs and the actual time on 100 runs, where (actual time) (total time) x (average early stopping epoch)/600. Not surprisingly, GNet took the least amount of time with 3096 seconds since the objective of GNet is the simplest. However, it is very interesting that the proposed methods, DGNet and DINet, are the first two to early stop. As a result, though DGNet and DINet have multi-objectives, they spent less actual training time but achieved better ITE estimation compared to GNet and INet. Since GNet and INet are actually DGNet and DINet with PPBR ablated, we find that PPBR component can help a model achieve better ITE estimates with less time. In addition, we find that DIGNet spent the longest time optimize since it has the most complex objective. To further study the stability of the model training, we also plot the metries VeF. Wass, dr, and VEPERE for the first 100 epochs of each model on the first IHDP dataset in Figure 8. We find that the training process of DIGNet is stable, even steadier than GNet and INet. From this perspective, we haven't seen a difficulty of optimizing DIGNet.
Table 10: Training time records on 100 IHDP datasets.
Model
Time for 600 epochs
Avg early stopping
Actual time
EPEEE test set
GNet
3096s
240.61
1241s
0.77=10118
INet
4042s
254.19
1712s
0.72++0.11
DGNet
3775s
169.17
1064s
0.60+0.009
DINet
3212s
157.98
846s
0.60++0.11
DIGNet
4984s
226.76
1884s
0.45-10.04
We also provide the ITE and ATE estimation results on 100 IHDP datasets when the combination of (an, 02) in DIGNet objective varies in (0.1,0.5, 1}. The relevant results are reported in Table 11, indicating our DIGNet model is robust to the hyperparameters varying.
Table 11: The results on 100 IHDP datascts with different combinations of (a1, 02) in DIGNet objective.
Training set
Test set
(ar, (2)
VEPEHE
6ATE
VEPEHE
EATE
(0.1,0.1)
0.407 + 0.018
0.125 + 0.015
0.434 + 0.022
0.138 + 0.016
(0.1,0.5)
0.414 + 0.026
0.120 0.015
0.434 + 0.028
0.123 0.015
(0.1,1)
0.416 0.019
0.116
0.452 + 0.026
0.121 0.015
00.5.0.1)
0.417 0.023
0.006
0.440 0.026
0.137
(0.5,0.5)
0.407 + 0.021
0.125+0.000
0.416
0.124
10.5,1)
0.126 0.014
0.455 0.028
0.133 0.016
(1,0.1)
0.411 0.021
0.119 + 0.015
0.439 + 0.027
0.118 0.015
(1,0.5)
0.020
0.118 0.015
0.430
0.128 0.016
(1,
0.402+0.00
0.112 0.014
0.437 0.027
0.121 0.015
33

---

Published in Transactions on Machine Learning Research (06/2024)
GNet
GNet
INet
DGNet
DGNet
DINet
DINet
DIGNet
DIGNet
2,00
3.00
1.50
2,00
CPENK
1,00
20
40
60
20
GNet
INet
DGNNtt DIGiNet
DINet
0.06
DIGINet
1.04
0.04
0.02
1.02
Wass
2x
0.00
1.01
20
40
60
20
Figure 8: Training loss plots for the first 100 epochs on the first IHDP dataset.
A.6 Objectives of Different Models
Objective of GNet.
min kg.ht
0,, + (x, 1E).
Objective of INet.
max (x; min
(;
Objective of DINet. Note that similar to DIGNet, the pre-balancing patterns are preserved by only updating 4 but fixing gg in the second step.
max
a2Cr (F,T)
min 4r
0, 0g
min R. Pr
L,(x,t,y; @ (0) 1E),
Objective of DGNet. Note that similar to DIGNet, the pre-balancing patterns preserved by only updating &G but fixing og in the first step.
min tx
da
min Ba.
C,(x,t,y; @ (40 1²),
34

---

Published in Transactions on Machine Learning Research (06/2024)
Objective of DIGNet.
min 4c 0
a; 9p),
max
min 4,
min 2e
Ly(x,t,y: Dg 2 (1 22) @ ( o &g), 1²).
35


Reviews:
[{"review_id": "3oYgwyGOgn", "reviewer": "Reviewer_y2MN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The idea of the proposed method is easy to follow\n2. The related works have been reviewed thoroughly \n3. The simulation results are intensive and evident\n\nWeaknesses:\n1. The theory is incremental given the existing works in the literature.\n2. I don't think the theory is sufficient to explain the improvements made by the method.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new deep learning method and some theoretical bounds for the problem of individual treatment effect estimation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "I think the writing of the paper can be improved. Many sentences are hard to read and follow.", "broader_impact_concerns": "N.A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pexdFgk3n1", "reviewer": "Reviewer_eHay", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- empirical results are strong - demonstrates wins over baselines\n- some useful theoretical results extending work of Shalit et al (2017) and Ben-David (2006)\n\nWeaknesses & Suggestions:\n- I find the proposed method somewhat confusing and the motivation a little unclear. It seems to me that the adversarial approach (minimizing treatment predictability) and the distributional discrepancy approach are somewhat redundant - it's been shown that one often implies the other, and so I don't see the benefit of concatenating both. Additionally, if in fact they do have different properties, it seems like concatenating them would be counterproductive; e.g. an adversary might perform poorly on the adversarially-optimized embedding, but well when we concatenate a non-adversarially-optimized embedding onto it. Because of this confusion I'm not able to gain much insight from the method or experimental results, which I feel like is particularly important in causal inference work; for instance, the theoretical results are nice but it's not so clear how the proposed setup helps optimize them\n- I think there should be more clarity off the top (intro and exposition) around the role of covariate shift here. For instance, the term \"covariate shift\" is ellided with \"selection bias\" at the bottom of page 1 and seemingly general distribution shift as well, when in fact there are many types of distribution shift that can present issues. A discussion of this would be helpful, at least for scoping purposes\n-I think the phrase \"over-balancing\" is not necessarily helpful - the issue is not always \"too much\" balancing, but rather the wrong kind of balancing when we lose accuracy\n- the motivating example on p2 is unclear: it seems to me like in this case, if T=age, then mapping these representations onto each other perfectly would be helpful, since then we can build an estimator for T=old and T=young separately. I'm not sure I understand what the issue is in this case.\n- the lack of insight around why this method might work makes me wonder if there is another reason, such as increased representational capacity achieved by concatenating representations. It would be good to either address some of these alternative reasons or clarify the motivations more \n\nSmaller points:\n- in 2.1, not sure why \\phi needs to be invertible - I know that's true in some works but it doesn't in general need to be the case for learning balanced representations right?\n- in 3, could use just a short explanation of what A and \\sigma^2 are intended to be - I understand but would help with readability\n- on bottom on p8, the authors claim that INet is novel - however there is quite a bit of work in the fair representation learning literature proposing very similar ideas, so I'm not sure I would agree with the claim. See \"Censoring Representations with an Adversary\" (Edwards & Storkey) and \"Learning Adversarially Fair and Transferable Representations\" (Madras et al) for two examples (should probably touch on this line of work in the related work section as well)\n- could use some more information in 5.1 on how IHDP is created and which treated samples are removed to create selection bias\n- I think the visualization in Fig 5 could be a bit stronger - no real reason AFAICT to lay out the x-axis over seeds like this. Also, would DGNet and DINet be useful to show in 5c and 5d respectively?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents an approach for estimating individual treatment effects - given an input X, the difference between the treatment and control potential (i.e. counterfactual) outcomes. The approach relies on two techniques for removing information from embeddings which have been explored in related literature: minimizing distributional distance between embeddings of treated and control units, and minimizing the ability of an adversary to predict from the embedding whether a unit is treated or control. They prove theoretical results similar to those in Shalit et al (2017) connecting these two properties to treatment effect estimation error (e.g. PEHE). Their proposed method includes two types of embeddings corresponding to these two properties, and concatenates them for the full model (DIGNet). They show that DIGNet outperforms baseline approaches on standard treatment effect estimation semi-synthetic setups like IHDP and Twins.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- my main point is around the method - it's not clear to me how the distributional distance minimizing method and the adversarial optimization method are either different, or complementary. This is key and I wouldn't accept this paper without understanding this\n- I would want to see more clarity in the intro around distribution shifts and the role of covariate shift in the scope\n- the motivating example on p2 needs clarification I think - this may help quite a bit with overall motivation\n- I'd like some extra thought put into empirical clarifications of a) why this approach works or b) why alternative explanations for the method working are not correct", "broader_impact_concerns": "no concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "J5M3yGiCVk", "reviewer": "Reviewer_2zS4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- This paper develops upper bounds for the counterfactual errors and ITE errors based on H-divergence.\n- The proposed framework DIGNet learns representations by adversarial learning using a set of objectives to balance factual prediction and causal effect estimation.\n\n**Weaknesses**\n- The motivation is not clear. It seems this paper targets treatment effect estimation using representation balancing. The underlying mechanism reveals the factual errors and the distribution discrepancy bound the counterfactual error. But the contradiction is unclear. I wonder whether one can directly optimize Eq (4).\n- The framework is based on potential outcomes and adopts the most common assumptions in potential outcomes. However, the framework implicitly assumes X->T->Y and X -> Y.\n- The proposed framework targets to mitigate the trade-off problem. However, the objective functions in Eq (20-23) are still difficult to balance. The adversarial games and the coefficients alpha are both hard to handle in optimization.\n- The method is only evaluated in synthetic data and semi-synthetic data. It is not evaluated in real-world datasets. The real-world datasets may not have the ground truth but there might be some challenges, such as causal graphs, and hyper-parameter tuning. I wonder whether the proposed method is sensitive or robust for misspecification.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies learning representation for treatment effect estimation under covariate shift. This paper proposes a representation balancing model, called DIGNet, for treatment effort estimation under covariate shift.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- This paper is about approximation rather than estimation. Given sufficient data samples, it is possible to estimate causal effects under a covariate shift. I assume this paper works on a numerical approximation from samples.\n- The motivation example is confusing as the age is identical to the treatment. It would be better to use a probabilistic relationship between - In Sec 2 preliminaries, there are the N iid random variables --> N randomly selected samples. The factual outcome and counterfactual come should be distinct.\n- In the discussion in Theorem 2, \"This new theoretical result provides a theoretical foundation for representation balancing models based on individual propensity confusion.\" Please explain either in this paragraph or in Sec 4.1.2.", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: In their final assessments, two reviewers voted "leaning reject" and one voted "leaning accept".  The core issues for the two who voted leaning reject were in the motivation of the approach and the associated clarity of how this is presented in the paper.

Also factoring into my decision is that this is a "major revision" of a previous TMLR submission (https://openreview.net/forum?id=uyp8eFbzzT) that was only relatively narrowly rejected with clear requested changes. Unfortunately, it was only possible to re-recruit one of the same reviewers from the previous submission (they changed their assessment from leaning reject to leaning accept). I generally feel that most of the issues raised from the previous submission have been corrected, with the exception of the overall clarity of the work (especially in motivating its core approach) which I still believe is problematic.  I do not feel that many substantial new issues have been raised in this revision round, but the two new revisions took a more critical stance on the clarity issues with the paper and lack of sufficient motivation for the core approach.

Overall, I think this is still quite a borderline decision.  However, ultimately I do believe that the vast majority of the core issues in the paper have been addressed over the two submission cycles and that the paper is now suitable for publication at TMLR.  In particular, while I think that many justified criticisms have been raised by reviewers that would likely prohibit publication at top ML conference venues, I do think it satisfies the criteria set out for acceptance by TMLR.  My recommendation is therefore that the paper should now be accepted.

That said, I still think the paper has noticeable room for improvement in the clarity of its key ideas.  I would therefore like to request minor changes based on further refinement of the introduction and general exposition of the core contributions.  I would like to see the current introduction made significantly more concise (potentially moving some content elsewhere if needed) so that it provides a more direct and focussed motivation and summary of the work, and more clearly explains the rationale behind the precise representation concatenation being used (i.e. why the approach being used achieves the two principles introduced, rather than just motivating them in isolation).
Decision: Accept (poster)