EconML 中的联邦学习

概述
将联邦学习整合到 EconML 库的动机
使用 EconML 进行联邦学习
- 介绍 FederatedEstimator
- 使用示例
理论

概述 

EconML 库中的联邦学习允许在单独的数据集上训练模型，然后将它们组合成一个统一的 CATE 模型，而无需将所有训练数据收集到一台机器上。

将联邦学习整合到 EconML 库的动机 

1. 大型数据集：对于大到无法放入单台机器的数据集，联邦学习允许您对数据进行分区，在每个分区上训练一个独立的因果模型，然后将这些模型组合成一个统一的模型。

2. 隐私保护：联邦学习使组织能够在不集中或共享敏感数据的情况下构建机器学习模型。这对于遵守数据隐私法规可能很重要，因为它能使数据保持本地化，并降低合规风险。

使用 EconML 进行联邦学习 

介绍 FederatedEstimator 

我们提供了 FederatedEstimator 类，用于聚合在不同数据子集上训练的独立估计器。这些独立估计器必须是同一类型，目前必须是 LinearDML、LinearDRLearner、LinearDRIV 或 LinearIntentToTreatDRIV 中的一种。

与其他估计器不同，您不应该在 FederatedEstimator 的实例上调用 fit 方法；相反，您应该单独训练您的独立估计器，然后将已训练的模型传递给 FederatedEstimator 的初始化函数。然后，FederatedEstimator 将把这些独立估计器聚合成一个统一模型。

请注意，由于使用联邦学习会带来内存开销，每个估计器都必须通过将 enable_federation 设置为 True 来选择启用它。

使用示例 

这里是对基本 FederatedEstimator API 的一个极其基础的演示

# Create individual LinearDML estimators
num_partitions = 3
estimators = []
for i in range(num_partitions):
    est = LinearDML(random_state=123, enable_federation=True)
    # Get the data for this partition
    X_part, y_part, t_part = (arr[i::num_partitions] for arr in (X, y, t))

    # In practice, each estimator could be trained in a distributed fashion
    # e.g. by using Spark
    est.fit(Y=y_part, T=t_part, X=X_part)
    estimators.append(est)

# Create a FederatedEstimator by providing a list of estimators
federated_estimator = FederatedEstimator(estimators)

# The federated estimator can now be used like a typical CATE estimator
cme = federated_estimator.const_marginal_effect(X)

实际上，如果您的所有数据都可以载入内存并集中到一台机器上（如上例所示），那么使用联邦学习而不是在整个数据上训练单个模型就没有优势。在更现实的场景中，将单个模型拟合到所有数据上是不切实际的，模型应该在不同的机器上并行训练，进行序列化，然后在能够创建联邦估计器的中心机器上进行反序列化。使用 Python 通过 pickle 模块进行的内置序列化，其过程大致如下：

# We're running this as a loop here, but in practice each of these iterations would be run on a different machine
for i in range(num_partitions):

    # Get the data for this partition, in practice probably by loading from disk
    X_part, y_part, t_part = (arr[i::num_partitions] for arr in (X, y, t))

    # The code to train a model and serialize it runs on each machine
    import pickle

    est = LinearDML(random_state=123, enable_federation=True)
    est.fit(Y=y_part, T=t_part, X=X_part)

    with open(f'model{i+1}.pkl', 'wb') as f:
        pickle.dump(est, f)


# On the central machine, deserialize the models and create the federated estimator
with open('model1.pkl', 'rb') as f1, open('model2.pkl', 'rb') as f2, open('model3.pkl', 'rb') as f3:
    est1 = pickle.load(f1)
    est2 = pickle.load(f2)
    est3 = pickle.load(f3)
federated_estimator = FederatedEstimator([est1, est2, est3])

# The federated estimator can now be used like a typical CATE estimator
cme = federated_estimator.const_marginal_effect(X)

理论 

许多估计器正在求解一个矩方程

\[\E[\psi(D; \theta; \eta)] = 0\]

其中 \(D\) 是数据，\(\theta\) 是参数，\(\eta\) 是 nuisance 参数（干扰参数）。通常，该矩是参数的线性函数，因此可以重写为

\[\E[\psi_a(D; \eta)\theta + \psi_b(D; \eta)] = 0\]

在这种情况下，使用经验期望求解方程得到

\[\begin{split}\begin{align*} \hat{\theta} &= -\E_n[\psi_a(D;\hat{\eta})]^{-1} \E_n[\psi_b(D;\hat{\eta})] \\ \sqrt{N}(\theta-\hat{\theta}) &\sim \mathcal{N}\left(0, \E_n[\psi_a(D;\hat{\eta})]^{-1} \E_n[\psi(D;\hat{\theta};\hat{\eta}) \psi(D;\hat{\theta};\hat{\eta})^\top] \E_n[\psi_a(D;\hat{\eta})^\top]^{-1}\right) \end{align*}\end{split}\]

方差计算中的中心项可以展开为

\begin{align*} \E_n[\psi(D;\hat\theta;\hat\eta) \psi(D;\hat\theta;\hat\eta)^\top] &= \E_n[(\psi_b(D;\hat\eta)+\psi_a(D;\hat\eta)\hat\theta) (\psi_b(D;\hat\eta)+\psi_a(D;\hat\eta)\hat\theta)^\top] \\ &= \E_n[\psi_b(D;\hat\eta) \psi_b(D;\hat\eta)^\top] + \E_n[\psi_a(D;\hat\eta)\hat\theta\psi_b(D;\hat\eta)^\top] \\ &+ \E_n[\psi_b(D;\hat\eta) \hat\theta^\top \psi_a(D;\hat\eta)^\top] + \E_n[\psi_a(D;\hat\eta) \hat\theta\hat\theta^\top\psi_a(D;\hat\eta)^\top ] \end{align*}

其中一些项包含 \(\hat\theta\) 出现在内部位置的乘积，但这可以等价地通过计算两侧矩阵的外积，然后与 \(\hat\theta\) 收缩（contracting）来计算。因此，我们可以分布式计算以下量：

\begin{align*} & \E_n[\psi_a(D;\hat\eta)] \\ & \E_n[\psi_b(D;\hat\eta)] \\ & \E_n[\psi_b(D;\hat\eta) \psi_b(D;\hat\eta)^\top] \\ & \E_n[\psi_b(D;\hat\eta) \otimes \psi_a(D;\hat\eta)] \\ & \E_n[\psi_a(D;\hat\eta) \otimes \psi_a(D;\hat\eta)] \\ \end{align*}

然后我们可以聚合这些分布式估计，使用前两个计算 \(\hat\theta\)，再使用 \(\hat\theta\) 和其余项计算解析方差。

例如，对于 \(y\) 关于 \(X\) 的线性回归，我们有

\[\begin{split}\psi_a(D;\eta) = X^\top X \\ \psi_b(D;\eta) = X^\top y\end{split}\]

因此我们需要分布计算的额外矩是

\[\begin{split}\begin{align*} & \E_n[X^\top y y^\top X] = \E_n[X^\top X y^2] = \E_n[X \otimes X \otimes y \otimes y] \\ & \E_n[X^\top y \otimes X^\top X] = \E_n[X \otimes X \otimes X \otimes y]\\ & \E_n[X^\top X \otimes X^\top X] = \E_n[X \otimes X \otimes X \otimes X] \\ \end{align*}\end{split}\]

因此，通过存储这三个额外矩的代价，我们可以分布式计算线性回归，并恢复与在完整数据集上计算完全相同的结果。

在联邦 CATE 估计的背景下，请注意实际上 nuisance 参数是在数据的子集上计算的，因此尽管聚合后的最终线性模型与在本地使用所有相同的 nuisance 参数计算的模型完全相同，但实际上，如果在所有数据上计算 nuisance 估计，它们可能会有所不同。在实践中，只要 nuisance 估计器以合理的速度收敛，这应该不是一个重大问题；例如，如果第一阶段模型足够准确，使得最终估计量以 \(O(1/\sqrt{n})\) 的速度收敛，那么将数据分割成 \(k\) 个分区应该只会将方差增加 \(\sqrt{k}\) 的因子。

EconML 中的联邦学习

概述

将联邦学习整合到 EconML 库的动机

使用 EconML 进行联邦学习

介绍 FederatedEstimator

使用示例

理论