美团履约平台技术部在因果推断领域持续的探索和实践中,自研了一系列分布式的工具。本文重点介绍了分布式因果树算法的实现,并系统地阐述如何设计实现一种分布式因果树算法,以及因果效应评估方面qini_curve/qini_score的不足与应对技巧。希望能为从事因果推断相关工作的同学们提供一些启发或帮助。
1. 业务背景
2. 分布式因果森林框架
2.1 技术选型与框架设计
2.2 性能优化
2.3 Serving实现
3. 分布式因果效应评估
3.1 无偏性校验
3.2 因果效应量级关系评估
3.3 分布式评估体系
4. 总结
近年来,因果推断在商品定价、补贴、营销等领域得到广泛应用并取得了显著的业务效果提升,例如用户增长、活动营销等业务场景。这些领域的共性是需要“反事实推断能力”,传统机器学习算法更关注预测问题,而因果推断提供了更佳的反事实推断能力。以营销活动为例,我们不仅需要知道当前优惠券金额下,订单数是多少(预测问题),还要知道在改变金额的情况下,订单数会发生怎样的变化(反事实问题)。
常见的因果建模方法主要包含Meta-Learner、深度表征学习和Tree-Base算法三大类。其中以因果树为代表的Tree-Base算法泛化性强,适用于多种业务场景。相较于Meta-Learner,树模型建模流程简单;相较于深度表征学习,树模型特征处理和调参过程简单并且具备极强的可解释性。
具体来说,XGBoost、LGBM、Spark Random Forest等树模型是为解决预测问题而提出的经典算法实现,而因果树算法引入了新的训练理论以及因果理论独有的干预变量、工具变量等概念。这意味着我们并不能通过对现有分布式树模型的简单改造,来实现因果理论下树模型的分布式训练,而是需要充分理解各类单机因果树算法的原理之后,选择合适的分布式编程范式高效地实现出来。
为了解决上述问题,美团履约平台技术部对开源项目进行了精细梳理,集各家之所长实现了一套高性能的分布式因果森林框架,在半小时内即可完成亿级样本100棵树的训练,突破了单机开源项目仅支持百万级样本的瓶颈。并经过复杂的抽象设计,最终实现通过自定义损失函数即可支持各类因果森林算法的能力,极大提升了框架的扩展性。
除此之外,美团履约平台技术部还在因果效应评估、观测数据去偏等方面建设了大量高效实用的分布式工具。本文将重点为大家分享如何设计实现一个分布式的因果森林算法,以及因果效应评估方面的经验技巧,将我们在分布式因果推断领域的一些探索和内部的实践经验分享给大家。
因果森林算法的提出引发了Tree-Base算法应用于因果建模的研究热潮,众多学者相继在因果森林的基础上提出了多种多样的改进算法。监督学习领域的树模型有众多优秀的开源分布式实现,例如Xgboost、LightGBM、SparkRandomForest等等。
但是开源的因果树模型分布式实现基本处于空白状态。因果树算法引入了新的训练理论(比如Honesty Tree)并且因果树的分裂还依赖于干预变量、工具变量,这导致我们无法通过对现有分布式树实现做简单来更改来实现。因此,我们立足于论文,充分调研并借鉴业内优秀的开源实现,最终设计实现了一套高性能的分布式框架,并能提供统一的Serving方案。
单机树模型的工程实现可以概括为:遍历所有潜在的切分点并计算分裂指标(损失函数),取指标指标最佳的分裂点分裂,不断分裂树节点直到满足退出条件。而分布式环境下每台机器只包含部分样本,分布式环境下任何全局指标计算都会带来极大的通讯成本,因此需要选择合适的分布式架构帮助我们计算分裂指标。
因此,对于分布式因果森林框架,我们关心三个问题:第一,如何计算因果树的分裂指标(损失函数);第二,如何求潜在分裂点;第三,选用何种分布式编程架构。在此基础上进一步抽象整合,就可以实现不同树模型共用一套分布式框架的目标。
从论文出发
为了深入了解因果森林类算法,我们仔细阅读了因果森林论文以及其作者Susan Athey的另一篇在因果领域有重要影响力的《Generalized Random Forests》论文。Susan Athey认为随机森林本质上是一种自适应的最近邻算法(KNN),也就是通过对样本空间的递归划分从而找到距离该样本点最近的K个点(落入同一个叶子节点)来表示该点的值。而因果森林算法本质上是随机森林算法在因果推断领域的一种特殊应用。
因果森林和传统分类、回归森林一样采用了二叉的CART树(Classification And Regression Tree)作为基模型。与分类和归回问题相同,特征值仅用于样本划分而不参与分裂指标的计算。不同之处在于,分类和回归问题仅研究预测观测值Y,而因果建模需要研究treatment、instrumental variable等变量与观测值Y之间的关联。此外,多维连续treatment是学界的热门研究方向。因此,相较于分类和回归问题,因果推断需要在样本表示上做出相应调整。
因果森林论文提出honestyTree的概念:将样本分成growSet和predictionSet两个部分,growSet用于树的生长,predictionSet用于prediction值的计算。在论文《Generalized Random Forests》中证明了最小化子节点评估值与真实值之间的误差等价于最大化左右节点间的异质性,并对CART树的生长过程做了更加广义的抽象,将其分解成labeling step和regression step两步。Susan Athey的单机C++开源项目grf-lab中将这两种观点融合在一起,把树的生长定义为relabeling/splitting/prediction三个步骤。
Pre-sorted Algorithm Or Histogram-based Algorithm ?
主流CART树模型求分裂点的实现有两种方式,以早期XGBoost为代表的预排序算法,以LightGBM和SparkRandomForest为代表的直方图算法(目前XGBoost也提供了直方图算法的实现)。
AllReduce Or MapReduce ?
因为XGBoost内建了一个AllReduce框架RABIT可以直接复用,因此我们迅速拟定了两个调研方向——复用XGBoost的AllReduce高性能实现和Spark MapReduce实现。
由于履使用的样本量在几千万级别,综合考虑开发测试成本和训练性能后,我们最终选择了MapReduce方案。
框架设计
在选定MapReduce+直方图的方案后,我们迅速将目光锁定在同样使用直方图算法的Spark RandomForest算法上(以下简称SparkRF)。我们在SparkRF上快速实现了一版分布式因果森林框架,并进一步实现了Generalized Causal Forests算法。
但是测试过程中我们发现,随着总节点数的增加,跨节点通信量(也就是Shuffle)剧增,同时还非常容易溢出。为了支持更大规模的模型训练,我们从跨节点通信、内存占用、计算复杂度、剪枝以及CPU缓存命中优化等多个方面优化了整个框架。为了讲清楚我们优化逻辑,大家先来看看SparkRF是如何实现的。
SparkRF的实现
SparkRF整个实现过程可以概括为如下几个步骤:
(1)将全量样本离散化并cache到内存,这一步包含三部分:
(2)树的生长
可以看到,Spark的实现除了直方图,还有不少精妙的地方。例如在每次可训练的总结点数有限的情况下,深度优先搜索相较于广度优先搜索更倾向于快速完成单棵树的训练,从而减少后续训练需要广播的树模型。篇幅所限,下面将主要为大家介绍分布式因果森林框架在内存占用方面的优化。
减少Cache体积
从上文可以看出,SparkRF使用int来表示最大分桶个数,而lightGBM使用无符号byte来存储,支持最多256个分桶。我们认为128个分桶足以支撑因果森林的业务需要,所以使用了有符号byte来表示分桶,相比int内存占用减少至1/4。
前文中提到,SparkRF为每个样本创建了一个标记数组。例如训练一个2棵树的森林,这个标记数组为[4,0],这表示此样本在tree0有放回采样4次,在tree1未被使用。此外,框架需要支持honestyTree,也就意味着需要另一个标记数组记录样本在growSet还是predictionSet。考虑到无放回采样足以覆盖绝大部分场景,并且为了不引入第二个标记数组,我们最终选择了BitSet实现。每棵树最多使用2个bit,1个bit表示是否是该树的样本,1个bit表示是否是honesty样本。当关闭honesty或者不使用下采样时,每棵树只需要1个bit,内存占用最多减少至1/32。
支持更大模型广播
上文中提到,SparkRF每一轮迭代调用reduceByKey之前都需要计算出哪些样本属于待分裂的节点,Spark通过useNodeIdCache参数提供了两种策略:
从源代码中我们发现,策略二每一轮迭代都会卸载上一轮持久化的nodeIdCache,再创建一个新的nodeIdCache持久化到内存。以1亿条样本100棵树的森林举例,每一轮迭代就是1亿个长度为100的int数组的创建与垃圾回收。实际测试中我们也发现策略二的效率不如方案一高。那么策略一又如何呢?
SparkRF在每一轮迭代中能够训练的最大节点数由maxMemoryInMB控制,我们希望通过增大这个参数来减少迭代次数。但随着树或树深的增加,往往陷入增大该参数就导致树模型广播到worker溢出的尴尬境地。经过对SparkRF源码分析,我们发现每个LearningNode都会存储当前节点、左子节点、右子节点的直方图,最终实现在一套通用框架下计算出每个节点的增益、纯度、预测值等等属性,但这导致了3倍的内存占用。
考虑到因果森林honestyTree原则,叶节点prediction值的计算使用predictionSet,因此生长过程中每个节点全都带着growSet的直方图是完全没有意义的。因此我们优化了树的生长逻辑,每个节点仅保留自身的直方图,对于已分裂的节点则清除直方图。以二叉满树为例,叶节点约占整棵树节点的1/2,结合直方图从3倍冗余到1倍存储,这一优化使树模型直方图的内存占用下降到原本的1/6,极大降低了模型体积。
BenchMark
经过一系列优化,最终实现了百棵树亿级样本小时级训练的目标。
备注:不同森林算法的复杂度不同,跨节点通讯量不同,总耗时会存在明显的差异。
因果森林本质上是随机森林算法的变种,由一棵棵彼此独立的二叉因果树构成,每棵树由innerNode和leafNode构成。其prediction的逻辑非常简单,每棵因果树单独predict获取leafOutput向量,森林中所有树预估的leafOutput向量求均值即可得到森林的输出值。因此,整个树模型的结构其实非常清晰,innerNode存储特征split信息,leafNode存储输出向量。除此之外还包含gain、impurity、count等属性用于计算特征重要性。
模型serving除了性能还需要考虑模型离线存储体积、模型的内存占用、模型字段的扩展性。结合因果树的特点,就需要特别注意leafOutput向量的实现。以下表中的场景为例,使用float数组大约就需要500409640* 4 byte / 1024/ 1024 = 312.5mb,而List则需要约4倍内存,正因如此我们快速放弃了简单快捷的Protobuff方案。
为什么要重视模型字段的扩展性呢?这是因为离线模型训练追求快速迭代而在线Serving追求稳定性。模型的扩展性好,不仅可以轻松做到新版本服务向下兼容老模型,还可以做到在不使用新特性的情况下,老版本服务向上兼容新模型,从而减少在线服务更新发版的次数。综合考虑以上因素以及对Spark的兼容性和对java serving生态的兼容性,我们设计了如下方案。
业内常见的因果效应评估手段主要评估ITE的序关系,例如qini score和auuc。但是存在如下三方面不足:
下文将为大家一一进行说明。
无偏性校验分为数据无偏性和模型无偏性。
数据无偏性校验可以通过X⊥T验证。首先可以训练一个X->T的倾向性得分模型,如果倾向性得分模型的auc在0.5附近则说明X无法正确地预测T,也就是说X⊥T,此时数据无偏。例如,使用了post-treatmen特征会导致特征穿越,最终导致数据是有偏的,这时候使用X⊥T的校验工具可以快速帮我们排查出这一类问题。
因果效应的序关系和量级关系同样重要,只是将弹性的序关系学习准确而没有将弹性的量级关系学习准确,决策者无法预估该treatment对用户的影响程度。例如,将量级错误的弹性应用到运筹优化决策中,可能会导致无法满足重要约束从而无法求得可行解。针对弹性量级无法评估的问题,我们在原有的qini_curve基础上增加了qini_pred_curve_counterfactual和qini_pred_curve。
qini_curve及其扩展
通过比较qini_pred_curve_counterfactual和qini_curve这两条曲线的重合程度和右端点纵坐标,我们可以观察出ITE的预估量级和真实量级是否一致。
根据这些曲线的形状、覆盖面积、重合程度,我们可以得到如下的判断:
avgITE和ATE的对比
上文中提到的三项指标都是累计因果效应的评估,我们还想更有针对性地观察每个弹性分桶下预估因果效应和真实因果效应量级的差异,所以开发了avgITE和CATE的对比工具。
早期我们也使用了pandas实现的单机评估算法,当样本量增加到400w条以上时遇到了严重的单机瓶颈。为此,我们对上述评估指标全部做了分布式改造。排序类指标的实现有分桶积分和逐条积分两种实现思路。考虑到逐条积分会有更高的精度,最终选择了分布式环境下逐条积分的方案。
---------- END ----------
美团履约平台诚招调度履约方向、LBS方向、机器学习平台、算法工程方向的技术专家和架构师,共建全行业最大的单一即时配送网络和平台,共同面对复杂业务和高并发流量的挑战,迎接履约业务全面智能化的时代。感兴趣同学可投递简历至:zhengzhenzhen02@meituan.com(邮件标题注明:美团履约技术团队)。
| 2023美团技术沙龙盘点:21个视频揭秘复杂业务场景背后的技术实践