前言
选自谷歌博客 机器之心编译 编辑：铭怿
大型机器学习模型在多个领域展现出具有变革性的结果，但在现实世界中，训练和推理效率正在成为关键需求。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 基准测试中最稳定的体验重放技术。
