文章名称

Transfer Learning for Estimating Causal Effects using Neural Networks

核心要点

文章主要关注CATE场景，把因果推断和迁移学习的概念结合，期望同时利用不同数据源的数据来学习他们底层共同的因果关系。作者首先提出了Y-learner的网络架构（这其实并不新鲜，因为已经有很多Multitask learning的方法了），随后作者利用了很多深度学习中的迁移学习技巧，来提升因果推断的能力，并有效的利用多个数据源的数据集，包括warn start，frozen-features, multi-head, and joint training。最后，作者提出了利用元学习的方法，快速学习多个数据源，并提升模型迁移到能力，使得模型能够更快的应用到新的场景。文章更像个工具箱，所有的方案是并行的，可以根据场景自己选择。

方法细节

问题引入

很多时候，我们可能有从不同的数据源收集到的数据，但是每个数据源的数据较少，并且即便总和起来也不够充分的训练神经网络。这是典型的小样本的场景，如果不采用合理的训练方式可能造成训练不充分（模型过于复杂），或者过拟合（模型过于简单）。因此，业界有各种各样的小样本学习方法提出，其中迁移学习是主流方法之一，元学习与两者有很大的交集，却又不尽相同。多场景的数据学习更适合元学习。而估计CATE本身就是数据缺失的，比如missing confounder是因果推断的根本问题之一，可以被看做某种数据缺失。这一问题可以利用随机实验收集数据来解决，但随机试验的数据量通常都很少。观测数据中selection bias的问题是另一个根本问题，样本较少会导致模型效果进一步变差（样本太少根本没有办法进行matching和balancing）。这些问题导致因果推断小样本学习其实是一个很重要的方向。

具体做法

作者利用Y-learner作为基础学习器来进行迁移学习，并估计causal effect。我们首先来看一下Y-learner的结构。其中，Y-learner会采用三个网络， $\pi_{\theta_0}(X_i)$ ， $\pi_{\theta_1}(X_i)$ 以及 $\pi_{\tau}(X_i)$ ，分别用来学习potential outcome（ $t=0$ ），potential outcome（ $t=1$ ）以及causal effect。如下图所示，Y-learner是交替进行训练的。

在 $t=0$ 的时候（图中是 $W_i=0$ ），会首先利用样本的实际观测结果（factual outcome） $Y_{i}^{obs}$ 来更新网络 $\pi_{\theta_0}(X_i)$ 的参数。随后利用实际观测结果加上估计的causal effect的结果 $\pi_{\tau}(X_i)$ 来训练估计网络 $\pi_{\theta_1}(X_i)$ ，这个类似于X-learner的方法。最后，利用 $\pi_{\theta_1}(X_i)$ 的估计值来代表 $t=1$ 时候的potential outcome，并减去实际的观测结果（实际观测时 $t=0$ ）来训练估计网络 $\pi_{\tau}(X_i)$ 。
在 $t=1$ 的时候（图中是 $W_i=1$ ），与 $t=0$ 的时候结构基本相同，只是插值的时候，实际的观测结果是 $t=1$ 的结果，而利用其他的网络的预测值作为估计量带入causal effect的计算公式来训练其他网络。

Y-learner structure

接来下，看一下transfer learning的部分。我们用户拥有多个场景的样本，不同的场景用上标 $i，j$ 代表。如 $\pi_{\theta_{0}^{i}}$ ， $\pi_{\theta_{0}^{j}}$ 分别代表两个不同场景下， $t = 0$ 时候的估计网络。迁移学习有多种处理方法，

warm start（也被成为fine-tuning）。这种情况下各个场景的输入特征空间是一样的，并且他们底层的内部因果关系是一样的，只是在同样的 $X$ 下，相同treatment下的potential outcome可能是不一样的，即 $\pi_{\theta_{0}^{0}}(x_i) \neq \pi_{\theta_{0}^{1}}(x_i)$ 。相比随机初始化，利用第一个任务训练好的权重初始化第二个任务并继续训练能够带来更快的收敛速度、更好的效果以及对样本的需求量更少（transfer learning的好处照搬过来...）。
Frozen-features。利用在场景1训练好的前几层网络的权重，来初始化场景2的网络。并且，在场景2训练的时候，冻结这部分权重，并不会训练这部分特征提取网络。
Multi-head，另一种是多头输出，这里的多头不只是outcome多头，还包括场景独立的多头部分。首先，仍然是逐个训练不同的场景网络。把第一个网络的参数用来初始化后边的场景的网络。网络存在不同任务间共享的网络，同时存在不同网络间独立的网络，并进行交替训练。
Joint training，本质上就是把所有场景的网络变成一个大的multitask learning，一次反向传播的时候会优化所有参数。损失函数是各个场景下样本的损失之和，即 $L = \sum_i || \hat{\mu}_{1}^{i} - \mu_{1}^{i} || + || \hat{\mu}_{0}^{i} - \mu_{0}^{i} ||$ 。

transfer learning

最后，作者利用元学习，来进行初始化权重的学习。其实就是利用Reptile学习初始化参数，分别利用multitask和S-learner的思路，来快速迁移。（具体明天再扯...）

代码实现

（留坑待填...）

心得体会

imputation like X-learner

Y-learner利用实际观测结果加上估计的causal effect的结果 $\pi_{\tau}(X_i)$ 来训练估计网络 $\pi_{\theta_1}(X_i)$ 。本质上是一种插值的训练方式，这样即分别利用了每一个potential outcome的数据信息，又充分利用了两者的潜在共有信息，比如共有的outcome的值结构（细节可以参见FlexTENet[1]）。

文章引用

[1] Curth, A., & van der Schaar, M. (2021). On Inductive Biases for Heterogeneous Treatment Effect Estimation. arXiv preprint arXiv:2106.03765.

因果推断深度学习工具箱 - Transfer Learning for Estimating Causal Effects using Neural Networks