如何提升深度学习算法效率，谷歌有这些绝招

2023-02-25 05:02

选自谷歌博客

机器之心编译

编辑：铭怿

大型机器学习模型在多个领域展现出具有变革性的结果，但在现实世界中，训练和推理效率正在成为关键需求。Google Research 一直在通过开发新的基础技术来提高大型 ML 模型的效率。本文中，谷歌以「提升模型效率」为主题，介绍如何使 ML 模型更加强韧、高效。

十年前，深度学习崛起的部分原因在于新算法和架构的融合、数据的显著增加以及计算能力方面的提升。过去 10 年里，AI 和 ML 模型更加深入、复杂、有了更多的参数和训练数据，也因而变得更大、更繁琐，这也为机器学习历史带来了最具变革性的成果。

这些模型越来越多地应用于生产和业务应用程序中，与此同时，其效率和成本已经从次要的考虑因素演变为主要的限制。为了应对在高效架构，训练效率，数据效率和推理效率四个层面的重大挑战，谷歌继续在 ML 效率上投入大量资金。除了效率之外，这些模型还面临着许多关于真实性、安全性、隐私性和新鲜度的挑战。接下来，本文将重点介绍 Google Research 为了应对上述挑战，在开发新算法方面所做的努力。

研究的基本问题在于「是否有更好的方法来参数化模型以提高效率？」2022 年，研究人员专注于通过检索上下文、混合专家系统，以及提升 Transformer（大型 ML 模型的心）的效率来研发通过增强模型来注入外部知识的新技术。

上下文增强模型

为了追求更高的质量和效率，可以使用来自大型数据库或可训练内存的外部上下文来增强神经模型。通过利用检索到的上下文，神经网络无需在其内部参数中广泛地储备知识，就能实现更好的参数效率、可解释性和真实性。

一篇名为《用于上下文增强语言建模的解耦上下文处理》（Decoupled Context Processing for Context Augmented Language Modeling）的文章中探索了一种基于解耦编码器 - 解码器架构的简单架构，用于将外部上下文合并到语言模型中。在自回归语言建模和开放领域问答任务中，这大大节省了计算量。然而，预训练大型语言模型（LLMs）通过在大型训练集上的自监督消耗了大量的信息。但是，目前还不清楚这些模型对世界的认识是如何与所呈现的上下文相互作用的。通过知识感知微调（KAFT），研究人员将反事实和不相关的上下文纳入标准监督数据集，这加强了 LLM 的可控性和鲁棒性。

论文地址：https://arxiv.org/abs/2210.05758

用于上下文合并的编码器 - 解码器交叉注意机制，允许上下文编码与语言模型推理解耦，进而提升上下文增强模型的效率。

在寻求模块化深度网络的过程中，其中一个问题是如何设计具有相应计算模块的概念数据库。研究人员提出了一种理论体系架构，将「记忆事件」（remember events）以 sketches 的形式存储在一个外部 LSH 表中，其中包括一个 pointers 模块来处理 sketches。

利用加速器从大型数据库中快速检索信息是上下文增强模型的另一大挑战。研究人员开发了一种基于 TPU 的相似度搜索算法，该算法与 TPU 的性能模型保持一致，并对预期召回率提供了分析保证，实现了峰值性能。搜索算法通常涉及大量的超参数和设计选择，这使得在执行新任务时很难对它们进行调整。研究人员提出了一种新的约束优化算法，用于自动化超参数调优。将期望的成本或召回率固定为输入，所提出的算法产生的调优在经验上非常接近速度 - 召回率帕累托边界（speed-recall Pareto frontier），并在标准基准测试中提供领先的性能。

混合专家模型

经证明，混合专家 (MoE）模型是在不过度增加计算成本的前提下，提高神经网络模型容量的有效手段。MoE 的基本思想是由许多专家子网络构建出一个统一网络，其中每个输入由一个合适的专家子集处理。因此，与标准神经网络相比，MoE 只调用了整个模型的一小部分，从而产生了如 GLaM 等语言模型应用程序所示的高效率。

GLaM 体系架构中的每个输入 token 被动态路由到 64 个专家网络中的两个进行预测。

对于给定的输入，路由函数负责决定应当激活哪些专家，该函数的设计具有挑战性，因为研究者希望避免对每个专家的利用不足和过度利用。最近的一项工作提出了专家选择路由，这是一种新的路由机制，它不是将每个输入 token 分配给 top-k 专家，而是将每个专家分配给 top-k token。这将自动确保专家的负载平衡，同时也自然地允许多个专家处理一个输入 token。

专家选择路由。具有预定缓冲区容量的专家被分配 top-k token，从而保证负载平衡。每个 token 可以由数量不定的专家处理。

有效的 Transformer

Transformer 是当下正火热的序列到序列模型，在从视觉到自然语言理解等一系列具有挑战性的问题上取得了显著的成功。这种模型的核心组成部分是注意力层，它识别查询和键之间的相似性，并使用这些相似性构造一个适当的值加权组合。虽然性能强，但注意力机制的计算效率却不高，复杂度通常为输入序列长度的二次方。

随着 Transformer 规模的不断增长，其中一个问题的研究十分有价值，即是否有任何自然发生的结构或模式的学习模型，可以破解注意力有效的原理。为此，研究人员研究了中间 MLP 层中的学习嵌入，并且发现它们非常稀疏 —— 例如，T5-Large 模型有 1% 的非零项。稀疏性进一步表明，人们可以在不影响模型性能的情况下潜在地降低 FLOPs。

论文地址：https://arxiv.org/pdf/2210.06313.pdf

最近，有研究推出 Treeformer—— 一种依赖于决策树的标准注意力计算的替代方案。简单来说，这可以快速识别与查询相关的键的一小部分，并且只对该集合执行注意力操作。根据经验，Treeformer 可以将注意力层的 FLOPs 降低 30 倍。除此之外还有序列注意力 —— 一种结合了注意力和贪婪算法的可微分特征选择方法。该技术对线性模型有很强的可证明保证，并可无缝扩展到大型嵌入模型。

另一种提高 Transformer 效率的方法是在注意力层中加速 softmax 计算。在「low-rank approximation of the softmax kernel」研究基础上，研究人员提出了一类新的随机特征，提供了 softmax kernel 的首个「positive and bounded」随机特征近似，并且在序列长度上的计算是线性的。

训练效率

高效的优化方法是现代 ML 应用程序的基石，在大规模设置中，这一点尤其重要。在这种设置下，即使是像 Adam 这样的一阶自适应方法通常也很昂贵，并且训练稳定性面临重重挑战。此外，这些方法通常对神经网络的架构不可知，从而忽略了架构的丰富性，导致训练效率低下。这也促使新技术不断被提出以更有效地优化现代神经网络模型。研究人员正在开发新的架构感知训练技术，例如，用于训练 Transformer 网络的一些研究，包括新的规模不变的 Transformer 网络和新的剪枝方法，并与随机梯度下降（SGD）结合使用，以加快训练进程。在该方法的帮助下，研究人员第一次能够使用简单的 SGD 有效地训练 BERT，而不需要自适应。

论文地址：https://arxiv.org/pdf/2210.05758.pdf

此外，研究人员在 LocoProp 的帮助下提出了一种新的方法 —— 在使用与一阶优化器相同的计算和内存资源的同时，实现与二阶优化器类似的性能。LocoProp 采用模块化的神经网络视图，将它们分解成层的组合。然后允许每一层都有自己的损失函数以及输出目标和权重正则化器。有了这个设置，经过适当的前后传递后，LocoProp 继续对每一层的局部损失进行并行更新。事实上，无论是理论上还是经验上，这些更新可以被证明类似于那些高阶优化器。在深度自编码器基准上，LocoProp 实现了与高阶优化器相当的性能，同时更具速度优势。

论文链接：https://proceedings.mlr.press/v151/amid22a.html

与反向传播类似，LocoProp 应用前向传递来计算激活。在向后传递中，LocoProp 为每一层设置每个神经元目标。最后，LocoProp 将模型训练拆分为跨层的独立问题，其中几个本地更新可以并行应用于每个层的权重。

SGD 等优化器的核心思想是，每个数据点从分布中独立且相同地采样。可惜这在实际环境中很难满足，比如强化学习，在这种情况下，模型（或智能体）必须从基于自己预测生成的数据中学习。研究人员提出了一种新的基于反向经验重放的 SGD 算法，该算法可以在线性动力系统、非线性动力系统和 Q-learning 中找到最优解。此外，已经有研究证明该方法的增强版本 IER 是目前最先进的版本，并且是在各种流行的 RL 基准测试中最稳定的体验重放技术。

论文地址：https://arxiv.org/pdf/2103.05896.pdf

数据效率

在许多任务中，深度神经网络严重依赖于大型数据集。除了大型数据集带来的存储成本和潜在的安全 / 隐私问题外，在此类数据集上训练现代深度神经网络还会产生很高的计算成本。解决这个问题的可行方法之一是对数据子集进行选择。

研究人员分析了子集选择框架，设计用于实际批处理设置中的任意模型族。在这种情况下，学习器可以一次采样一个示例，同时访问上下文和 true 标签，但为了限制开销，只有在选择了足够多的示例批次后才能更新其状态（即进一步的训练模型权重）。研究人员开发了一种算法，称为 IWeS，它通过重要性抽样来选择示例，其中分配给每个示例的抽样概率是基于在先前选择的批次上训练的模型的熵。研究所提供的理论分析证明了泛化和抽样率的界限。

论文地址：https://arxiv.org/pdf/2301.12052.pdf

训练大型网络的另一个问题是，它们可能对训练数据和部署时看到的数据之间的分布变化高度敏感，特别是在使用有限数量的训练数据时，这些数据可能无法覆盖所有部署时场景。最近的一项研究假设「极端简单性偏差」是神经网络这种脆弱性背后的关键问题，其使这一假设成为可行，导致了两种新的互补方法 ——DAFT 和 FRR—— 结合在一起时，可以提供明显更鲁棒的神经网络。特别是，这两种方法使用对抗性微调以及反向特征预测来加强学习网络。

论文地址：https://arxiv.org/pdf/2006.07710.pdf

推理效率

事实证明，增加神经网络的规模能够提高其预测准确性，然而，在现实世界中实现这些收益是具有挑战性的，因为大型模型的推理成本对于部署来说非常高。这促使策略提高服务效率，而不牺牲准确性。在 2022 年，专家研究了实现这一目标的不同策略，特别是基于知识蒸馏和自适应计算的策略。

蒸馏

蒸馏是一种简单而有效的模型压缩方法，极大地扩展了大型神经模型的潜在适用性。已有研究证明蒸馏能在广告推荐等一系列实际应用中发挥其作用。蒸馏的大多数用例涉及基本配方在给定领域的直接应用，对何时以及为什么这应该起作用的理解有限。谷歌今年的研究着眼于针对特定环境定制蒸馏，并正式研究了控制蒸馏成功的因素。

在算法方面，研究通过仔细建模教师标签中的噪声，开发出了一项重新加权训练示例的重要途径，以及一项有效措施来对数据子集进行采样以获得教师标签。谷歌在《Teacher Guided Training: An Efficient Framework for Knowledge Transfer》中表明：不是被动地使用教师来注释固定的数据集，而是主动地使用教师来指导选择要注释的信息性样本。这使得蒸馏过程在有限的数据或长尾设置中脱颖而出。

论文地址：https://arxiv.org/pdf/2208.06825.pdf

除此之外，谷歌还研究了从交叉编码器 (dual-encoder, 例如 BERT) 到因子双编码器（dual-encoder）的新方法，这也是对（query, document）对的相关性进行评分的重要设置。研究人员探讨了交叉编码器和双编码器之间性能差距的原因，注意到这可能是泛化的结果，而不是双编码器的容量限制。精馏损失函数的精心构造可以缓解这种情况，并减少交叉编码器和双编码器性能之间的差距。随后，在 embedtitil 中，研究通过匹配教师模型中的嵌入来进一步改进双编码器蒸馏。这种策略也可以用于从大到小的双编码器模型中提取信息，其中继承和冻结教师的文档嵌入可以证明是非常有效的。

论文地址：https://arxiv.org/pdf/2301.12005.pdf

理论方面，研究从监督复杂性角度出发，提供了一个关于蒸馏的新视角，这是一种衡量学生预测教师标签程度的方法。NTK（neural tangent kernel ）理论提供了概念性的见解。研究进一步证明，蒸馏会导致学生不适合教师模型认为难以建模的点。直观地说，这可以帮助学生将有限的能力集中在那些可以合理建模的样本上。

论文地址：https://arxiv.org/pdf/2301.12245.pdf

自适应计算

虽然蒸馏是一种降低推理成本的有效手段，但它在所有样本上都是一致的。然而，直观地说，一些简单（easy）样本可能天生就比困难（hard）样本需要更少的计算。自适应计算的目标是设计能够实现这种依赖于样本的计算的机制。

CALM（Confident Adaptive Language Modeling）为基于 Transformer 的文本生成器（如 T5）引入了受控的早期退出（early-exit）功能。

论文地址：https://arxiv.org/pdf/2207.07061.pdf

在这种形式的自适应计算中，模型动态地修改每个解码步骤使用的 Transformer 层数。早期退出门使用带有决策阈值的置信度测量，该决策阈值经过校准，以满足统计性能保证。通过这种方式，模型只需要为最具挑战性的预测计算解码器层的完整堆栈。更简单的预测只需要计算几个解码器层。在实践中，该模型平均使用大约三分之一的层进行预测，在保持相同水平的生成质量的同时，速度能提升 2-3 倍。

使用常规语言模型（top）和 CALM （bottom）生成文本。CALM 试图做出早期预测。一旦对生成的内容足够自信（深蓝色调），它就会跳过来节省时间。

一种流行的自适应计算机制是两个或多个基本模型的级联。使用级联的一个关键问题：是简单地使用当前模型的预测，还是将预测推迟到下游模型。学习何时延迟需要设计一个合适的损失函数，它可以利用适当的信号作为延迟决策的监督。为了实现这一目标，研究人员正式研究了现有的损失函数，证明由于隐含的标签平滑应用，它们可能不适合训练样本。研究展示了可以通过延迟规则的 post-hoc 训练来缓解这种情况，这并不需要以任何方式修改模型内部。

论文地址：https://openreview.net/pdf?id=_jg6Sf6tuF7

对于检索应用，标准语义搜索技术对大型模型生成的每个嵌入使用固定表示。也就是说，无论下游任务及其相关的计算环境或约束如何，表示的大小和能力基本上是固定的。MRL（Matryoshka representation learning）引入了根据部署环境调整表示法的灵活性。当与标准的近似最近邻搜索技术（如 ScaNN）结合使用时，MRL 能够提供高达 16 倍的低计算，同时具有相同的召回率和准确率指标。