Google Research 示例#
使用 Flax 完成的 Google Research 研究集锦。
注意#
快速注意 (FAVOR+) 和用 Performer 重新思考注意#
GitHub 上的代码
研究论文
用 Performer 重新思考注意 (Choromanski 等人，2020)
介绍了 *“Performer，可以以可证明的准确性估计常规 (softmax) 全秩注意变换器，但仅使用线性（而不是二次）空间和时间复杂度，而无需依赖于任何先验知识，例如稀疏性或低秩性。为了近似 softmax 注意内核，Performer 使用了一种新颖的通过正交随机特征的快速注意方法 (FAVOR+)，该方法可能对可扩展的内核方法具有独立的意义。FAVOR+ 也可用于有效地建模超出 softmax 的可内核化的注意机制。”*。
自注意不需要 O(n^2) 内存#
研究论文
自注意不需要 O(n^2) 内存 (Rabe 和 Staats，2021)
“我们提出了一种非常简单的注意算法，该算法需要关于序列长度的 O(1) 内存，以及需要 O(log n) 内存的自注意扩展。这与经常提出的关于自注意需要 O(n^2) 内存的说法形成对比。虽然时间复杂度仍然是 O(n^2)，但设备内存而不是计算能力往往是现代加速器的限制因素。因此，降低注意的内存要求可以处理比本来可能实现的更长的序列……”
计算机视觉#
着色变换器 (ColTran)#
研究论文
着色变换器 (Kumar 等人，2020)
“我们介绍了着色变换器 (ColTran)，这是一种完全依赖于自注意进行图像着色的架构。我们引入了条件变换器层，这是一种基于自注意的条件生成模型的新颖构建块。我们的消融研究表明，在许多不同的基线中使用这种机制的优越性。最后，我们证明 ColTran 可以为 ImageNet 生成多样化、高保真度的着色，即使对于人类评分者而言，这些着色也与真实情况基本无法区分。”
视觉变换器 (ViT)、MLP 混合器架构 *以及* Big Vision#
GitHub 上的代码
-
“这个代码库旨在使用 Cloud TPU VM 或 GPU 机器训练大规模视觉模型。它基于 Jax/Flax 库，并使用 tf.data 和 TensorFlow Datasets 来实现可扩展且可重复的输入管道。”
-
视觉变换器和 MLP 混合器的 JAX 代码
超过 50,000 个视觉变换器和混合检查点，用于生成“如何训练你的 ViT？”的数据。
研究论文
一张图片胜过 16x16 个词：用于大规模图像识别的变换器 (Dosovitskiy 等人，2020)
“在视觉中，注意要么与卷积网络一起使用，要么用于替换卷积网络的某些组件，同时保持其整体结构。我们表明，这种对 CNN 的依赖并不是必需的，直接应用于图像块序列的纯变换器在图像分类任务中可以表现非常出色。当在大规模数据上进行预训练并迁移到多个中等规模或小型图像识别基准（ImageNet、CIFAR-100、VTAB 等）时，视觉变换器 (ViT) 在与最先进的卷积网络相比时取得了出色的结果，同时训练所需的计算资源大大减少。”
MLP 混合器：用于视觉的全 MLP 架构 (Tolstikhin 等人，2021)
“在本文中，我们表明，虽然卷积和注意都是足以获得良好性能的，但它们都不是必需的。我们提出了 MLP 混合器，这是一种完全基于多层感知器 (MLP) 的架构。MLP 混合器包含两种类型的层：一种在图像块上独立应用 MLP（即“混合”每个位置的特征），另一种跨块应用 MLP（即“混合”空间信息）。当在大型数据集上训练或使用现代正则化方案时，MLP 混合器在图像分类基准上取得了有竞争力的分数，其预训练和推理成本与最先进的模型相当。”
如何训练你的 ViT？视觉变换器中的数据、增强和正则化 (Steiner 等人，2021)
“视觉变换器 (ViT) 已被证明在各种视觉应用中取得了极具竞争力的性能，例如图像分类、目标检测和语义图像分割。与卷积神经网络相比，视觉变换器的较弱归纳偏差通常会导致在较小的训练数据集上训练时更加依赖模型正则化或数据增强（简称“AugReg”）。我们进行了一项系统的实证研究，以更好地理解训练数据量、AugReg、模型大小和计算预算之间的相互作用。”
当视觉变换器在没有预训练或强数据增强的情况下优于 ResNet 时 (X. Chen 等人，2021)
“视觉变换器 (ViT) 和 MLP 信号表明，用通用神经架构取代人工特征或归纳偏差的进一步努力。现有的作品通过海量数据赋予模型能力，例如大规模预训练和/或重复的强数据增强，但仍然报告了与优化相关的問題（例如，对初始化和学习率的敏感性）。因此，本文从损失几何的角度研究了 ViT 和 MLP 混合器，旨在提高模型在训练时的数据效率和在推理时的泛化能力。”
LiT：使用锁定图像文本调优的零样本迁移 (X. Zhai 等人，2021)
“本文介绍了对比调优，这是一种简单的方法，它使用对比训练来对齐图像和文本模型，同时仍然利用它们的预训练。在我们的实证研究中，我们发现锁定预训练的图像模型和解锁文本模型效果最好。我们将对比调优的这种实例称为“锁定图像调优”（LiT），它只是教文本模型从预训练的图像模型中读取出新的任务的良好表示。LiT 模型获得了对新视觉任务进行零样本迁移的能力，例如图像分类或检索。提出的 LiT 具有广泛的适用性；它使用三个不同的图像文本数据集，可以可靠地与多种预训练方法（监督和无监督）以及不同的架构（ResNet、视觉变换器和 MLP 混合器）一起使用。”
使用稀疏专家混合 (MoE) 扩展视觉#
研究论文
使用稀疏专家混合扩展视觉 (Riquelme 等人，2021)
“稀疏门控专家混合网络 (MoE) 已在自然语言处理中展现出极佳的可扩展性。然而，在计算机视觉中，几乎所有性能良好的网络都是“密集的”，即每个输入都由每个参数处理。我们提出了视觉 MoE (V-MoE)，这是视觉变换器的稀疏版本，它具有可扩展性，并且与最大的密集网络相比具有竞争力…… 我们展示了 V-MoE 在扩展视觉模型方面的潜力，并训练了一个 150 亿参数的模型，该模型在 ImageNet 上获得了 90.35% 的准确率……”
扩散#
变分扩散模型#
研究论文
变分扩散模型 (Kingma 等人，2021)
“基于扩散的生成模型在感知上展现出令人印象深刻的合成能力，但它们也能成为出色的基于似然的模型吗？我们对这个问题做出了肯定的回答，并介绍了一系列基于扩散的生成模型，它们在标准图像密度估计基准测试中获得了最先进的似然性。与其他基于扩散的模型不同，我们的方法允许有效地联合优化噪声调度和模型的其余部分。我们表明，变分下界 (VLB) 简化为一个关于扩散数据的信噪比的非常简洁的表达式，从而加深了我们对该模型类别的理论理解。利用这一见解，我们证明了文献中提出的几个模型之间的等价性。此外，我们证明了连续时间 VLB 对噪声调度是不变的，除了其端点的信噪比之外。这使我们能够学习一个噪声调度，最大程度地减少所得到的 VLB 估计量的方差，从而加速优化……”
域适应#
GIFT（面向目标的特征逐步插值）#
研究论文
野外逐步域适应：当中间分布不存在时 (Abnar 等人，2021)
“我们关注域适应问题，其中目标是将模型朝向目标分布进行迁移，而不是学习域不变表示。已经证明，在以下两个假设下：(a) 能够访问来自中间分布的样本，以及 (b) 样本被标注了与源分布的变化量，自训练可以成功地应用于逐步迁移的样本以使模型适应目标分布。我们假设 (a) 足以使迭代自训练能够通过利用隐式课程缓慢地将模型适应到目标分布。在 (a) 不成立的情况下，我们观察到迭代自训练效果不佳。我们提出了 GIFT，这是一种通过插值来自源域和目标域的示例表示来创建来自中间分布的虚拟样本的方法……”
泛化#
代理差距最小化改善了锐度感知训练#
研究论文
代理差距最小化改善了锐度感知训练 (J. Zhuang 等人，2022)
“最近提出的锐度感知最小化 (SAM) 通过最小化一个在参数空间邻域内定义的扰动损失来改善泛化能力。然而，我们表明尖锐和扁平的最小值都可能具有较低的扰动损失，这意味着 SAM 并不总是偏好扁平的最小值。相反，我们定义了一个代理差距，一个在邻域半径（用于导出扰动损失）较小时等效于局部最小值处海森矩阵的主特征值的度量。代理差距易于计算，并且在训练期间可以轻松地直接最小化。基于以上观察，我们提出了代理差距引导的锐度感知最小化 (GSAM)，这是一种对 SAM 的改进，其计算开销可忽略不计……”
元学习#
learned_optimization#
GitHub 代码：learned_optimization
研究论文
持久进化策略的展开计算图中的无偏梯度估计 (Vicol 等人，2021)
“我们介绍了一种称为持久进化策略 (PES) 的方法，它将计算图划分为一系列截断展开，并在每次展开后执行基于进化策略的更新步骤。PES 通过在整个展开序列中累积校正项来消除这些截断带来的偏差。PES 允许快速参数更新，具有较低的内存使用量，是无偏的，并且具有合理的方差特性。”
梯度并非全部 (Metz 等人，2021)
“……在本简短报告中，我们讨论了一种常见的基于混沌的故障模式，这种模式出现在各种可微分情况下，从循环神经网络和数值物理模拟到训练学习优化器。我们追溯了这种故障到所研究系统的雅可比矩阵的频谱，并提供了从业人员何时可能会预期这种故障破坏其基于微分的优化算法的标准。”
模型效率#
高效扩展 Transformer 推理#
GitHub 上的代码
研究论文
高效扩展 Transformer 推理 (Pope 等人，2022)
“我们开发了一个用于推理效率的简单分析模型，以根据应用程序需求选择针对 TPU v4 切片优化的最佳多维分区技术。我们将这些技术与一系列低级优化相结合，在 500B+ 参数模型的延迟和模型 FLOPS 利用率 (MFU) 权衡之间实现了新的帕累托边界，性能优于 FasterTransformer 基准测试套件。我们进一步表明，通过适当的分区，多查询注意（即多个查询头共享单个键/值头）的较低内存需求能够扩展到 32 倍更大的上下文长度。”
神经渲染 / NeRF#
可泛化的基于补丁的神经渲染#
研究论文
可泛化的基于补丁的神经渲染 (Suhail 等人，2022)
“……我们提出了一种不同的范式，其中不需要深度特征和 NeRF 式体积渲染。我们的方法能够直接从场景中采样的补丁集合中预测新场景中目标射线の色，而无需进行深度特征和 NeRF 式体积渲染。”
优化#
Amos 优化器 *和* JEstimator#
GitHub 上的代码
-
“……实现了 Amos，一个与 optax 库兼容的优化器，以及 JEstimator，一个轻量级库，它具有类似 tf.Estimator 的接口，用于管理 JAX 中机器学习程序的 T5X 兼容检查点，我们使用它来运行论文中的实验。”
-
研究论文
Amos：一种具有自适应权重衰减的 Adam 式优化器，用于面向模型的尺度 (Tian 和 Parikh，2022)
介绍了 *“Amos，一个与 optax 库兼容的优化器，以及 JEstimator，一个轻量级库，它具有类似 tf.Estimator 的接口，用于管理 JAX 中机器学习程序的 T5X 兼容检查点。”* *“当用于预训练 BERT 变体和 T5 时，Amos 始终比 AdamW 的最先进设置收敛更快，在 <=70% 的训练步骤和时间内实现了更好的验证损失，同时所需的槽变量内存 <=51%。”*
量化#
帕累托最优量化 ResNet 主要为 4 位 *和* AQT：准确量化训练#
GitHub 上的代码
研究论文
帕累托最优量化 ResNet 主要为 4 位 (Abdolrashidi 等人，2021)
“在这项工作中，我们使用 ResNet 作为案例研究，系统地研究了量化对推理计算成本-质量权衡曲线的影响。我们的结果表明，对于每个 bfloat16 ResNet 模型，都存在具有更低成本和更高精度的量化模型；换句话说，bfloat16 计算成本-质量权衡曲线被 4 位和 8 位曲线帕累托支配，主要量化为 4 位的模型产生了最佳的帕累托曲线……我们使用的量化方法针对实用性进行了优化：它只需要很少的调整，并且是根据硬件能力设计的……作为这项工作的一部分，我们贡献了一个用 JAX 编写的量化库……”
强化学习#
来自演示的动作量化的连续控制 (AQuaDem)#
研究论文
来自演示的动作量化的连续控制 (Dadashi 等人，2021)
提出了 *“一个针对具有连续动作空间的问题的新型强化学习 (RL) 框架：来自演示的动作量化 (AQuaDem)。所提出的方法包括从人类演示中学习连续动作空间的离散化。这种离散化方法为每个输入状态返回一组合理的动作（根据演示），从而捕获演示者的先验及其多峰行为。通过离散化动作空间，任何离散动作深度 RL 技术都可以很容易地应用于连续控制问题。实验表明，所提出的方法在 RL 设置中优于最先进的方法，如 SAC，以及在模仿学习设置中优于 GAIL。”*
序列模型 / 模型并行#
T5X：使用 t5x 和 seqio 扩展模型和数据#
-
“T5X 是一个模块化、可组合、研究友好的框架，用于在多个规模上进行高性能、可配置、自助式序列模型（从语言开始）的训练、评估和推理。”
研究论文
T5X：使用 t5x 和 seqio 扩展模型和数据 (Roberts 等人，2022)
“最近基于神经网络的语言模型从扩展训练数据集的大小和模型本身的参数数量中获益匪浅。由于各种因素，扩展可能会很复杂，包括需要在超级计算机集群（例如 TPU）上分布计算、防止数据输入时出现瓶颈，并确保结果可重复。在这项工作中，我们介绍了两个简化这些问题的软件库：t5x 简化了大规模构建和训练大型语言模型的过程，同时保持易用性，而 seqio 为简单创建快速且可重复的训练数据和评估管道提供了基于任务的 API。这些开源库已用于在具有多个 TB 训练数据的训练集上训练具有数百亿个参数的模型。除了这些库之外，我们还发布了用于 T5 类编码器-解码器模型以及 GPT 类仅解码器架构的配置和说明。”
模拟#
Brax - 用于大规模刚体模拟的可微物理引擎#
研究论文
Brax - 用于大规模刚体模拟的可微物理引擎 (Freeman 等人，2021)
“我们展示了 Brax，一个开源库，用于刚体模拟，重点关注加速器上的性能和并行性，用 JAX 编写。我们在受现有强化学习文献启发的一组任务上展示了结果，但在我们的引擎中重新制作了这些任务。此外，我们提供了 PPO、SAC、ES 和直接策略优化的 JAX 重新实现，这些实现与我们的环境一起编译，允许学习算法和环境处理在同一设备上发生，并且在加速器上无缝扩展。”