1. 引言

在许多决策问题中，从决策制定的角度来看，估计治疗的异质性结果并不足够。神经网络可能能够很好地近似结果和相应的反事实情况，但由于神经模型缺乏透明度和可解释性，其在现实世界决策制定中的部署将受到限制。异质效应混合模型（HEMM）方法最初是在论文《在阿片类药物处方指南中应用的治疗效果估计中的可解释子群发现》中提出的，目的是缓解这一挑战，使决策者获得更多洞见。

1.1 子群发现和异质性治疗效果问题

异质效应混合模型（HEMM）背后的思想涉及假设数据集中每个个体都有一个低维聚类或潜在的分组。下面图A直观地展示了以下现象：注意，在Z1中几乎所有接受治疗的实例都有积极的结果，而在Z3中很少。我们感兴趣的是恢复这样的潜在子群。

1.2 HEMM 描述和板式符号

在HEMM模型中，这是通过使用有限（K）正态和伯努利混合来实现的。执行治疗的贡献随后由属于K个子群之一来介导。

2. 合成数据示例

本节将展示HEMM在合成数据集上的能力。

2.1 数据描述和生成过程

我们将使用包含在causallib中的HEMM模块的合成数据集来演示HEMM发现子群以及异质性效果的能力。合成数据集是按照以下方案生成的。

1. 我们取 $(X=(X_0,X_1)\in {\mathbb{R}}^2)$ 并从 $(X=[0,1{]}^2)$ 上的均匀分布中采样它。

2. 为了模拟观察性研究中固有的选择偏差，治疗变量依赖于 $(X)$ 作为：
   $[T\sim \text{Bernoulli}(0.4)\text{ for }{X}_0<0.5]$
   $[T\sim \text{Bernoulli}(0.6)\text{ for }{X}_0>0.5]$

3. 潜在结果 $(Y(0))$ 和$(Y(1))$ 也是伯努利分布，其均值由图示的 $(X,T)$ 函数给出。图示显示 $(p(Y(1)=1|X)>p(Y(0)=1|X))$，即治疗增加了接受治疗人群有更高概率获得积极结果。注意，在条件可交换性假设下我们有：
   $[p(Y(t)=1|X)=p(Y=1|T=t,X)]$

4. 我们通过将更高的概率分配给 ( X ) 的上三角区域来模拟混杂因素 $(X)$ 的效应。这与 $(T)$ 的分布一起意味着，无论治疗如何，更有可能获得积极结果的个体（上三角）也更有可能接受治疗（右半方形）。

5. 最后，我们将增强治疗效果组建模为一个圆形区域 ( S = { a : ∥ a − c ∥ 2 < r } ) ( S = \{a : \|a - c\|_2 < r\} ) (S={a:∥a−c∥2​<r})，其中$[p(Y(1)=1|S)>p(Y(1)=1|X)]$我们设置 $(c=(0.5,0.5))$ 和 $(r=1)$。总共生成了 2,500 个样本 $((x_i,t_i,y_i))$，如上所述。

causallib.contrib.hemm.gen_synthetic_data 中的 gen_montecarlo 函数允许我们随机采样多个合成数据集的实例化版本。

from causallib.contrib.hemm.gen_synthetic_data import gen_montecarlo

syn_data = gen_montecarlo(5000, 2, 100)

2.2 反事实结果估计和 PEHE 估计

我们将针对合成数据集的每次实现拟合一个 HEMM 模型，并使用精确度衡量异质性效应指标 (PEHE) 评估 HEMM 对 ITE（也称为条件平均治疗效应 (CATE) 的估计能力。PEHE 更正式定义为围绕 ITE 的均方误差。

$[PEHE=\frac{1}{n_1}{\sum }_{i=1}^{n_1}(f_i(x_i)-f_0(x_i)-\mathbb{E}[Y(1)-Y(0)]|X=x_i{)}^2]$

在这里，$(f_0)$ 和 $(f_1)$ 分别代表用于估计具有特征 $(X_i)$ 的观测的反事实结果的函数。括号内的第二项是真实的条件平均治疗效果（CATE）。请注意，真实的CATE只能在合成设置中观察到，而在现实世界的观察数据中则无法观察到，因此使用PEHE作为估计CATE拟合优度的指标仅限于合成数据集。

from causallib.contrib.hemm import HEMM
#Leave 30% of the Training Data as Validation
vsize = int(0.3*syn_data['TRAIN']['x'].shape[0])
#setting the Training Dataset
Xtr  = syn_data['TRAIN']['x'] [:-vsize,:, 0]
Ytr  = syn_data['TRAIN']['yf'][:-vsize  , 0]
Ttr  = syn_data['TRAIN']['t'] [:-vsize  , 0]

#setting the Dev/Val Dataset
Xdev = syn_data['TRAIN']['x'] [-vsize:,:, 0]
Ydev = syn_data['TRAIN']['yf'][-vsize:  , 0]
Tdev = syn_data['TRAIN']['t'] [-vsize:  , 0]

#setting the Test Dataset
Xte  = syn_data['TEST']['x'] [:,:, 0]
Yte  = syn_data['TEST']['yf'][:  , 0]
Tte  = syn_data['TEST']['t'] [:  , 0]

#Feature size
Xdim = Xtr.shape[1]

#Number of Components to Discover.
K = 3

#Initialize the model with Population means and Std. Devs.
#Empirically results in faster, and better convergence.
mu  = Xtr.mean(axis=0).reshape(1,-1)
std = Xtr.std(axis=0).reshape(1,-1)

#Set the Learning Rate for Adam Optimizer and Batch Size
learning_rate = 1e-4
batch_size    = 100

#Indicate the Outcome Distribution (Y), could be 'bin' for 
#Binary Outcomes or 'cont' for Continuous Outcomes.
response = 'bin'

#Indicate what kind of a model to be used to adjust for 
#confounding. You can also pass your own PyTorch Model!!!!
outcome_model='linear'

#Instantiate an HEMM model
model = HEMM(Xdim, K, homo=True, mu=mu, std=std, bc=2, lamb=0.0000,\
                spread=.1,outcome_model=outcome_model,sep_heads=True,epochs=10,\
                 learning_rate=learning_rate,weight_decay=0.0001,metric='LL', use_p_correction=False,\
                 response=response,imb_fun=None,batch_size=batch_size )


#Fit the HEMM model on the Training Data, with Early Stopping
#on the Dev Set.
cd = model.fit(Xtr, Ttr,Ytr, validation_data=(Xdev, Tdev, Ydev))
#Extract the True Potential outcomes for the Training 
#and the Test Dataset.
mu1tr = syn_data['TRAIN']['mu1'][:-vsize, 0] 
mu0tr = syn_data['TRAIN']['mu0'][:-vsize, 0]

mu1te = syn_data['TEST']['mu1'][:, 0] 
mu0te = syn_data['TEST']['mu0'][:, 0]
#Estimate the Individual Potential Outcomes using HEMM
inSampleCFoutcomes  = model.estimate_individual_outcome(Xtr, Ttr)
outSampleCFoutcomes = model.estimate_individual_outcome(Xte, Tte)
#Compute the True In-sample Individual Treatment Effect (ITE)
trueCATE = mu1tr - mu0tr

#Compute the In-Sample  ITEs
inSampleITE = inSampleCFoutcomes[1]-inSampleCFoutcomes[0]

#Compute the In-Sample PEHE
inSamplePEHE = ((inSampleITE - trueCATE)**2).mean()
print("In Sample PEHE:", inSamplePEHE)
In Sample PEHE: 0.011010778786300154
#Compute the True Out of Sample Individual Treatment Effect (ITE)
trueCATE = mu1te - mu0te

#Compute the Out of Sample  ITEs
outSampleITE = outSampleCFoutcomes[1]-outSampleCFoutcomes[0]

#Compute the Out of Sample PEHE
outSamplePEHE = ((outSampleITE - trueCATE)**2).mean()

print("In Sample PEHE:", outSamplePEHE)
In Sample PEHE: 0.010915930183215048

如您所见，HEMM 返回了一个较低的 PEHE 分数，这证明了它在估计治疗的异质性效应方面的能力。