让大模型的训练和推理，比更快还更快！谷歌2022年终总结第四弹

虽然谷歌的Bard翻车了，但谷歌的AI实力仍然不容小觑。

自开年以来，由Jeff Dean领衔的Google Research年终总结系列「Google Research, 2022 & beyond」一直在持续更新，最近也是更新到了第四期。

本期以「提升模型效率」为主题，一起看看谷歌工程师们都想出了哪些ideas！

运行效率成关键

在过去十年里，深度学习呈现爆炸式发展，很大程度上是由于新算法和体系结构的融合、数据量的显著增加以及计算能力的提高。

相比十年前，人工智能和机器学习模型变得更加巨大、更复杂，具有更深且更复杂的网络结构、更多的参数、训练时用到了更多的数据，共同促进了机器学习历史上一些最具变革性的成果。

随着这些模型越来越多地部署在生产和业务应用程序中，模型的推理效率和运行成本已经从一个次要因素变成了一个主要的限制因素。

Google在这方面的应对措施就是继续在机器学习效率方面投入巨资，主要解决以下四个难题：

1、高效的模型架构（Efficient Architecture）

2、数据效率（Data Efficiency）

3、训练效率（Training Efficiency）

4、推理效率（Inference Efficiency）

除了效率之外，模型还面临着围绕真实性、安全性、隐私性和时效性（freshness）等诸多难题。

这篇文章将重点介绍一系列Google Research研究公司为应对上述挑战而开发的新算法。

高效的模型架构

一个最基本的问题是：是否有更好的方法来参数化模型以提高效率？

2022年，Google Research专注于新技术，通过检索上下文增强模型，混合专家，使变压器(位于大多数大型机器学习模型的核心)更有效，从而注入外部知识。

上下文增强模型（Context-augmented models）

为了追求更高的质量和效率，神经模型可以从大型数据库或可训练的记忆库中的外部context中得到增强。

通过利用检索到的上下文，神经网络不需要在其内部参数中记忆大量的世界知识，就能够获得更好的参数效率、可解释性和真实性。

在「上下文增强语言建模的解耦上下文处理」一文中，研究人员探索了一种基于解耦编解码器架构将外部上下文合并到语言模型中的简单架构。

论文链接：​https://arxiv.org/abs/2210.05758​

该模型能够显著地节省计算量，同时在自回归语言建模和开放领域问题回答任务方面给出了有竞争力的结果。

预训练的大语言模型(LLM)通过对大训练集的自监督消耗了大量的信息，但目前还不清楚这些模型的「世界知识」是如何与输入的context相互作用的。

通过知识感知微调(knowledge aware fine-tuning, KAFT) ，研究人员将反事实上下文和无关上下文结合到标准的监督数据集中，增强了 LLM 的可控性和鲁棒性。

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

上下文增强模型的另一个难题是从大型数据库中快速检索信息的加速器。

研究人员开发了一个基于 TPU 的最近邻搜索算法，该算法与 TPU 的性能模型保持一致，并为预期召回提供分析保证，从而达到最佳性能。

搜索算法通常涉及大量的超参数和设计选择，这使得它们很难在新任务上进行优化，所以研究人员进一步提出了一种新的约束优化算法来自动调整超参数，将期望的成本或召回作为输入，该算法产生的调整在经验上非常接近速度召回的Pareto frontier，并在标准基准上给出领先的性能。

混合专家模型

混合专家(MoE, Mixture-of-experts)模型已被证明是增加神经网络模型容量而不会过度增加其计算成本的有效手段。MoEs 的基本思想是从多个专家子网络中构建一个网络，其中每个输入由一个合适的专家子网络处理。

因此，与标准的神经网络相比，MoEs 只调用整个模型的一小部分，从而提高了语言模型应用程序(如 GLaM)的效率。

决定哪些专家应该被激活（be active）以参与到某一特定输入取决于路由功能（routing function），路由的设计非常具有挑战性，因为开发者的预期是每个专家模型都是恰如其分的，不会出现利用不足和过度的情况。

在最近的一项工作中，研究人员提出了专家选择路由(Expert Choice Routing) ，一种新的路由机制，它不是将每个输入token分配给 top-k 专家，而是反过来将每个专家分配给 top-k token，可以自动调节专家的负载平衡，同时自然地允许多个专家处理输入token

论文链接：​https://openreview.net/pdf?id=jdJo1HIVinI​

高效的Transformers

Transformer是当下最流行的序列到序列的模型，从视觉到自然语言的理解，在一系列具有挑战性的任务中已经展现了非常强大的性能。

这类模型的一个核心组成部分是注意力层（attention layer），通过计算「查询」（query）和「键」（key）之间的相似度来构建一个适当的「值」（value）加权组合。虽然性能强，但注意力机制的计算效率却不高，复杂度通常为输入序列长度的二次方。

随着Transformer规模的不断扩大，其中一个重要问题的研究十分有价值：是否有任何自然发生的结构或模式的学习模型，可以解了注意力有效的原理。

对此，Google Research研究了中间 MLP 层的学习嵌入，发现它们非常稀疏，比如，T5-大模型只有 < 1% 的非零项，其稀疏性进一步表明，可以在不影响模型性能的情况下降低FLOP

论文链接：​https://arxiv.org/pdf/2210.06313.pdf​

研究人员最近提出了 Treeformer模型，一种依赖于决策树的标准注意力计算的替代方案，直观地说，它可以快速识别与查询相关的一小部分键，并且只在这个集合上执行注意力操作。根据经验，Treeformer 可以使注意力层的 FLOP 减少30倍。

同时研究人员还引入了序贯注意（ sequential attention），一种将注意和贪婪算法相结合的可微特征选择方法，该技术经过证明，可以直接无缝迁移到大规模模型。

另一种提升Transformer效率的方法是在注意层加速softmax的计算。

在low-rank approximation of the softmax kernel研究基础上，研究人员提出了一类新的随机特征，提供了softmax kernel的首个「positive and bounded」随机特征近似，并且在序列长度上是计算线性的。

论文链接：​https://arxiv.org/abs/2205.15317​

并且还提出了首个涵盖多个注意力masking的机制，如因果编码和相对位置编码。

训练效率

有效的优化方法是现代机器学习应用的基石，在大规模环境中尤为重要。

在这种情况下，即使是像Adam这样的一阶自适应方法也往往需要大量的计算量，而且训练的稳定性也会变得非常困难。

此外，这些方法往往与神经网络的体系结构无关，没有考虑模型架构内的结构信息，导致训练效率低下，这也促使新技术能够更有效地优化现代神经网络模型。

Google Research开发了一些新的基于模型架构的训练技术，例如，用于训练Transofmre网络，包括新的规模不变的Transofmre网络和新的clipping方法，当与原版随机梯度下降(SGD)结合时，可以加快训练速度。

论文链接：​https://arxiv.org/pdf/2202.00980.pdf​

使用这种方法，研究人员首次实现了仅使用简单的 SGD 就能有效地训练 BERT，而无需自适应（adaptivity）。

此外，研究人员提出了一种新的方法LocoProp，在使用与一阶优化器相同的计算和内存资源的情况下，获得类似于二阶优化器的性能。

LocoProp 通过将神经网络分解为一个层的组合，从而获得了神经网络的模块化视图。然后允许每一层都有自己的损失函数以及输出目标和重量调整器。通过这种设置，经过适当的向前向后梯度传递后，LocoProp 继续对每个层的「局部损失」执行并行更新。

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

事实上，这些更新在理论上和经验上都类似于高阶优化器，在深度自编码器基准测试中，LocoProp 获得了与高阶优化器相当的性能，同时速度明显更快。

像 SGD 这样的优化器的一个关键假设是，每个数据点都是从一个分布中独立和完全相同地采样的，不过在强化学习等实际应用环境中，很难满足该假设，因为模型(或智能体)必须从根据自己的预测生成的数据中学习。

研究人员提出了一种新的算法方法，称为带有反向经验重放的 SGD，可以在线性动力系统（linear dynamical system）、非线性动力系统和强化学习的 Q 学习等几种情况下找到最优解。

论文链接：​https://arxiv.org/abs/2103.05896​

此外，这种方法的一个改进版本IER ，经过实验被证明是最先进的，并且是各种流行的 RL 基准测试上最稳定的经验重播技术。

数据效率

深度神经网络严重依赖于大型数据集，随之而来的还有存储成本和潜在的安全/隐私问题，在这些数据集上训练现代深度神经网络也会带来很高的计算成本。

解决这个问题的一个有前景的方法是数据子集选择，其中学习器的目标是从大量的训练样本中找到信息量最大的子集，以接近(甚至改进)整个训练集的训练。

研究人员分析了一个子集选择框架设计工作与任意模型族在一个实际的batch设置，在这种设置下，学习器可以一次采样一个样本，访问上下文和真实标签，但为了限制开销，一旦选择了足够大的一批样本，只能更新其状态，即，进一步训练模型权重。

并基于此开发了一种称为 IWeS 的算法，该算法通过重要抽样来选择样本，其中分配给每个样本的抽样概率是基于对先前选定的批次训练的模型的熵。文中提供了一个理论分析，证明了泛化和抽样率的界限。

论文链接：​https://arxiv.org/pdf/2301.12052.pdf​

训练大型网络的另一个问题是，它们可能对部署时看到的训练数据和数据之间的分布变化高度敏感，特别是在处理数量有限的训练数据时，这些数据可能不包括所有的部署时间情景。

最近的一项研究假设「极端简单性偏差」是神经网络这种脆弱性背后的关键问题，其最新工作使这一假设成为可行，导致了两种新的互补方法DAFT 和 FRR ，结合起来提供了显著更强大的神经网络。特别是，这两种方法使用对抗微调以及逆特征预测，提升了学习网络的健壮性。

论文链接：​https://arxiv.org/pdf/2006.07710.pdf​

推理效率

增加神经网络的规模已被证明在提高其预测准确性方面有惊人的效果，然而，在现实世界中利用这些优势是具有挑战性的，因为大型模型的推理成本可能过高了，该问题也促使策略提高服务效率，而不牺牲准确性。

研究人员提出了实现该目标的不同策略，特别是基于知识蒸馏和自适应计算的策略。

蒸馏

蒸馏是一种简单而有效的模型压缩方法，极大地扩展了大型神经模型的潜在适用性，并且在广告推荐等一系列实际应用中被证明是非常有效的。

大多数蒸馏的用例都涉及到将基本模型直接应用到给定的领域，对于何时以及为什么应该这样做只有有限的理解。Google的研究着眼于根据特定的环境对蒸馏进行裁剪，并系统性地研究了决定蒸馏成功的因素。

在算法方面，通过仔细建模教师模型提供的标签中的噪声，研究人员开发了一个原则性的方法来重新加权的训练例子，和一个稳健的方法来抽样的数据子集有教师标签。

论文链接：https://arxiv.org/abs/2210.06711

在「教师指导训练」的过程中，研究人员提出了一个新的蒸馏框架: 不是被动地使用教师来标注一个固定的数据集，而是主动地使用教师来指导信息样本的选择来进行标注，这使得蒸馏过程在有限的数据或长尾设置中更有效。

论文链接：​https://arxiv.org/abs/2208.06825​

同时还研究了从交叉编码器(dual-encoder, 例如 BERT)到因子双编码器（dual-encoder）的新方法，这也是对（query, document）对的相关性进行评分的重要设置。

论文链接：https://proceedings.mlr.press/v162/menon22a/menon22a.pdf

文中研究了交叉编码器和双重编码器之间性能差距的原因，指出这可能是泛化的结果，而非双编码器的容量限制。

精心构造的蒸馏损耗函数可以缓解这种情况，缩小交叉编码器和双编码器之间的性能差距。

随后，在 EmbedDistil进一步研究了通过匹配来自教师模型的嵌入来进一步改进双编码器精馏。该策略还可用于从大到小的双编码器模型中提取信息，继承和冻结教师的文档嵌入被证明是非常有效的。

论文链接：​https://arxiv.org/abs/2301.12005​

在理论方面提供了一个新的视角，通过监督复杂性的蒸馏，衡量学生如何能够预测教师的标签。

论文链接：​https://arxiv.org/abs/2301.12245​

利用神经正切核(NTK)理论得出了一些概念性的结论，比如，能力差距可能会影响蒸馏，因为这样的教师的标签可能出现类似于纯粹的随机标签的学生。

论文链接：​https://arxiv.org/abs/2301.12923​

同时进一步证明，蒸馏导致学生欠拟合的点也是对于教师模型比较难的问题，从直觉上讲，这可能有助于学生将其有限的能力集中在那些可以合理建模的样本上。

自适应计算

虽然蒸馏是降低推理成本的有效方法，但它在所有样本中的作用是一致的（uniformly），直观地说，一些「简单」的样本可能本质上需要比相对更「难」样本更少的计算。

自适应计算的目标是设计能够进行这种依赖于样本的计算的机制。

Confident Adaptive Language Modeling(CALM)为 T5等基于Transformer的文本生成器引入了受控的早期退出功能。

论文链接：​https://arxiv.org/abs/2207.07061​

在这种形式的自适应计算中，模型动态修改每个解码步骤使用的Transformer层数，其中early-exit gate使用带有决策阈值的置信度量，该置信度量经过校准以满足统计性能保证。

这样，模型只需要为最具挑战性的预测计算完整的解码器层堆栈，更简单的预测只需要计算几个解码器层。在实践中，该模型平均使用约三分之一的层进行预测，在保持同等水平的生成质量的同时，产生2-3倍的加速。

一个常用的自适应计算机制包括两个或多个基本模型的级联，其关键问题是决定是简单地使用当前模型的预测，还是将预测推迟到下游模型，学习何时推迟需要设计一个合适的损失函数，它可以利用适当的信号作为推迟决策的监督。

Google Research系统性地研究了现有的损失函数，证明由于标签平滑的隐式应用，它们可能不适合训练样本，文中还展示了可以通过延迟规则的事后训练来缓解这种情况，这种训练不需要以任何方式修改模型内部。

论文链接：​https://openreview.net/pdf?id=_jg6Sf6tuF7​

对于检索应用程序，标准的语义搜索技术对大型模型生成的每个嵌入使用固定的表示，也就是说，不管下游任务及其相关的计算环境或约束，表示的大小和能力大多是固定的。

Matryoshka 表示学习引入了根据部署环境调整表示的灵活性，强制表示在它的坐标中有一个自然的排序，这样对于资源受限的环境，只使用表示的最高的几个坐标；而对于更丰富和精度关键的设置，可以使用表示的更多坐标。

论文链接：​https://openreview.net/pdf?id=9njZa1fm35​

当结合标准近似最近邻搜索技术，如扫描神经网络，MRL 能够提供多达16倍的低计算相同的召回率和准确度量。

总结

大型机器学习模型在多个领域都展现出具有变革性的结果，但训练和推理的效率正在成为使这些模型在现实世界中切实可行的关键需求。

通过开发新的基础技术，Google Research已经在使大型机器学习模型高效方面进行了大量投资，这也需要持续性的努力，未来将继续探索核心挑战，使机器学习模型更加健壮和高效。