介绍

ali在2020《XDL: An Industrial Deep Learning Framework for High-dimensional Sparse Data》提出了它们的XDL平台。我们看下它的实现:

摘要

随着数据和计算能力的快速增长,基于深度学习的方法已成为许多人工智能问题的主流解决方案,如图像分类、语音识别和计算机视觉。包括Tensorflow、MxNet和PyTorch在内的几个优秀的深度学习(DL)框架已经开源,进一步加速了社区的进步。然而,现有的DL框架并不针对涉及高维稀疏数据的应用程序设计,这种数据在许多成功的在线业务中广泛存在,如搜索引擎、推荐系统和在线广告。在这些工业场景中,深度模型通常在大规模数据集上进行训练,这些数据集包含高达数十亿的稀疏特征和数千亿的样本,给DL框架带来了巨大挑战。

在本文中,我们介绍了一个高性能、大规模且分布式的DL框架XDL,它提供了一个优雅的解决方案,填补了现有DL框架的通用设计与高维稀疏数据引起的工业需求之间的差距。自2016年以来,XDL已在阿里巴巴成功部署,服务于许多生产环境,如在线广告和推荐系统。在数百个GPU卡上并行运行,XDL可以在仅几个小时内训练具有数千亿参数的深度模型。除了其卓越的性能和灵活性,XDL对开发人员也非常友好。阿里巴巴的算法科学家可以用几行简单的代码开发和部署新的深度模型。XDL API和参考实现作为开源包在2018年12月发布,遵循Apache 2.0许可,并可在https://github.com/alibaba/xdeeplearning找到。

1 引言

在过去的十年中,作为人工智能最令人兴奋和强大的分支之一,深度学习在语音识别、计算机视觉、自然语言处理和医学诊断等领域取得了重要的突破。得益于许多现实世界应用中产生的天文数量级的数据,以及图形处理单元(GPU)等基础设施带来的前所未有的计算能力,以及TensorFlow[1]、MxNet[6]、Caffe[18]等复杂的开源深度学习框架,深度学习技术已被广泛用于解决现实世界问题。

尽管现有的深度学习框架在许多领域取得了巨大成功,但它们并不针对涉及高维稀疏数据的应用程序设计友好。高维稀疏数据在许多互联网规模的应用中广泛存在,如搜索引擎、推荐系统和在线广告。例如,在我们的展示广告系统中,我们每天生成的用户行为日志数据量达到PB。从这些数据中提取的训练样本包含数十亿个特征,而每个样本只有少数几个维度是非零的

与语音识别、计算机视觉和自然语言处理等数据密集型应用不同,这些互联网规模的在线业务,数据稀疏,对深度学习框架提出了独特的挑战。

  • 首先,问题的规模如此之大,以至于无法在单个节点上解决。需要并行化来解决这种规模的问题。
  • 其次,数据极其稀疏。如果处理不当,稀疏性可能导致效率低下。
  • 第三,这些数据的表述并不规范。样本之间可能有重复的特征,给带宽带来高压。

此外,在快速模型演化中,需要添加或删除特征以测试其与目标的相关性。现有的深度学习框架并不设计友好地解决这些挑战。这种差距使得这些业务难以充分从蓬勃发展的深度学习技术中受益。

通过抽象我们在构建在线广告系统的深度学习解决方案中的实践经验,涉及高维稀疏数据的深度模型通常包括两部分:sparse特征学习和dense模型学习。稀疏特征学习是将大量高维稀疏数据转换为密集特征的过程,而密集模型学习则负责基于密集特征学习最优模型结构。这两部分还需要智能连接,以便这两个学习任务可以相互交互以实现最佳结果。

在本文中,我们介绍了XDL,这是一个为涉及高维稀疏数据的学习任务设计的高性能、大规模且分布式的深度学习框架。在稀疏特征学习部分,XDL提供了一个精心设计的分布式系统,对I/O、数据管道、通信和GPU进行了深入优化,提供了极高的效率和可扩展性。用户可以使用embedding字典或像CNN/RNN这样的深度模型将稀疏项目特征或图像/文本特征映射到密集表示。在密集模型学习部分,XDL采用任何开源深度学习框架作为其后端,并采用一种名为桥接技术的新技术支持。有了这种设计,XDL能够加速互联网规模学习问题的培训,并在端到端生产系统中提供服务。

在对XDL进行真实世界数据集评估时,我们发现它至少可以比Tensorflow/MxNet的原生分布式版本快5倍。

自2016年以来,XDL已在我们公司的核心业务中部署,如电子商务推荐和在线广告。在数百台服务器上运行,XDL在仅几个小时内训练了数十亿参数的多个模型。除了其性能和灵活性,XDL致力于隐藏复杂的工程细节。使用XDL,我们的算法科学家可以用几行简单的代码开发和部署新模型。这种系统设计带来了灵活性,并解锁了许多算法创新。

2 相关工作

随着问题规模和复杂性的增加,我们面临着深度学习框架性能和可扩展性的更严重挑战。传统的解决方案如TensorFlow[1]、MxNet[6]和Caffe[18]通常在单机上运行,并为分布式训练提供有限支持。即使有多个GPU,这些单机解决方案也难以处理具有1011样本和1010参数的模型训练。因此,分布式训练平台似乎是一个有希望的解决方案,许多学者在这方面做出了巨大贡献。

MPI[14]定义了一组接口,用于机器之间的通信和协调。由于其兼容性和可扩展性,MPI在分布式内存应用的并行计算中被广泛使用。许多机器学习框架都建立在MPI之上。

  • Chen等人提出了RABIT[5],这是一个AllReduce库,通过增加容错属性改进了OpenMPI。
  • Adam Coates等人提出了COTS HPC技术[3, 9]。他们在InfiniBand互连上用MPI构建了一个GPU服务器集群。仅用16台机器,这个系统就可以有效地训练具有110亿参数的模型。

然而,这些框架的数据处理部分过于简化。因此,它们既不能充分利用数据特性,也不能处理大规模稀疏数据的工作。此外,基于MPI的系统在容错方面支持不足,因此在一些工业场景中导致了可用性问题。

另一个广泛使用的范式是参数服务器(PS),它可以利用键值存储适当处理稀疏数据。

  • Dean等人[11]首次引入了PS,并使用downpour SGD在谷歌训练大规模深度网络。
  • Petuum[30]引入了有界延迟模型到PS架构[23],实现了异步计算。

参数服务器在工业中也广泛使用。

  • DSSTNE[2]是亚马逊设计的DL框架,特别适用于大型稀疏数据集。它提供了极其高效的自动模型并行多GPU支持和100%确定性执行。模型并行和数据并行同时支持。
  • Kunpeng[34]是阿里巴巴基于PS的分布式学习系统,支持像稀疏逻辑回归和多项式加性回归树这样的模型。

这些基于PS的框架通常设计用来支持传统模型,如逻辑回归或卷积神经网络,它们只包含稀疏部分或密集部分。这些框架很难处理涉及大规模稀疏特征和定制深度网络的模型,如[12, 31, 32]中所述。

另一方面,机器翻译[29]和语音识别[3]等自然语言处理系统使用类似的SparseNet + DenseNet模型结构。然而,它们通常只有多达百万个单词,而我们通常有数十亿个项目ID或图像作为特征。

另一个活跃的研究领域是自动寻找分布式模型训练的最佳运行时配置

  • [27]使用强化学习来寻找最优设备放置。
  • [19, 20]基于Legion[4]搜索混合并行策略的最佳配置,Legion是一种用于异构和并行机器的高性能编程系统。

然而,这些算法都没有考虑到稀疏性,也没有为如此规模的问题提供解决方案。

这些相关工作为构建高效的模型训练系统提供了许多见解。通过学习它们的优缺点,我们开发了XDL以支持涉及高维稀疏数据的工业问题。

3 XDL架构

在本节中,我们介绍XDL的架构和设计哲学。在讨论之前,有必要描述我们设计的动力。

3.1 涉及高维稀疏数据的深度模型的网络特性

近年来,基于深度学习的方法成功地革新了在线业务中应用的算法,并取得了最先进的性能。作为先驱工作,DSSM[17]提出在网络搜索引擎的场景中用深度网络建模查询和文档的相关性。DSSM模型的核心架构遵循embedding&多层感知器(MLP)范式:首先将高维sparse输入映射到低维embedding空间,然后使用多层感知器(MLP)拟合标签。这种embedding&MLP架构激发了大多数后续工作设计具有高维sparse输入的深度模型,这些模型来自视频推荐、网络搜索和电子商务网站广告等各种应用。代表性网络包括Wide & Deep Learning[8]、Youtube深度模型[10]、DeepFM[15]、深度兴趣网络(DIN)[32, 33]和CrossMedia[12]等。关于推荐系统深度模型的详细调查可以在[31]中找到。

从算法的角度来看,基于embedding&MLP的网络抽象了这些工业模型家族,这些模型通常将从稀疏数据中学习分为两个步骤:

  • i)表示学习(representation learning),从高维sparse输入中捕获信息并将它们embedding到低维空间,
  • ii)函数拟合(function fitting),模型denseembedding表示和监督label之间的关系。

为了简单起见,我们称第一步的网络为SparseNet,第二步的网络为DenseNet。图1说明了这种抽象。

图片名称

图1 embedding和多层感知器(MLP)网络架构的抽象,主要包含SparseNet和DenseNet。SparseNet将原始的sparse输入(维度为D,其中D可扩展至数十亿)映射到低维表示(维度为d,传统上扩展至数千)。DenseNet学习拟合数据的函数

在工业应用场景中,新的深度模型应该快速发展以提高业务利润。因此,算法科学家希望训练系统易于理解和使用:新模型应该在脚本代码中以独立视角设计,并并行运行,将复杂的分布式训练细节隐藏在后台。基于这些考虑,我们设计了XDL,这是一个工业深度学习框架,旨在构建一个高性能系统来训练涉及高维稀疏数据的深度模型。我们将展示我们如何应对高维稀疏数据带来的挑战,并提出一个遵循桥接策略的优雅设计。

3.2 XDL的设计哲学和桥接方法论

从上述内容中,我们可以看到涉及高维稀疏数据的模型中的SparseNet和DenseNet需要不同的系统功能。DenseNet由几个密集层组成,需要在本地机器上进行高计算密度。包括Tensorflow、MxNet和Pytorch在内的几个DL框架已经开源,这些框架可以很好地处理这类DenseNet。然而,SparseNet包含数千亿个特征,这些特征来自原始样本。因此,能够处理这种情况的成功DL框架必须具有出色的分布式特性和足够的稀疏数据计算能力。上述提到的现有DL框架并不针对涉及高维稀疏数据的网络设计良好。

为了更好地理解训练涉及高维稀疏数据的深度模型的挑战,让我们详细看看SparseNet和DenseNet的规模。首先展示我们数据集的规模。表1显示了我们在展示广告系统中一天使用的典型生产数据量。

表1

表2显示了我们日常任务中使用的一些模型的网络参数统计。显然,SparseNet是这些模型中贡献最大困难的部分。DenseNet遵循现有深度学习框架所持有的传统设置。此外,SparseNet需要处理输入数据的实际I/O问题,其体积为数TB,以及复杂和繁重的并行问题。这些挑战使得SparseNet的训练与DenseNet相比截然不同且至关重要。

模型 输入维度 网络参数 输出维度
MLP 10亿 18亿 1440
DIN 10亿 18亿 33438
CrossMedia 2.5亿 5亿 1464
模型 输入维度 网络参数 输出维度
MLP 1440 1.2万 1
DIN 33438 1.7万 1
CrossMedia 1464 1万 1

表2: 多层感知器、深度兴趣网络[32]和Crossmedia[12]模型的网络统计。

综上所述,我们设计了XDL,其架构如图2所示。遵循分而治之的策略,我们提出了一种全新的桥接架构,其中训练系统由两个主要子系统构成:

图片名称

图2 XDL架构。服务器和worker不需要具有相同的等级。它们可以部署在同一物理节点上,以提高通信效率。

  • 高级模型服务器(AMS)。AMS提供了一个精心设计的分布式系统,为训练SparseNet提供极高的效率和可扩展性。AMS处理快速演变的表示学习算法,这些算法具有大规模sparse输入。支持embedding字典和像CNN/RNN这样的模型,将大型sparse输入映射到密集向量。

  • 后端worker(BW)。BW遵循常见的深度学习设置,从低维输入中学习,并允许采用任何开源DL框架作为其后端。有了这种灵活性,我们的算法科学家可以轻松尝试涉及DNN[17]、注意力机制[32]或门控循环单元[33]的新模型结构,以模拟复杂且不断演变的用户行为。

在前向传递中,后端worker首先从I/O模块读取数据,然后向AMS发送特征ID请求。AMS将根据ID计算embedding向量并将其发送回后端worker。worker还将接收到上一次迭代的更新后的DenseNet模型参数。在收集完所有dense输入向量和模型参数后,worker将对输入向量执行sum pooling,并将结果输入到Dense中。在反向传递中,数据流被反转,worker计算梯度。AMS将收集DenseNet和SparseNet的梯度。参数更新将在服务器端使用像Adam[22]或Momentum SGD这样的求解器进行。关于AMS模块的实现和优化的更多细节将在下一节中描述。

XDL具有高性能和良好的可扩展性,因此能够支持工业生产。从2016年开始,XDL被用于训练目标广告部门的CTR预测模型。在数百台机器上运行,XDL训练了一系列具有数十亿参数的DNN模型,支持数百种业务场景

为了便于在各种计算平台上部署,XDL可以被多种资源管理平台调度,如Yarn,并为各种数据存储系统提供数据I/O接口,如HDFS和Kafka。XDL的良好兼容性使得业务应用能够快速升级。

3.3 高级模型服务器

从参数服务器的架构设计中学习,我们设计了高级模型服务器(AMS),这是一个管理SparseNet和DenseNet训练的分布式系统。AMS负责存储和更新所有参数,包括SparseNet和DenseNet。因此,参数放置(Parameter Placement)成为高性能设计的关键时刻。考虑到在线业务的稳定性和可用性,容错也应得到充分考虑。

参数放置(Parameter Placement)

由于sparse参数和dense参数都存储在AMS上,参数放置(Parameter Placement)算法应处理这两种参数的不同特性:

  • 一方面,sparse参数占用大量内存,因为稀疏特征的维度可能高达数千亿。然而,每个mini-batch中,只有少数特征ID被worker请求,这表明sparse参数的I/O压力很低。此外,随着训练过程的进行,sparse参数的内存使用量会发生变化,因为样本会不断将新的ID带入sparse参数中。
  • 另一方面,dense参数占用的内存较少,但I/O压力较大,因为worker在每个mini-batch中都请求整个dense参数。

根据上述分析,我们的kv存储使用hash mapping作为AMS中sparse参数的低级数据结构,这也存在于传统的参数服务器中。因此,我们可以在每个服务器上平均分配哈希表桶,减轻服务器上的内存存储压力。不幸的是,这个过程会导致worker请求数量大幅增加。为了解决这个问题,我们在embedding字典查找之前,在worker端合并worker请求。这种优化在实践中为我们提供了巨大的加速。至于dense参数,每个参数的I/O压力将在内存限制内平均分配在每个服务器上。

容错

随着样本数量和模型复杂性的快速增长,特定深度模型的离线训练需要越来越多的时间。当DL系统中的某些角色失败时,从一开始就重新训练在时间和资源方面都非常昂贵。此外,在一些在线学习场景中,稳定性和可用性非常重要。因此,XDL中精心设计了容错机制。AMS由调度器和多个服务器组成。服务器使用定期心跳与调度器保持同步。在训练过程中,整个模型的快照将存储在某个位置。如果任何服务器由于某种原因失败,调度器将注意到异常状态,并设置整个AMS系统为不可用状态,直到服务器被Yarn或其他调度系统重启。如果AMS准备好了,调度器将通知服务器从最新的快照恢复并继续训练过程。

同时,由于worker负责读取样本,每个worker的读取状态也作为特殊参数存储在AMS中。当worker失败并重新启动时,它们可以通过从AMS拉取参数来恢复其读取状态。通过这种努力,worker在训练过程中不会丢失或重复读取任何样本。

异步更新

XDL支持同步模式和异步模式。

  • 在同步模式下,worker在每次迭代中拉取参数和推送更新各一次。每次迭代直到所有AMS从所有worker接收更新并对参数应用平均梯度后才完成
  • 在异步模式下,没有这样的迭代概念。worker独立操作。因此,每个worker可以在将更新推送到AMS后立即进入下一步,而不需要等待其他worker完成他们的更新。同时,在服务器端,来自worker的更新以无锁风格应用。我们使用这种方法,因为embedding表的更新非常稀疏,冲突很少发生。对于密集DNN,尽管每个更新可能不会成功应用于整个模型,但总是只有一个全局模型版本。因此,无锁风格更新可以被视为在模型上添加一个掩码以随机丢弃部分梯度

在实践中,我们发现异步模式可以实现更大的系统吞吐量,同时几乎不会损失模型性能。

4 系统实现和优化

4.1 I/O

在生产场景中,我们训练系统的输入,如样本和初始模型,通常存储在不同的位置。例如,一些样本从在线队列系统中流式传输,而其他样本则从离线分布式文件系统中检索。I/O模块设计的目标是:最大化来自数据源(如HDFS、Kafka等)的输入带宽使用

众所周知,分布式ML系统中的I/O带宽已成为最大的瓶颈之一。由于我们使用现有的DL框架作为计算后端,如何更快地吞吐数据成为一个更关键的问题。因此,我们实现了各种优化以最大化I/O吞吐量。

首先,不同样本之间存在大量重复特征。这些重复来自我们数据源的最开始:用户行为。一个用户点击两个不同的链接将产生两个包含相同用户重复特征的不同样本。图3说明了重复的结构。压缩可以优化存储、通信和计算效率

图片名称

图3 重复

分层样本压缩:

XDL充分利用样本之间存在ID重复的事实。在预处理阶段,原始样本被组织成多前缀树(multi-prefix trees),这大大减少了存储空间和通信成本。在训练阶段,许多前缀树被组合成用户定义的batch大小的mini-batch。这种压缩也为计算带来好处,因为它减少了重复的embedding向量计算。如果前缀树的总样本大小超过batch size,最后一个前缀树将被分割。

图4展示了分层样本压缩的一个例子。原始样本被组织成两个不同的3层前缀树。树的第一层、第二层和第三层分别代表用户特征、广告特征和创意特征。在训练阶段,最后一个前缀树根据batch size=6进行分割。构建了两个辅助张量(指示器),以指示相邻层之间的关系。

图片名称

图4 I/O数据和计算压缩。在数据预处理阶段,生成多前缀树以用于训练阶段的mini-batch生成。

4.2 workflow流水线

workflow流水线是加速程序运行过程的常用技术。通过多线程技术使程序中的阶段在时间线上重叠,可以节省相当的运行时间。XDL中也实现了几种workflow流水线。

为了最大化I/O吞吐量,我们用线程流水线处理过程。一个完整的训练迭代被分为三个阶段:

  • 1)读取样本并将它们组合成mini-batch;
  • 2)从键值存储或输入数据中预提取mini-batch的参数索引,以在服务器上进行计算;
  • 3)拉取模型参数并进行前向/反向传播。

流水线利用线程并行性通过重叠执行这三个阶段。

如图5所示,这三个阶段的工作被调度到三个不同的线程池。这三个线程池通过一个无锁队列进行协调。如果使用GPU,除了I/O实现中的流水线外,还设计了SparseNet和DenseNet之间的workflow流水线,以进一步提高XDL的性能。在这种情况下,由于DenseNet的输入是SparseNet的输出,SparseNet的结果可以在DenseNet的前一个训练过程结束之前预计算。假设SparseNet和DenseNet的训练时间分别为TsparseTdense总运行时间将从Tsparse+Tdense减少到maxTsparse,Tdense。在这种情况下,训练过程的效率显著提高,模型性能损失很小,因为当SparseNet的输出被预提取时,可能存在一些未更新的稀疏ID。因此,当相邻mini-batch之间有很少共同稀疏ID时,这种流水线通常被适应。

图片名称

图5 训练迭代由三个阶段组成,使用流水线来最大化并行性。

XDL还允许算法科学家用几行代码构建连接训练过程中任何阶段的流水线。深度模型可以被划分为任意数量的阶段(stage)。算法科学家在考虑性能和准确性时有多种选择,以在训练过程中构建流水线。

4.3 高级模型服务器的优化

为了提高AMS的效率,在各个方面都采用了优化。

为了加速embedding过程中的K-V查询,XDL利用了GPU的能力。GPU有两个优势:

  • 首先,GPU的内存带宽更高(Nvidia P100为700GBps,Intel E5-2699 v4为100GBps)。
  • 其次,它在数千个核心上并行运行(Nvidia P100有3584个核心)。

因此,我们可以在GPU上实现更快的embedding字典查找。GPU的缺点是GRAM大小有限。为了解决这个问题,只有embedding字典的索引将被预提取到GRAM中。

由于AMS在大多数情况下扮演参数服务器的角色,因此AMS和后端worker之间存在大量的网络通信。当并行度很高时,网络通信量可能非常大。在这种情况下,网络通信消耗的运行时间成为总运行时间的关键部分。XDL采用Seastar[28]作为其通信库,并在Seastar框架的基础上进行了许多有效的优化。涉及零拷贝和CPU绑定等技术,以确保网络通信不会成为整个训练过程的瓶颈。

在深度学习中,batch size是一个非常重要的超参数。一些模型在使用mini-batch size时表现出良好的性能,而一些模型更喜欢大batch。这要求DL框架能够处理不同的样本batch size。正如我们之前提到的,CPU绑定(cpu-binding)是AMS中的一项重要技术。AMS中的每个线程都绑定到一个CPU核心,因此节省了不同核心之间的线程切换时间。

  • 当样本batch size较小时,这种CPU绑定机制使AMS上的计算性能高。
  • 然而,如果样本batch size很大,将有太多的ID需要由单个线程处理。

在这种情况下,AMS有一个自适应机制,其中创建了一个额外的线程池,用于处理大量的ID。有了这种优化,AMS在mini-batch和大batch样本上都表现出色。

4.4 使用XDL进行在线学习

在线学习方法最近在工业中被广泛使用,因为它能够实时捕捉客户行为的变化。一个有效的在线学习模型在许多业务场景中非常有价值。XDL为在线学习提供了一系列机制。

特征入口过滤器(Feature Entry Filter)

正如我们之前提到的,由于样本将新的ID带入sparse参数,sparse参数的内存使用量随着训练过程而变化。因此,模型存储在AMS上不断增加。因此,必须将ID的规模控制在一定水平。为了解决这个问题,XDL提供了几种特征入口过滤器机制,如基于概率的过滤器和计数布隆过滤器。在这种情况下,低频ID将在训练过程之前被丢弃,因为它们对模型的贡献较小。同时,在线模型的规模和内存存储得到很好的控制,这保证了在线系统的稳定性。

增量模型导出

随着持续的训练过程,当模型存储达到数百GB甚至更高时,全模型导出在时间和计算资源方面将变得极其昂贵。因此,XDL为在线训练和推理系统提供了增量模型导出。每个增量模型只需要少量的存储空间,并且非常容易部署在在线系统上。

特征过期

与特征入口过滤器具有相同目的,特征过期机制也控制模型的大小。在在线训练过程中,长时间未更新的特征将变得无用,可以被删除。XDL允许算法专家编写自己的特征过期函数,并定制特征过期计划。

用户界面

XDL为算法科学家提供了Python API,以便开发新的深度学习模型。这些API让用户专注于模型本身,而不必担心I/O效率、分布式策略、系统一致性和低级通信。在并行化模型上的编码和调试几乎与在串行模型上相同。通过隐藏并行化和优化的细节,用户可以非常容易地将新开发的模型部署到沙箱和生产环境中的大量GPU集群上。大多数情况下,科学家只需要修改DenseNet结构来测试新想法。用户可以通过运行用户定义的代码完全控制后端worker。修改数据流和SparseNet组件以尝试更激进的算法创新也非常方便。

5 XDL生态系统

除了XDL,阿里巴巴还开源了XDL算法解决方案和Blaze推理引擎,它们与XDL框架一起构成了一个生态系统。所有这些基于XDL的工具都可以在https://github.com/alibaba/x-deeplearning找到。

XDL算法解决方案包括几个在阿里巴巴广告业务中成功采用的有效模型,如深度兴趣网络(DIN)[32]、深度兴趣演化网络(DIEN)[33]、基于树的深度匹配(TDM)[35]等。

Blaze推理引擎是一个高性能的推理引擎,它利用内核融合和低精度推理等技术,为在线广告业务启用了大规模实时计算。

XDL生态系统为广告行业树立了示例和指南。越来越多的公司采用XDL作为他们的离线或在线深度学习工具。

6 评估

在本节中,我们提供了前一节描述的优化策略的详细评估。实验在通过25Gbps以太网互联的节点上进行。每个节点配备了2个Intel(R) Xeon(R) Platinum 8163 CPU @ 2.50GHz(96核)和512GB RAM。在编译时,我们使用GCC 4.8.5,带有-O2 -fopenmp -mavx2标志。

在我们的实践中,当选择后端时,MxNet或TensorFlow的性能几乎相同,因为它们在DenseNet计算方面具有相同的性能。在以下评估中,我们选择TensorFlow作为XDL的计算后端。

6.1 数据集

在我们的实验中,我们使用[25]发布的开放数据集,训练了一个用于点击率(CTR)预测问题的深度神经网络(DNN)模型。该数据集是从淘宝的电子商务推荐系统中以1%的比例抽样的。该数据集的统计信息显示在表3中。

开放数据集包括三个方面:用户兴趣、广告特征和上下文。它们中的每一个都由高维稀疏ID表示。使用XDL的Python接口,我们基于SparseNet + DenseNet抽象构建了模型

  • 在SparseNet中,我们使用embedding字典将每个ID转换为维度为18的密集向量。
  • 随后是一个分组sum pooling,将这些向量减少到23个向量。
  • 在DenseNet中,我们使用一个3层全连接DNN,包含约100K参数。(在这个DNN中,输入层有23*18个节点,第一层有200个节点,第二层有80个节点,输出层有1个节点。)

模型使用点击label来估计CTR。

6.2 子系统评估

为了了解XDL的性能,我们将逐一评估优化策略,看看我们能从每一个策略中获得多少好处。以下实验在上述开放数据集上使用相同的环境运行。所有实验在200个worker(除了可扩展性实验)上运行,使用80个CPU模式的AMS。我们为每个worker分配8个CPU核心,每个AMS将占用一台机器中的所有96个CPU核心。XDL在异步模式下运行,更新是无锁的,如第4.3节所述。

分层样本压缩:

I/O模块中的分层样本压缩对通信和计算都带来了巨大的加速。平均来说,每个用户ID在开放数据集中可以重复57次,这意味着那些用户共享的公共特征的压缩比可以达到57。为了评估使用样本压缩可以获得多少性能提升,我们重新抽样数据集以构建具有不同平均压缩比的新数据集。然后我们比较XDL在这些数据集上运行时的吞吐量。结果可以在图6中看到。

图片名称

图6 压缩比与吞吐量。运行在200个工作进程上,batch size为5000。

大batch size

使用大batch size将减少迭代次数,因此可以大幅降低数据复制和通信的开销。尽管Keskar等人[21]发现使用大batch训练可能导致尖锐的最小值并损害泛化。已有几种技术[13, 16]被提出来缓解这个问题。因此,我们倾向于使用更大的batch size来提高吞吐量。不同batch size下AMS的吞吐量显示在图7中。

图片名称

图7 批量大小与吞吐量。压缩比为57,使用200个工作进程。

可扩展性

我们还测试了XDL的可扩展性,并将其与不同数量的worker上的原生TensorFlow进行了比较,使用了80个AMS。同样,每个worker使用8个CPU核心,每个AMS使用96个CPU核心。XDL和TensorFlow都使用无锁风格的异步更新。结果显示在图8中。由于其对通信和重复特征的优化,XDL始终表现更好。当我们固定AMS的数量时,在异步模式下运行时,我们几乎实现了线性可扩展性。

图片名称

图8

7 结论

我们介绍了一个名为XDL的深度学习框架,它对于具有SparseNet + DenseNet抽象的高维稀疏数据集非常强大。凭借面向高维度的设计和系统优化,XDL在真实世界的数据集上可以比原生TensorFlow快得多。我们的算法科学家可以轻松地使用XDL灵活的API开发新模型,以支持阿里巴巴业务的快速发展。

随着XDL生态系统的开源进程,XDL可以在未来进一步发展。首先,像TVM[7]、深度梯度压缩[24]和混合精度训练[26]这样的高级加速技术正在测试中,并将很快部署。其次,作为参数存储和更新的开放解决方案,AMS将需要更灵活的数据流和模型结构,以满足最新模型如记忆网络和动态计算图的要求。

参考

介绍

baidu在2020《AIBox: CTR Prediction Model Training on a Single Node》提出了它们的AIBox平台。我们看下它的实现:

摘要

作为世界上主要的搜索引擎之一,百度的赞助搜索(Sponsored Search)自2013年起就采用了深度神经网络(DNN)模型来预测广告点击率(CTR)。百度在线广告系统(即“凤巢(Phoenix Nest)”)使用的输入特征极其高维(例如,成百上千的十亿个特征)且极其稀疏。百度生产系统中使用的CTR模型的大小可以轻松超过10TB。这给在生产环境中对这些模型进行训练、更新、使用带来了巨大的挑战。

对于百度的广告系统来说,保持模型训练过程的高效率显然是重要的,以便工程师(和研究人员)能够快速完善(refine)和测试他们的新模型或新功能。此外,由于每天有数十亿的用户广告点击历史记录条目到达,模型必须快速重新训练,因为CTR预测是一项极其时间敏感的任务。百度当前的CTR模型是在MPI(消息传递接口)集群上训练的,这需要高容错性和同步,从而产生昂贵的通信和计算成本。当然,集群的维护成本也是巨大的。

本文介绍了AIBox,这是一个集中式系统,通过使用固态硬盘(SSD)和GPU来训练具有数十TB规模参数的CTR模型。由于GPU上的内存限制,我们仔细地将CTR模型划分为两部分:一部分适合CPU,另一部分适合GPU。我们进一步引入了一个双层缓存管理系统,用于在SSD上存储10TB参数的同时提供低延迟访问。在生产数据上的广泛实验显示了新系统的有效性。AIBox具有与大型MPI集群相当的训练性能,同时只需要一小部分集群的成本。

1 引言

作为世界上领先的搜索引擎之一,百度的赞助搜索系统(即“凤巢”)[20] 自2013年起就采用了深度神经网络(DNN)模型来预测广告点击率(CTR)。CTR预测[19, 22]在决定最佳广告空间分配中起着关键作用,因为它直接影响用户体验和广告盈利能力。通常,CTR预测以多个资源作为输入,例如查询-广告相关性、广告特征和用户画像。然后,它估计用户点击给定广告的概率。最近,深度学习在计算机视觉[29]和自然语言处理[5]中取得了巨大成功。受此启发,提出了用于CTR预测任务的学习方法[13, 15, 20, 61, 65, 66]。与常用的逻辑回归[40]相比,深度学习模型可以显著提高准确性,但训练成本也显著增加。

在百度搜索广告的当前生产系统中,我们模型的训练过程既耗费资源又耗时。模型是在MPI(消息传递接口)[52]集群中的参数服务器[24, 32]上训练的,该集群拥有数百个CPU-only节点。生产中使用的主要模型大小超过10TB,存储/管理在特殊硬件上。参数服务器解决方案受到百级节点环境中节点故障和网络故障的影响。更糟糕的是,参数服务器中的同步阻塞了训练计算,并产生了大量的网络通信开销,而异步训练框架由于每个工作节点上的模型过时[48, 62],存在模型收敛问题

改进赞助搜索的生产系统存在许多不同的方向上的迷人机会和挑战。一个活跃的研究领域是在调用CTR模型之前提高“召回”(广告)的质量。例如,百度最近向社区分享了这样的技术论文[20],该论文建立在快速近邻搜索算法[55, 64]和最大内积搜索技术[51, 56]之上。

在本文中,我们介绍了百度为改进在线广告系统的另一项并行主要努力,即将CTR模型训练从MPI集群转移到GPU上。虽然使用GPU进行机器学习和科学计算已经成为一种常见做法,但使用GPU训练商业CTR模型仍然带来许多实质性挑战。最值得注意的挑战是训练数据规模达到PB(PeteByte)级别,训练出的模型大小超过10TB。训练样本的数量可以达到数千亿,特征的数量可以达到数千亿(我们通常使用2的64次方作为一个方便的特征空间大小的替代)。输入模型的数据也非常稀疏,每个特征向量仅有几百个非零项。

作为商业赞助搜索系统,任何模型压缩技术都不应妥协预测性能(收入)。实际上,即使是微小的(例如,0.1%)预测准确性下降也会导致不可接受的收入损失。事实上,整个系统已经高度优化,几乎没有冗余(例如,参数已经被仔细量化为整数)。看来改进的空间非常小。流行的模型压缩技术,如下采样[8, 30, 46]和哈希[12, 27, 33, 34, 57],对于训练数据极其高维(例如,数百亿个特征)和极其稀疏(例如,每个特征仅有几百个非零项)的商业CTR模型来说效果较差。研究论文中常见的论点,如“以仅损失0.3%的准确性为代价减少一半的训练成本”,在这个行业不再适用。另一方面,在百度训练DNN CTR模型是一项日常任务。工程师和数据科学家必须尝试许多不同的模型/特征/策略/参数,并且必须非常频繁地训练和重新训练CTR模型。硬件(如MPI集群)和能源消耗的成本可能非常高。

为了应对这些挑战,我们提出了AIBox,这是一个新颖的集中式系统,可以在单个节点上高效地训练这个庞大的机器学习模型。AIBox采用新兴硬件,SSD(固态硬盘)和GPU,存储大量参数并加速繁重的神经网络训练计算。作为一个集中式系统,AIBox直接消除了分布式系统中由网络通信引起的缺点。此外,与大型计算集群中的数千个节点相比,单个节点的AIBox硬件故障数量要少几个数量级[49]。由于只需要内存中的锁和GPU片上通信,单个节点中的同步成本也大大降低。与分布式环境相比,没有通过网络传输数据。

然而,在设计AIBox时存在两个主要挑战。

第一个挑战是:在单个节点上存储10TB规模的模型参数。当容量超过1TB时,内存价格飙升。随着模型未来变得更大,这在可扩展性上是不切实际的,并且在实际应用中大规模生产也不实用。由于成本高昂,我们不能将整个10TB参数存储在主内存中。新兴的PCIe总线上的非易失性内存高速(NVMe)SSD比硬盘驱动器的延迟低50倍以上[60]。我们利用SSD作为二级存储来保存参数。然而,SSD有两个缺点:

  • 首先,在延迟方面,SSD仍然比主内存慢两个数量级,导致训练过程中参数的访问和更新速度慢。
  • SSD的另一个缺点是SSD中的存储单元仅能持续几千次写入周期。

因此,我们必须维护一个有效的内存缓存,以隐藏SSD延迟并减少对SSD的磁盘写入。

第二个挑战是:在单个节点上使用多个GPU来加速训练计算。最近,配备32GB高带宽内存(HBM)的Nvidia Tesla V100[28]的单精度性能达到15.7 TFLOPS,比顶级服务器CPU节点(Intel Xeon系列)在深度学习推理上快47倍[38]。这为设计一个多GPU计算节点提供了独特的机会,该节点的计算性能与集群相当。然而,当前现成的GPU没有TB级HBM。我们不能将整个CTR预测神经网络保持在GPU HBM中。在这项工作中,我们提出了一种新颖的方法(第2节)将神经网络划分为两部分:

  • 第一部分是内存密集型的,并且在CPU上训练。
  • 网络的另一部分是计算密集型的,同时输入特征数量有限,我们在GPU上训练它。

训练数据和模型参数在主内存和多GPU的HBM之间传输。然而,主内存和GPU HBM之间的通信受到PCIe总线带宽的限制。当通信带宽成为瓶颈时,高GPU数值计算性能被阻塞。新兴的NVLink[21]和NVSwitch[44]技术使GPU之间的直接通信无需涉及PCIe总线。我们采用NVLink并设计了一个HBM内的参数服务器,以减少GPU数据传输。

总结来说,本工作的主要贡献包括:

  • 我们介绍了AIBox,这是一个由SSD和GPU加速的单节点系统,用于训练具有10TB参数的CTR预测模型。单节点设计范式消除了分布式系统的昂贵网络通信和同步成本。据我们所知,AIBox是第一个为如此大规模的实际机器学习应用设计的集中式系统。
  • 我们展示了一种新颖的方法,将大型CTR预测模型划分为两部分。划分后,我们能够在CPU上保留内存密集型训练部分,并利用内存受限的GPU来加速计算密集部分的训练
  • 我们提出了稀疏表(Sparse Table)来通过在SSD上存储模型参数并利用内存作为快速缓存来减少SSD I/O延迟。此外,我们实现了一个3阶段流水线,将网络、稀疏表和CPU-GPU学习阶段的执行重叠起来。
  • 我们通过将AIBox与由75个节点组成的分布式集群解决方案进行比较,在包含10PB样本的实际CTR预测数据上进行了广泛的实验评估。结果表明AIBox的有效性——AIBox具有与集群解决方案相当的训练性能,而其成本仅为我们为集群支付成本的不到1/10。

本文的其余部分组织如下。第2节介绍了我们的CTR预测神经网络模型。第3节,我们介绍了AIBox的高层设计。我们在第4节探讨了在稀疏表上大型模型存储的挑战,并提出了采用SSD上的双层缓存管理系统的解决方案。第5节展示了实验结果。第6节讨论了相关工作,第7节总结了本文。

2 CTR预测神经网络

工业深度网络被设计和训练以处理大规模数据样本,以帮助准确、快速、可靠地预测广告的点击率(CTR)。百度CTR预测模型中的特征通常是极其稀疏的特征(例如,成百上千的十亿个特征),每个向量中仅有极少数非零值(例如,几百个)。这个庞大的DNN模型在仅存储非零参数并进行仔细量化后,参数大小超过10TB。由于GPU的HBM容量有限,显然不切实际将整个模型的10TB参数保持在GPU的HBM中

在本文中,我们提出了一个双模块架构,用于在CPU+GPU上训练庞大的DNN CTR模型。第一个模块专注于高维&稀疏特征的嵌入学习,第二个模块用于与第一个模块产生的密集特征进行联合学习。嵌入学习在CPU上处理,以帮助学习低维密集嵌入表示。由于10TB参数的内存密集问题使得在训练期间无法将整个模型保持在内存中,我们利用SSD存储模型参数。参数可以从SSD快速访问到CPU。

通过将从CPU学习到的嵌入向量传输到GPU,计算密集型的联合学习模块可以充分利用强大的GPU进行CTR预测。在联合学习模块中,几个全连接神经网络被建模,以嵌入作为输入。这些神经网络的最后一层被连接在一起,用于最终的CTR预测。图1显示了设计的CTR神经网络模型的概览。我们在以下小节中介绍模型的细节。

图片名称

图1: 设计的CTR预测神经网络模型概览。嵌入学习输入层的节点:代表高维稀疏特征。联合学习合并层中没有入箭头的节点是:除了嵌入之外的密集个性化输入特征。

为了方便讨论,在本文的其余部分,我们将使用1012(千亿)作为CTR模型中极其稀疏特征的维度。读者应记住,实际生产系统中的特征数量可能远远超过1012

2.1 在CPU上的嵌入学习

嵌入学习模块旨在将高维稀疏向量(例如,1012维)映射到低维密集表示。如图1所示,嵌入学习模块包括高维稀疏特征的输入层和输出嵌入层。使用ReLU作为激活函数。这个模块主要是内存密集型的,因为1012特征导致10TB规模的模型参数,并且无法将所有参数加载到主内存中。为了学习嵌入,我们将10TB参数存储到SSD中。由于SSD和CPU之间的高效访问速度,我们可以轻松地从SSD加载参数并在CPU上学习embedding。

2.2 在GPU上的联合学习

在CPU上计算高维稀疏特征的嵌入后,我们将嵌入从CPU传输到GPU进行CTR预测过程。联合学习的输入包括:dense个性化特征和学习到的embedding。个性化特征通常来自多种来源,包括广告创意的文本、用户个性化行为和各种与广告相关的元数据。如果我们直接将这些特征进行concate并输入到神经网络中,可能无法充分探索个性化特征中的信息,导致CTR预测结果不准确。因此,我们设计了几个深度神经网络,并联合学习有意义的表示,以进行最终的CTR预测。如图1所示,联合学习模块包含几个(图中为两个)深度神经网络。每个网络将学习到的嵌入和一种个性化信息一起作为输入层。然后应用几个全连接层,以自动方式帮助捕获特征的交互。这些网络的最后一层隐藏层被组合用于softmax层和CTR预测的输出层。我们使用负对数似然作为目标函数:

L=1|S|(x,y)S(ylogp(x)+(1y)log(1p(x)))

这里:

  • S是训练集,
  • S是训练样本的大小。
  • x代表我们CTR预测模型的输入特征,
  • y{0,1}是标签,表示用户是否点击了广告。
  • p(x)是softmax层之后的输出,表示数据样本被点击的预测概率。

为了有效地从前一个神经网络中学习,表示从第一层和最后一层隐藏层提取出来,然后与当前神经网络的输入层连接,用于联合学习。具体来说,第一隐藏层代表低级特征学习,从输入层提取最相关的信息。最后一层隐藏层显示高级特征学习,并检测对最终CTR预测最抽象但最有帮助的信息。我们结合从前一个网络中最有意义的低级和最强大的高级信息,以获得更准确的CTR预测结果。

这个模块主要是计算密集型的,因为有多个神经网络。我们在GPU上训练这个模块以加速。完成一轮或几轮联合学习后,使用反向传播首先更新GPU上的神经网络参数和传输的嵌入。然后更新的嵌入被返回以更新CPU上的模型参数。通过这种方式,两个模块共同工作进行CTR模型训练。

3 AIBox系统概览

在本节中,我们介绍AIBox系统的概览,并从高层次描述其主要模块。图2描绘了AIBox的架构。它包含三个组件:CPU模块、稀疏表模块和GPU模块。

图片名称

图2 AIBox架构

CPU模块:协调和嵌入学习

CPU模块协调整个训练工作流程。首先,它通过高速网络从分布式文件系统(例如,HDFS [7])读取训练样本,并构建后续并行处理的mini-batches。我们称这些mini-batches为一个pass。然后,CPU模块通过与稀疏表模块交互来计算mini-batches的特征嵌入,以获取引用的参数。之后,嵌入被传输到GPU进行神经网络的联合学习部分。在GPU上完成联合学习后,GPU模块内联合神经网络的反向传播梯度返回到CPU模块。它通过稀疏表的push操作更新CPU侧的参数。除了主工作流程外,CPU模块还会定期将当前训练快照保存为检查点到分布式文件系统,以便作业失败恢复。

稀疏表:关键值存储系统

稀疏表是一个键值存储系统,存储SSD上的1012离散特征的值(例如,模型权重和点击信息)。这些特征消耗超过10TB的空间。特征到文件映射的键哈希索引位于内存中进行键查找。尽管SSD的延迟比硬盘低得多,但它们仍然无法与纳秒级内存延迟相媲美。为了解决这一关键问题,我们为SSD访问设计了一个显式的内存缓存策略以隐藏延迟。同时,缓存还充当缓冲区,避免了频繁写入SSD(第4节)。

GPU模块:联合学习

GPU模块接收从CPU模块传输的sparse embedding,并将它们存储在HBMs中。然后,embedding被送入dense联合学习网络。通过PCI-E总线传输的embedding和GPU核心上的联合学习计算是时序重叠的。我们为数据传输创建一个单独的CUDA流[43],另一个CUDA流用于学习计算。GPU HBMs充当片上参数服务器(on-chip parameter server)。dense神经网络的参数在所有GPU之间复制,并存储在片上GPU HBM中。对于每个mini-batches的pass,每个GPU根据其本地副本计算新参数。之后,AIBox中的所有GPU通过NVLink高速互连进行集体通信以同步参数。

3阶段流水线

训练工作流程主要涉及3个耗时任务:训练样本读取、sparse表操作和神经网络训练(嵌入学习+联合学习)。这3个任务对应于独立的硬件资源:网络、SSD和CPU+GPU。如图3所示,我们构建了一个3阶段流水线来隐藏这些任务的延迟,并为每个阶段维护一个预取队列。此外,预取还隐藏了阶段之间的通信时间。对于每个阶段,我们创建一个工作线程,从预取队列中取作业并提供给相应的资源。然后,处理结果被推入下一阶段的预取队列。当下一阶段的预取队列已满时,即下一阶段已经有太多未处理的作业时,每个阶段的线程将被暂停。预取队列的容量根据每个阶段的执行时间预设。对于图3中的RE阶段,我们没有任何执行依赖。因此,我们通过网络从分布式文件系统中提取样本,直到我们在队列中有足够的未处理样本为止。ST阶段仅依赖于RE——通过我们的缓存管理系统消除了对EL和JL的依赖。ST阶段从SSD中加载RE中引用的参数。由于我们的缓存内存足够大,可以缓存过去几次pass中引用的参数,我们不需要等待EL和JL的完成。我们只需要确保加载的参数在被EL + JL阶段的相应pass使用之前不能从缓存中移除。EL计算依赖于ST阶段,因为我们只能在将参数加载到内存后开始计算嵌入。而JL训练必须等待从EL部分计算出的嵌入。由于数据依赖性,JL和EL处于一个流水线中。配备多个高端GPU的AIBox具有与数十个仅CPU节点的分布式集群相当的计算能力。我们可以为神经网络模型的训练实现类似的执行时间。此外,RE是最简单的和最快的阶段。其延迟在流水线中被隐藏。然而,与所有参数都在内存中的分布式集群解决方案相比,与SSD交互的稀疏表操作纯粹是开销。我们需要积极优化ST,以确保ST比EL + JL更快。在这种情况下,ST的延迟在流水线中被隐藏,以便我们可以拥有相当的训练速度。因此,我们关注以下文章中稀疏表的设计细节。

图片名称

图3 设计的CTR预测神经网络模型包括一个3阶段流水线,该流水线对于网络、SSD和神经网络训练(CPU + GPU)上的执行操作是时序重叠的。

4 稀疏表

稀疏表旨在高效地在SSD上存储模型参数。它利用内存作为SSD的快速缓存,同时减少SSD I/O并提供低延迟的SSD访问。它由两个主要组件组成:键哈希索引和双层缓存管理。

4.1 键哈希索引

为了通过特征键访问SSD上的参数文件,我们必须为CTR预测模型中的1012个参数存储1012个键到文件的映射。在内存中将每个键到文件的映射存储为64位值对需要1.6TB = (8字节键 + 8字节SSD上的偏移量) × 1012,这超出了1TB的内存预算。我们必须仔细设计键哈希索引和SSD上的文件结构以减少内存占用。

我们引入了一个分组函数,将键映射到组ID,使得每个组包含m个键,即group(key) → {0, 1, ···, 1012/m1}。这里1012个键被划分为1012/m组。在对键进行分组后,我们能够保持组到文件的映射在内存中,因为内存消耗仅为原始键到文件映射的1/m。由于键从1到1012是连续的,分组函数可以通过均匀划分键空间轻松获得,例如,group(key) → key mod 1012/m。我们设置m=BLOCK/(8+sizeof(value)),其中BLOCK是SSD的I/O单元,它由SSD块大小(通常是4096)决定,8代表键占用的字节,sizeof(value)是值(模型参数)的大小,字节为单位,在我们的CTR预测模型中大约是50字节。m绝不能设置为小于BLOCK/(8+sizeof(value))的值,因为SSD访问必须从磁盘获取BLOCK字节。设置一个太小的m是次优的。另一方面,我们选择的m越大,键哈希索引的内存占用就越小。然而,一个较大的m导致一个较大的组,因为我们必须从SSD获取多个页面以获得一个组。因此,我们设置的m值在组到文件映射占用内存的可接受空间时是最优的。当块大小远大于值大小时,这是真的。

作为内存占用的权衡,分组策略的缺点是同一组中的值从磁盘获取,即使它们在当前mini-batch中未被引用——I/O带宽被浪费。一个可能的优化是将具有高共现特征的组在一起,例如,预训练一个学习哈希函数[26]以最大化特征共现。这属于垂直分区[42, 63]的另一个研究领域,超出了本文的范围。此外,这个缺点被缓存管理组件减少,我们跳过缓存键的组从磁盘读取。

4.2 双层缓存管理

缓存管理设计受到以下两个挑战的指导:访问性能和SSD的寿命。

首先,内存访问延迟是纳秒级,而SSD需要微秒级才能查看数据,因为SSD比内存慢大约1,000倍。然而,CTR中的参数如此稀疏和偏斜,以至于在mini-batch的一个pass中少于1%的参数被引用。这为我们提供了一个机会来构建一个内存中的缓存系统,保存有限内存预算中的频繁使用的“热参数”

第二个挑战是:SSD的物理属性只允许每个存储单元进行数千次写入。参数在训练的每次迭代中更新。如果参数立即更新,将显著缩短SSD的寿命。缓存管理还充当参数缓冲区。缓冲的参数在内存中更新,不涉及SSD I/O。当缓冲区达到其容量并且缓存替换策略将其交换出缓存时,它们才被懒洋洋地物化到SSD

我们提出了一个双层缓存管理来应对这些挑战(图4)。与通过链表链接处理冲突的经典哈希表相比,我们的系统在链接之前有一个额外的哈希级别。此外,我们引入了两个单独的链表来优化缓存探测性能。我们还为每个SSD文件附加了一个布隆过滤器,以减少不必要的SSD读取。一个后台工作线程运行以维护缓存策略,即执行缓存替换并将更新的参数写回SSD。

图片名称

图4 二级cache管理架构

第一层,图4中的Level 1,将键组ID哈希到缓存插槽sisi=hash1(gid),其中:gid是在键分组中计算的组ID(g_id = group(key))。每个si都有一个指向SSD中的文件的指针,该文件存储所有处理过的键的序列化模型参数,哈希1(g_id) = si。由于我们模型的极端稀疏性,许多参数在训练的早期阶段没有被触及。用默认值初始化所有参数并将其保存到SSD上是低效的。对于第一次引用的键,我们必须从SSD读取一个数据块,它只返回默认值。有1012个参数在我们的模型中,它可能产生1012个不必要的SSD读取。因此,我们执行延迟参数初始化,即我们直到它被mini-batch引用才物化一个参数。此外,我们为每个物化文件附加了一个布隆过滤器[9, 41]以过滤不必要的磁盘读取。布隆过滤器是一个低内存占用的概率数据结构,用于测试一个元素是否存在于一个集合中。它保证没有假阴性,因为当布隆过滤器返回假时,被测试的元素肯定不是集合的成员。此外,查询布隆过滤器是快速且内存带宽效率高的。只访问一个常数位数。

由于我们在布隆过滤器返回假时跳过探测SSD文件,布隆过滤器可能减少了SSD文件中潜在键缺失的数量。

第二层中的哈希码(图4中的Level 2)由hash2(g_id, bucket)计算,其中bucket是一个可调参数,控制每个缓存插槽si中的桶数。我们引入第二层哈希以缩短链表的链长。探测链表需要遍历链表直到我们找到探测键。在最坏情况下,它可能需要链表长度的步骤,而其他哈希表操作只需要恒定时间。因此,遍历链表是哈希表探测中最耗时的步骤。第二层哈希将si的长链链表分成多个桶,每个桶都有短链表以减少低效的迭代。参数bucket用空间换取高效探测。桶越大,我们期望中的链表就越短。

对于所有哈希到第二层同一个桶的键值对,我们构建两个单独的链表,最近最少使用[45](LRU)列表和最少频繁使用[53](LFU)列表,以在内存中缓存这些对(图4中的分离链表)。LRU列表存储当前pass的mini-batch中使用的参数。它将最近使用的键值对保持在链表的前面,从而减少遍历链表以定位数据的步骤。另一方面,LFU列表充当缓冲区,包含缓存替换的候选。最频繁使用的值位于链表的头部,因此可以高效访问。而最少频繁使用的值被冲入SSD并从缓存内存中移除。一个后台运行的工作线程定期将不在当前pass的mini-batch中的数据从LRU列表移动到LFU列表。链表导致许多小链表节点的分配和释放。我们维护一个内存池,采用Slab内存分配机制[6]来管理链表中节点的创建和删除。

4.3 稀疏表操作符

稀疏表的拉取(pull)操作符接收当前pass中引用的一组键作为输入,并将这些键的值加载到内存中。图5说明了如下工作流程:首先,我们通过键分组哈希函数计算组ID,并在双层缓存管理的第一层定位缓存插槽si。然后,我们查询每个定位缓存插槽的布隆过滤器,以测试这些引用的键可能存在(当键第一次被引用时不存在)。对于每个未通过布隆过滤器的键,我们在内存中创建一个带有默认值的链表节点,并将键插入布隆过滤器。对于通过布隆过滤器的键(它们可能存在于缓存或SSD文件中),我们通过它们的键和组ID在第二层定位相应的缓存桶bi,j,并遍历两个链表(LRU和LFU)以查找引用的键。如果在链表中找到键,则为缓存命中,因为我们不需要访问SSD。另一方面,对于未在链表中定位的键,我们必须读取SSD上的文件。以布隆过滤器理论上保证的低概率,引用的键在文件中未找到。我们为每个未找到的键创建一个默认节点,并像对未通过布隆过滤器的键一样更新相应的布隆过滤器。最后,我们将定位/创建的链表节点移动到LRU列表的头部。拉取操作后,所有引用键的参数都存储在内存缓存中,并且它们已准备好使用。

图片名称

图5

稀疏表的推送(push)操作符通过给定的一组键更新模型参数。更新通过学习率和反向传播梯度的组合值来增加引用参数。我们的工作流程(第3节)确保所有给定键的参数已通过相同键的拉取操作符加载到内存中。因此,我们可以通过双层缓存管理定位这些参数的值,并在不涉及写入SSD数据的情况下高效地增加它们。磁盘写入被推迟,直到更新的缓存节点从缓存内存中移除。当缓存内存预算已满时,即LFU列表的长度大于预设阈值MAX_LFU,我们从LFU列表中移除最后的FLUSH_STEP个节点,其中FLUSH_STEP是一个可调常数。由于节点按频率降序排序在LFU列表中,最后的FLUSH_STEP个节点的使用频率最低。被移除的参数必须被刷新到SSD,因为它们已被过去的推送操作更新。我们通过缓存管理组件的第一层定位包含被移除键的文件,从SSD读取它们,并将它们与更新的值合并到新文件中。

4.4 文件管理

推送稀疏表操作符不断在SSD上创建小文件,而现代操作系统上的本地文件系统,例如ext4[39]和XFS[54],并不是为处理许多小文件而设计的。大量的文件元数据压倒了这些文件系统。我们通过利用小文件创建总是成批进行的事实,为这项任务设计了一个特定的轻量级文件管理系统。与键分组类似,我们将F个小文件分组到一个大文件中以创建更少的文件。此外,当将一批独立的小文件磁盘写入分组为顺序写入时,它充分利用了磁盘带宽。对于每个小文件,其在大文件中的偏移量被保存。偏移量和大文件的名称的组合作为小文件的“文件名”。这种组合在缓存管理系统的第一层缓存插槽中维护。当我们将更新刷新到SSD时,它会更新。在我们创建越来越多的新文件后,一些旧文件没有被任何缓存插槽引用,就会过时。我们在内存中为每个大文件维护一个引用计数器。当一个“小文件”在其中更新并引用一个新文件时,计数器减少。我们一旦其计数器达到零就删除大文件。此外,我们定期监控磁盘使用情况。我们合并引用计数低的文件,并在磁盘使用量高于模型大小MAX_REPLICATION倍时删除它们。这里MAX_REPLICATION由SSD容量 ∗ (85%+超额配置)/模型大小计算,其中超额配置是制造商指定的SSD的保留存储空间。为了最大化SSD寿命,SSD控制器采用磨损均衡算法[11],这需要空闲空间来分散磁盘写入。当空闲空间少于85%时,SSD的性能会下降[2]。

5.实验

参考

Pinterest在《Graph Convolutional Neural Networks for Web-Scale Recommender Systems》中提出了PinSage:

介绍

深度学习方法已经在推荐系统应用中非常重要,被用于学习关于图片、文本、单个用户的有用高维embeddings。使用deep models学到的representations可以被用于补全、或者替换传统的推荐算法(比如:CF)。并且这些学到的representations具有高可用性,因为他们可以被用在许多推荐任务中。例如,使用一个deep model学到的item embeddings可以被用于item-item推荐,并且可以推荐有主题的集合(例如:playlists,或者“feed”内容)

最近,在该领域内有较大发展——特别是能够基于在图结构化数据上学习的新deep learning方法的发展,这对于推荐应用来说很基础(例如:可以利用user-to-item交互图,也可以使用社交网络图)。

在这些成就中,最著名的是Graph Convolutional Networks(GCNs)深度学习结构的成功。GCNs背后的核心思想是:学习如何从局部图邻居(local graph neighborhoods)的feature信息中进行迭代式聚合(aggregate)(图1)。这样的一个“卷积(convolution)”操作会从一个node的一跳图邻居的feature信息进行转换和聚合,并且通过将多个这样的convolutions信息进行stacking可以传播到图的远方。不同于纯content-based deep models(例如:RNNs),GCNs会同时利用内容信息和图结构。GCN-based方法已经在无数推荐系统benchmarks上设置了一个新的标准。然而,在benchmark任务上的收益还没有转换到真实生产环境中。

同时将GCN-based node embeddings的training和inference扩展到具有数十亿nodes和数百亿edges的图上是一个巨大的挑战。当在一个大数据环境中,对GCNs进行扩展很难,因为在设计底下的许多核心假设是冲突的。例如,在训练期间,所有已存在的GCN-based推荐系统需要在整个图拉普拉斯算子(full graph Laplacian)上操作,因而是不可行:当底层图(underlying graph)具有数十亿nodes,它们的结构通常是演化的。

当前工作

这里我们提出了一个高度可扩展的GCN框架,我们在Pinterest的生产环境中开发和部署过。我们的框架,是一个random-walk-based GCN 称为“PinSage”,它在30亿nodes和180亿edges的一个大型图(massive graph)上操作——该graph是GCNs的常见应用的10000倍以上大。

PinSage会利用许多关键insights来弹性提升GCNs的可扩展性:

  • on-the-fly convolutions:传统的GCN算法通过将feature matrics乘以full graph Laplacian的幂来执行图卷积。相反的,PinSage算法执行很高效,通过对在一个node周围的邻居进行抽样,并且从该抽样后的邻居来动态构建一个计算图。这些动态构建的计算图(图1)指定了如何在一个特定node周围执行一个局部卷积(localized convolution),并且缓和了在训练期间对整个graph操作的需求。
  • Producer-consumer minibatch construction:我们开发了一个producer-consumer架构来构建minibatches,它能确保在模型训练期间最大化GPU利用率。一个大内存、CPU-bound producer可以有效抽样node network邻居,并且获取必要。
  • 有效的MapReduce inference:给定一个fully-trained GCN模型,我们会设计一个有效的MapReduce pipeline,可以将训练好的模型分散来生成数十亿节点的embeddings,可以最小化重复计算。

图片名称

图1

除了在可扩展性上的这些基础优点外,我们也会引入新的训练技术和算法创新点。这些创新点会改善由PinSage学到的representations的质量,从而能导致在下游推荐系统任务上获得大的提升:

  • 通过random walks来构建convolutions:通过采用节点的所有邻居来执行convolutions(图1)会导致很大的计算图,因此,我们会采用sampling。然而,random sampling是次优的,我们会开发一种新的技术,它使用short random walks来采样计算图(computation graph)。一个额外的好处是:每个node具有一个importance score,我们会在pooling/aggregation step中使用。
  • Importance pooling:graph convolutions的一个核心组件是:在graph中的局部邻居的feature信息的aggregation。我们会引入一个方法来对在该aggregation中的node features的importance进行权衡,它基于random-walk相似度measures,在离线评估指标中会有一个46%的效果增益。
  • 课程培训(Curriculum training):我们设计了一个Curriculum training的scheme,其中该算法会在训练期间feed越来越hard的样本,从而产生一个12%的效果增益。

对于在Pinterest中的多个推荐任务,我们已经部署了PinSage,它是一个用户可以对交互的pins进行流行内容发现和管理的应用,其中:pins是在线内容的一个可视化书签。用户会将这些pins组织成boards,它包含了许多相似pins的collections。总之,Pinterest是世界上关于图片的最大user-curated graph(用户组织图),具有20亿唯一的pins,被收集到超过10亿的boards上。

通过大量离线metrics、A/B tests,我们展示了:对比起其它可扩展的deep content-based推荐算法,我们的方法在一个item-item推荐任务上(例如:相关pins推荐)以及一个”homefeed”推荐任务上可以达到SOTA的效果。在离线ranking指标上,我们对比baseline获得了超过40%的提升,在head-to-head人工评估上,我们的推荐要好60%。。

据我们所知,这是最大的deep graph embeddings,为新一代基于graph convolutional结构的推荐系统铺平了道路。

2.相关工作

3.方法

在本节中,我们会描述PinSage结构、training、以及MapReduce pipeline的技术细节,可以使用一个训练好的PinSage模型来有效生成embeddings。

我们的方法的关键计算重任是:局部图卷积(localized graph convolutions)。为了为一个node(例如:一个item)生成embedding,我们会应用多个convolutional模块,它们会聚合来自该node的局部图邻居(local graph neighborhood)(图1)的feature信息(例如:可视化features、文本features)。每个module会学习如何从一个小的图邻居(graph neighborhood)来聚合信息,通过将多个这样的模型进行stacking,我们的方法可以获得关于局部网络拓朴的信息。重要的是,这些localizied convolutional modules的参数会跨所有nodes进行共享,使得我们的方法的参数复杂度完全取决于input graph size。

3.1 问题设定

Pinterest是一个内容发现应用,其中:

  • 用户会将这些pins组织成boards,
  • 用户会将比较相关的pins包成集合

总之,Pinterest graph包含了20亿pins,10亿boards,并且超过180亿的edges(例如:pins的成员和它们相应的boards)。

我们的任务是,生成pins的high-quality embeddings 或者 representations(例如:通过最近邻查询相关的pin推荐,或者在下游的reranking系统中使用)。为了学习这些embeddings,我们会将Pinterest环境建模成一个二部图,它包含了两个不相互交集合I(包含了pins)和C(包含了boards)的结点。注意,我们的方法也是天然泛化的,其中I被看成是一个items集合,C被看成是user-defined contexts和collections。

除了图结构外,我们也会假设:pins/items uI会与real-valued属性相关联,xuRd。总之,这些属性可以指定关于一个item的metadata或content信息,另外在Pinterest的case中,我们有:具有丰富文本信息和图片features有关的pins。我们的目标是,利用这些input属性,以及二部图结构来生成高质量embeddings。这些embeddings接着会被用于通过最近邻lookup的推荐系统的候选生成(例如:给定一个pin,发现相关pins)或者在对候选排序时作为features使用。

出于便利性和泛化性,当我们描述PinSage算法时,我们会简单地将完整graph的node set使用 V=IC来表示,不会显式对pin和board nodes进行显式区别(除非严格必要),会统一使用更通用的术语“node”.

3.2 模型结构

我们会使用localized convolutinal模块来生成nodes的embeddings。我们从input node features开始,接着开始学习nueral networks,它会将通过graph来对features进行transform和aggregate来计算node embeddings(图1)。

前向传播算法

我们会考虑生成一个embedding的任务, 对于一个node u的zu,它取决于node的input features和围绕该node的图结构。

图片名称

算法1

PinSage算法的核心是一个localized convolution操作,其中,我们会学习如何从u的邻居(图1)的信息聚合。该过程在算法1 CONVOLVE中有详解。基本思想是:我们将关于u的邻居通过一个dense neural network进行转换成representations zvvN(u),接着在vectors的结果集(第一行)上应用一个aggregator/pooling function(例如:一个element-wise mean或weighted sum,表示为γ)。该aggregation step会提供一个关于u的局部邻居N(u)的vector representation nu。我们接着将aggretated neighborhood vector nu与u的当前representation hu进行拼接,并将contantenated vector nu与u的当前representation hu进行concatenate在一起,并通过另一个dense neural network layer(第2行)将concatenated vector进行转换。经验上,我们观察到:当使用concatenation operation来替代average operation时会取得极大的效果提升。另外,在第3行中的normalization会使得训练更稳定,它对于normalized embeddings来说执行近似最近邻搜索更高效(第3.5节)。该算法的output是一个关于u的representation,同时包含了它自己以及它的局部图邻居间的信息。

Importance-based neighborhoods

在我们的方案中,一个重要概念是:如何定义节点邻居N(u),例如:如何在算法1中选择要convolve的邻居集合。然而,之前的GCN方法会简单检查k-hop的图邻居,在PinSage中,我们定义了importance-based 邻居,其中:一个node u的邻居会被定义成T nodes会对node u施加最可能的影响。具体的,我们会从node u开始模拟random walks,并且计算出由random walk访问到的L1-normalized的节点访问数。u的邻居接着被定义成:对于node u具有最高的normalized visit counts的T个nodes。

该importnace-based邻居定义的优点是两面的。

  • 首先,选择一个固定数目的节点来聚合,允许我们控制着在训练期间该算法的memory footprint。
  • 第二,当聚合邻居的vector representations时,它会使算法1考虑邻居的importance

特别的,我们会在算法1中实现γ作为一个weighted-mean,它使用根据L1 normalized visit counts定义的weights。我们将该方法称为importance pooling

Stacking convolutions

每个时间,我们会使用CONVOLVE操作(算法1),我们会从一个node处获得一个新的representation,并且将多个这样的convolutions相互进行stack,以便获得围绕node u的局部图结构的更多信息。特别的,我们会使用多个关于convolutions的layers,其中:在layer k上的convolutions的inputs依赖于来自layer k-1(图1)的representations output,其中:intial(例如:“layer 0”)representations等于input node features。注意,在算法1(Q,q,W,w)中的模型参数会跨nodes共享,但在layers间的参数则不同。

算法2会详细说明:stacked convolutions 是如何为一个nodes的minibatch set M生成embeddings。我们首先计算每个node的邻居,接着使用K个convolutional迭代来生成关于target nodes的layer-K的representations。final convolutional layer的output接着通过一个fully-connected neural network进行feed来生成最终的output embeddings zu,uM

我们接着学习模型的full set参数:对于每个convolutional layer (Qk,q(k),W(k),w(k),k{1,,K})的weight和bias参数,以及最终dense neural network layer G1,G2,g的参数。在算法1中的Line 1的输出唯度(例如:Q的column-space维度)会被设置为,在所有layers上均是m。出于简洁性,我们将所有convolutional layers的输出维度设置为相当(算法1第3行),接着我们将该size参数通过d进行表示。该模型的最终output维度(算法2)也被设置为d。

算法2

3.3 模型训练

我们会以一个监督学习的方式,使用max-margin ranking loss来训练PinSage。在该setup中,我们假设:我们已经访问了一个关于items L的labeled pairs集合,其中:在该set中的pairs (q,i)L会被假设是相关的——例如:我们假设:如果 (q,i)L,那么item i是对于query item q的一个好推荐侯选。training阶段的目标是,最优化PinSage参数,以便在labeled set中的pairs (q,i)L的output embedding 是非常接近的。

我们首先详细描述了margin-based loss function。根据此,我们给出了一个关于我们开发的多个技术的总览,它会导致PinSage的计算效率以及快速收敛,允许我们可以在数十亿节点的graphs和数十亿训练样本上进行训练。最终,我们描述了我们的curriculum-training scheme,它会提升推荐的总体质量。

Loss function

为了训练模型的参数,我们会使用一个Max-margin-based loss function。基本思想是:我们希望最大化正样本的inner product,例如:query item的embedding以及相应的相关item。同时,我们希望确认,负样本的inner product(例如:在query item的embedding和一个unrelated item间的inner product)要小于正样本由一些pre-defined margin给出的正样本。对于 (zq,zi)L的node embeddings的单个pair的loss function:

JG(zqzi)=EnkPn(q)max{0,zqznkzqzi+Δ}

…(1)

其中:

  • Pn(q)表示了对于item q的负样本的分布
  • Δ表示了margin hyper-parameter

我们会在下面解释负样本的采样。

使用large minibatches的Multi-GPU训练

为了在单机训练上充分利用多个GPUs,我们会以一个multi-tower的形式来运行forward和backward propagation。有了多个GPUs后,我们首先会将每个minibatch(图1底部)划分成等size的比例。每个GPU会获取minibatch的一部分,并使用参数的相同集合进行计算。在backward propagation后,在所有GPUs上每个参数的gradients会聚合到一起,会执行单频synchronous SGD。由于需要在非常大数目的样本上(数十亿规模)训练,我们会使用large batch sizes进行运营,范围从512到4096.

我们使用由Goyal[16]提出的相似技术来确保快速收敛,当处理大的batch sizes时,并保持训练和泛化的准确率。我们使用一个渐近的热启过程(gradual warmup produre),在第一个epoch,根据线性scaling rule从小到大增加learning rate。之后,learning rate会指数递减。

Producer-consumer minibatch构建

在训练期间,由于size很大,对于数十亿nodes的feature matrix和adjacency list会在CPU memory中放置。然而,在PinSage的CONVOLVE step,每个GPU过程需要访问在邻居中节点的邻居和feature信息。从CPU内容中访问来自GPU的数据并不高效。为了解决该问题,我使用一个re-indexing技术来创建一个sub-graph G=(V,E)包含了nodes和它的邻居,它会在当前minibatch的计算中被涉及。一个小的feature matrix会只包含了与当前minibatch的计算相关的node features,会被抽取出来,以便顺序与在G’中的nodes的index相一致。G’的adjacency list和small feature matrix会在每个minibatch迭代中被feed到GPUs中,因此,在GPU和CPU间的通信在CONVOLVE step期间不需要,极大提升了GPU利用率。

训练过程会交替使用CPUs和GPUs。模型会在GPUs中计算,然而:抽取features、re-indexing、negative sampling会在CPUs中计算。除了使用multi-tower training的并行GPU计算之外,CPU计算会使用OpenMP,我们会设计一个 producer-consumer pattern来在当前迭代中运行GPU计算,在下一迭代上并行使用CPU计算。这会进一步减少一半的训练时间。

对negative items进行采样

negative sampling会在我们的loss function中会被使用,作为edge likelihood的normalization factor的一个近似。当使用大的batch sizes训练时会提升效率,我们抽样一个关于500 negative items的集合,会在每个minibatch中与所有训练样本进行共享。对比起对于每个node独立运行负样本来说,这会极大节约需要在每个training step期间被计算的embeddings的数目。经验上,我们不会观察到在两个sampling schemes间效果的一个不同之处。

在最简单的case中,我们会从items的整个集合中均匀抽样负样本。然而,确保正样本(items(q,i)的pair)的内积大于q,500 negative items的每个都太“easy”,对于要学习的系统来说,不会提供足够好的“resolution”。特别的,我们的推荐算法可以发现在20亿items的catalog间与q 最相关的1000个relevant items. 换句话说,我们的模型会在20亿items上能够区分/标识 1个item。但在500个random negaive items,模型的resolution只有1/500. 因此,如果我们从20亿items中抽样出500个随机负样本items,任何这些与query items更轻微相关的items的机会更小。因此,该学习大概率不会做出好的参数更新,不能区分非常相关items和轻微相关items。

为了解决上述问题,对于每个正的训练样本(例如:item pair (q,i)),我们会增加“hard” negative examples,例如:items一定程度上与query item q相关,但与正样本item i不相关。我们称这些为“hard negative items”。他们会根据query item q的Personlized PageRank score通过在graph中的items进行ranking来生成。在2000-5000上排序的items会被随机抽样作为hard negative items。如图2所示,hard negative examples对比起random negative examples与query更相似,对于模型排序来说挑战更大,强制模型去学习以细粒度的方式区分items。

图片名称

图2

通过training produre使用hard negative items,是需要训练至收敛的epochs数目的两倍。为了帮助收敛,我们开发了一个curriculum training scheme【4】。在训练的第一个epoch,不会使用hard negative items,以便算法快速发现在参数空间中的一个区域,其中:loss是相对小的。我们接着在后续epochs中添加hard negative items,聚焦于模型去学习如何区分高度相关pins和轻微相关pins。在训练的epoch n,我们会为每个item添加n-1个hard negative items到negative items集合。

3.4 通过MapReduce的Node embeddings

在模型被训练之后,直接使用训练好的模型来为所有items(包含了在训练期间没有见过的items)来生成embeddings仍然是很具挑战的。使用算法2以naive的方式计算nodes的embeddings,会导致重复计算,这是由nodes的K-hop邻居间的重合引起的。如图1所示,当为不同的target nodes生成embeddings时,许多nodes会在多个layers上重复计算。为了确保有效的inference,我们开发一个MapReduce方法,它无需重复计算即可运行model inference。

图片名称

图3

我们观察到:node embeddings的inference可以非常好地借助于MapReduce的计算模型。图3详述了在二部图上的data flow,pin-to-board Pinterest graph,其中,我们假设:input(例如:“layer-0”) nodes是pins/items(以及layer-1 nodes是boards/contexts)。MapReduce pipeline具有两个关键部分:

  • (1) 一个MapReduce 工作可以用于将所有pins投影到一个低维latent space中,其中:aggregation oepration会被执行(算法1,第一行)
  • (2) 另一个MapReduce job接着用于将产生的pin representations,与他们出现的boards的ids进行联合,board embedding的计算会通过它的邻居的features进行pooling。

注意:我们的方法会避免冗余计算,每个node的latent vector只会计算一次。在boards的embedding会被包含后,我们会使用两个多的MapReduce jobs来计算关于pins的第二层的embeddings,以一个相似的方式,该过程可以尽快迭代(直到K个convolutional layers)。

3.5 高效地最近邻lookups

由PinSage生成的embeddings可以被用于许多下游推荐任务中,在许多settings中,我们可以直接使用在学到的embedding空间上通过最近邻查询这些embeddings来做出推荐。也就是说,给定一个query item q,我们可以推荐那些关于query item的embedding的embedding是K个最近邻。ANN可以通过locality sensitive hashing有效获得。在hash function被计算后,items的检索可以使用一个two-level检索过程,基于Weak AND操作符来实现。假设PingSage model通过离线训练给出,所有node embeddings会通过MapReduce进行计算,并在database中保存,有效的最近邻lookup operation可以确保系统以一个在线的方式提供服务

standford在《Inductive Representation Learning on Large Graphs》中提出了GraphSage:

介绍

在large graphs中的节点(nodes)的低维vector embedding,已经被证明对于许多预估和图分析任务来说作为feature inputs是很有用的。在node embedding方法的基本思想是:使用降维技术将关于一个node的图邻居的高维信息蒸馏(distill)到一个dense vector embedding中。这些node embeddings可以接着被feed到下游的机器学习系统中,并用于分类、聚类、连接预测等任务中。

然而,之前的工作主要关注于来自单个fixed graph的embedding nodes,许多真实应用,会对于未知nodes、或者全新的subgraphs也能快速生成embeddings。这些归纳能力对于高吞吐、生产环境机器学习系统来说很重要,它会在演进的图上操作、并能总是遇到未知的nodes(例如:在Reddit上的posts、在Youtube上的users和videos)。对于生成node embeddings的归纳法来说,会面临着在具有相同形式features的各种图上的泛化(generalization):例如:来自一个模式生物的在蛋白质的相互作用图上,训练一个embedding genreator,接着可以很容易使用训练好的模型,来对于新的生物体(organisms)收集来的数据生成node embeddings。

对比起直推式setting(transductive setting),inductive node embedding问题是特别难的,因为泛化到unseen nodes需要将新观察到的subgraphs“安排(aligning)”到已经通过算法最优化好的node embeddings中。一个inductive framework必须学习去认识一个node的邻居的结构化属性,它能表明节点的在图中的局部角色(local role),以及全局位置(global position)

对于生成node embeddings的大多数已经存在的方法,是天然直推式的(transductive)。绝大多数方法是直接使用MF-based目标来最优化每个节点的embeddings,不会天然生成unseen data,因为他们会在一个单一fixed graph上做出预估。这些方法可以在一个inductive setting环境中被修改来执行,但这些修改版本的计算开销都很大,需要在做出新预估之前额外迭代好几轮gradient descent。一些最近的方法在图结构上使用卷积操作(convolutional operators)来进行学习,能提供一个embedding方法。因此,GCNs(graph convolutional networks)已经被成功应用到在fixed graph的直推式setting(transductive setting)上。而在本工作中,我们同时将GCNs扩展到归纳式无监督学习(inductive unsupervised learning)任务上,并提出一个framework来生成GCN方法,它使用trainable aggregation function(而不是简单的convolutions).

我们提出了一个关于inductive node embedding的general framework,称为GraphSage(抽样和聚合:SAmple and aggreGatE)。不同于基于MF的embedding方法,我们会利用node features(例如:文本属性、node profile信息、node degrees)来学习一个embedding function,它会泛化到unseen nodes上。通过在学习算法中包含node features,我们会同时学习每个node邻居的拓朴结构,以及在邻居上的node features分布。当我们关注feature-rich graphs(例如:具有文本属性的引文数据、功能/分子标记的生物学数据),我们的方法也会利用出现在所有graphs(例如:node degrees)中的结构化features,因而,我们的算法也会被应用于没有node features的graphs中。

不同于为每个node训练一个不同的embedding vector的方式,我们的方法会训练一个关于aggregator functions的集合,它们会从一个node的局部邻居(local neighborhood)(图1)中学习到聚合特征信息(aggregates feature information)。每个aggregator function会从一个远离一个给定结点的不同跳数、搜索深度的信息进行聚合。在测试(test)或推断(inference time)时,我们使用已训练系统来为整个unseen nodes通过使用学到的aggregation functions来生成embeddings。根据之前在node embeddings生成方面的工作,我们设计了一个无监督loss function,它允许GraphSage使用task-specific supervision来进行训练。我们也表明了:GraphSage可以以一个完全监督的方式进行训练。

图片名称

图1 GraphSAGE抽样和聚合方法的可视化演示

我们会在三个node分类benchmarks上评估我们的算法,它们会测试GraphSAGE的能力来在unseen data上生成有用的embeddings。我们会使用两个演进的document graphs,它们基于citation data和Reddit post data(分别预估paper和post类目),以及基于一个一个蛋白质的相互作用的multi-graph生成实验。使用这些benchmarks,我们展示了我们的方法能够有效生成unseen nodes的表示,效果对比baseline有一个大的提升。。。

2.相关工作

3.GraphSAGE

我们方法的关键思想是,我们会学习如何从一个节点的局部节点(local neighborhood)中聚合特征信息(例如:邻近节点的degrees和文本属性)。我们首先描述了GraphSAGE的embedding生成算法(例如:forward propagation),它会假设:为GraphSAGE模型参数已经被学习的节点生成embeddings。我们接着描述:GraphSAGE模型参数可以使用标准的SGD和BP技术被学到。

3.1 Embedding生成算法(例如:forward propagation)

在本节中,我们描述了embedding生成,或者forward propagation算法(算法1),它会假设:模型已经被训练过,并且参数是固定的。特别的,我们会假设:我们已经学到了关于K个aggregator functions的参数(表示为:AGGREGATEk,k{1,,K}),它会被用于在模型的不同layers间、或者“搜索深度”上传播信息。第3.2节描述了我们是如何训练这些参数的。

  • G(V,E):图
  • {xv,vV}:输入特征
  • K:深度
  • Wk,k{1,,K}: 权重矩阵
  • σ:非线性激活
  • AGGREGATEk,k{1,,K}:可微聚合函数
  • N:v2v:邻居函数(neighborhood function)

图片名称

算法1

算法1的背后意图是:

  • 在每个迭代中,或者搜索深度上,节点会聚合来自它们的local neighbors的信息;
  • 并且随着该过程迭代,nodes随着graph的更进一步触达逐渐获得越来越多的信息。

算法1描述了case中的embedding生成过程,其中:

  • G=(V,ϵ):表示整个graph
  • {xv,vV}:表示graph中的所有node的input features
  • K:表示深度
  • Wk,{1,,K}:表示weight matrics

我们在下面描述了如何生成该过程到minibatch setting中。算法1的外循环中的每个step过程如下,其中:

  • k:表示在外循环(或者搜索深度)中的当前step
  • hk:表示在该step中一个node的representation

首先,每个node vV会将在它的立即邻居(immediate neighborhood){hk1u,uN(v)}上的nodes 的representations聚合到单个vector hk1N(v)。注意,该聚合step依赖于:在outer loop的前一迭代生成的representations,并且k=0(”bad case”)representations会被定义成input node features。

接着,在聚合了邻近的feature vectors之后,GraphSAGE会将该node的当前representation hk1v与聚合的邻近vector hk1N(v)进行拼接(concatenates),该concatenated vector会通过一个具有非线性activation function σ的fully connected layer进行feed,它会将representations转换成在算法的下一step进行使用(例如:hkv,vV)。neighbor representations的聚合可以通过多个aggregator结构来完成,并且我们会在第3.3节中讨论不同的结构选择。

为了扩展算法1到minibatch setting中,给定一个关于input nodes的集合,我们首先对所需要的neighborhood sets(到深度K)进行forward sample,接着我们在inner loop进行运行,通过替代迭代所有nodes,我们只会计算:必须满足在每个depth上的递归的representations。

与 Weisfeiler-Lehman Isomorphism Test的关系

GraphSAGE算法在概念上受testing graph isomorphism的经典算法的启发。在算法1中,我们:

  • (i) 设置K=V
  • (ii) 设置weight矩阵作为identity
  • (iii) 使用一个合适的hash function作为一个aggregator(无非线性),

接着算法1是一个关于WL isomorphism test的实例,也被称为“naive vertex refinement”。如果由算法1输出的representations {zv,vV}。该test在一些情况下会失败,但是对于许多图是合法的。GraphSAGE是对WL test的连续近似,其中,我们将hash function替代成trainable neural network aggregators。当然,我们使用GraphSAGE来生成有用的节点表示(node representations)——而非test graph isomorphism。然而,在GraphSAGE和经典的WL test间的连接,为我们的算法设计提供了理论context来学习关于节点邻居的拓朴结构。

Neighborhood定义

在本工作中,我们会均匀地抽样一个关于节点邻居的fixed-size集合,而非使用算法1中完整的邻居集合,是便保持每个batch的计算开销是固定的。也就是说,使用过载的概念,在算法1中,我们将N(v)定义为一个从集合{uV:(u,v)ϵ}中fixed-size的均匀抽取,在每个迭代k上我们会从不同的均匀样本中进行抽样。如果没有这种抽样,单个batch的内存和runtime会不可预知,最坏情况下为O(V)。作为对比,对于GraphSAGE的每个batch space和时间复杂度被固定在O(Ki=1Si),其中Si,i{},K是user-specified常数。实际来说,我们发现我们的方法可以达到具有K=2的高效果,并且S1S2500

3.2 学习GraphSAGE的参数

为了以一个完全无监督setting方式学习有用的、可预测的表示(representations),我们会通过SGD使用一个graph-based loss function给output representations zu,uV,并且调节weight矩阵 Wk,k{1,,K}, 以及得到的aggregator functions的参数。graph-based loss function会鼓励邻近节点具有相似的表示(representations),从而强制分离的nodes的representations具有高度的不同:

JG(zu)=log(σ(zTuzv))QEvnPn(v)log(σ(zTuzvn))

…(1)

其中:

  • v:是一个node,它会与u在一个fixed-length的random walk中共现
  • σ:是sigmoid function
  • Pn:是一个negative sampling分布
  • Q:定义了negative samples的数目

重要的是,不同于之前的embedding方法为每个node(通过一个embedding look-up)训练一个唯一的embedding,我们feed到该loss function的representations zu会由在一个node的局部邻居(local neighborhood)中包含的features来生成

该无监督setting会模拟以下情况:其中node features会提供到下游的机器学习应用,或者作为一个服务 或者 在一个静态仓库中。在本case中,其中representations只会在一个指定的下游任务中,无监督loss(等式1)可以简单地被一个task-specific目标(比如:cross-entropy loss)替换。

3.3 Aggregator结构

在N-D lattices(例如:句子、图片、或3D空间),一个node的邻居没有自然序:在算法1的aggregator function必须在一个关于vectors的无序集合上操作。理由的,一个aggregator function必须是对称的(例如:它的inputs排列是不变的),同时是可训练的,并且保持一个高度表达的能力。aggregation function的对称性确保了我们的neural network model是可训练的,并且可以被用于随意顺序的节点邻居的feature sets上。我们会检查三个候选aggregator function:

Mean aggregator

我们的第一个候选 aggregator function是mean aggregator,其中:我们简单地对在{hk1u,uN(v)}中的vectors做elementwise平均。该mean aggregator接近等于在transductive GCN network中的convolutional propagation rule。特别的,我们可以通过在算法1中的第4和第5行使用下面进行替换,派生一个关于GCN方法的inductive variant:

hkvσ(WMEAN({hk1v}{hk1u,uN(v)})

…(2)

我们将该修改版称为mean-based aggregator convolutional,因为它是一个关于一个localized spectral convolution的不平滑的线性近似。在该convolutional aggregator和其它aggregators间的一个重要区别是:在算法1的第5行,不会执行concatenation操作。例如:convolutional aggregator不会将node的前一layer representation hk1v与he aggregated neighborhood vector hkN(v)进行concatenate。该concatenate可以被看成是关于一个关于在不同“search depths”或GraphSAGE算法中“layers”间的”skip connection”的简单形式,并且它会产生在效果上的巨大收益。

LSTM aggregator

我们也会检查一个更复杂的基于一个LSTM结构的aggregator。对比起mean aggregator,LSTMs的优点是:具有更强的表达能力。然而,需要注意的是,LSTMs没有天然的对称性(例如:它们没有排列不变性),因为他们以序列方式处理它们的inputs。我们采用LSTMs用来操作一个无序集合,通过简单地将LSTMs应用到一个关于节点邻居的随机排列上。

Pooling aggregator

该aggregator具有对称性,同时可训练。在pooling方法中,每个邻居的vector会独立地通过一个fully-connected neural network进行feed;根据该转换,一个elementwise max-pooling操作会被应用来对跨邻居集合的信息进行聚合:

AGGREGATEpoolk=max({σ(Wpoolhkui+b),uiN(v)})

…(3)

其中:

  • max表示element-wise max操作符
  • σ是一个非线性activation function

原则上,该函数会在max pooling之前被使用,可以作为一个随意的deep multi-layer perceptron,但我们关注于简单的single layer结构。该方法受最近研究【29】的一些启发。直觉上,multi-layer perceptron可以被认为是:作为在邻居集合中的每个节点表示的feature计算的一个函数集合,该模型可以有效地捕获邻居集合的不同方面。注意,原则上,任务对称vector function可以被用来替代max operator(例如:一个element-wise mean)。我们发现:关于developments test在max-pooling和mean-pooling间没有大差别,因此后续主要关注max-pooling。

实验

facebook在《MEMORY NETWORKS》中提出了memory networks,在另一文《End-To-End Memory Networks》提出了end2end的实现:

摘要

我们会介绍一种使用recurrent attention model的neural network,它构建在一个可能很大的外部memory之上。该架构与Memory network的形式一样,但与其中的model不同的是:它的训练是end-to-end的,因而在训练期间几乎是无监督学习,这使得它很容易应用到实际环境中。它可以看成是关于RNNsearch[2]的一个扩展:其中在执行每次output symbol时会执行多个计算steps(hops)。该模型的灵活性允许我们将它应用到像QA等不同的任务以及语言建模上。对比起之前的Memory Networks的研究,它几乎是无监督的。我们在Penn TreeBank和Text8 datasets上验证了我们的方法要好于RNNs和LSTMs。

1.介绍

在AI 研究领域,在QA或补全任务上,构建模型会涉及多个计算步骤,模型可以描述序列数据中的long-term dependencies。

最近的一些工作,在建模时会使用显式存储和attention;维持这样的一个storeage对于解决这样的挑战提供了一种方法。在[23]中,存储会通过一个continuous representation给出;从该存储进行读取和写入,可以通过neural networks的actions进行建模。

在本工作中,我们提出了一种新的RNN结构,其中:在输出一个symbol之前,recurrence会从一个可能很大的external memory中读取多次。我们的模型可以被看成是在[23]中的Memory Network的一个continuous form。本工作中的模型通过BP训练并不容易,需要在该network的每个layer上进行监督。模型的连续性意味着:它可以从input-output pairs通过end-to-end的方式进行训练,因此很容易应用到许多任务上,例如:语言建模或者真实的QA任务。我们的模型可以看成是RNNsearch[2]的一个版本,它在每个output symbol上具有多个计算步骤。我们会通过实验展示:在long-term memory上的多跳对于我们的模型的效果来说很重要,训练memory representation可以以可扩展的方式进行集成到end-to-end neural network model上。

2.方法

我们的模型会采用:

  • 一个关于inputs x1,,xn的离散集合,它们会被存储到memory中
  • 一个query q
  • 输出一个answer a

xi,q,a的每一个都包含了来自具有V个词的字典的symbols。该模型会将所有x写到memory中,直到达到一个确定的buffer size,接着我们会为x和q寻找一个连续的representation。这会允许在训练期间,error signal的BP通过多个memory accesses回到input。

2.1 Single Layer

我们先在single layer的case中开始描述我们的模型,它会实现一个单个memory hop操作。后续我们会展示可以对它进行stack来给出在memory上的多跳。

Input memory representation

假设我们给定一个input set x1,,xi被存到memory中。{xi}的整个集合会被转到d维memory vectors {mi}中,它通过在一个连续空间上嵌入每个xi计算得到,在最简单的case中,使用一个embedding matrix A(其中:size=d×V)。query q也会被嵌入来获得一个internal state u。在embedding space中,我们会计算在u和每个memory mi间的match程度,通过以下公式对内积采用softmax得到:

pi=Softmax(uTmi)

…(1)

其中,Softmax(zi)=ezijezj。在该方式中,p是在inputs上的一个概率向量(probability vector)。

Output memory representation

每个xi都具有一个相应的output vector ci。memory o的response vector是一个在transformed input ci与来自input的probability vector进行加权求和:

o=ipici

…(2)

由于从input到output的函数是smooth的,我们可以轻易地计算gradients以及BP。其它最近提出的memory或attention形式也采用该方法【2】【8】【9】。

生成最终的prediction

在single layer case中,output vector o和input embedding u接着通过一个最终的weight matrix W(size为V×d)和一个softmax来生成predicted label:

ˆa=Softmax(W(o+u))

…(3)

整体模型如图1(a)所示。在训练期间,所有三个embedding matrics A, B, C,以及W都通过最小化ˆa和 true label a间的一个标准的cross-entropy loss进行联合学习。训练会使用SGD进行执行。

图片名称

图1

2.2 Multiple Layers

我们现在扩展我们的模型来处理K跳的操作。memory layers会以如下方式进行stack:

  • 在第一个之上的layers的input,是从layer k的output ok和input uk的求和(后续会有不同组合):
uk+1=uk+ok

…(4)

  • 每个layer会具有它自己的embedding matrics Ak,Ck,用于嵌入到inputs {xi}中。然而,如下所示,他们会被限制以便减轻训练、减少参数数目.

  • 在network的顶层,W的input也会将top memory layer的input和output进行组合:

ˆa=Softmax(WuK+1)=Softmax(W(oK+uK))

我们探索了在模型中两种类型的weight tying机制:

  • 1.Adjacent:一个layer的output embedding是下一个的input embedding,例如:Ak+1=Ck。我们也会限制:a) answer prediction matrix会与最终的output embedding相似,例如:WT=CK, b) question embedding会与第一层的input embedding相匹配,例如:B=A1
  • 2.Layer-wise(RNN-like):input和output embeddings对于不同的layers是相同的,例如:A1=A2==AK以及C1=C2==CK。我们已经发现:添加一个线性映射H到在hops间的u的update上是有用的;也就是说:uk+1=Huk+ok。该mapping会随着剩余参数学习,并在我们的实验上用于layer-wise weight tying。

图1(b)展示了一个3-layer版本。总体上,它与[23]中的Memory Network相似,除了每一layer中的hard max操作已经使用了一个来自softmax的continuous weighting替换外。

注意,如果我们使用layer-wise weight tying scheme,我们的模型可以被转成一个传统的RNN,其中我们会将RNN的outputs分割成internal和external outputs。触发一个internal output可以与考虑一个memory相对应,触发一个external output对应于预测一个label。从RNN的角度,图1(b)和等式(4)的u是一个hidden state,该模型会使用A生成一个internal output p(图1(a)中的attention weights)。该模型接着会使用C来吸收p,并更新hidden state等。这里,与一个标准RNN不同的是,我们会在K hops期间,显式的基于在memory中储存的outputs作为条件,我们会采用soft的方式来保存这些outputs,而非对它们采样。这样,我们的模型会在生成一个output之前做出一些计算step,这意味着被“外部世界”见过。

其它