WSDM 2023 | GNN竟是动态规划？轻量级特征交互图网络它来了！
© 作者｜田震
机构｜中国人民大学高瓴人工智能学院
研究方向｜推荐系统
本文主要提出了一个轻量级的图神经网络模型，其通过知识蒸馏技术对目前主流点击率预估模型进行轻量化压缩，在4个点击率预估数据集上的性能与效率均取得了较好效果。文章也同步发布在AI Box知乎专栏（知乎搜索 AI Box专栏），欢迎大家在知乎专栏的文章下方评论留言，交流探讨！
论文题目：Directed Acyclic Graph Factorization Machines for CTR Prediction via Knowledge Distillation
论文下载地址：https://arxiv.org/pdf/2211.11159.pdf
论文开源代码：https://github.com/RUCAIBox/DAGFM
Directed Acyclic Graph Factorization Machines for CTR Prediction via Knowledge Distillation
Zhen Tian, Ting Bai, Zibin Zhang, Zhiyuan Xu, Kangyi Lin, Ji-Rong Wen, Wayne Xin Zhao
一. 前言
点击率预测任务旨在预测用户点击某个商品的概率，是推荐系统中一类非常重要的任务。而点击率预测任务的核心环节是建模特征交互。如何建模高阶交互一直是学术研究的热点，并且有诸多学者设计出了有效的模型结构来显式地建模高阶特征。然而，这些方法通常具有非常高的时间复杂度，往往无法满足大规模工业级场景下的效率要求。因此，点击率预测模型的轻量化成为一个亟待解决的重要课题。
知识蒸馏是一类十分有效的轻量化技术，被广泛地应用到计算机视觉以及自然语言处理等任务中。它的主要思想是将复杂模型作为教师模型，将其知识迁移到一个轻量级的学生模型中。然而，由于学生模型结构的局限性，知识蒸馏技术很难将教师模型的特征交互知识完全迁移到学生模型中，因此会产生较大的精度损失。
为了解决上述问题，我们想要设计一个能够完美拟合各种特征交互模型的轻量级学生模型。我们注意到目前大部分特征交互模型的参数低秩性，即真实场景中往往只有一小部分交互起到重要作用，其特征交互的语义权重具有极大的分解潜力。为了提取教师模型全部特征交互知识，我们设计了一个轻量级图神经网络模型(DAGFM)，其中图信息传播过程与一个构造任意阶交互的动态规划算法严格对应。因此，任意一个特征交互都对应着动态规划图中唯一一条转移路径，特征交互的权重为转移路径中边权重之积。DAGFM对任意复杂的特征交互模型都达到了几乎无损的压缩效果，并且极少的参数量使得它在微调阶段可以缓解教师模型面临的过拟合效应，在不同的推荐场景下的性能与效率都有较大提升。
二. KD-DAGFM模型
1. 拓扑结构
模型结构如图所示，我们首先将用户和物品特征 建模为一个 阶无环竞赛图 。其中 ， 每个图节点有一个状态向量 ，其初始时刻状态为特征表示 。
2. 特征交互
特征交互过程建模为图节点的邻域聚合过程。在 时刻，节点 聚合其邻居节点的状态信息，并更新其自身状态向量：
其中 为交互函数，其参数为边权重 ，其具体可以设计为：
Inner Product：
Kernel Product:
其中，Inner Product 复杂度较低( )，但建模域交互能力较差，代表模型有FwFM，FvFM以及XDeepFM等。而Kernel Product 建模域交互能力较强，但复杂度较高( )，其代表模型有FmFM，DCNV2等。
为了构建一个具有较低时间复杂度并且具有较强域交互建模能力的交互函数，我们利用Kernel Product的低秩特性，将其参数矩阵进行谱分解，并保留主特征值，即分解为： 。基于此，我们设计了具有线性复杂度的交互函数：
Outer Product:
在 时刻，我们对节点状态作sum pooling：
得到图状态： ，特征交互信息定义为： 。网络给出的预测为： 。
3. 动态规划与任意阶交互
我们指出， 阶无环竞赛图上的邻域聚合过程与一个构建任意阶交互的动态规划算法对应。
我们定义交互后缀为其最大索引特征，即 的后缀为 。其次，我们定义 为所有以 为后缀的 阶交互的集合，例如 。定义状态变量 ，其初始状态 。显然，所有以 为后缀的 阶交互可以表示为所有以 ( )为后缀的 阶交互与 构成的交互之和。因此，其状态转移方程可以表示为：
所有特征的全阶交互可以表示为所有状态向量之和。注意到在DAGFM中，如果将Inner Product的参数置为全1向量(或者将Kernel Product的参数置为单位矩阵)，其交互方程可以表示为：
因此，其图聚合过程的本质为动态规划的转移过程，图节点的状态变量与动态规划算法定义的状态完全等价。图中的边权重刻画了动态规划的转移权重，亦即特征域间的交互权重。
如图所示为图聚合展开的DAG，同样为DP算法的状态转移图。其中，任意一个 阶交互都对应着图中唯一一条长度为 的路径，其特征交互的语义权重为路径中边权重之积。例如，节点 分别代表特征male, weekday, Canada 与 sunny，特征交互male x weekday x Canada 对应的路径为 ，权重为 。事实上，边权重 控制特征 与 在第 阶的交互权重。
综上，DAGFM可以无重复无遗漏地表达出任意阶交互，并且在任意时刻每个节点状态，每条边的权重均有着确切的语义信息，保证所有交互语义的无重无耦性。
4. 知识蒸馏
知识蒸馏过程如图1(b)所示，在蒸馏过程中，我们共享教师与学生网络间的特征表示层，主要进行教师-学生之间特征交互语义权重间的拟合。由之前阐述，学生模型具有任意阶交互的拟合能力，并且每个交互的均有其特有的参数组合。因此这个过程可以理解为对教师的参数进行因子分解，将其知识提取到DAGFM的带权有向边中。蒸馏过程采用MSE Loss：
同时，结合BCE Loss保证学习方向的正确性： 最终的学习Loss为二者加权：
由于在蒸馏过程中，学生模型的特征表示层继承自教师模型，在蒸馏过程被冻结，导致其与最终学习到的主干网络参数存在bias。于是，我们最后微调学生网络得到最终用于上线的推理模型。
三. 实验结果
1. 数据集
我们使用了3个公开数据集：Criteo，Avazu以及MovieLens-1M。以及一个工业级数据集：微信公众平台数据集。
2. 蒸馏实验
我们首先将CrossNet与CIN作为教师模型，并将FwFM，DAGFM，FmFM以及Tiny MLP等四个轻量级网络作为学生模型进行蒸馏实验。可以看到，由于FwFM以及FmFM只能学习到二阶交互知识，它们在蒸馏过程无法达到完全收敛，在精度方面有较大损失。而MLP只能拟合隐式交互信息，无法对Inner Product或者Kernel Product这类显式运算进行有效学习。相比之下，DAGFM对于CIN以及CrossNet都达到了近乎无损的压缩效果。
3. 预估效果实验
我们与目前CTR主流baseline在预测效果与效率方面进行对比。可以看到，我们提出的蒸馏模型最多能够压缩教师模型132倍的参数量以及1200倍的浮点运算量，并且达到了接近无损的压缩效果，这也说明了教师模型大部分的参数计算都是冗余的。同时，极低的参数量使得DAGFM能够有效地避免教师模型面临的过拟合现象，我们通过微调可以进一步提升学生模型的预测表现，其推荐性能优于绝大部分baseline。
4. 消融实验
我们首先我们选择层数为3的CIN（构造4阶交互）模型作为教师模型来探究不同聚合层数的学生模型对于蒸馏表现的影响。可以看到，层数为1或2的DAGFM无法学习到4阶交互的知识，无法达到完全收敛，精度有一定损失。而层数为3或4的DAGFM可以涵盖到教师模型中的全部特征交互，达到无损的压缩效果。
其次我们选择CIN与CrossNet作为教师模型来探究不同交互函数下的学生模型对于蒸馏表现的影响。由于CIN和CrossNet分别采用inner product和kernel product构造特征交互，而kernel product可以退化到inner product，因此DAGFM-kernel可以完全拟合CIN和CrossNet，而DAGFM-inner无法拟合CrossNet。而我们设计的outer product退化自kernel product，对于CIN和CrossNet均有极佳的拟合效果。
5. 在线A/B Test
我们在微信公众号文章推荐场景下对算法进行在线测试，可以看到对比之前的state-of-the-art方法，文章的点击用户数以及CTR均有所提升，并且推荐计算效率提升效果明显。
6. 通用蒸馏模型
DAGFM能够对各种显式特征交互模型进行近乎无损的压缩，然而其无法有效地对MLP等隐式特征交互结构进行知识提取。因此，我们将DAGFM与MLP进行结合得到一个通用的蒸馏模型。结果显示，其能够对xDeepFM，DCNV2，AutoInt+以及FiBiNet等兼具显式与隐式的交互模型进行无损压缩，显著提高计算效率。
四. 总结
我们提出了一个轻量级图神经网络模型以及点击率预测模型轻量化的知识蒸馏算法。图网络的建模过程与一个构造任意阶交互的动态规划算法相对应，证明了其能够有效地建模任意阶交互。图聚合过程可以集成不同的特征交互函数，使得它能够接近无损地提取任意复杂的特征交互模型的知识，并且可以极大地提高计算效率。我们通过在线和离线实验验证了算法的有效性，体现了DAGFM在处理大规模工业级推荐场景下的优越性。
更多推荐
800篇顶会论文纵览推荐系统的前沿进展
PLM是做题家吗？一文速览预训练语言模型数学推理能力新进展
ICLR2023推荐系统投稿论文集锦