# Learning to Infer Counterfactuals: Meta-Learning for Estimating Multiple
Imbalanced Treatment Effects

Guanglin Zhou1, Lina Yao1, Xiwei Xu2, Chen Wang2, Liming Zhu2 1 University of
New South Wales, Sydney, Australia, 2 Data61, CSIRO, Sydney, Australia
guanglin.zhou<EMAIL_ADDRESS>xiwei.xu, chen.wang,
<EMAIL_ADDRESS>
###### Abstract.

We regularly consider answering counterfactual questions in practice, such as
”Would people with diabetes take a turn for the better had they choose another
medication?”. Observational studies are growing in significance in answering
such questions due to their widespread accumulation and comparatively easier
acquisition than Randomized Control Trials (RCTs). Recently, some works have
introduced representation learning and domain adaptation into counterfactual
inference. However, most current works focus on the setting of binary
treatments. None of them considers that different treatments’ sample sizes are
imbalanced, especially data examples in some treatment groups are relatively
limited due to inherent user preference. In this paper, we design a new
algorithmic framework for counterfactual inference, which brings an idea from
Meta-learning for Estimating Individual Treatment Effects (MetaITE) to fill
the above research gaps, especially considering multiple imbalanced
treatments. Specifically, we regard data episodes among treatment groups in
counterfactual inference as meta-learning tasks. We train a meta-learner from
a set of source treatment groups with sufficient samples and update the model
by gradient descent with limited samples in target treatment. Moreover, we
introduce two complementary losses. One is the supervised loss on multiple
source treatments. The other loss which aligns latent distributions among
various treatment groups is proposed to reduce the discrepancy. We perform
experiments on two real-world datasets to evaluate inference accuracy and
generalization ability. Experimental results demonstrate that the model
MetaITE matches/outperforms state-of-the-art methods.

Counterfactual inference; Meta-learning; Multiple imbalanced treatments

††copyright: none

##  1\. Introduction

Counterfactual questions are widespread (Liu et al., 2020; Chernozhukov et
al., 2013; Bottou et al., 2013; Alaa and van der Schaar, 2017; Glass et al.,
2013), i.e., ”Would people with diabetes take a turn for the better had they
choose another medication”. Randomized Control Trials (RCTs) are regarded as
the golden standard for estimating counterfactual effects due to the
randomness that the assigned treatments are not dependent on users.
Nevertheless, it is expensive and time-consuming to control randomness
strictly. Besides, RCTs are sometimes immoral and can not be performed in the
filed of healthcare (Khang, 2015). In contrast, observational data are
pervasive and relatively easier to acquire. For instance, electronic health
records (EHRs), which store patients’ information and doctors’ disease
diagnosis, are typical examples of observational data. This type of data is
composed of co-variables related to patients, different medications as
treatments for patients, and outcomes that refer to final response effects,
i.e., whether the patient recovers or not.

Figure 1. Illustration of effect estimation on two imbalanced treatments from
observational data. Consider that people with diabetes choose treatments from
two drugs: the control treatment is relatively cheaper and works well, while
the treated group refers to a new drug that is expensive but probably more
performance-enhancing. We aim to estimate the effects for patients in test
sets on these two treatments. The problem is that the control treatment seems
a better choice (simply $\frac{4}{6}$>$\frac{1}{2}$) due to the limited
samples for the treated group, opposite to the ground truth in the right
graph. For simplicity, we only illustrate two imbalanced treatments.

Recently, machine learning, especially representation learning and domain
adaptation, have been demonstrated as an effective mechanism to infer
counterfactual effects from observational data. The literature (Johansson et
al., 2016) is the first to connect counterfactual inference with domain
adaptation and introduces a regularization item that enforces similarity in
the distributions between binary treatments. (Shalit et al., 2017) extends
(Johansson et al., 2016) and adopts Integral Probability Metrics (IPMs) to
estimate distances between binary distributions, then formulates the model as
CFR-Wass. (Yoon et al., 2018) explores generative adversarial nets (GANs) to
infer treatment effects. Considering that previous works ignore local
similarity information, (Yao et al., 2018) approaches to preserve local
similarity and infer effects based on deep representation learning. (Du et
al., 2021) employs mutual information to reduce information loss, and (Zhou et
al., 2021) constructs an information loop to improve inference performance
further.

However, the methods mentioned above are all related to binary treatments
except that the GANITE (Yoon et al., 2018) is naively suitable for multiple
treatments. Moreover, none of them explicitly considers the problem of
imbalanced treatments. On one hand, patients face multiple choices. For
example, diabetes is a disease that damages the health of many people, and it
is characterized as an increase in blood sugar (Kooti et al., 2016). There are
mainly three treatments for diabetes: drug, diet and exercise treatment.
Furthermore, the drug treatment consists of oral medicine and injection
medicine. Technically speaking, only modelling two treatments in many
applications is not rigorous enough. On the other hand, imbalance is expected
due to inherent self-selection in observational studies. Sample sizes among
different treatment groups present significant variations and thus covariate
distribution changes across treatment groups, which causes biased estimation
for treatment effects. For example, in Figure 1, samples in the control group
are sufficient, while patients in the treated group are limited, probably
because the new drugs are expensive and hard to afford for most patients. The
imbalance causes wrong estimation for the true responses. (Shalit et al.,
2017) also mentions imbalance issue and presents that in a real-world program
called Infant Health and Development Program (IHDP), the size in the treated
group is much smaller than the control group (139 treated, 608 control).
Apparently, the setting of multiple imbalanced treatments is more in line with
the actual situation and research on counterfactual inference in this setting
is more supportive for practical applications.

Figure 2. The overview of our model MetaITE. Data collection
$\\{\mathcal{D}_{S^{j}},\mathcal{D}_{T}\\}$ from one randomly chosen source
treatment group $S^{j}$ and target treatment group $T$. Model parameters are
updated locally on the support set within the inner loop and updated globally
via three losses within the outer loop. Besides, we adopt MMD discrepancy to
regularize distribution disparity among treatment groups.

In order to cope with the above commonplace but ignored problems, in this
paper, we propose MetaITE, a meta-learning framework for inferring individual
treatment effects from observational data in the setting of multiple
imbalanced treatments. Meta-learning is an efficient mechanism towards few-
shot scenario and domain generalization (Li et al., 2018). Some recent works
have applied meta-learning into causality (Sharma et al., 2019; Bengio et al.,
2019; Ton et al., 2021). The works (Bengio et al., 2019; Ton et al., 2021)
deal with causal directionality to distinguish the cause from effect, in
contrast to counterfactual causality in our work. MetaCI is the first work
that employs meta-learning for counterfactual inference (Sharma et al., 2019).
It incorporates CFRNet (Johansson et al., 2016) into a meta-learning based
reptile framework. Nevertheless, it still focuses on binary treatments and
does not study the problem of imbalanced treatments.

In this work, we focus on the setting of multiple imbalanced treatments, where
samples in some groups are sufficient but limited in others. In particular, we
cast the estimation of causal effects of multiple imbalanced treatments in a
novel meta-learning framework, in which data episodes among treatment groups
are regarded as meta-learning tasks. The treatment groups with plenty of
samples are considered as source domains, and the treated groups with limited
examples are regarded as target domains. We develop a model agnostic gradient
descent framework to optimize the proposed method. We train a base meta-
learner on a set of source tasks and utilize tasks in the target domain to
fine-tune the learned model. We design the meta-learner with feature extractor
and inference network. In addition, we make use of supervised information of
the source samples. Moreover, considering distribution disparity among
different treatment groups, we conceive discrepancy loss to reduce disparity
via maximum mean discrepancy (MMD). As far as we know, this is the first work
to apply meta-learning for estimating multiple imbalanced treatment effects,
and this competitive framework is flexible and straightforward.

The key contributions are summarized as follows.

  * •

To the best of our knowledge, this is the first work that studies the problem
of multiple imbalanced treatments in counterfactual inference.

  * •

We propose a novel framework MetaITE for estimating individual treatment
effects by incorporating inferring counterfactuals into the meta-learning
paradigm. Moreover, we utilize the supervised information with source samples
and attach discrepancy measures to reduce distribution disparity among
treatment groups.

  * •

Extensive experiments on two real-world datasets demonstrate MetaITE
outperforms the state-of-the-arts.

##  2\. Problem Setup

Table 1. Notations. Symbol | Description
---|---
$\mathbf{x}_{i}$ | covariates of the $i$-th unit
$\mathbf{t}_{i}$ | treatment assignment for $i$-th unit
$\mathbf{y}_{i}$ or $\mathbf{y}_{i}^{F}$ | observed/factual outcome for $i$-th unit
$\mathbf{y}_{i}^{CF}$ | counterfactual outcome for $i$-th unit
$\hat{\mathbf{y}}_{i}^{F}$ | estimated factual outcome for $i$-th unit
$\hat{\mathbf{y}}_{i}^{CF}$ | estimated counterfactual outcome for $i$-th unit
$n$ | number of units
$p$ | dimension of raw data
$k$ | number of treatments
$\\{{S}^{j}\\}_{j=1}^{k-1}$ | source treatment groups
${T}$ | target treatment group
$\mathcal{D}_{S^{j}},\mathcal{D}_{T}$ | sample data collection in meta-update

Let $\mathcal{X}$ denote the $p$-dimensional feature space, and $\mathcal{Y}$
represent potential outcomes. Consider obtaining a set of observational data
$\\{\mathbf{x}_{i},\mathbf{t}_{i},\mathbf{y}_{i}\\}_{i=1}^{n}$, where
$\mathbf{x}_{i}\in\mathcal{X}\subseteq\mathbb{R}^{p}$ represents covariates
matrix related to users or patients. The $\mathbf{t}_{i}$ signifies the
treatment assigned to the unit $i$ in fact, and $\mathbf{t}_{i}$ is from the
set of $\mathcal{T}\in\\{0,1\\}^{k}$, while
$\mathbf{y}_{i}\in\mathcal{Y}\subseteq\mathbb{R}$ indicates the factual
outcome we observe.

For each unit $i$, we get covariates
$\mathbf{x}_{i}^{(0)},\mathbf{x}_{i}^{(1)},...,\mathbf{x}_{i}^{(p-1)}$, along
with the treatment $\mathbf{t}_{i}$ and factual outcome $\mathbf{y}_{i}$.
Unlike prior works, we extend binary treatments to multiple settings, where
$k\geq 2$. Following the potential outcome framework (Rubin, 1974, 2005),
there are $k$ potential outcomes for each sample. We clarify the observed
outcome as factual outcome $\mathbf{y}_{i}^{F}$ and unobserved outcome as
counterfactual outcome $\mathbf{y}_{i}^{CF}$. Note that users select only one
treatment for the factual outcome, and a set of outcomes for each unit
associated with the remaining $k-1$ treatments are all called
$\mathbf{y}_{i}^{CF}$. We aim to precisely estimate $\mathbf{y}_{i}^{F}$ and
$\mathbf{y}_{i}^{CF}$ for each sample, given $\mathbf{x}_{i}$ and
$\mathbf{t}_{i}$.

Following the potential outcome framework (Rosenbaum and Rubin, 1983), we make
two assumptions of unconfoundedness and overlap in this work. Strong
Ignorability consists of the above two assumptions. Within the strong
ignorability assumption, we can tackle the problem of approximating potential
outcomes using a machine learning model
$f:\mathcal{X}\times\mathcal{T}\rightarrow$ $\mathcal{Y}$.

Problem Definition Now, we describe formal problem definition. Suppose that we
obtain a set of observational data divided by multiple treatments, where
sample sizes among different groups are various . Furthermore, considering
that covariate distribution discrepancy among various treatments exists, we
regard each treatment group as one domain (Shalit et al., 2017). We first
select one treatment with limited samples as a target domain ${T}$. Then we
view the remaining multiple treatments as source domains
${S}^{1},\dots,{S}^{k-1}$, where some treatments have sufficient samples and
others have relatively fewer samples. The number of domains equals the number
of treatments.

For simplicity, we refer to domains as assigned treatments subsequently. We
build data collection for each domain which represented as
$\\{\mathbf{X},\mathbf{y}\\}_{j},j=1,2,\dots,k$. A base model is trained with
data collection comprised of one source domain $\mathcal{D}_{S^{j}}$ and
target domain $\mathcal{D}_{T}$ in meta update. Our goal is to learn the model
$f$ so that it generalizes well on several domains. In that manner, via the
predictive model $f$, we obtain response outcome $\mathbf{y}_{i}$ related to
each sample $\mathbf{x}_{i}$ and treatment $\mathbf{t}_{i}$, a vector made up
of factual outcome $\mathbf{y}_{i}^{F}$ and counterfactual outcomes
$\mathbf{y}_{i}^{CF}$.

##  3\. The proposed method

On the basis of problem setup, we present a meta-learning based method for
counterfactual inference in this section. In particular, we focus on multiple
treatments and analyze the situation where sample sizes among treatments are
imbalanced. The overview of our model is shown in Figure 2. First, we
introduce the base inference model and then present our MetaITE architecture
for counterfactual inference.

###  3.1. Base Model

Similar to previous works, we build a prediction model by feeding covariates
and treatments. The model’s outcomes refer to factual and counterfactual
effects when different treatments are fed into the model. The base prediction
model $f$ in this work includes two parts with feature extractor and inference
network, which is $f=g\circ h$. The feature extractor
$g(\psi):\mathcal{X}\rightarrow\mathcal{Z}$ stands for the process that embeds
raw data from input space $\mathcal{X}$ into latent representation space
$\mathcal{Z}$. And $h(\theta):\mathcal{Z}\rightarrow\mathcal{Y}$ predicts the
outcome with latent features.

####  3.1.1. Feature Extractor

Due to inherent selection bias in observational studies, users who choose
various treatments are in discrepancy. In other words, distributions among
different treatment groups deviate from each other. Therefore, the base model
aims to learn a treatment-invariant model using data samples from multiple
treatment groups. Specifically, in each iteration, the base model is fed with
data from one source treatment $\mathcal{D}_{S^{j}}$ that is selected randomly
and target treatment $\mathcal{D}_{T}$, and we expect the base model
generalizes well on these two domains. In such a manner, we obtain the
treatment-invariant model on all domains after the outer-loop is completed.

To a certain extent, the covariates matrices $\mathbf{X}$ can be represented
as user preference as it generally involves user profiles and contextual
semantics when the user makes a choice. We construct a multi-layers fully
connected neural network to embed raw covariates into latent feature space,
denoted as $g(\mathbf{X};\psi)$. $\mathbf{Z}=g(\mathbf{X};\psi)$, where
$\mathbf{Z}\in\mathbb{R}^{d_{z}}$; $d_{z}$ is the embedding size; $g$ denotes
the fully connected layers; $\psi$ is the parameter of the fully connected
layers in the feature extractor.

In this section, we consider attaching discrepancy measures between two
domains to enhance the latent feature transferability and improve the
generalization ability of inference in multiple domains (Kumagai and Iwata,
2019; Johansson et al., 2016; Long et al., 2015). The maximum mean discrepancy
(MMD) is an integral probability metric that is a class of distances on
probability measures (Gretton et al., 2012). MMD is a powerful paradigm that
compares two distributions via transforming each distribution into a
Reproducing Kernel Hilbert Space (PKHS). In our scenario, given two
distributions from one source domain ${S^{j}}$ and the target domain ${T}$, we
obtain corresponding latent representations $\mathbf{Z}_{{S^{j}}}$ and
$\mathbf{Z}_{{T}}$. Then an empirical estimate of the squared MMD with two
datasets is computed by

(1)
$\displaystyle\widehat{MMD}^{2}(\mathcal{D}_{S^{j}},\mathcal{D}_{T})=\left\|\sum_{n=1}^{N}\frac{\mathit{k}(\cdot,\mathbf{Z}_{{S^{j}}})}{N}-\sum_{m=1}^{M}\frac{\mathit{k}(\cdot,\mathbf{Z}_{{T}})}{M}\right\|_{\mathcal{H}_{k}}^{2}$
$\displaystyle=\sum_{n=1}^{N}\sum_{n=1}^{N}\frac{\mathit{k}(\mathbf{Z}_{{S^{j}}},\mathbf{Z}_{{S^{j}}}^{{}^{\prime}})}{N^{2}}-2\sum_{n=1}^{N}\sum_{m=1}^{M}\frac{\mathit{k}(\mathbf{Z}_{{S^{j}}},\mathbf{Z}_{{T}})}{NM}$
$\displaystyle+\sum_{m=1}^{M}\sum_{m=1}^{M}\frac{\mathit{k}(\mathbf{Z}_{{T}},\mathbf{Z}_{T}^{{}^{\prime}})}{M^{2}}$

where $\mathit{k(\cdot,\cdot)}$ measures distance in Hilbert space and we use
Gaussian kernel (Sriperumbudur et al., 2011) for distance calculation. Through
MMD, we formulate
$\mathcal{L}_{disc}=\widehat{MMD}^{2}(g(\mathbf{X}_{S^{j}};\psi),g(\mathbf{X}_{T};\psi))$
to measure the distribution discrepancy between each source domain $S^{j}$ and
target domain $T$. By minimizing $\mathcal{L}_{disc}$, we ensure that
$g(\mathbf{X};\psi)$ generates more balanced embeddings, even though
$\mathbf{X}$ come from various domains. Furthermore, $g(\cdot\ ;\psi)$ can be
optimized by stochastic gradient descent (SGD).

####  3.1.2. Inference Network

Given the embedding feature vector $\mathbf{z}_{i}$ for each unit, we obtain
the outcome prediction

(2) $\hat{\mathbf{y}_{i}}=h(\mathbf{z}_{i};\theta)$

where $h$ is implemented by a fully-connected layers neural network with the
network parameter $\theta$, as shown in blue blocks of Figure
LABEL:fig:basemodel.

For the scenario of classification, i.e., Twins dataset, a cross-entropy loss
is used to optimize the inference network:

(3)
$\mathcal{L}_{inf}=\frac{1}{n}\sum_{i=1}^{n}(\mathbf{y}_{i}log(\hat{\mathbf{y}}_{i})+(1-\mathbf{y}_{i})log(1-\hat{\mathbf{y}}_{i}))$

On the other hand, in response to the regression problem like the News
dataset, the network can be optimized by minimizing the objective of the mean-
squared error (MSE):

(4)
$\mathcal{L}_{inf}=\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{y}_{i}-\hat{\mathbf{y}}_{i}\right\|_{2}^{2}$

where $n$ denotes the number of samples in the training data. In the next
section, we will introduce how to utilize and optimize these two losses
$\mathcal{L}_{disc}$ and $\mathcal{L}_{inf}$ in the meta adaptation.

###  3.2. MetaITE Architecture

Generally, meta-learning aims to generalize well over various learning tasks
from numerous domains. The learning procedure consists of three parts:

  * •

episodic training scheme that splits available data into support sets and
query sets for meta-train and meta-test respectively

  * •

local update that optimizes model parameters on support sets within the inner-
loop

  * •

global update that updates parameters on the query sets within the outer-loop

####  3.2.1. Episodic Training

As Figure 2 shows, given training data from a set of domains
$\mathcal{D}_{S}=\\{\mathcal{D}_{S^{j}}\\}_{j=1}^{k-1}$ and one target domain
$\mathcal{D}_{T}$, we split training data into support sets for meta-train and
query sets for meta-test. In particular, we uniformly sample one source domain
${S^{j}}$ from the source sets for meta-train and choose data collection
$\mathcal{D}_{T}$ from the target domain for meta-test. As Algorithm 1
illustrates, we construct a meta batch in each outer loop. Each batch includes
two domains, one for the source and the other for the target. We randomly
choose $K$ samples in selected source domain as support set
$\\{\mathbf{X}_{i}^{Sup},\mathbf{y}_{i}^{Sup}\\}_{i=1}^{K}$ for the local
update, and $K$ identities from target domain as query set
$\\{\mathbf{X}_{i}^{Que},\mathbf{y}_{i}^{Que}\\}_{i=1}^{K}$ for global update,
where $K$ is a hyper-parameter.

####  3.2.2. Local Update

We follow the model-agnostic meta-learning (MAML) that learns transferrable
internal representations among different domains (Finn et al., 2017). The
local update corresponds to the inner-loop in Algorithm 1. Note that we aim to
ensure the base model generalizes well on multiple domains. We firstly let the
model learn well on the source domain with support set
$\\{\mathbf{X}_{i}^{Sup},\mathbf{y}_{i}^{Sup}\\}_{i=1}^{K}$. The corresponding
loss can be formulated as:

(5)
$\mathcal{L}_{Sup}=\sum_{i=1}^{K}\mathcal{L}_{inf}(\mathbf{y}_{i}^{Sup},h(g(\mathbf{X}_{i}^{Sup};\psi);\theta))$

And we use stochastic gradient descent (SGD) to optimize the parameters $\psi$
and $\theta$.

(6)
$(\psi^{{}^{\prime}},\theta^{{}^{\prime}})\leftarrow(\psi,\theta)-\alpha\nabla_{\psi,\theta}\mathcal{L}_{Sup}$

where $\alpha$ is a fixed hyper-parameter controlling the update rate. In
practice, the updated parameters $\psi^{{}^{\prime}}$ and
$\theta^{{}^{\prime}}$ are computed using several gradient descent updates. We
illustrate one gradient descent update for simplicity in Equation 6. The model
parameters are trained and optimized with the performance of $f(\cdot\
;\psi^{{}^{\prime}},\theta^{{}^{\prime}})$. Correspondingly, the meta-
objective can be formulated as :

(7)
$\mathop{min}\limits_{\psi}\mathcal{L}_{Sup}(f_{\psi^{{}^{\prime}}})=\mathop{min}\limits_{\psi}\sum_{\\{\mathbf{X},\mathbf{y}\\}\sim
p(\mathcal{D}_{S^{j}})}\mathcal{L}_{Sup}(\mathbf{y},f(\mathbf{X};\psi^{{}^{\prime}}))$

(8)
$\mathop{min}\limits_{\theta}\mathcal{L}_{Sup}(f_{\theta^{{}^{\prime}}})=\mathop{min}\limits_{\theta}\sum_{\\{\mathbf{X},\mathbf{y}\\}\sim
p(\mathcal{D}_{S^{j}})}\mathcal{L}_{Sup}(\mathbf{y},f(\mathbf{X};\theta^{{}^{\prime}}))$

where the number of sums equals to the number of gradient update in the inner
loop.

####  3.2.3. Global Update

When we obtain $\psi^{{}^{\prime}}$ and $\theta^{{}^{\prime}}$, the loss is
optimized with query set:

(9)
$\mathcal{L}_{Que}=\sum_{i=1}^{K}\mathcal{L}_{inf}(\mathbf{y}_{i}^{Que},h(g(\mathbf{X}_{i}^{Que};\psi^{{}^{\prime}});\theta^{{}^{\prime}}))$

Then, meta-optimization is performed:

(10)
$(\psi,\theta)\leftarrow(\psi,\theta)-\beta\nabla_{\psi,\theta}\mathcal{L}_{Que}$

where $\beta$ is the meta-learning rate and set as a hyper-parameter. The
global update corresponds to a meta-gradient update that involves a gradient
through a gradient in the outer-loop of Algorithm 1.

With the aid of local and global updates, the base model learns from source
domains with sufficient data and is able to fast adapt to the target domain.
In this way, we learn a model with great generalization on multiple
treatments.

####  3.2.4. Meta optimization

Although the MAML paradigm is capable of learning transferable information
among multiple domains via the above three steps, in the context of
counterfactual inference, we are required to provide precise outcomes related
to each treatment domain. Furthermore, domain discrepancy affects inference
performance. Therefore, in addition to combining the above losses on the
support set and query set, we consider incorporating $\mathcal{L}_{disc}$ in
Eq. (LABEL:eq_mmd) to optimize the MetaTIE jointly. The $\mathcal{L}_{disc}$
in the MetaITE criterion is formulated as:

(11)
$\mathcal{L}_{disc}=\widehat{MMD}^{2}(g(\mathbf{X}^{Sup};\psi),g(\mathbf{X}^{Que};\psi))$

where $\widehat{MMD}^{2}$ equals to Eq. (LABEL:eq_mmd). By enforcing latent
distributions in the source and target domains, discrepancy among different
treatment groups is reduced after the training iterations. The ultimate
objective loss for meta optimization is concluded, and Adam (Kingma and Ba,
2014) is used.

(12)
$\mathcal{L}_{obj}=\mu\mathcal{L}_{Que}+\epsilon\mathcal{L}_{Sup}+\gamma\mathcal{L}_{disc}+\left\|\omega\right\|_{2}$

The final term $\left\|\cdot\right\|_{2}$ is $l_{2}$ regularization for model
complexity. The detailed pseudo code of the training process is shown in
Algorithm 1. Regarding the testing section, we select test data to construct
query sets and sample support sets from each domain iteratively . In this way,
we gain all estimated effect outcomes coordinating with each treatment, i.e.,
$\\{\hat{\mathbf{y}}_{i}^{\mathbf{t}_{i}=t}\\}_{t=1}^{k}$. The estimate
process of MetaITE is shown in Algorithm 2.

Algorithm 1 : Training process of MetaITE

1:Input: Set of source domains
$\mathcal{D}_{S}=\\{\mathcal{D}_{S^{j}}\\}_{j=1}^{k-1}$; Target domain
$\mathcal{D}_{T}=\\{\mathcal{D}_{T}\\}$

2:Require: Hyperparameters $\alpha,\beta,\mu,\epsilon,\gamma$

3:Output: Feature Extractor $g(\psi)$, Inference Network $h(\theta)$

4:$\triangleright$ learning $\psi,\theta$ from training data

5:randomly initialize model parameters $\psi,\theta$

6:while Outer-loop not done do

7: Uniformly select one source domain ${S^{j}}$

8: Sample a batch from $\\{\mathcal{D}_{S^{j}},\mathcal{D}_{T}\\}$

9: while Inner-loop not done do

10: $K$ samples in support set from $\mathcal{D}_{S^{j}}$ :
$\\{\mathbf{X}_{i}^{Sup},\mathbf{y}_{i}^{Sup}\\}_{i=1}^{K}$

11: $K$ samples in the query set from $\mathcal{D}_{T}$ :
$\\{\mathbf{X}_{i}^{Que},\mathbf{y}_{i}^{Que}\\}_{i=1}^{K}$

12: Compute $\mathcal{L}_{Sup}$ using Eq.(5)

13: Compute adapted parameters $(\psi^{{}^{\prime}},\theta^{{}^{\prime}})$ via
Eq.(6)

14: end while

15: Compute $\mathcal{L}_{Que}$ with Eq.(9)

16: Compute discrepancy $\mathcal{L}_{disc}$ by using Eq.(11)

17: Update
$(\psi,\theta)\leftarrow(\psi,\theta)-\beta\nabla_{\psi,\theta}\mathcal{L}_{obj}$

18:end while

##  4\. Experiments

In this section, we evaluate model performance on two real-world datasets,
Twins and News. These two datasets are specified as benchmarks in baselines,
and we verify MetaITE and baselines in binary and multiple treatments. In
addition to performance evaluation, we also analyze robustness on imbalance
among treatment groups.

###  4.1. Datasets

Twins. This dataset is created from all twins
birth111http://data.nber.org/data/linked-birth-infant-death-data-vital-
statistics-data.html in the USA between 1989-1991, and we only pay attention
to the twins weighing less than 2kg and without missing features (Louizos et
al., 2017). The dataset is the same as (Yoon et al., 2018; Zhou et al., 2021)
for binary treatments. Significantly, there are 30 covariates related to the
parents, the pregnancy and the birth. We use the treatment $t=1$ as being the
heavier twin, and $t=0$ is expressed as the lighter twin. The outcome is
mortality after one year. The final dataset includes 11,400 pairs of twins.
The selection bias is introduced as that we selectively choose one of the
twins as the observation and hide the other: $t_{i}|x_{i}\sim
Bern(Sigmoid(w^{\mathrm{T}}x_{i}+n))$, where
$w^{T}\sim\mathcal{U}((-0.1,0.1)^{30\times 1})$ and $n\sim\mathcal{N}(0,01)$.
Moreover, the mortality rate for the lighter twin is 17.7%, and that in the
heavier twin is 16.1%. Concerning multiple treatments, we follow the procedure
in (Yoon et al., 2018; Louizos et al., 2017) and four treatments are shaped as
$(1)\ t=0$: lower weight and female sex, $(2)\ t=1$: lower weight and male
sex, $(3)\ t=2$: higher weight and female sex, $(4)\ t=3$: higher weight and
male sex. There are 11984 samples and 50 covariates. The outcome is the same
as that in binary treatments.

News. (Johansson et al., 2016) introduces the News dataset that simulates the
opinions of a media consumer exposed to multiple news items, which is
originated from the NY Times
corpus222https://archive.ics.uci.edu/ml/datasets/bag+of+words. The News
dataset is introduced for binary treatments in (Johansson et al., 2016) and
(Schwab et al., 2020) extends the original dataset to the multiple setting for
estimating average dose-response curves. We assume that readers prefer to read
certain content on some specific devices. For example, readers may like to
read brief news items like hot topics on mobile phones but prefer to read news
with professional content and longer paragraphs via computers, i.e., The
Economist. Furthermore, we aim to infer the individual treatment effects of
obtaining more content from some specific devices on the reader’s opinions. In
particular, each sample $\mathbf{x}_{i}$ refers to news items represented by
word counts, and simulated outcome $\mathbf{y}_{i}\in\mathbb{R}$ represents
the reader’s opinions of the news. Regarding the intervention
$\mathbf{t}_{i}$, we follow (Schwab et al., 2020) and construct multiple
treatments which correspond to various devices used to view the news items,
such as desktop, smartphone, newspaper, and tablet. We train an LDA topic
model (Blei et al., 2003) that characterizes the topic distribution of each
news item. We select 5000 samples, and 50 LDA topics are learned from the
corpus. We define $z(\mathbf{X})$ as the topic distribution of news items, and
randomly choose $k$ centroids in topic space $\\{z^{0},\cdots,z^{k-1}\\}$ as
different devices and set $z^{m}$ as the mean centroid of the whole topic
space. Random Gaussian distribution is derived with mean
$\mu\sim\mathcal{N}(0.45,0.15)$ and standard deviation
$\sigma\sim\mathcal{N}(0.1,0.05)$. The outcomes are firstly deduced with
$\tilde{{y}}\sim\mathcal{N}(\mu,\sigma)+\epsilon$ with
$\epsilon\sim\mathcal{N}(0,0.15)$. Then for each sample, we provide $k$
potential outcomes, which is $\mathbf{y}_{i}^{j}\sim C(\tilde{y}\
\mathbf{D}(z(\mathbf{x}_{i}),z^{j})+\mathbf{D}(z(\mathbf{x}_{i}),z^{m})),j=0,\cdots,k-1$,
where $C$ is a scaling factor and $\mathbf{D}$ is the Euclidean distance
metric. The observed treatment is regarded as
$\mathbf{t}_{i}|\mathbf{x}_{i}\sim$ softmax$(\kappa\mathbf{y}_{i})$ with a
treatment assignment bias $\kappa$. We set $C=50$ and $\kappa=10$ by following
previous work (Schwab et al., 2020). There are two different variants of this
dataset with $k=\\{2,4\\}$ viewing devices.

We randomly separate both datasets into training and testing sets with a ratio
of $80\%:20\%$.

###  4.2. Baselines

We adopt the following state-of-the-art models that achieve the best
performance in counterfactual inference.

  * •

OLS/LR1: least square regression using the treatment as a feature.

  * •

OLS/LR2: separate least square regressions for each treatment group.

  * •

k-NN: (Crump et al., 2008) k-nearest neighbour method for inferring individual
treatment effects.

  * •

BNN: (Johansson et al., 2016) refers to the balancing neural networks that
firstly incorporates representation learning into counterfactual inference.

  * •

CFR-Wass333https://github.com/clinicalml/cfrnet : (Shalit et al., 2017) builds
on and extends the work in (Johansson et al., 2016). We add a regularization
and learn a balanced representation between binary treatment groups
considering selection bias in observational data. They use Wasserstein
distance to measure and control discrepancy.

  * •

CFR-MMD : (Shalit et al., 2017) is a variant of CFR-Wass, which uses MMD
measures instead of Wasserstein distance.

  * •

TARNet: (Shalit et al., 2017) is a variant of CFR-Wass without balancing
regularization.

  * •

GANITE444https://github.com/jsyoon0823/GANITE: (Yoon et al., 2018) uses a
generative adversarial network to capture and learn counterfactual
distributions. The generator generates the proxies of counterfactual outcomes
and passes to infer individual treatment effects. GANITE is the only work that
can be applied to multiple treatments naively.

  * •

SITE555https://github.com/Osier-Yi/SITE: (Yao et al., 2018) considers that
existing estimation methods almost care about balancing the distributions of
control and treated groups, but none of them notes similarity information.
SITE selects triple samples and preserves local similarity information by
position-dependent deep metric.

  * •

ABCEI666https://github.com/octeufer/Adversarial-Balancing-based-
representation-learning-for-Causal-Effect-Inference: (Du et al., 2021) adopts
a mutual information estimator as the regularization item to preserve helpful
information for predicting counterfactual effects.

  * •

CBRE777https://github.com/jameszhou-gl/CBRE: (Zhou et al., 2021) utilizes
adversarial training to balance distributions between treated and control
groups. And CBRE builds an information loop to further reduce highly
predictive information loss.

###  4.3. Experimental Settings

####  4.3.1. Evaluation Metrics

We use the Rooted Precision in Estimation of Heterogeneous Effect
($\sqrt{\epsilon_{PEHE}}$) and Mean Absolute Error on ATE ($\epsilon_{ATE}$)
as performance metrics for binary treatments. The smaller the two metrics are,
the better the performance is.

(13)
$\sqrt{\epsilon_{PEHE}}=\sqrt{\frac{1}{n}\sum_{i=1}^{n}((\mathbf{y}_{i}^{t=1}-\mathbf{y}_{i}^{t=0})-(\hat{\mathbf{y}}_{i}^{t=1}-\hat{\mathbf{y}}_{i}^{t=0}))^{2}}$

(14)
$\epsilon_{ATE}=\left|\frac{1}{n}\sum_{i=1}^{n}(\mathbf{y}_{i}^{t=1}-\mathbf{y}_{i}^{t=0})-\frac{1}{n}\sum_{i=1}^{n}(\hat{\mathbf{y}}_{i}^{t=1}-\hat{\mathbf{y}}_{i}^{t=0})\right|$

In the context of multiple treatments, we use the root mean square error
($RMSE$) (Yoon et al., 2018) to evaluate counterfactual effects because
$\sqrt{\epsilon_{PEHE}}$ cannot be extended to the multiple treatments
setting.

(15)
$RMSE=\sqrt{\frac{1}{n\times\left|\mathcal{T}\right|}\sum_{i=1}^{n}\sum_{t_{i}\in\mathcal{T}}(\mathbf{y}^{t=t_{i}}(\mathbf{x}_{i})-\hat{\mathbf{y}}^{t=t_{i}}(\mathbf{x}_{i}))^{2}}$

Table 2. Performance Evaluation of MetaITE with other state-of-the-art methods on Twins and News datasets. Bold indicates the method with the best performance. The lower is the better. n.r. = not reported for reasons (null value or no convergence).  Methods | Twins_bin | News_2 | Twins_4 | News_4
---|---|---|---|---
$\sqrt{\epsilon_{PEHE}}$ | $\epsilon_{ATE}$ | $\sqrt{\epsilon_{PEHE}}$ | $\epsilon_{ATE}$ | $RMSE$ | $RMSE$
OLS/LR1 | $0.3097\pm 0.000$ | $0.0079\pm 0.000$ | $16.1785\pm 0.000$ | $16.9774\pm 0.000$ | $0.2267\pm 0.000$ | $11.04542\pm 0.000$
OLS/LR2 | $0.3175\pm 0.000$ | ${0.0099}\pm{0.000}$ | $14.1958\pm 0.000$ | $12.6928\pm 0.000$ | n.r. | $16.3344\pm 0.000$
K-NN | $0.3101\pm 0.000$ | $0.0077\pm 0.000$ | $27.3473\pm 0.000$ | $26.3602\pm 0.000$ | $0.2359\pm 0.000$ | $8.75516\pm 0.000$
BNN | $0.3179\pm 0.000$ | $0.0102\pm 0.000$ | $\textbf{13.6390}\pm\textbf{0.000}$ | $12.1902\pm 0.000$ | $1.4939\pm 0.000$ | $9.45723\pm 0.000$
TARNet | $0.3355\pm 0.000$ | $0.0101\pm 0.000$ | $19.0558\pm 0.000$ | $19.0550\pm 0.000$ | $0.1925\pm 0.000$ | $9.2592\pm 0.000$
CFR-Wass | $0.3112\pm 0.000$ | $0.0078\pm 0.000$ | $19.2389\pm 0.000$ | $19.2375\pm 0.000$ | ${0.1947}\pm{0.000}$ | $9.4580\pm 0.000$
CFR-MMD | $0.3355\pm 0.000$ | $0.0101\pm 0.000$ | $19.2642\pm 0.000$ | $19.2628\pm 0.000$ | $0.1968\pm 0.000$ | $10.7841\pm 0.000$
GANITE | $0.3228\pm 0.008$ | $0.0690\pm 0.039$ | n.r. | n.r. | $0.2277\pm 0.040$ | n.r.
SITE | $0.3180\pm 0.002$ | $0.0069\pm 0.001$ | $20.5922\pm 0.003$ | $16.7775\pm 0.000$ | $0.2004\pm 0.000$ | $9.2265\pm 0.000$
ABCEI | n.r. | n.r. | n.r. | n.r. | $0.2090\pm 0.000$ | n.r.
CBRE | $0.3302\pm 0.000$ | $0.0107\pm 0.000$ | $19.5432\pm 0.000$ | $19.542\pm 0.000$ | $0.2093\pm 0.000$ | $11.9447\pm 0.001$
MetaITE | $\textbf{0.3093}\pm\textbf{0.000}$ | $\textbf{0.0062}\pm\textbf{0.000}$ | ${17.1726}\pm{0.359}$ | $\textbf{9.2552}\pm\textbf{0.737}$ | $\textbf{0.1921}\pm\textbf{0.000}$ | $\textbf{8.7303 }\pm\textbf{0.100}$

Considering all of these baselines except the GANITE (Yoon et al., 2018) are
tailored explicitly to binary treatments, we extend these models for multiple
treatments as follows: we select one treatment as the control group and then
choose one as the treated group to perform binary treatment effects
estimation. Note that there are $2^{k}$ distinct treatment subsets for $k$
treatments, so we need to infer $2^{k}$ binary treatment-effects to sum up the
final multiple treatment effects. All baseline models are performed with
parameters as mentioned in their papers or code repositories.

####  4.3.2. Parameter Setting

The parameter setting for our model is included in appendix A.2. All the
models are trained by Adam optimizer, and the maximum iteration of MetaITE is
set as 15000.

###  4.4. Performance Evaluation

In this section, we evaluate our model MetaITE and compare it with baselines
on two benchmarks. The model is implemented with Tensorflow
1.5.1888https://www.tensorflow.org/ in Python 3.6. We perform all of the
experiments on one computing device with 16-core Intel Core i7-10700 CPUs and
a total of 500 GB Memory RAM. We repeat ten times and report mean and std
results.

Generally speaking, we derive datasets with two and four treatments from raw
data, i.e., Twins_bin, Twins_4, News_2, News_4. We believe evaluation on these
four subsets of data collection can demonstrate the inference performance and
generalization ability of MetaITE for estimating multiple treatments effects.

Below we give a detailed breakdown of the results of our experiments. We first
perform experiments on the Twins dataset with binary and four treatments
separately. For the Twins_bin dataset, there are $4594$ samples in the control
group for $t=0$, and we configure the sample size as 80 for the treated group.
As for Twins_4 dataset, there are $6058$ and $5926$ samples in $t=0$ and $t=2$
correspondingly. The sample number of $t=1$ and sample number of $t=3$ are
both 160. In this way, we simulate datasets for both binary and multiple
imbalanced treatments. As shown in Table 2, MetaITE outperforms other state-
of-the-art models for both binary and four treatments. The results in the
Twins dataset, regardless of binary treatments or four treatments, show that
models based on representation learning almost outperform regression-based
methods, especially for multiple treatments. So we believe a learning-based
method is more suitable for counterfactual inference. On the other hand,
relatively simpler models like CFR-Wass perform better than complex models
such as GANITE and CBRE. We consider that some treated groups have limited
samples in the imbalanced setting. In this way, model complexity in these
models does not match the sample size, which leads to poor generalization. The
results and analyses also indicate that the MetaITE is a promising framework
for estimating multiple imbalanced treatments. Note that the result of ABCEI
on Twins_bin contains null values, so we record n.r. for that.

Then we perform experiments on the News dataset. On the News_2 dataset, there
are 1634 samples and 160 samples for $t=0$ and $t=1$. And the News_4 contains
$\\{860,80,80,80\\}$ samples for treatment set $\\{0,1,2,3\\}$. The test
procedure is the same as that on Twins dataset. As results in Table 2 show,
MetaITE realizes the best performance on the metric $\sqrt{\epsilon_{ATE}}$ of
News_2 and $RMSE$ of News_4, and the performance is close to the best result
on the metric $\sqrt{\epsilon_{PEHE}}$. Because the code of GANITE is designed
for classification, it cannot be performed for regression tasks on the News
dataset so that we record n.r. for results of GANITE.

From the results on two benchmarks, We can observe that our method MetaITE
demonstrates the stable and nearly best performance in various scenarios,
including classification and regression tasks. While SITE and CBRE have better
performance than CFR-Wass on other regular benchmarks (Yao et al., 2018; Zhou
et al., 2021), they perform worse in the situation of multiple imbalanced
treatments, as shown in Table 2. Because these two models design extra modules
to preserve similarity and reduce information loss, both models have higher
complexity than CFR-Wass and TARNet. We guess that is why they are not
suitable for the small data setting. Meta-learning is proposed to consider a
few-shot setting, and therefore it is an appropriate mechanism for multiple
imbalanced treatments. Experimental results also prove that it generalizes
well. Moreover, by comparing CFR-WAss with TARNet, we can see that discrepancy
measure is helpful for the inference performance. We derive MMD distance as a
discrepancy measure and incorporate it in MetaITE, which improves the
inference performance further. We evaluate the function of the discrepancy
measure in section 4.6.

###  4.5. Robustness Study of Unbalanced Treatments

In order to examine the ability of our method to handle the problem of
imbalanced treatments, in this section, we perform a robustness study for
imbalanced treatments. Specifically, we control the imbalance among treatment
groups and assess the performance of MetaITE and baselines when imbalance
increases. We use Twins_bin and News_4 datasets for the evaluation.

As for the original Twins_bin dataset, there are 4561 samples in the control
group with $t=0$ and 4464 samples in the treated group with $t=1$. we keep the
control group constant, and we configure the sample size of the treated group
as a percentage of the initial size. We reduce the sample size from $100\%$ to
$5\%$, and the smaller the percentage is, the greater the imbalance is.
Correspondingly, the imbalance increases from 1 to 20. We still use
$\sqrt{PEHE}$ as the performance metric for the Twins_bin. From the results in
Figure 3, MetaITE outperforms four other models. As the imbalance goes larger,
MetaITE has stable performance and generalizes well.

Figure 3. Robustness Study of imbalanced treatments on Twins_bin. The
imbalance means that we configure this proportion of control sizes to treated
samples in the experiment.

Regarding to the raw News_4 dataset, there are $\\{$925, 1073, 1021, 957$\\}$
samples for $t=\\{$0,1,2,3$\\}$. As with the process on Twins_bin, we keep the
sample size in the control group with $t=0$ and reduce the sample size from
$100\%$ to $5\%$ in the other three treated groups. We adopt $RMSE$ to measure
the performance of News_4. From the results in Figure 4, MetaITE shows a
stable performance when the imbalance increases. Especially when the imbalance
is large, our model still presents competitive performance compared to
baselines. We find that when the imbalance decreases from $5$ to $10$, the
inference performance of CBRE degrades by $39\%$. It proves that the existing
complex models are not suitable for imbalanced treatments. Note that when the
samples of all treatment groups are sufficient, our method is not as good as
the baseline models on the News_4 dataset. Nevertheless, treatment imbalance
is more in line with real-world datasets, i.e., IHDP (139 treated, 608
control). In this work, we focus on the problem of multiple imbalanced
treatments. We believe MetaITE has great potential as the applicable framework
for estimating multiple imbalanced treatment effects.

Figure 4. Robustness Study of imbalanced treatments on News_4. The imbalance
means that we configure this proportion of control sizes to treated samples in
the experiment. TARNet and CFR-Wass cannot solve the situation where the
imbalance is greater than 10, so we record n.r. for these.

###  4.6. Ablation Study

The objective of MetaITE includes several loss functions. In this section, we
discuss the impacts on model performance and perform experiments on the News_2
and News_4 datasets. The $l_{2}$ regularization is set with a weight decay of
0.05. The standard meta-learning mechanism utilizes the loss on the query set
to update model parameters within the outer loop and adds weight
regularization. Furthermore, considering the situation of multiple imbalanced
treatments, we incorporate the predictive results on source domains as extra
supervision signals and combine the discrepancy measures. We configure that
$\\{\mu,\epsilon,\gamma\\}$ are all from $0$ to $1$ with step $0.1$, and each
parameter combination is evaluated 10 times. For News_2 dataset, we first
report the top three best results of $\sqrt{PEHE}$ and related hyper-
parameters: $17.1726$ with $\\{1.0,0.9,1.0\\}$, $18.9048$ with
$\\{1.0,0.8,1.0\\}$, $19.1451$ with $\\{1.0,0.7,1.0\\}$. As for News_4, we
report top three best results of RMSE and their corresponding hyper-
parameters, which are 8.7303 with $\\{1.0,0.0,1.0\\}$, 8.9489 with
$\\{0.3,0.0,0.8\\}$, 8.9807 with $\\{0.8,0.0,1.0\\}$. We see that the weight
of discrepancy loss is large, which means that balancing distributions among
treatment groups improves performance. On News_4, we find that the inference
loss on source treatments is not really helpful for total inference accuracy.
As Eq.(9) shows, the calculation of $\mathcal{L}_{Que}$ is related to
parameters updated by $\mathcal{L}_{Sup}$. We guess these two loss functions
are redundant to a certain extent on the News_4 and on the other hand, the
supervision signal of source treatments works well on the News_2. To sum up,
we obtain the optimal hyper-parameters $\\{\mu,\epsilon,\gamma\\}$ of
$\\{1.0,0.0,1.0\\}$ for the News_4 dataset and $\\{1.0,0.9,1.0\\}$ for the
News_2.

##  5\. Related work

###  5.1. Counterfactual Inference

Answering counterfactual questions is a fundamental problem in many
applications. Researchers focus on matching methods and regression at first,
such as OLS/LR and k-NN (Crump et al., 2008). The performance of inference
models is greatly improved when incorporated with representation learning or
neural networks. In this work, we focus on representation learning-based
works. (Johansson et al., 2016) is the first framework that brings domain
adaptation and representation learning for counterfactual inference. The study
(Shalit et al., 2017) learns balanced representation that enforces treated and
control distributions similar via Integral Probability Metrics. (Yoon et al.,
2018) proposes GANITE that utilizes two generative adversarial networks to
infer individual treatment effects. Moreover, GANITE is suitable to be applied
in the setting of multiple treatments. Comparing previous works, Yao et al.
(Yao et al., 2018) consider the local similarity information that provides
meaningful regularization for inference but is ignored by most methods. (Du et
al., 2021) utilizes adversarial training to balance distribution between two
groups and applies mutual information metrics to reduce information loss.
(Zhou et al., 2021) extends (Du et al., 2021) and constructs an information
loop to preserve original data properties via two supplementary losses.
However, only one of them can naively infer counterfactual outcomes in the
setting of multiple treatments, and none of them considers the situation of
imbalanced treatments.

###  5.2. Meta-Learning

Meta-learning has been widely used in various fields, such as few-show
learning, including regression and classification (Finn et al., 2017) and
domain generalization (Li et al., 2018). The model-agnostic meta-learning
framework is a widely adopted method called MAML (Finn et al., 2017). In
recent years, some works related to meta-learning has been utilized in
causality (Sharma et al., 2019; Bengio et al., 2019; Ton et al., 2021).
(Sharma et al., 2019) incorporates CFRNet (Shalit et al., 2017) into a meta
learning-based reptile framework in order to infer counterfactual outcomes.
The works (Bengio et al., 2019; Ton et al., 2021) focus on causal discovery
with meta-learning. Moreover, (Ton et al., 2021) pay attention to the small
data setting and utilize a meta-learning framework, which is consistent with
ours. However, they mainly deal with causal directionality to distinguish the
cause from effect rather than counterfactual causality.

##  6\. Conclusion

In this paper, we introduce MetaITE, a meta-learning algorithm for estimating
individual treatment effects, primarily concentrating on multiple imbalanced
treatments. MetaITE leverages predictive neural networks to express the
feature extractor and the inference network as a base model. Moreover, it
further incorporates the base model in a meta-learning framework and provides
two extra regularizations: distribution discrepancy measures and the
supervision of domains with high-volume samples to improve inference
performance. We perform experiments on two benchmarks with real-world data,
and our results empirically demonstrate that our model matches/outperforms the
eleven baselines.

In future work, we consider the existence of hidden confounders. We look
forward to incorporating variational autoencoders (VAEs) into the MetaITE
framework (Louizos et al., 2017).

## References

  * (1)
  * Alaa and van der Schaar (2017) Ahmed M Alaa and Mihaela van der Schaar. 2017.  Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes.  _Advances in Neural Information Processing Systems_ 30 (2017).
  * Bengio et al. (2019) Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. 2019\.  A meta-transfer objective for learning to disentangle causal mechanisms.  _arXiv preprint arXiv:1901.10912_ (2019).
  * Blei et al. (2003) David M Blei, Andrew Y Ng, and Michael I Jordan. 2003\.  Latent dirichlet allocation.  _the Journal of machine Learning research_ 3 (2003), 993–1022.
  * Bottou et al. (2013) Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candela, Denis X Charles, D Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and Ed Snelson. 2013\.  Counterfactual Reasoning and Learning Systems: The Example of Computational Advertising.  _Journal of Machine Learning Research_ 14, 11 (2013).
  * Chernozhukov et al. (2013) Victor Chernozhukov, Iván Fernández-Val, and Blaise Melly. 2013.  Inference on counterfactual distributions.  _Econometrica_ 81, 6 (2013), 2205–2268.
  * Crump et al. (2008) Richard K Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. 2008.  Nonparametric tests for treatment effect heterogeneity.  _The Review of Economics and Statistics_ 90, 3 (2008), 389–405.
  * Du et al. (2021) Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy. 2021.  Adversarial balancing-based representation learning for causal effect inference with observational data.  _Data Mining and Knowledge Discovery_ (2021), 1–26.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine. 2017\.  Model-agnostic meta-learning for fast adaptation of deep networks. In _International Conference on Machine Learning_. PMLR, 1126–1135.
  * Glass et al. (2013) Thomas A Glass, Steven N Goodman, Miguel A Hernán, and Jonathan M Samet. 2013.  Causal inference in public health.  _Annual review of public health_ 34 (2013), 61–75.
  * Gretton et al. (2012) Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. 2012\.  A kernel two-sample test.  _The Journal of Machine Learning Research_ 13, 1 (2012), 723–773.
  * Johansson et al. (2016) Fredrik Johansson, Uri Shalit, and David Sontag. 2016\.  Learning representations for counterfactual inference. In _International conference on machine learning_. PMLR, 3020–3029.
  * Khang (2015) Young-Ho Khang. 2015\.  The causality between smoking and lung cancer among groups and individuals: addressing issues in tobacco litigation in South Korea.  _Epidemiology and health_ 37 (2015).
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ (2014).
  * Kooti et al. (2016) Wesam Kooti, Maryam Farokhipour, Zahra Asadzadeh, Damoon Ashtary-Larky, and Majid Asadi-Samani. 2016\.  The role of medicinal plants in the treatment of diabetes: a systematic review.  _Electronic physician_ 8, 1 (2016), 1832.
  * Kumagai and Iwata (2019) Atsutoshi Kumagai and Tomoharu Iwata. 2019.  Unsupervised domain adaptation by matching distributions based on the maximum mean discrepancy via unilateral transformations. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 33. 4106–4113.
  * Li et al. (2018) Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. 2018.  Learning to Generalize: Meta-Learning for Domain Generalization. In _AAAI_.
  * Liu et al. (2020) Dugang Liu, Pengxiang Cheng, Zhenhua Dong, Xiuqiang He, Weike Pan, and Zhong Ming. 2020\.  A general knowledge distillation framework for counterfactual recommendation via uniform data. In _Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval_. 831–840.
  * Long et al. (2015) Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. 2015\.  Learning transferable features with deep adaptation networks. In _International conference on machine learning_. PMLR, 97–105.
  * Louizos et al. (2017) Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling. 2017\.  Causal effect inference with deep latent-variable models.  _arXiv preprint arXiv:1705.08821_ (2017).
  * Rosenbaum and Rubin (1983) Paul R Rosenbaum and Donald B Rubin. 1983.  The central role of the propensity score in observational studies for causal effects.  _Biometrika_ 70, 1 (1983), 41–55.
  * Rubin (1974) Donald B Rubin. 1974\.  Estimating causal effects of treatments in randomized and nonrandomized studies.  _Journal of educational Psychology_ 66, 5 (1974), 688\.
  * Rubin (2005) Donald B Rubin. 2005\.  Causal inference using potential outcomes: Design, modeling, decisions.  _J. Amer. Statist. Assoc._ 100, 469 (2005), 322–331.
  * Schwab et al. (2020) Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M Buhmann, and Walter Karlen. 2020\.  Learning counterfactual representations for estimating individual dose-response curves. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 34. 5612–5619.
  * Shalit et al. (2017) Uri Shalit, Fredrik D Johansson, and David Sontag. 2017\.  Estimating individual treatment effect: generalization bounds and algorithms. In _International Conference on Machine Learning_. PMLR, 3076–3085.
  * Sharma et al. (2019) Ankit Sharma, Garima Gupta, Ranjitha Prasad, Arnab Chatterjee, Lovekesh Vig, and Gautam Shroff. 2019.  Metaci: Meta-learning for causal inference in a heterogeneous population.  _arXiv preprint arXiv:1912.03960_ (2019).
  * Sriperumbudur et al. (2011) Bharath K Sriperumbudur, Kenji Fukumizu, and Gert RG Lanckriet. 2011\.  Universality, Characteristic Kernels and RKHS Embedding of Measures.  _Journal of Machine Learning Research_ 12, 7 (2011).
  * Ton et al. (2021) Jean-François Ton, Dino Sejdinovic, and Kenji Fukumizu. 2021\.  Meta Learning for Causal Direction. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 35. 9897–9905.
  * Yao et al. (2018) Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. 2018\.  Representation learning for treatment effect estimation from observational data.  _Advances in Neural Information Processing Systems_ 31 (2018).
  * Yoon et al. (2018) Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. 2018\.  GANITE: Estimation of individualized treatment effects using generative adversarial nets. In _International Conference on Learning Representations_.
  * Zhou et al. (2021) Guanglin Zhou, Lina Yao, Xiwei Xu, Chen Wang, and Liming Zhu. 2021.  Cycle-Balanced Representation Learning For Counterfactual Inference.  _arXiv preprint arXiv:2110.15484_ (2021).

##  Appendix A APPENDIX

###  A.1. Estimate Process

Algorithm 2 : Estimate process of MetaITE

1:Input: Feature Extractor $g(\psi)$, Inference Network $h(\theta)$; Training
Set and Testing Set in all domains; Hyperparameters
$\alpha,\beta,\mu,\epsilon,\gamma$.

2:$\triangleright$ estimate

3:while estimate not done do

4: construct query set $\\{\mathbf{X}^{Que},\mathbf{y}^{Que}\\}$ with $K$
samples from

5: testing data

6: for $t\in\mathcal{T}$ do

7: randomly sample $K$ units from training data with

8: treatment that equals to $t$ for $\\{\mathbf{X}^{Sup},\mathbf{y}^{Sup}\\}$

9: feed
$\\{\mathbf{X}^{Sup},\mathbf{y}^{Sup},\mathbf{X}^{Que},\mathbf{y}^{Que}\\}$
into trained model, and

10: get prediction $\\{\hat{\mathbf{y}}_{i}^{\mathbf{t}_{i}=t}\\}_{i=1}^{K}$
End for

11: concatenate predictions in for-loop and obtain outcome for

12: each treatment $\\{\hat{\mathbf{y}}_{i}^{\mathbf{t}_{i}=t}\\}_{t=1}^{k}$

13:end while

###  A.2. Parameter Setting

There are three types of parameters: meta setup, model structure and hyper
parameters.

  1. (1)

Meta setup parameters: we set meta batch size as 5, which means the number of
tasks sampled in each meta-update. The number of samples in the support and
query sets are both set as 8. Note that the product of meta batch size and the
samples number in support set or query set corresponds to the value of $K$ in
Algorithm 1. The number of inner gradient updates during training is 4.

  2. (2)

Model structure parameters: these parameters consist of the number of layers
and the embedding size, which are related to both feature extractor and
inference network. We build two-layers fully connected networks for the
feature extractor with embedding size $[256,128]$ and four-layers fully
connected networks for the inference network with embedding size
$[128,128,64,64]$.

  3. (3)

Hyper-parameters: we mainly have five hyper-parameters
$\\{\alpha,\beta,\mu,\epsilon,\gamma\\}$ in Algorithm 1. For the meta-learning
rate $\beta$ and update learning rate $\alpha$ in the inner loop, we set both
$\alpha$ and $\beta$ as $1e-3$. As for three coefficients of loss functions,
we set the parameters $\\{\mu,\epsilon,\gamma\\}$ to range from $0$ to $1$
with step $0.1$ and obtain the optimal parameters.

###  A.3. Parameter Sensitivity

As introduced in the parameter setting, MetaITE contains meta setup
parameters, model structure parameters and hyper-parameters. In this
subsection, we evaluate all parameters except $\\{\mu,\epsilon,\gamma\\}$.
Meta bach size and the number of samples in the support and the query set can
influence the number of data samples fed into the local and global updates. As
Figure 5 (LABEL:sub@fig:e) shows, smaller batch size is helpful for the model
performance where the size is between 5 to 10. As for the sample size in
support or query set, we compare the results with size in $\\{3,5,8\\}$. In
Figure 5 (LABEL:sub@fig:a), we can find that the median value may lead to a
bad result. The number of updates in the inner loop is configured from 4 to
12, and Figure 5 (LABEL:sub@fig:f) illustrates that more minor gradient
updates provide better results, which is related to data size. According to
Figure 5 (LABEL:sub@fig:b)-(LABEL:sub@fig:c), which shows the impact of the
number of layers and embedding size on model performance, we can see that
shallower layers in feature extractor and deeper layers in inference network
lead to better results. Moreover, the median value of embedding size is more
suitable due to the comparatively small data setting. Similarly, two learning
rates $\alpha,\beta$ are chosen from $\\{1e-4,1e-3,1e-2\\}$ and median value
$1e-3$ leads to better results, as seen in Figure 5 (LABEL:sub@fig:d).

(a) Support sample size

(b) Number of layers in feature extractor and inference network

(c) Embedding size in feature extractor and inference network

(d) Learning rate $\alpha,\beta$

(e) Meta batch size

(f) Number of gradient updates

Figure 5. Parameter sensitivity in different settings.