summaries: This paper proposed a CATE estimation approach when the treatment variable $T$ is a graph. Under a separable assumption (A3) between confounders $X$ and the treatment $T$ in the outcome structural equation ($Y$), Robinson's decomposition of CATE for binary $T$ can be applied. The performance of the authors' proposal was evaluated in both synthetic experiments and real data analysis. The paper proposed a CATE estimation method for graph-structured treatment by Robinson Decomposition. It is an extension of Robinson Decomposition, and the experiment shows the effectiveness of GIN on two small world datasets. The writing is good, but some contents are missing, so that the paper is not that easy to read. My detailed comments and concerns are described in the following. This paper proposes graph intervention network (GIN) to estimate conditional average causal effects when treatments are graph-structured. GIN builds on an extension of traditional Robinson decomposition to arbitrary treatment types. In particular, GIN factorizes the outcome Y into an inner product of two terms $g(X)^\top h(T)$. During training, GIN first estimates a mean outcome model $m(X)$ and then estimate $g,h$ using the estimated $m(X)$ function. The method is evaluated on two datasets with synthetic outcome functions. The paper considers the problem of estimating conditional average treatment effect when the treatment is a graph. The paper generalizes the Robinson decomposition, adapts the R-learner to this graph-treatment setting and introduces a two-stage algorithm. The performance of the proposed algorithm is compared to a few benchmarks through numerical experiments.  This paper studies the estimation of conditional average treatment effect (CATE) with multiple treatments that are graph structured. The proposed method extends the R-learner (Nie and Wager, 2020), which is designed for binary treatment, by learning a representation that reduces the dimension of the treatments. The representation is parametrized by a graph neural network that takes the graph structure into account. Unlike the standard R-learner that learns the nuisance components $E[Y|X]$ and $E[T|X]$ separately from the CATE, the proposed method only estimates $E[Y|X]$ separately while learns the representation $h(T)$, the propensity features $E[h(T)|X]$, and CATE jointly via gradient-based methods. 
abstractText: We address the estimation of conditional average treatment effects (CATEs) for structured treatments (e.g., graphs, images, texts). Given a weak condition on the effect, we propose the generalized Robinson decomposition, which (i) isolates the causal estimand (reducing regularization bias), (ii) allows one to plug in arbitrary models for learning, and (iii) possesses a quasi-oracle convergence guarantee under mild assumptions. In experiments with small-world and molecular graphs we demonstrate that our approach outperforms prior work in CATE estimation.
authors.name: Jean Kaddour
authors.name: Yuchen Zhu
authors.name: Matt J. Kusner
references.authors: A. Alaa M. van der Schaar
references.title: Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design
references.venue: In Dy,
references.authors: A.M. Alaa M. van der Schaar
references.title: Bayesian inference of individualized treatment effects using multi-task gaussian processes
references.venue: Advances in Neural Information Processing Systems,
references.authors: S. Athey G. Imbens
references.title: Recursive partitioning for heterogeneous causal effects
references.venue: Proceedings of the National Academy of Sciences,
references.authors: S. Athey S. Wager
references.title: Estimating treatment effects with causal forests: An application
references.venue: arXiv preprint arXiv:1902.07409,
references.authors: S. Athey J. Tibshirani S. Wager
references.title: Generalized random forests
references.venue: The Annals of Statistics,
references.authors: I. Bica J. Jordon M. van der Schaar
references.title: Estimating the effects of continuous-valued interventions using generative adversarial networks
references.venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems
references.authors: A. Caron I. Manolopoulou G. Baio
references.title: Estimating individual treatment effects using non-parametric regression models: a review
references.venue: arXiv preprint arXiv:2009.06472,
references.authors: V. Chernozhukov D. Chetverikov M. Demirer E. Duflo C. Hansen W. Newey J. Robins
references.title: Double/debiased machine learning for treatment and structural parameters
references.venue: The Econometrics Journal, 21(1):C1–C68,
references.authors: A. Curth M. van der Schaar
references.title: Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms
references.venue: The 24th International Conference on Artificial Intelligence and Statistics,
references.authors: M. Fey J.E. Lenssen
references.title: Fast graph representation learning with PyTorch Geometric
references.venue: In ICLR Workshop on Representation Learning on Graphs and Manifolds,
references.authors: S. Flaxman S. Mishra A. Gandy H.J.T. Unwin T.A. Mellan H. Coupland C. Whittaker H. Zhu T. Berah Eaton J. W
references.title: Estimating the effects of non-pharmaceutical interventions on covid-19 in europe
references.authors: Health Statistics N. C
references.title: 2007–2008 national health and nutrition examination survey (nhanes). US Department of Health and Human Services, Centers for Disease Control and Prevention
references.authors: A. Gretton K. Fukumizu C.H. Teo L. Song B. Schölkopf A.J. Smola
references.title: A kernel statistical test of independence
references.venue: Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems,
references.authors: P.R. Hahn J.S. Murray Carvalho C. M
references.title: Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion)
references.venue: Bayesian Analysis,
references.authors: S. Harada H. Kashima
references.title: Graphite: Estimating individual effects of graph-structured treatments
references.venue: arXiv preprint arXiv:2009.14061,
references.authors: J. Hartford G. Lewis K. Leyton-Brown M. Taddy
references.title: Deep IV: A flexible approach for counterfactual prediction
references.venue: Proceedings of the 34th International Conference on Machine Learning,
references.authors: T. Hatt S. Feuerriegel
references.title: Estimating average treatment effects via orthogonal regularization
references.venue: arXiv preprint arXiv:2101.08490,
references.authors: J.L. Hill
references.title: Bayesian nonparametric modeling for causal inference
references.venue: Journal of Computational and Graphical Statistics,
references.authors: G. Imbens D. Rubin
references.title: Causal Inference for Statistics, Social, and Biomedical Sciences
references.authors: G.W. Imbens
references.title: Nonparametric estimation of average treatment effects under exogeneity: A review
references.venue: The Review of Economics and Statistics,
references.authors: A. Jesson S. Mindermann U. Shalit Y. Gal
references.title: Identifying causal-effect inference failure with uncertainty-aware models
references.venue: Advances in Neural Information Processing Systems,
references.authors: A. Jesson S. Mindermann Y. Gal U. Shalit
references.title: Quantifying ignorance in individual-level causal-effect estimates under hidden confounding
references.venue: In International Conference on Machine Learning. PMLR,
references.authors: F. Johansson U. Shalit D. Sontag
references.title: Learning representations for counterfactual inference
references.venue: In International conference on machine learning,
references.authors: N. Kallus
references.title: DeepMatch: Balancing deep covariate representations for causal inference using adversarial training
references.venue: Proceedings of the 37th International Conference on Machine Learning,
references.authors: E.H. Kennedy
references.title: Optimal doubly robust estimation of heterogeneous causal effects
references.venue: arXiv preprint arXiv:2004.14497,
references.authors: E.H. Kennedy Z. Ma M.D. McHugh D.S. Small
references.title: Nonparametric methods for doubly robust estimation of continuous treatment effects
references.venue: Journal of the Royal Statistical Society. Series B, Statistical Methodology,
references.authors: D.P. Kingma J. Ba
references.title: Adam: A method for stochastic optimization
references.venue: 3rd International Conference on Learning Representations,
references.authors: S.R. Künzel J.S. Sekhon P.J. Bickel B. Yu
references.title: Metalearners for estimating heterogeneous treatment effects using machine learning
references.venue: Proceedings of the National Academy of Sciences,
references.authors: F Li
references.title: Propensity score weighting for causal inference with multiple treatments
references.venue: Annals of Applied Statistics,
references.authors: M.J. Lopez R. Gutman
references.title: Estimation of causal effects with multiple treatments: a review and new ideas
references.venue: Statistical Science,
references.authors: C. Louizos U. Shalit J.M. Mooij D. Sontag R. Zemel M. Welling
references.title: Causal effect inference with deep latent-variable models
references.venue: Advances in Neural Information Processing Systems,
references.authors: D. Lu C. Tao J. Chen F. Li F. Guo L. Carin
references.title: Reconsidering generative objectives for counterfactual reasoning
references.venue: Advances in Neural Information Processing Systems,
references.authors: K.W. Ma J.P. Lewis W.B. Kleijn
references.title: The HSIC bottleneck: Deep learning without back-propagation
references.venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,
references.authors: A. Mastouri Y. Zhu L. Gultchin A. Korba R. Silva M.J. Kusner A. Gretton K. Muandet
references.title: Proximal causal learning with kernels: Two-stage estimation and moment restriction
references.venue: In International Conference on Machine Learning. PMLR,
references.authors: S. Mendelson J. Neeman
references.title: Regularization in kernel learning
references.venue: The Annals of Statistics,
references.authors: W. Miao Z. Geng E.T. Tchetgen
references.title: Identifying causal effects with proxy variables of an unmeasured confounder
references.venue: arXiv preprint arXiv:1609.08816,
references.authors: C. Morris M. Ritzert M. Fey W.L. Hamilton J.E. Lenssen G. Rattan M. Grohe
references.title: Weisfeiler and leman go neural: Higher-order graph neural networks
references.venue: Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: R. Nabi T. McNutt I. Shpitser
references.title: Semiparametric causal sufficient dimension reduction of high dimensional treatments
references.venue: arXiv preprint arXiv:1710.06727,
references.authors: B. Neal Huang C.-W S. Raghupathi
references.title: Realcause: Realistic causal inference benchmarking
references.venue: arXiv preprint arXiv:2011.15007,
references.authors: L. Nie M. Ye qiang liu D. Nicolae
references.title: VCN}et and functional targeted regularization for learning causal effects of continuous treatments
references.venue: In International Conference on Learning Representations,
references.authors: X. Nie S. Wager
references.title: Quasi-oracle estimation of heterogeneous treatment
references.venue: effects. Biometrika,
references.authors: A. Paszke S. Gross F. Massa A. Lerer J. Bradbury G. Chanan T. Killeen Z. Lin N. Gimelshein L. Antiga A. Desmaison A. Köpf E. Yang Z. DeVito M. Raison A. Tejani S. Chilamkurthy B. Steiner L. Fang J. Bai S. Chintala
references.title: Pytorch: An imperative style, high-performance deep learning library
references.venue: Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems
references.authors: J. Pearl
references.title: Causality: models, reasoning, and inference
references.authors: M. Pollmann
references.title: Causal inference for spatial treatments
references.venue: arXiv preprint arXiv:2011.00373,
references.authors: R. Ramakrishnan P.O. Dral M. Rupp O.A. von Lilienfeld
references.title: Quantum chemistry structures and properties of 134 kilo molecules
references.venue: Scientific Data,
references.authors: P.M. Robinson
references.title: Root-n-consistent semiparametric regression
references.venue: Econometrica, 56(4):931–954,
references.authors: P.R. Rosenbaum D.B. Rubin
references.title: The central role of the propensity score in observational studies for causal effects
references.authors: L. Ruddigkeit R. van Deursen L.C. Blum Reymond J.-L
references.title: Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17
references.venue: Journal of Chemical Information and Modeling,
references.authors: M.S. Schlichtkrull T.N. Kipf P. Bloem R. van den Berg I. Titov M. Welling
references.title: Modeling relational data with graph convolutional networks
references.venue: The Semantic Web - 15th International Conference,
references.authors: P. Schwab L. Linhardt W. Karlen
references.title: Perfect match: A simple method for learning representations for counterfactual inference with neural networks
references.venue: arXiv preprint arXiv:1810.00656,
references.authors: P. Schwab L. Linhardt S. Bauer J.M. Buhmann W. Karlen
references.title: Learning Counterfactual Representations for Estimating Individual Dose-Response Curves
references.venue: In AAAI Conference on Artificial Intelligence,
references.authors: U. Shalit F.D. Johansson D. Sontag
references.title: Estimating individual treatment effect: generalization bounds and algorithms
references.venue: In International Conference on Machine Learning,
references.authors: C. Shi D. Blei V. Veitch
references.title: Adapting neural networks for the estimation of treatment effects
references.venue: Advances in Neural Information Processing Systems,
references.authors: C. Shi V. Veitch D. Blei
references.title: Invariant representation learning for treatment effect estimation
references.venue: arXiv preprint arXiv:2011.12379,
references.authors: R. Silva
references.title: Observational-interventional priors for dose-response learning
references.venue: Advances in Neural Information Processing Systems,
references.authors: R. Singh M. Sahani A. Gretton
references.title: Kernel instrumental variable regression
references.venue: Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems
references.authors: E.J.T. Tchetgen A. Ying Y. Cui X. Shi W. Miao
references.title: An introduction to proximal causal learning
references.venue: arXiv preprint arXiv:2009.10982,
references.authors: S. Wager S. Athey
references.title: Estimation and inference of heterogeneous treatment effects using random forests
references.venue: Journal of the American Statistical Association,
references.authors: J.N. Weinstein E.A. Collisson G.B. Mills K.R.M. Shaw B.A. Ozenberger K. Ellrott I. Shmulevich C. Sander J.M. Stuart
references.title: The cancer genome atlas pan-cancer analysis project
references.venue: Nature genetics,
references.authors: L. Xu Y. Chen S. Srinivasan N. de Freitas A. Doucet A. Gretton
references.title: Learning deep features in instrumental variable regression
references.venue: In International Conference on Learning Representations,
references.authors: L. Yao S. Li Y. Li M. Huai J. Gao A. Zhang
references.title: Representation learning for treatment effect estimation from observational data
references.venue: Advances in Neural Information Processing Systems,
references.authors: D.W. Zhang G.J. Burghouts C.G.M. Snoek
references.title: Set prediction without imposing structure as conditional density estimation
references.venue: In International Conference on Learning Representations,
references.authors: Z. Zhou S. Athey S. Wager
references.title: Offline multi-action policy learning: Generalization and optimization
references.venue: arXiv preprint arXiv:1810.04778,
references.authors: Y. Zhu D.L. Coffman D. Ghosh
references.title: A boosting algorithm for estimating generalized propensity scores with continuous treatments
references.venue: Journal of Causal Inference,
sections.heading: 1 Introduction
sections.text: Estimating feature-level causal effects, so-called conditional average treatment effects (CATEs), from observational data is a fundamental problem across many domains. Examples include understanding the effects of non-pharmaceutical interventions on the transmission of COVID-19 in a specific region [12], how school meal programs impact child health [13], and the effects of chemotherapy drugs on cancer patients [52]. Supervised learning methods face two challenges in such settings: (i) missing interventions, the fact that we only observe one treatment for each individual means models must extrapolate to new treatments without access to ground truth, and (ii) confounding factors that affect both treatment assignment and the outcome means that extrapolation from observation to intervention requires assumptions. Many approaches have been proposed to overcome these issues [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 18, 19, 21, 22, 23, 25, 27, 29, 33, 39, 41, 42, 45, 52, 56, 57, 60, 64, 67].
  In many cases, treatments are naturally structured. For instance, a drug is commonly represented by its molecular structure (graph), the nutritional content of a meal as a food label (text), and geographic regions affected by a new policy as a map (image). Taking this structure into account can provide several advantages: (i) higher data-efficiency, (ii) capability to work with many treatments, and (iii) generalizing to unseen treatments during test time. However, the vast majority of prior work operates on either binary or continuous scalar treatments (structured treatments are rarely considered, a notable exception to this trend is Harada & Kashima [16] which we describe in Section 2).
  To estimate CATEs with structured interventions, our contributions include:
  • Generalized Robinson decomposition (GRD): A generalization of the Robinson decomposition [47] to treatments that can be vectorized as a continuous embedding. This GRD reveals a learnable ∗Correspondence to jean.kaddour.20@ucl.ac.uk
  35th Conference on Neural Information Processing Systems (NeurIPS 2021).
  pseudo-outcome target that isolates the causal component of the observed signal by partialling out confounding associations. Further, it allows one to learn the nuisance and target functions using any supervised learning method, thus extending recent work on plug-in estimators [42, 29].
  • Quasi-oracle convergence guarantee: A result that shows that given access to estimators of certain nuisance functions, as long as the estimates converge at an O(n−1/4) rate, the target estimator for the CATE achieves the same error bounds as an oracle who has ground-truth knowledge of both nuisance components, the propensity features, and conditional mean outcome.
  • Structured Intervention Networks (SIN): A practical algorithm using GRD, representation learning, and alternating gradient descent. Our PyTorch [43] implementation is online.2
  • Evaluation metrics designed for structured treatments. Since previous evaluation protocols of CATE estimators have mostly focused on binary or scalar-continuous treatment settings, we believe that our proposed evaluation metrics can be useful for comparing future work.
  • Experimental results with graph treatments in which SIN outperforms previous approaches.
sections.heading: 2 Related Work
sections.text: Closest to our work is GraphITE [16], a method that learns representations of graph interventions for CATE estimation. They propose to minimize prediction loss plus a regularization term that aims to control for confounding based on the Hilbert-Schmidt Independence Criterion (HSIC) [14]. This technique suffers from two drawbacks: (i) the HSIC requires multiplication of kernel matrices and scales quadratically in the batch size; (ii) selecting the HSIC kernel hyper-parameter is not straightforward, as ground-truth CATEs are never observed, and empirical loss does not bound CATE estimation error [1]. We discuss other related work not on structured treatments in Appendix A.
sections.heading: 3 Preliminaries
sections.heading: 3.1 Conditional Average Treatment Effects (CATEs)
sections.text: Imagine a dataset where each example (xi, ti, yi) ∈ D represents a hospital patient’s medical history record xi, prescribed drug treatment ti, and health outcome yi, as illustrated in Figure 1 (Center). Further, we wish to understand how changing the treatment changes a patient’s health outcome. The CATE, τ ( t′, ti,xi ) , describes the expected change in outcome for individuals with history xi, when treatment ti is replaced by t′, depicted in Figure 1 (Right). In real-world scenarios, we only observe one outcome for each patient at one treatment level. Further, the patient’s pre-treatment health conditions xi influence both the doctor’s treatment prescription and outcome, thereby confounding the effect of the treatment on the outcome.
  Formally, we have the dataset D= { (xi, ti, yi) }n i=1
  sampled from a joint distribution p (X,T, Y ), where Y = f (X,T) + ε, as depicted in Figure 1 (Left). We define the causal effect of fixing
  2https://github.com/JeanKaddour/SIN
  treatment variable T ∈ T to a value t on outcome variable Y ∈ R using the do-operator [44] as E [ Y | do(T = t) ] . Crucially, this estimate differs from the conditional expectation E [ Y | T = t ] in that it describes the effect of an external entity intervening on T by fixing it to a value t (removing the edge X → T). We further condition on pre-treatment covariates X to define the conditional causal estimand E [ Y | X = x, do(T = t) ] . The conditional average treatment effect (CATE) is the difference between expected outcomes at different treatment values t, t′ for given covariates x,
  τ(t′, t,x) , E [ Y | X = x, do(T = t′) ]︸ ︷︷ ︸ =:µt′ (x) −E [ Y | X = x, do(T = t) ]︸ ︷︷ ︸ =:µt(x) , (1)
  where µt (x) is defined as the expected outcome for a covariate vector x under treatment t.
  Because we do not observe both treatments t, t′ for a single covariate x, we need to make assumptions that allow us to identify the CATE from observational data.
  Assumption 1. (Unconfoundedness) There are no confounders of the effect between T and Y beyond X. Therefore, Pr ( Y ≤ y | x, do(t) ) = Pr ( Y ≤ y | x, t ) , for all (x, t, y).
  Assumption 2. (Overlap) It holds that 0 < p ( t | x ) < 1, for all (x, t).
  Assumption 2 means that all sub-populations have some probability of receiving any value of treatment (otherwise, some τ(t′, t,x) may be undefined or impossible to estimate.) These assumptions allow us to estimate the causal quantity τ(t′, t,x) through statistical estimands:
  τ ( t′, t,x ) = µt′ (x)− µt (x) = E [ Y | X = x,T = t′ ] − E [ Y | X = x,T = t ] . (2)
  While one can model µt(x) with regression models, such approaches suffer from bias [9, 26, 29] due to two factors: (i) associations between X and T, due to confounding, makes it hard to identify the distinct contributions of X and T on Y , and (ii) regularization for predictive performance can harm effect estimation. Mitigating these biases relies on exposing and removing nuisance components. This transforms the optimization into a (regularized) regression problem that isolates the causal effect.
sections.heading: 3.2 Robinson Decomposition
sections.text: One way to formulate such nuisance components is via the Robinson decomposition [47]. Originally a reformulation of the CATE for binary treatments, it was used by the R-learner [42] to construct a plug-in estimator. The R-learner exploits the decomposition by partialling out the confounding of X on T and Y . It also isolates the CATE, thereby removing regularization bias.
  Let the treatment variable be T ∈ {0, 1} and the outcome model p ( y | x, t ) parameterized as
  Y = f(X, T ) + ε ≡ µ0(X) + T × τb(X) + ε, (3)
  where we define error term ε such that E [ ε | x, t ] = E [ ε | x ] = 0, and τb (x) , τ (1, 0,x).
  Define the propensity score [48] e (x) , p ( T = 1 | x ) and the conditional mean outcome as
  m (x) , E [ Y | x ] = µ0 (x) + e (x) τb (x) . (4)
  From model (3) and the previous definitions, it follows that Y −m (X) = ( T − e (X) ) τb (X) + ε, (5)
  allowing us to define the estimator
  τ̂b (·) = arg min τb
  { 1
  n n∑ i=1 ( ỹi − t̃i × τb (xi) )2 + Λ ( τb (·) )} , (6)
  where ỹi , yi−m̂ (xi) and t̃i , ti−ê (xi) are pseudo-data points defined through estimated nuisance functions m̂(·), ê(·), which can be learned separately with any supervised learning algorithm.
sections.heading: 4 The Generalized Robinson Decomposition
sections.text: Our goal is to estimate the CATE τ(t′, t,x) for structured interventions t′, t (e.g., graphs, images, text) while accounting for the confounding of X on T and Y . Inspired by the Robinson decomposition, which has enabled flexible CATE estimation for binary treatments [6, 9, 33, 42], we propose the Generalized Robinson Decomposition from which we extract a pseudo-outcome that targets the causal effect. We demonstrate the usefulness of this decomposition from both a theoretical view (quasi-oracle convergence rate in Section 4.2) and practical view (Structured Intervention Networks in Section 5). For details on its motivation and derivation, we refer the reader to Appendix B.
sections.heading: 4.1 Generalizing the Robinson Decomposition
sections.text: To generalize the Robinson decomposition to structured treatments, we introduce two concepts: (a) we assume that the causal effect is a product effect: the outcome function f∗ (X,T) can be written as an inner product of two separate functionals, one over the covariates and one over the treatment, and (b) propensity features, which partial out the effects from the covariates on the treatment features. Similar techniques have been previously shown to add to the robustness of estimation [9, 42]. Assumption 3. (Product effect) We consider the following partial parameterization of p(y | x, t),
  Y = g (X) > h (T) + ε, (7)
  where g : X → Rd, h : T → Rd and E[ε | x, t] = E [ ε | x ] = 0, for all (x, t) ∈ X × T .
  This assumption is mild, as we can formally justify its universality. The following asserts that provided we allow the dimensionality of g and h to grow, we may approximate any arbitrary bounded continuous functions in C (X × T ) where X × T is compact. Proposition 1. (Universality of product effect) LetHX×T be a Reproducing Kernel Hilbert Space (RKHS) on the set X × T with universal kernel k. For any δ > 0, and any f ∈ HX×T , there is a d ∈ N such that there exist two d-dimensional vector fields g : X → Rd and h : T → Rd, where ‖f − g>h‖L2(PX×T ) ≤ δ. (Proof in Appendix C)
  This assumption allows us to simplify the expression of the CATE for treatments t′, t, given x,
  τ ( t′, t,x ) = g (x) > ( h ( t′ ) − h (t) ) . (8)
  Define propensity features eh (x) , E [ h (T) | x ] and m (x) , E [ Y | x ] = g (x) > eh (x) .
  Following the same steps as in Section 3.2, the Generalized Robinson Decomposition for eq. (7) is Y −m (X) = g (X)> ( h (T)− eh (X) ) + ε. (9)
  Given nuisance estimates m̂(·), êh(·), we can use this decomposition to derive an optimization problem for h(·), g(·) (note êh(·) implicitly depends on h(·), we address this dependence in Section 5).
  ĝ (·) , ĥ (·) , arg min g,h  1n n∑ i=1 ( Yi − m̂ (Xi)− g (Xi)> ( h (Ti)− êh (Xi) ))2 + Λ ( g (·) ) (10)
sections.heading: 4.2 Quasi-oracle error bound of Generalized Robinson Decomposition
sections.text: We establish the main theoretical result of our paper: a quasi-oracle convergence guarantee for the Generalized Robinson Decomposition under a finite-basis representation of the outcome function. This result is analogous to the R-learner for binary CATEs [42]: when the true e (·) ,m (·) are unknown, and we only have access to the estimators ê (·) , m̂ (·), then as long as the estimates converge at n−1/4 rate, the estimator τ̂b (·) achieves the same error bounds as an oracle who has ground-truth knowledge of these two nuisance components.
  More formally, provided the nuisance estimators m̂(·) and êh(·) converge at an O ( n−1/4 ) rate, our
  CATE estimator will converge at an Õ(n− 1
  2(1+p) ) rate for arbitrarily small p > 0, recovering the parametric convergence rate for when the true m(·) and eh(·) are provided as oracle quantities.
  Our analysis assumes that the outcome E [ Y | X = x,T = t ] can be written as a linear combination of fixed basis functions. By Proposition 1, as long as we have enough basis functions, this representation is flexible enough to capture the true outcome function. Assumption 4. Let α(X) ∈ Rdα , β(T) ∈ Rdβ be fixed, known orthonormal basis features on X ∈ Rdx , T ∈ Rdt , respectively. The true outcome function f∗(x, t) = E[Y | X = x,T = t] can be written as f∗(x, t) = α>(x)Θ∗β(t) for some (unknown) matrix of coefficients Θ∗.
  Note that by setting g = α>Θ∗ and h = β, we recover eq. (7). Additionally, we will need overlap in the basis features α(X ),β(T ). Assumption 5 (Overlap in features). The marginal distribution of features Pα(X )×β(T ) is positive, i.e. supp[Pα(X )×β(T )] = α(X )× β(T ).
  Assumption 5 is typically weaker than requiring overlap in X and T, i.e., when dα, dβ dx, dt. With further technical assumptions specified in Appendix F, we establish the following theorem. Theorem 2. Let Θ∗ denote the representer of the true outcome function. Suppose Assumptions 5, 6, and 4 hold. Moreover, suppose that the propensity estimate êh is uniformly consistent,
  sup x∈X ‖êh(x)− eh(x)‖ →p 0 (11)
  and the L2 errors converge at rate E [{ m̂(X)−m∗(X) }2] ,E [ ‖êh (X)− eh (X) ‖2 ] = O(a2n) (12)
  for some sequence an → 0, where (an) is such that an = O(n−κ) with κ > 14 . Further, we define the regret as the excess risk
  R ( Θ̂n ) , L ( Θ̂n ) − L (Θ∗) , L (Θ) , E [{( Y −m∗ (X) ) −α (X) Θ ( β (T)− eh (X) )}2] .
  (13)
  Suppose that we obtain Θ̂n via a penalized basis function regression variant of the Generalized Robinson Decomposition, with a properly chosen penalty Λn ( ‖Θ̂n‖2 ) (specified in the proof). Then,
  Θ̂n satisfies the regret bound: R ( Θ̂n ) = Õ(r2n) with rn = n − 1 2(1+p) for arbitrarily small p > 0.
sections.heading: 5 Structured Intervention Networks
sections.text: We introduce Structured Intervention Networks (SIN), a two-stage training algorithm for neural networks, which enables flexibility in learning complex causal relationships, and scalability to large data-sets. This implementation of GRD strikes a balance between theory and practice: while we assumed fixed basis-functions in Section 4.2, in practice, we often need to learn the feature maps from data. We leave the convergence analysis of this representation learning setting for future work.
sections.heading: 5.1 Training Algorithm
sections.text: We propose to simultaneously learn feature maps ĝ (X) , ĥ (T) using alternating gradient descent, so that they can adapt to each other. A remaining challenge is that learning êh(X) is now entangled with learning ĥ (T). While the R-learner is based on the idea of cross-fitting, where at each data point i we pick estimates of the nuisances that do not use that data point, we introduce a pragmatic representation learning approach for (ĝ, ĥ) that does not use cross-fitting3.
  3We could in principle use cross-fitting for êh, although the loop between fitting ĥ alternating with êh would break the overall independence between êhi (X) and data point i. While it is possible that cross-fitting for ê
  h is still beneficial in this case, for simplicity and for computational savings, we did not implement it.
  a SIN Training. Input: Stage 1 data D1 := {(xi, yi)}mi=1, Stage 2 data D2 := {(xi, ti, yi)}ni=1 Step sizes λθ, λη, λψ, λφ. Number of update steps K. Mini-batch sizes B1, B2.
  1: Initialize parameters: θ,η,ψ,φ 2: while not converged do . Stage 1 3: Sample mini-batch {(xb, yb)} mB1 b=1 4: Evaluate Jm (θ) 5: Update θ ← θ − λθ∇̂θJ(θ) 6: end while 7: while not converged do . Stage 2 8: Sample mini-batch {(xb, tb, yb)} nB2 b=1
  9: Evaluate Jg,h (ψ,φ) , Jeh (η) 10: for k = 1 to K do 11: Update φ← φ− λφ∇̂φJg,h (ψ,φ) 12: Updateψ ← ψ−λψ∇̂ψJg,h (ψ,φ) 13: end for 14: Update η ← η − λη∇̂ηJeh (η) 15: end while
  b Pseudocode in a PyTorch-like style. # Initialize submodels and optimizers m, e, g, h = MLP(...), MLP(...), MLP(...),
  GNN(...) m_opt, e_opt, g_opt, h_opt = Adam(m.params(),
  m_lr), Adam(e.params(), e_lr), ...
  # Stage 1 for batch in train_loader:
  X, Y = batch.X, batch.Y m_opt.zero_grad() F.mse_loss(m(X), Y).backward() m_opt.step()
  # Stage 2 for batch in train_loader:
  X, T, Y = batch.X, batch.T, batch.Y for _ in range(num_update_steps):
  g_opt.zero_grad() h_opt.zero_grad() F.mse_loss((g(X)*(h(T) - e(X))).sum
  (-1), (Y-m(X))).backward() g_opt.step() h_opt.step()
  e_opt.zero_grad() F.mse_loss(e(X), h(T)).backward() e_opt.step()
  Figure 2: The two-stage algorithm for training SIN.
  We learn surrogate models for the mean outcome and propensity features m̂θ(X) and êhη(X) with parameters θ ∈ Rdθ ,η ∈ Rdη , as well as feature maps for covariates and treatments ĝψ(X), ĥφ(T), parameterized by ψ ∈ Rdψ ,φ ∈ Rdφ . We denote regularizers by Λ (·). Figure 2 summarizes the algorithm. As the mean outcome model m̂θ (X) does not depend on the other components, we learn it separately in Stage 1. In Stage 2, we alternate between learning ψ,φ,η.
  Stage 1: Learn parameters θ of the mean outcome model m̂θ (X) based on the objective
  Jm (θ) = 1
  m m∑ i=1 ( yi − m̂θ (xi) )2 + Λ (θ) , (14)
  which relies only on covariates and outcome data D1 := { (xi, yi) }m i=1 .
  Stage 2: Learn parameters ψ,φ for the covariates and treatments feature maps ĝψ (X) , ĥφ (T), as well as parameters η for the propensity features êhη (X).
  Jg,h (φ,ψ) = 1
  n n∑ i=1
  ( yi − { m̂θ (xi) + ĝψ (xi) > ( ĥφ (ti)− êhη (xi) )})2 + Λ (ψ) + Λ (φ) .
  (15)
  This loss hinges on êhη (X), which needs to be learned by
  Jeh (η) = n∑ i=1 ∥∥∥ĥφ (ti)− êhη (xi)∥∥∥2 2 + Λ (η) , (16)
  note again the dependence on ĥφ (T). While it may be tempting to learn ψ,φ and η jointly, they have fundamentally different objectives (êhη (X) is defined as an estimate of the expectation E [ h (T) | x ] ). Therefore, we employ an alternating optimization procedure, where we take k ∈ {1, . . . ,K} optimization steps for ψ,φ towards Jg,h (ψ,φ) and one step for learning η. We observe that setting K > 1, i.e. updating ψ,φ more frequently than η, stabilizes the training process.
sections.heading: 5.2 Advantages of SIN
sections.text: We conclude by describing the beneficial properties of SIN, particularly in finite-sample regimes:
  1. Targeted regularization: Regularizing ĝψ (X) , ĥφ (T) in eq. (15) after partialing out confounding is a type of targeted regularization of the isolated causal effect. In contrast, outcome estimation methods can suffer from regularization-induced confounding, e.g., regularizing the effect estimate away from zero in the service of trying to improve predictive performance [29].
  2. Propensity features: Learning propensity features can help us to (i) partial out parts of X that cause the treatment but not the outcome, and (ii) dispose unnecessary components of T.
  3. Data-efficiency: In contrast to methods that split the data into disjoint models for each treatment group (known as T-learners for binary treatments [8, 10]), sharing causal effect parameters between all covariates regardless of their assigned treatment increases data-efficiency.
  4. Partial data: In settings without access to both the treatment assignment and the outcome but only access to one of them, one can leverage that data to improve the (nuisance) estimator further, e.g., when a patient’s recovery is observed one year after a drug was administered [33].
sections.heading: 6 Experiments
sections.text: Here we evaluate how CATE estimation with our proposed model SIN compares with prior methods.
sections.heading: 6.1 Experimental Setup
sections.text: Datasets. To be able to compute CATE estimation error w.r.t. a ground truth, we design two causal models: a simpler synthetic model with small-world graph treatments and a more complex model with real-world molecular graph treatments and gene expression covariates. The Small-World (SW) simulation contains 1,000 uniformly sampled covariates and 200 randomly generated Watts–Strogatz small-world graphs [61] as treatments. The Cancer Genomic Atlas (TCGA) simulation uses 9,659 gene expression measurements of cancer patients for covariates [62] and 10,000 sampled molecules from the QM9 dataset [46] as treatments. Appendix D details the data-generating schemes.
  Baselines. We compare our method to (1) Zero, a sanity-check baseline that consistently predicts zero treatment effect and equals the mean squared treatment effect (poorly regularized models may perform worse than that due to confounding), (2) CAT, a categorical treatment variable model using one-hot encoded treatment indicator vectors, (3) GNN, a model that first encodes treatments with a GNN and then concatenates treatment and individual features for regression, (4) GraphITE [16], a CATE estimation method designed for graph treatments (more details in Section 2). GNN and CAT reflect the performance of standard regression models. The contrast between these two provides insight into whether the additional graph structure of the treatment improves CATE estimation. To deal with unseen treatments during CATE evaluation, we map such to the most similar ones seen during training based on their Euclidean distance in the embedding space of the GNN baseline.
  Graph models. For small-world networks, we use k-dimensional GNNs [38], as to distinguish graphs they take higher-order structures into account. To model molecular graphs, we use Relational Graph Convolutional Networks [50], where the nodes are atoms and each edge type corresponds to a specific bond type. We use the implementations of PyTorch Geometric [11].
  Evaluation metrics. We extend the expected Precision in Estimation of Heterogeneous Effect (PEHE) commonly used in binary treatment settings [19] to arbitrary pairs of treatments (t, t′) as follows. We denote the Unweighted PEHE (UPEHE) and the Weighted PEHE (WPEHE) as
  UPEHE(WPEHE) , ∫ X ( τ̂ ( t′, t,x ) − τ ( t′, t,x ))2 p ( t | x ) p ( t′ | x ) p (x) dx, (17)
  where the weighted version gives less importance to treatment pairs that are less likely; to account for the fact that such pairs will have higher estimation errors. In fact, as the reliability of estimated effects decreases by how likely they are in the observational study, we evaluate all methods on U/WPEHE truncated to the top K treatments, which we call U/WPEHE@K. To compute this, for each x, we rank all treatments by their propensity p ( t | x ) (given by the causal model) in descending order. We
  take the top K treatments and compute the U/WPEHE for all ( K 2 ) treatment pairs.
  In-sample vs. out-sample. A common benchmark for causal inference methods is the in-sample task, which we include here for completeness: estimating CATEs for covariate values x found in the training set. This task is still non-trivial, as the outcome of only one treatment is observed during training 4. In contrast, and arguably of more relevance to decision making, the goal of the out-sample task is to estimate CATEs for completely unseen covariate realizations x′.
  Hyper-parameter tuning. To ensure a fair comparison, we perform hyper-parameter optimization with random search for all models on held-out data and select the best hyper-parameters over 10 runs.
  Propensity. We define the propensity (or treatment selection bias) as p ( T | x ) =
  softmax ( κW>X ) , where W ∈ R|T |×d,∀i, j : Wij ∼ U [0, 1] is a random matrix (sampled then fixed for each run). Recall |T | is the number of available treatments and let d be the dimensionality of the covariates. Here the bias strength κ is a temperature parameter that determines the flatness of the propensity (the lower the flatter, i.e., κ=0 corresponds to the uniform distribution).
sections.heading: 6.2 Comparison of Performances on different K Treatments
sections.text: Figure 3 shows the UPEHE@K of all methods for K ∈ {2, . . . , 10}. We also report the WPEHE@6 of all methods in Table 1. Unless stated otherwise, we report results for bias strengths κ = 10 and κ = 0.1 in the SW and TCGA datasets, respectively across 10 random trials.
  The results indicate that the relative performance of each method, for both the in-sample and outsample estimation tasks, is consistent. Further, they suggest that, overall, the performance of SIN is best due to a better isolation of the causal effect from the observed data compared to other methods. The performance difference between CAT and GNN across all results indicate that accounting for graph information significantly improves the estimates. We observe from the SW experiments that GraphITE [16] performs slightly better than GNN, while it is nearly the same as GNN on TCGA.
  Surprisingly, the results of the TCGA experiments with low bias strength κ = 0.1 expose that all models but SIN fail to isolate causal effects better than the Zero baseline. These results confirm that confounding effects of X on Y combined with moderate causal effects can cause severe regularization bias for black-box regression models, while SIN partials these out from the outcome by m̂θ (X). We include additional results on convergence and larger values of K in Appendix E.1.
  4The original motivation comes from Fisherian designs where the only source of randomness is on the treatment assignment [20]. Our motivation is simpler: rule out the extra variability from different covariates, highlighting the difference between methods due to different loss functions and less due to smoothing abilities.
sections.heading: 6.3 Comparison of Robustness to different Bias Strengths κ
sections.text: A strong selection bias (i.e. large κ) in the observed data makes CATE estimation more difficult, as it becomes unlikely to see certain treatments t ∈ T for particular covariates x. Here, we assess each model’s robustness to varying levels of selection bias, determined by κ, across 5 random seeds. In Figure 4, we see that SIN outperforms the baselines across the entire range of considered biases. Interestingly, SIN performs competitively even in a case with no selection bias (κ=0, which corresponds to a randomized experiment). Importantly, all performances seem to either stagnate (SW) or to increase (TCGA) with increasing biases. Notably, the poor performance of CAT suddenly improves on datasets with high bias. We believe this is because, in high bias regimes, we see fewer distinct treatments overall, which allows the CAT model to approach the performance of GNN.
sections.heading: 7 Limitations, Future Work and Potential Negative Societal Impacts
sections.text: Limitations and future work. Firstly, in some real-life domains, Assumption 1 (Unconfoundedness) can be too strong, as there may exist hidden confounders. There are two common strategies to deal with them: utilizing instrumental variables [17, 58, 63] or proxy variables [35, 37, 59]. Developing new approaches for structured interventions in such settings is a promising future direction. Secondly, SIN is based on neural networks; however, neural network initialization can impact final estimates. To obtain consistency guarantees, GRD can be combined with kernel methods [35, 58].
  Potential negative societal impacts. Because causal inference methods make recommendations about interventions to apply in real-world settings, misapplying them can have a negative real-world impact. It is crucial to thoroughly test these methods on realistic simulations and alter aspects of them to understand how violations of assumptions impact estimation. We have aimed to provide a comprehensive evaluation of structured treatment methods by showing how estimation degrades as less likely treatments are considered (Figure 3) and as treatment bias increases (Figure 4).
sections.heading: 8 Conclusion
sections.text: The main contributions of this paper are two-fold: (i) the generalized Robinson decomposition that yields a pseudo-outcome targeting the causal effect while possessing a quasi-oracle convergence guarantee under mild assumptions, and (ii) Structured Intervention Networks, a practical algorithm using representation learning that outperforms prior approaches in experiments with graph treatments.
sections.heading: Acknowledgements
sections.text: We thank Antonin Schrab, David Watson, Jakob Zeitler, Limor Gultchin, Marc Deisenroth and Shonosuke Harada for useful discussions and constructive feedback on the paper. JK and YZ acknowledge support by the Engineering and Physical Sciences Research Council with grant number EP/S021566/1. This work was partially supported by an ONR grant number N62909-19-1-2096 to RS. We thank the Alan Turing Institute for the provision of Azure cloud computing resources.
title: Causal Effect Inference for Structured Treatments