xNN：支付宝端侧深度学习框架

编者荐语：
比较新颖的一点，可伸缩建模，实现对终端设备算力的精细化挖掘，即实现更细粒度的机型定制化适配，不过看来目前还处于探索之中，类似的，文末有字节和快手的相关方案文章。
另外是与淘系开源框架MNN的关系，xNN是支付宝闭源的，二者无关。
🙋🏻‍♀️ 编者按：
蚂蚁端智能诉求可以追溯到 2017 年的支付宝新春扫福活动。2017 年是支付宝第一次引入 AR 实景扫福，通过扫描任意“福”字帮助大家去集收集福卡。当时的福字识别模型选择服务端服务部署的技术方案，为了在活动期间识别福字，需要调用大量的服务端部署资源来部署识别模型；另一方面，DL 在云端则意味着数据必须上传。即使不考虑计算压力，从网络延时、流量、隐私保护等角度，也给用户体验带来种种限制。因此，对相当多的应用来说，DL 模型前移到移动端部署可以看作是一种刚需。
以下文章来源于支付宝体验科技，作者世豪

背景介绍

在上述背景下，支付宝开始了端侧 AI 的探索，并与 2017 年八月在支付宝APP 中首次上线 xNN 端侧推理引擎。在同一个时间段前后，我们也看到业界许多公司都开始布局并落地端上 AI 相关能力，蚂蚁也是业界较早布局端智能技术方向的公司之一。

xNN 引擎上线率先通过图像识别能力支撑“AR 扫花识花”业务，并逐步扩展了包括图像检测、文字识别、关键点回归等认知智能业务场景，及以搜索推荐营销、安全风控为主的数据智能业务场景。同时，端侧 AI 能力也在线下 IoT 场景的支付宝刷脸、智能货柜等中广泛应用，并通过支付宝小程序、扫一扫等入口将部分成熟的 AI 能力开放到生态支持。

在广泛的业务支持背后，我们需要首先定义端侧 AI 的一个核心技术课题。总结来讲端侧 AI 核心问题是模型优化的问题。我们知道在服务端部署 AI 模型时，通常整个部署环境的资源无论从算力、存储资源来讲，都是相对来讲比较丰富的。然而反观终端移动端设备以及 IOT 设备，无论是算力还是存储等指标上有若干数量级下的的严苛约束。因此，核心技术挑战在于如何通过可能的模型优化技术手段来降低端上部署模型的性能、尺寸、功耗等等。

作为终端 APP 应用开发团队，我们把能够触及到的模型优化空间抽象为四层：

模型层：通过设计终端友好的模型结构及其相应的参数来实现在模型任务精度和计算复杂度间的一个权衡。例如目前广泛应用的 MobileNets、ShuffleNets 系列的模型结构均属于这个层面上的优化。
计算图层：通过对固化的推理 DAG 计算图进行算子调度、融合、拆解等手段，保证模型计算依赖关系不变的前提下，提高我们推理效能。社区和学术界都会有相关的工作。类似 TorchScript、onnx、TVM-Relay 等工作均在这个维度有非常多的能力建设。
算子层：这个维度的优化也是端侧推理引擎核心关注的部分，即能否通过更为高效的算法、kernel 实现来达到对硬件算力的使用。相关工作开源社区中各家端侧推理引擎，以及服务端 NVIDIA 的 cuDNN 等等加速 Library。
设备层：硬件设备是我们端侧推理所依赖的基础设施，硬件厂商通过提供高效的指令集、AI 专用芯片硬件加速来提高算力峰值。

在此基础上，蚂蚁端侧 AI 框架 xNN 定位于覆盖上述四层优化空间。xNN 目标在于构建可以被不同业务算法在研发阶段复用的端 AI 基础能力。从模型结构设计、训练效率优化、模型参数压缩、端设备的部署，我们希望构建端到端的技术能力来提高整体体验。在这样的定位之下，既需要对各个业务算法领域异同有深入理解、了解其背后的核心诉求，同时也要紧跟整个业界的 AI 硬件发展趋势变化。xNN 团队也持续与华为、高通等各厂商保持紧密的合作与沟通。

xNN 框架 1.0：轻量化建模

首先跟大家分享我们在前期所建设的第一代框架的整体设计思路：轻量化建模。

轻量化建模和主流的模型优化的思路相对一致。我们基于刚才所定义的四个层级的优化主体逐层建设相应的能力，整体上形成一个自上而下的逐层优化。在这个过程中，我们沉淀了四个方面的技术能力：结构搜索，模型压缩，模型转换，计算引擎。通四个维度逐层向下优化，我们希望能够为我们内部不同的业务算法需求提供层次化的解决方案。接下来我会对这四个模块整体做一个简单的介绍。

首先是模型压缩和结构搜索，这两个技术模块是与端模型算法研发过程紧密相关，且需要与算法同学深度协同的一些技术能力。

模型压缩：

早在 xNN 计算引擎首次上线时，我们便配套布局建设了模型压缩的核心技术。它包括了模型的剪枝，参数的量化，以及数据编码。

剪枝（Pruning）：模型剪枝可以简单的分为三个维度。（1）通道剪枝（Channel Pruning）虽然能够有效的降低模型的尺寸以及计算的性能，但在实际的使用过程中调整通道间的压缩比例往往需要大量调参。（2）模版化剪枝（Pattern Pruning）及突触剪枝（Synapse Pruning）：这种方式通可以在引擎侧设计特殊的稀疏计算 kernel 来实现加速，但往往依赖较高的压比率才能获得较好的加速效果。因此想要把这个能力充分的利用起来，往往依赖较大的优化精力投入。整体来看，我们认为剪枝相关的技术在逐步的被结构搜索（NAS）的方式所去替代，并逐步将更多的技术资源迁移至 NAS 方向的建设中。
量化（Quantization）：量化是广泛应用于我们实际的业务场景中的一项模型压缩技术。它一方面能够帮助我们压缩模型尺寸以减小模型下发的流量消耗和设备端存储消耗，比如说我们可以用量化的方式对搜索推荐模型中的 Embedding 的词表做模型的压缩；另一方面，我们也可以用量化的方式去做加速。因为很多硬件厂商所提供的指令集都支持了 INT8 量化加速指令。所以量化是一个投入产出性价比较高的模型压缩手段。
编码（Coding）：利用数据编码方式来进一步的减小模型的尺寸，是在 xNN 引擎上线第一个模型时便引入的核心能力。我们可以简单的通过稀疏的表达存储方式来去做一个信息的存储；与此同时，xNN 设计了有效的编码算法，通过对AI的数据分布来实现更高的压缩比率来逼近信息熵的极限。

结构搜索：

网络结构搜索（Network Architecture Search， NAS）正在 xNN 体系中逐步取代很多单点模型压缩能力。因为上述的压缩算法往往会依赖于算法研发同学优化经验和反复尝试迭代才能获得较优的效果，那么通过自动化的方式来去替代人工的优化和搜索，对研发的效能帮助很大。

在这个背景下，我们也持续跟进并沉淀了非常多的能力。比如说我们通过 Black-box Optimization 实现自动化的超参搜索，比如通过 Differiential NAS 这种可微的方式，或者是通过 one-shot NAS 的方式来去做模型压缩。这些方案都广泛应用在我们的业务场景中，并取得了很不错的实践效果，绝大多数场景都进一步超过了以往经过专家经验深度调优的模型指标。

模型转换：

这个阶段主要的工作是把我们主流的深度学习训练框架所产出的模型转换为实际部署的 xNN 模型格式。整个的转换涉及的各种图优化的过程对相关领域同学都比较熟悉，我们不做过多的介绍。

在此基础之上，围绕着蚂蚁业务中实际需求，也逐步建设了非常多的增量功能。例如模型串联的需求，大家做真正的做端智能的业务落地，会发现往往一个业务所需要的端智能的算法通常并不是由一个单一的模型就能够解决的，往往是由多个 AI 模型通过一定的算法逻辑串联而成。因此有没有办法把多个AI模型有效的串联到一个模型之中？如果可以，那我们能够极大的减少算法研发同学和实际的工程部署同学之间的对接成本，以及整个开发的生命周期。

因此我们在编译阶段提供了相关的工具：模型串联工具帮助把一些包括循环分支；动态化工具能够实现引擎功能、资源调度逻辑的动态发布；另外包括模型保护工具、端侧训练工具等等相关的一些能力，都伴随着业务的发展逐步建设和完善起来。

计算引擎：

最后一块是计算引擎，这个也是我们通常所称的端侧 AI 计算相关的一个核心模块。计算引擎在我们 APP 的开发者的视角上来讲，我们更多的是聚焦于如何寻找高效的 kernel 实现，在访存密集型和计算密集型的计算任务上去做合理的优化和调整。

在整个建设过程中，我们针对不同的部署平台，会有不同的差异化的一个优化策略。比如说在 CPU 和 GPU 上，我们能够通过自己去实现定制化的高效 kernel 实现，来达到更优的效果。因此我们会在这些平台上去利用我们的专家经验，针对核心算子做深度的优化，并会逐步的以更自动化的算子编译的方式去做演进。

面对现在逐渐发展起来的 NPU 等 AI 芯片，我们 APP 开发者往往不太能够具备非常深入底层的优化方式。因此这一块我们会和各个芯片及硬件厂商会做较多协同，以期能够利用最大化的利用到端侧 AI 芯片的硬件加速算力。

以上是我们在轻量化建模这个技术建设阶段下几个核心模块相关的简单介绍，下面是相关技术在业务场景中的效果评估，我们主要关注在模型性能、模型尺寸及精度表现。

轻量化建模总结 – 优化表现

性能：xNN 计算引擎的性能持续对标业界多个优秀的开源引擎框架，同时我们也围绕着自己的场景做定制油画，比如在定点 kernel 配合定点压缩算法的定制方案能够在许多定点场景下，能够达到更具有竞争力的性能表现。
尺寸：模型尺寸极大的影响到了模型下发的成功率，以及在端侧存储限制。我们在实际的业务场景中，通过刚刚介绍的各种压缩优化能力，线上模型普遍在百 KB 的量级。
精度：模型压缩往往会影响原始模型的精度，我们也在持续关注和并在如量化算法上面的一些技术创新，并持续保持着在定点优化上面的优势。

以上是我们整个对 xNN 轻量化建模阶段工作的简单总结。当然轻量化建模这件事情是需要持续优化投入的课题，我们也会持续的在不同的核心技术点上去深度打磨。

端智能面临的问题及挑战

端 AI 整体发展趋势及诉求

在构建 xNN 轻量化建模技术的过程中，我们持续关注到整个端智能领域的发展趋势。我们去看当下的端智能领域，和四五年前去刚开始做端智能的时候相比，已经有了一些相当大的变化。我们会把整体端 AI 的一个发展诉求和趋势总结为下面三个角度：

从场景角度：我们看到即使是在公司内部已经有越来越多样化的业务场景希望尝试端智能来提升整个业务的一个效果。研发场景、研发人数及模型生产数量都相比于两三年前会有数倍数的增长。
从 AI 算法角度：这是一个蓬勃发展和高速演进领域。相比最初的 DNN、CNN 网络结构，我们看到了如 RNN、LSTM、Transformer 等新的网络结构设计的演进。这意味着端智能技术侧需要持续的理解以支撑好新的算法架构相关所产生的需求。
从 AI 芯片角度：在四五年前开始做端侧 AI 时还是以 CPU 为主的算力支撑。但今天我们再看终端算力环境，主流的旗舰机型往往都标配了像 NPU 的 AI 硬件加速器。这意味着对于支付宝这样的国民级 APP，我们需要覆盖的用户的手机算力会从 GFlops 到 TFlops、 3 到 5 个数量级的变化，算力的碎片化严重也意味着我们很难达到一致性的业务表达。

领域内技术的发展也带来短端侧 AI 核心诉求的变化。早期我们会投入较多的资源来利用端智能能力实现业务价值。今天在越来越多的场景和算法研发已经证明了它的价值的基础之上，我们要考虑的是如何构建我们的技术能力，让更多的业务能够高效的接入，并让已有的业务能够实现更为极致的优化。

这其中在关注算法研发部署效率的基础之上，一个核心的挑战是算力的碎片化的问题。如何把我们端上目前已经有的更高的算力利用起来。是要去深入去探讨的问题。

轻量化建模：算力资源利用率低

我们从算力资源利用率的角度来看一下轻量化建模主要的应用方式。在以往，轻量化建模希望产出极致轻量的一单一的模型。它其实意味着是什么呢？我们要么放弃低端机的用户覆盖，要么我们牺牲了高端机用户的体验。下面这张图是一个比较直观的例子。

我们从 AI benchmark 网站上抽取部分高通平台随着年份迭代的芯片算力的增长。我们可以看到如果以注重机型覆盖为例，使得这个模型在不同的手机上都能够被它的算力所支撑，达到一个实时的体验效果。但也就意味着，对于许多新的主流机型，它明明有更高的算力，明明可以用更大模型来达到更好的体验，但却不可以做到。

如何最大化算力资源利用率？

因此我们去考虑如何最大化资源算力的一个利用率呢？那我们考虑如果我们能够在感知到要部署的具体设备时，能否定制化的够构建这个设备最优的模型。整体来看，我们如果能通过多个模型做部署，就能够实现对整体算力利用率的提升。这是我们最开始的动机，在这个动机之下，引出了我们目前正在做的第二代 xNN 框架，并称之为可伸缩建模。

xNN 2.0 可伸缩建模

轻量化建模 -> 可伸缩建模

相比于第一代的框架，我们首先讨论下这两个思路本质的技术区别是什么。在以往我们追求一个单模型，希望在很多的可能的模型结构中去寻找到一个最优解，因此在这个过程中，我们的技术建设思路。始终围绕着对整个优化的路径上极致提升每一个优化工具的功能，来实现一个逐步的逼近最优解的方式。

但是如果我们想要把这样的模式应用在多模型的研发生产过程中，会存在很多问题。单模型的研发模式需要大量的人力资源的投入来做端到端的调优门槛，如果把这套模式复制到生产多个模型，意味着整个研发成本会有O（N）的线性复杂度的提升。因此我们希望一种新的研发部署方式，他能够有高效的生产多个模型。意味着我们每次的生产过程不是单一的生产最优的模型，而是生产一组在不同部署设备上分别最优的一个最优模型组的集合，即一个在多目标优化下的帕雷托前沿模型解几何。如果我们能够实现这一步，也就意味着来说我们能够一次性的产出一批模型，而这些模型能够分别针对我们所要部署的设备来寻找到它最优的模型结构。

基于这样一个可伸缩建模的思路，我们进一步思考如何生产高质量的、贴近实际生产环境所需的多模型。我们认为多模型之间它不应该是一个无序的存在，它应该是在某些方面具备了一定的伸缩性的。也就是说我可能在某指标上通过调整某一些核心的因子就能够实现。比如说我们希望在不同的设备上具备性能上的变化。我调整这些因子就能够实现不同模型的产出，来达到性能上差异化。只有这些有序的伸缩因子的挖掘，才能有持续指导我们做更高效的模型研发部署。所以可伸缩建模其中的核心是可伸缩变换因子的概念。

核心挑战：算法研发门槛高

影响多模型指标的核心因子是不是一个相对清晰、确定的答案呢？很明显它是否定的。学术界的论文往往基于许多理想假设、代理指标来做问题的建模，但是当地迁移到实际工业场景中，问题就会变得相当复杂，对于一个模型研发同学来讲的门槛是相当高的。

这里我会举两个实际的例子，也是我们在实际的业务支撑过程中非常有体感的一些例子。比如说刚才我们提到了就是端侧 AI 算力的差异化，算力的变化可以通过计算精度的表达来去伸缩。比如说高通的一些新款的设备已经能够支持到 INT4 等混合精度的表达。

如果作为一个算法研发同学，想要把这个能力应用起来，实现不同性能的表达。那它整个过程中可能会遇到什么阻碍呢？我们会发现，首先我们需要底层的部署设备具备这样的一个能力，这本身就是设备差异极大的；同时向上整个的推理引擎侧要能够非常好的兼容设备的能力；同时我们在整个模型编译的过程中，对设备所依赖的计算图表达做定制优化；我们也依赖定点量化算法上去做一些新的设计；最后我们更需要在整个的模型结构上对设备的功能有相应的理解，才能设计出设备友好的模型结构。也就意味着，如果我想要真正的把推理表达精度的因子广泛应用起来的话，需要从上到下的一个各个层面上的协同，只有有充分的协同，才能够把这个能力挖掘到最大化。

另一个例子我们从算法的视角来看，上文也提到了整个 AI 算法领域也是在高速的发展的，很多任务都被 Transformer-based 模型结构刷到了一个更高的一个精度。我们在自己的业务场景中会发现如果不针对 Transformer 结构做针对的优化，效果是不尽如人意的。那我们再深入的去看这里面为什么呢？例如我们去看 Transformer 里面的 MHSA 核心结构，想要把这个结构给充分的应用起来：在做结构搜索时我们要考虑怎么针对这么多头注意力机制去搜索可变的参数；同时思考 MHSA 是否有能够全定点压缩算法来达到落地时的性能加速；考虑我们的硬件厂商所提供编译工具，是不是能够友好的兼容 Transformer 结构等等；另一方面，如果我们去自己去实现相关 kernel，我们以往的实现有没有真正关注到这些结构所需要的一些核心算子。只有把这些东西全部的综合的考虑来看，我们才能得到一个结论，说 Transformer 是不是更适合在某一款具体的 CPU/GPU/NPU 上去执行。否则，任何一环有了缺失，都会影响到最终结论的判断，进一步影响到算法研发效果天花板。

因此，如果我们从算法研发的视角来看，想要让算法同学去理解这中间的每一环的领域知识及技术依赖，成本是相当高的。但如果我们能够以因子的视角把这些影响因子所关联的各个层级的技术栈充分的打通，并以一种高效的方式提供给到我们的算法同学，实现面向不同设备的多模型建模，便是我们希望建设的可伸缩建模。

可伸缩建模 – 接入体验

可伸缩建模，是 xNN 面向下一代建模的建设思路。我们希望做到的是什么呢？就是在以往单模型的算法研发的代码基础之上，仅仅通过少量的代码修改，就能够完成多模型建模的迁移。举例来看，比如说我们想要去构建多个模型以覆盖从 10ms 到 50ms 的设计目标，我们只需要在原先的模型定义基础之上，增加一行语法糖去描述 latency 目标的伸缩范围。接下来，利用我们所沉淀的可伸缩因子、或者是算法同学自研的自己的一些优化策略，将原先的静态模型结构变换为一个可伸缩的搜索空间描述，最后利用框架的空间-算法解耦设计，使用合适的搜索算法完成最终模型的生产。