TensorFlow实战：5步搞定因果推断模型TARNet（附完整代码）

TensorFlow实战：5步搞定因果推断模型TARNet（附完整代码）

如果你正在处理营销效果评估、药物疗效分析或者任何需要回答“如果...会怎样”的业务问题，那么因果推断就是你工具箱里不可或缺的利器。传统的机器学习模型擅长预测相关性，但要剥离出真正的因果效应，尤其是在高维、非随机化的观测数据中，常常力不从心。TARNet（Treatment-Agnostic Representation Network）作为深度因果推断领域的经典模型，提供了一种优雅的解决方案：它通过一个共享的表示层学习协变量特征，再针对不同干预（Treatment）分支进行预测，巧妙地平衡了偏差与方差。今天，我们就抛开复杂的理论推导，直接上手TensorFlow，用五个清晰的步骤，从零搭建一个可运行的TARNet模型，并解决实际工程中令人头疼的维度灾难和样本不平衡问题。

1. 环境准备与数据理解

在开始敲代码之前，确保你的工作环境已经就绪。我们将使用TensorFlow 2.x，这是目前的主流选择，其Keras API能让我们像搭积木一样构建模型。此外，一些数据处理和可视化的库也会用到。

pip install tensorflow==2.10.0 pandas numpy scikit-learn matplotlib seaborn

接下来，理解我们要处理的数据结构至关重要。因果推断的数据通常包含三部分：

• 协变量 (X): 描述样本特征的多维向量，例如用户的年龄、历史行为、设备信息等。
• 干预/处理 (T): 一个二元或多元的指示变量，表示样本接受了哪种处理（如：广告A=1，广告B=0；或用药=1，不用药=0）。
• 结果 (Y): 我们关心的观测结果，比如点击率、康复指标等。

核心挑战在于反事实的缺失：对于一个给定的用户，我们只能观察到他在一种处理下的结果，而无法知道如果给他另一种处理，结果会怎样。TARNet的目标，正是从观测数据中学习，去估计这个“缺失”的反事实结果。

一个典型的数据集可能长这样（以Pandas DataFrame为例）：

注意：在实际业务数据中，处理组（T=1）和对照组（T=0）的样本分布很可能是不平衡的，例如高收入用户更可能看到高价广告。这种“选择偏差”会直接污染效应估计，是TARNet需要克服的关键。

2. 数据预处理与特征工程

拿到原始数据后，直接丢给模型往往效果不佳。我们需要进行一系列预处理，为模型训练打下坚实基础。这一步虽然繁琐，但很大程度上决定了模型的上限。

首先，处理缺失值和异常值。对于连续型协变量，常见的做法是用中位数或均值填充缺失值；对于类别型变量，可以单独设置一个“缺失”类别。异常值则可以根据业务逻辑或统计方法（如IQR法则）进行截断或转换。

import pandas as pd import numpy as np from sklearn.impute import SimpleImputer from sklearn.preprocessing import StandardScaler, OneHotEncoder from sklearn.compose import ColumnTransformer # 假设df是我们的DataFrame # 分离特征、处理和结果 X = df.drop([‘treatment‘, ‘outcome‘], axis=1) T = df[‘treatment‘].values Y = df[‘outcome‘].values # 区分数值型和类别型特征 numeric_features = X.select_dtypes(include=[‘int64‘, ‘float64‘]).columns categorical_features = X.select_dtypes(include=[‘object‘, ‘category‘]).columns # 构建预处理管道 numeric_transformer = Pipeline(steps=[ (‘imputer‘, SimpleImputer(strategy=‘median‘)), (‘scaler‘, StandardScaler()) ]) categorical_transformer = Pipeline(steps=[ (‘imputer‘, SimpleImputer(strategy=‘constant‘, fill_value=‘missing‘)), (‘onehot‘, OneHotEncoder(handle_unknown=‘ignore‘, sparse=False)) ]) preprocessor = ColumnTransformer( transformers=[ (‘num‘, numeric_transformer, numeric_features), (‘cat‘, categorical_transformer, categorical_features) ]) # 拟合并转换训练数据 X_processed = preprocessor.fit_transform(X)

其次，样本权重调整。为了缓解处理组和对照组因样本量差异或倾向性得分不同带来的偏差，我们经常需要为每个样本计算一个权重。逆概率加权（IPW）是一种常用方法，权重 ( w_i = \frac{T_i}{e(X_i)} + \frac{1-T_i}{1-e(X_i)} )，其中 ( e(X_i) ) 是倾向性得分，即给定特征X下接受处理T=1的概率。我们可以先训练一个简单的逻辑回归模型来估计倾向性得分。

from sklearn.linear_model import LogisticRegression from sklearn.model_selection import train_test_split # 估计倾向性得分 ps_model = LogisticRegression(max_iter=1000).fit(X_processed, T) propensity_scores = ps_model.predict_proba(X_processed)[:, 1] # 计算IPW权重，加入小常数防止除零 epsilon = 1e-6 weights = T / (propensity_scores + epsilon) + (1 - T) / (1 - propensity_scores + epsilon) # 对权重进行归一化，使其在训练中更稳定 weights = weights / np.mean(weights)

最后，将数据划分为训练集、验证集和测试集。切记，划分必须基于样本ID或时间戳进行，绝不能随机打散所有数据，因为我们需要评估模型在未见过的协变量分布上的泛化能力，随机划分会带来数据泄露，高估模型性能。

# 假设我们有时间戳列 ‘timestamp‘ df_sorted = df.sort_values(‘timestamp‘) split_idx = int(0.7 * len(df_sorted)) val_idx = int(0.85 * len(df_sorted)) train_data = df_sorted.iloc[:split_idx] val_data = df_sorted.iloc[split_idx:val_idx] test_data = df_sorted.iloc[val_idx:] # 分别对三个数据集应用相同的预处理流程 X_train_processed = preprocessor.transform(train_data.drop([‘treatment‘, ‘outcome‘], axis=1)) T_train = train_data[‘treatment‘].values Y_train = train_data[‘outcome‘].values W_train = weights[train_data.index] # 获取对应的权重 # 同理处理验证集和测试集...

3. 构建TARNet模型核心架构

现在进入最核心的部分：用TensorFlow构建TARNet模型。其设计哲学非常直观：一个共享的表示层（Φ网络）学习从高维协变量X到一个低维、平衡的表示空间Z的映射；然后，针对每一种处理T，都有一个独立的头网络（Y网络）来预测结果。这种结构迫使共享层提取与处理无关的通用特征，而头网络专注于学习特定处理下的效应。

我们将以面向对象的方式，继承tf.keras.Model类来构建模型，这样能更灵活地定义前向传播和自定义损失函数。

import tensorflow as tf from tensorflow.keras import layers, regularizers class TARNet(tf.keras.Model): """ TARNet 模型实现。 参数: input_dim: 输入特征维度。 representation_dim: 共享表示层的维度。 hidden_units_rep: 共享层各隐藏层的神经元数量列表。 hidden_units_out: 每个输出头网络的隐藏层神经元数量列表。 n_treatments: 处理的数量（默认为2，即二元处理）。 dropout_rate: Dropout比率，用于防止过拟合。 l2_reg: L2正则化系数。 """ def __init__(self, input_dim, representation_dim=64, hidden_units_rep=[100, 80], hidden_units_out=[60, 40], n_treatments=2, dropout_rate=0.2, l2_reg=1e-4): super(TARNet, self).__init__() # 1. 共享表示层 (Phi网络) self.representation_layers = [] prev_units = input_dim for i, units in enumerate(hidden_units_rep): self.representation_layers.append( layers.Dense(units, activation=‘relu‘, kernel_regularizer=regularizers.l2(l2_reg), name=f‘rep_dense_{i}‘) ) self.representation_layers.append( layers.Dropout(dropout_rate, name=f‘rep_dropout_{i}‘) ) prev_units = units # 最终表示层，输出到指定维度 self.representation_layers.append( layers.Dense(representation_dim, activation=None, kernel_regularizer=regularizers.l2(l2_reg), name=‘representation_output‘) ) # 2. 处理特定的输出头网络 (Y网络) self.output_heads = [] for t in range(n_treatments): head_layers = [] prev_units = representation_dim for i, units in enumerate(hidden_units_out): head_layers.append( layers.Dense(units, activation=‘relu‘, kernel_regularizer=regularizers.l2(l2_reg), name=f‘head_{t}_dense_{i}‘) ) head_layers.append( layers.Dropout(dropout_rate, name=f‘head_{t}_dropout_{i}‘) ) prev_units = units # 最终输出层，根据任务选择激活函数（回归为None，分类为sigmoid/softmax） head_layers.append( layers.Dense(1, activation=None, name=f‘head_{t}_output‘) ) self.output_heads.append(tf.keras.Sequential(head_layers, name=f‘y_head_{t}‘)) self.n_treatments = n_treatments def call(self, inputs, training=False): """ 前向传播。 输入: 一个元组 (X_features, T_indices) X_features: 协变量特征张量。 T_indices: 处理索引张量，形状为(batch_size,)，值为0或1（对于二元处理）。 输出: 预测的结果张量，形状为(batch_size, 1)。 """ x, t = inputs # 通过共享表示层 for layer in self.representation_layers: x = layer(x, training=training) # 初始化一个全零张量用于收集批处理结果（更高效） batch_size = tf.shape(x)[0] outputs = tf.zeros((batch_size, 1)) # 为每一种处理动态选择对应的头网络进行计算 for treatment_idx in range(self.n_treatments): # 创建当前处理对应的掩码 mask = tf.cast(tf.equal(t, treatment_idx), tf.float32) mask_expanded = tf.expand_dims(mask, axis=-1) # 从 (batch,) 变为 (batch, 1) # 通过对应的头网络得到预测 head_prediction = self.output_heads[treatment_idx](x, training=training) # 将预测结果累加到outputs中，非当前处理的样本对应位置贡献为0 outputs += mask_expanded * head_prediction return outputs

这个实现的关键在于call方法中动态路由的机制。我们根据输入的T_indices，将经过共享层后的表示x，分别送入对应的头网络，并只选取对应处理的预测结果进行输出。这种方式比分别对两个群体前向传播两次要高效。

4. 实现反事实损失函数与模型训练

模型结构搭建好了，但如何训练它才能学到真正的因果效应？这就需要引入反事实损失函数。TARNet的损失函数通常由三部分组成：

1. 预测损失 (Prediction Loss): 衡量模型对观测结果的拟合程度，通常是均方误差（MSE）或交叉熵。
2. 表示平衡正则项 (Representation Balancing Regularizer): 这是TARNet及其变体CFR的精髓。它通过最小化处理组和对照组在表示空间Z上的分布距离，来减少混杂偏差。常用最大均值差异（MMD）或Wasserstein距离。
3. 权重项 (Weighting): 结合之前计算的样本权重（如IPW权重），让模型更关注那些对平衡分布贡献大的样本。

我们将实现一个结合了加权MSE和MMD正则化的损失函数。

def mmd_loss(rep_t1, rep_t0, sigma=1.0): """ 计算处理组和对照组表示之间的最大均值差异（MMD）损失。 使用高斯核函数。 """ # 合并所有表示 X = tf.concat([rep_t1, rep_t0], axis=0) # 计算两两之间的核矩阵 pairwise_dists = tf.reduce_sum(tf.square(X[:, tf.newaxis] - X[tf.newaxis, :]), axis=-1) K = tf.exp(-pairwise_dists / (2 * sigma ** 2)) # 分割核矩阵 n1 = tf.shape(rep_t1)[0] n0 = tf.shape(rep_t0)[0] K_tt = K[:n1, :n1] K_cc = K[n1:, n1:] K_tc = K[:n1, n1:] # 计算MMD^2的无偏估计 mmd2 = (tf.reduce_sum(K_tt) / (n1 * (n1 - 1)) + tf.reduce_sum(K_cc) / (n0 * (n0 - 1)) - 2 * tf.reduce_sum(K_tc) / (n1 * n0)) # 确保非负 return tf.maximum(mmd2, 0.0) def tarnet_loss_fn(y_true, y_pred, sample_weight, representation, treatment, alpha=1.0): """ 自定义TARNet损失函数。 y_true: 真实结果。 y_pred: 模型预测结果。 sample_weight: 样本权重。 representation: 共享层的输出表示。 treatment: 处理指示。 alpha: 平衡正则项的强度系数。alpha=0即为原始TARNet，alpha>0则为CFR。 """ # 1. 加权预测损失 mse = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE) prediction_loss = mse(y_true, y_pred) weighted_prediction_loss = tf.reduce_mean(prediction_loss * sample_weight) # 2. 表示平衡正则项 (MMD) # 根据处理标签分离表示 mask_t1 = tf.cast(tf.equal(treatment, 1), tf.bool) rep_t1 = tf.boolean_mask(representation, mask_t1) rep_t0 = tf.boolean_mask(representation, tf.logical_not(mask_t1)) # 只有当两组都有样本时才计算MMD balancing_loss = tf.cond( tf.logical_and(tf.shape(rep_t1)[0] > 1, tf.shape(rep_t0)[0] > 1), lambda: mmd_loss(rep_t1, rep_t0), lambda: tf.constant(0.0, dtype=tf.float32) ) # 总损失 total_loss = weighted_prediction_loss + alpha * balancing_loss return total_loss

现在，我们可以组装训练流程。由于使用了自定义损失，我们需要在训练循环中手动计算损失和梯度。

# 初始化模型、优化器和损失跟踪器 model = TARNet(input_dim=X_train_processed.shape[1]) optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3) train_loss_metric = tf.keras.metrics.Mean(‘train_loss‘) val_loss_metric = tf.keras.metrics.Mean(‘val_loss‘) # 准备TensorFlow Dataset batch_size = 256 train_dataset = tf.data.Dataset.from_tensor_slices( ((X_train_processed, T_train), Y_train, W_train) ).shuffle(buffer_size=10000).batch(batch_size).prefetch(tf.data.AUTOTUNE) val_dataset = tf.data.Dataset.from_tensor_slices( ((X_val_processed, T_val), Y_val, W_val) ).batch(batch_size).prefetch(tf.data.AUTOTUNE) @tf.function def train_step(x_batch, t_batch, y_batch, w_batch): with tf.GradientTape() as tape: # 获取中间表示：需要修改call方法以返回表示，或使用子模型 # 这里我们创建一个返回表示和输出的模型 representation = model.get_representation(x_batch, training=True) y_pred = model((x_batch, t_batch), training=True) loss = tarnet_loss_fn(y_batch, y_pred, w_batch, representation, t_batch, alpha=0.5) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) train_loss_metric.update_state(loss) # 需要在TARNet类中添加一个方法来获取表示 class TARNet(tf.keras.Model): # ... __init__ 部分同上 ... def get_representation(self, x, training=False): for layer in self.representation_layers: x = layer(x, training=training) return x # ... call 部分同上 ... # 训练循环 epochs = 100 for epoch in range(epochs): # 重置指标 train_loss_metric.reset_states() val_loss_metric.reset_states() # 训练 for (x_batch, t_batch), y_batch, w_batch in train_dataset: train_step(x_batch, t_batch, y_batch, w_batch) # 验证（可选，在验证集上计算损失但不更新梯度） for (x_batch, t_batch), y_batch, w_batch in val_dataset: representation = model.get_representation(x_batch, training=False) y_pred = model((x_batch, t_batch), training=False) val_loss = tarnet_loss_fn(y_batch, y_pred, w_batch, representation, t_batch, alpha=0.5) val_loss_metric.update_state(val_loss) # 打印日志 print(f‘Epoch {epoch+1}, Train Loss: {train_loss_metric.result():.4f}, Val Loss: {val_loss_metric.result():.4f}‘)

训练过程中，密切关注训练损失和验证损失的变化。如果验证损失很早就开始上升，可能是过拟合的迹象，需要增加Dropout率、L2正则化强度或减少网络复杂度。调整alpha参数可以控制表示平衡的强度，通常需要通过交叉验证来寻找最佳值。

5. 模型评估、效果可视化与结果解读

模型训练完成后，我们不能只看损失函数，必须评估其因果效应估计的准确性。由于反事实的真实值无法获得，我们通常采用以下策略进行评估：

1. 模拟数据验证：在已知真实数据生成过程（例如，基于合成数据）的情况下，我们可以直接计算估计的个体处理效应（ITE）与真实ITE之间的误差，如精确匹配误差（PEHE）。

def evaluate_on_synthetic(model, X_test, t_test, y_test, true_ite): """ 在模拟数据上评估模型。 true_ite: 模拟数据中已知的真实个体处理效应。 """ # 为所有样本计算反事实预测 # 预测当T=0时的结果 y_pred_t0 = model.predict((X_test, np.zeros_like(t_test))) # 预测当T=1时的结果 y_pred_t1 = model.predict((X_test, np.ones_like(t_test))) # 计算估计的ITE estimated_ite = y_pred_t1 - y_pred_t0 # 计算PEHE (sqrt[mean((estimated_ite - true_ite)^2)]) pehe = np.sqrt(np.mean((estimated_ite.squeeze() - true_ite) ** 2)) print(f‘PEHE on test set: {pehe:.4f}‘) # 计算ATE误差 true_ate = np.mean(true_ite) estimated_ate = np.mean(estimated_ite) ate_error = np.abs(true_ate - estimated_ate) print(f‘True ATE: {true_ate:.4f}, Estimated ATE: {estimated_ate:.4f}, Absolute Error: {ate_error:.4f}‘) return estimated_ite, pehe, ate_error

2. 在真实数据上的评估：对于真实数据，我们可以采用倾向性得分分层或基于匹配的评估。基本思想是，在倾向性得分相近的样本子群内，处理分配可以近似看作是随机的。我们可以检查在这些子群内，模型预测的处理效应是否一致，或者计算分组后的平均处理效应（ATE）并与已知的领域知识或小规模随机试验结果对比。

3. 可视化分析：可视化是理解模型行为和结果的有力工具。

• 效应异质性分析：将估计的ITE按照协变量的某个维度（如用户价值）排序并绘图，观察效应如何随特征变化。

import matplotlib.pyplot as plt # 假设我们估计出了ITE，并有一个重要的连续特征‘user_value‘ ite_estimates = estimated_ite.squeeze() user_value = X_test[‘user_value‘].values # 排序 sort_idx = np.argsort(user_value) sorted_ite = ite_estimates[sort_idx] sorted_value = user_value[sort_idx] # 绘制散点图与局部平滑曲线 plt.figure(figsize=(10, 6)) plt.scatter(sorted_value, sorted_ite, alpha=0.5, s=10, label=‘Estimated ITE‘) # 使用移动平均或LOESS平滑 window_size = 50 smoothed_ite = np.convolve(sorted_ite, np.ones(window_size)/window_size, mode=‘valid‘) smoothed_value = np.convolve(sorted_value, np.ones(window_size)/window_size, mode=‘valid‘) plt.plot(smoothed_value, smoothed_ite, color=‘red‘, linewidth=3, label=‘Smoothed Trend‘) plt.xlabel(‘User Value‘) plt.ylabel(‘Estimated Individual Treatment Effect (ITE)‘) plt.title(‘Heterogeneity of Treatment Effect across User Value‘) plt.legend() plt.grid(True, alpha=0.3) plt.show()

• 表示空间可视化：使用t-SNE或UMAP将共享表示层学到的特征Z降维到2D，并用不同颜色标注处理组和对照组。一个好的表示应该让两组样本的分布高度重叠。

from sklearn.manifold import TSNE import seaborn as sns # 获取测试集样本的表示 Z_test = model.get_representation(X_test_processed, training=False).numpy() # t-SNE降维 tsne = TSNE(n_components=2, random_state=42, perplexity=30) Z_2d = tsne.fit_transform(Z_test) # 绘制散点图 plt.figure(figsize=(10, 8)) scatter = plt.scatter(Z_2d[:, 0], Z_2d[:, 1], c=t_test, cmap=‘coolwarm‘, alpha=0.6, s=20) plt.colorbar(scatter, label=‘Treatment (T=1)‘) plt.xlabel(‘t-SNE dimension 1‘) plt.ylabel(‘t-SNE dimension 2‘) plt.title(‘2D Visualization of Learned Representation (Colored by Treatment)‘) plt.show()

• 校准图：对于二元结果，可以绘制预测概率与实际观测频率的校准图，评估模型预测的准确性。

结果解读与业务应用： 最终，模型输出的不只是预测值，更是每个个体在两种（或多种）处理下的潜在结果差异。业务方可以据此进行精细化决策：

• 个性化策略：对ITE为正且值很大的用户（即对处理反应积极的用户），优先施加处理（如投放广告、给予优惠）。
• 避免伤害：对ITE为负的用户（处理可能产生负面效果），避免施加处理，节省资源并提升用户体验。
• 资源优化：将处理资源集中在效应最显著的群体上，最大化整体收益。

例如，在一个营销场景中，模型可能告诉你，对于高活跃度、但近期消费频率下降的用户（特征X），发送一张特定品类的优惠券（T=1）比发送通用红包（T=0）能带来高出15%的转化概率提升。这个结论可以直接指导运营人员制定精准的召回策略。

整个流程走下来，从数据准备到模型部署，最深的体会是因果推断模型对数据质量和特征工程的依赖远超普通预测模型。一个微小的混淆变量就可能导致效应估计严重偏误。因此，在应用TARNet或任何因果模型时，与业务专家紧密合作，深入理解数据生成过程，和花在调参上的时间同等重要，甚至更为关键。代码和模型是骨架，业务逻辑和领域知识才是赋予其生命的灵魂。