本周论文包括谷歌放出下一代 AI 架构 Pathways 论文;何恺明组最新论文等研究。
目录
Training-free Transformer Architecture Search
PATHWAYS: ASYNCHRONOUS DISTRIBUTED DATAFLOW FOR ML
Autoregressive Image Generation using Residual Quantization
Exploring Plain Vision Transformer Backbones for Object Detection
FewNLU: Benchmarking State-of-the-Art Methods for Few-Shot Natural Language Understanding
LANGUAGE MODELING VIA STOCHASTIC PROCESSES
DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery
ArXiv Weekly Radiostation:NLP、CV、ML 更多精选论文(附音频)
论文 1:Training-free Transformer Architecture Search
- 作者:Qinqin Zhou 、 Kekai Sheng 等
- 论文链接:https://arxiv.org/pdf/2203.12217.pdf
摘要:在近期的一篇论文《Training-free Transformer Architecture Search》中,来自腾讯优图实验室、厦门大学、鹏城实验室等结构的研究者回顾近些年 NAS 领域的进展,并注意到:为了提高搜索效率,研究社区提出了若干零成本代理(zero-cost proxy)的评估指标(如 GraSP、TE-score 和 NASWOT)。这些方法让我们能够在无需训练的条件下就能评估出不同 CNN 结构的排序关系,从而极大程度上节省计算成本。
从技术上来说,一个典型的 CNN 模型主要由卷积模块组成,而一个 ViT 模型主要由多头注意力模块(MSA)和多层感知机模块(MLP)组成。这种网络结构上的差异会让现有的、在 CNN 搜索空间上验证有效的零成本代理无法保证其在 ViT 搜索空间上模型评估效果(见下图 1)。
因此,研究一种更适合 ViT 结构评估、有利于 TAS 训练效率的零成本代理指标是有必要且值得探索的。这一问题也将促使研究者进一步研究和更好地理解 ViT 结构,从而设计一种有效的、无需训练的 TAS 搜索算法。
图 1. (a)研究者采样的 1000 个 ViT 模型的参数量和效果分布。(b-e)在 CNN 搜索空间效果好的 zero-cost proxy 方法并不适用于 ViT 搜索空间。(f)他们的 DSS-indicator 更适合用来评估不同的 ViT 模型。
研究者设计了一个有效且高效的零代价代理评估指标 DSS-indicator(下图 2),并基于此设计了一个包含模块化策略的无训练 Transformer 结构搜索算法(Transformer Architecture Search,TF-TAS)。
图 2. 方法的整体框架图。
具体来说,DSS-indicator 通过计算 MSA 的突触多样性和 MLP 的突触显著性来得到 ViT 结构的评价分数。这是学术界首次提出基于 MSA 的突触多样性和 MLP 的突触显著性来作为评价 ViT 结构的代理评估指标。而且需要注意的是,TF-TAS 与搜索空间设计和权值共享策略是正交的。因此,可以灵活地将 TF-TAS 与其他 ViT 搜索空间或 TAS 方法相结合,进一步提高搜索效率。与人工设计的 ViT 和自动搜索的 ViT 相比,研究者设计的 TF-TAS 实现了具有竞争力的效果,将搜索过程从 24 GPU 天数缩短到不到 0.5 GPU 天数,大约快 48 倍。
推荐:CVPR 2022 Oral ,腾讯优图 & 厦门大学提出无需训练的 ViT 结构搜索算法。
论文 2:PATHWAYS: ASYNCHRONOUS DISTRIBUTED DATAFLOW FOR ML
- 作者:Paul Barham 、 Aakanksha Chowdhery 等
- 论文链接:https://arxiv.org/pdf/2203.12533.pdf
摘要:Jeff Dean 等人去年提出了一种名叫「Pathways」的通用 AI 架构。他介绍说,Pathways 旨在用一个架构同时处理多项任务,并且拥有快速学习新任务、更好地理解世界的能力。
该架构的特点可以概括为:
- 能够训练一个模型来做成千上万件事情;
- 当前模型只注重一种感官,Pathways 可做到多种;
- 当前模型密集且效率低下,Pathways 会把模型变得稀疏而高效。
在发布想法大半年之后,Jeff Dean 终于公布了 Pathways 的论文,其中包含很多技术细节。
论文写道,PATHWAYS 使用了异步算子的一个分片数据流图(sharded dataflow graph),这些算子消耗并产生 futures,并在数千个加速器上高效地对异构并行计算进行 gang-schedule,同时在它们专用的 interconnect 上协调数据传输。PATHWAYS 使用了一种新的异步分布式数据流设计,它允许控制平面并行执行,尽管数据平面中存在依赖关系。这种设计允许 PATHWAYS 采用单控制器模型,从而更容易表达复杂的新并行模式。
实验结果表明,当在 2048 个 TPU 上运行 SPMD(single program multiple data)计算时,PATHWAYS 的性能(加速器利用率接近 100%)可以媲美 SOTA 系统,同时吞吐量可媲美跨越 16 个 stage 或者被分割成两个通过数据中心网络连接的加速器岛的 Transformer 模型的 SPMD 案例。
PATHWAYS 构建在先前的系统的基础上,包括用于表征和执行 TPU 计算的 XLA (TensorFlow, 2019)、用于表征和执行分布式 CPU 计算的 TensorFlow 图和执行器 (Abadi et al., 2016),以及包括 JAX (Bradbury et al., 2016) 在内的 Python 编程框架 (Bradbury et al., 2018) 和 TensorFlow API。利用这些构建块,PATHWAYS 在兼顾协调性的同时,仅用最少的代码更改就能运行现有的 ML 模型。
推荐:谷歌下一代 AI 架构 Pathways 论文解读。
论文 3:Autoregressive Image Generation using Residual Quantization
- 作者:Doyup Lee、 Chiheon Kim 等
- 论文链接:https://arxiv.org/pdf/2203.01941.pdf
摘要:来自 Kakao Brain 、韩国浦项科技大学的研究者提出了一种残差量化 VAE (RQ-VAE) 方法,它使用残差量化 (RQ) 来精确逼近特征图并降低其空间分辨率。RQ 没有增加编码簿大小,而是使用固定大小的编码簿以从粗到细的方式递归量化特征图。在 RQ 的 D 次迭代之后,特征图表示为 D 个离散编码的堆叠图。由于 RQ 可以组成与编码簿大小一样多的向量,因此 RQ-VAE 可以精确地逼近特征图,同时保留编码图像的信息,而无需庞大的编码簿。由于精确的近似,RQ-VAE 可以比以前的研究 [14,37,45] 进一步降低量化特征图的空间分辨率。例如, RQ-VAE 可以使用 8×8 分辨率的特征图对 256×256 图像进行 AR 建模。该论文已被 CVPR'22 接收。
此外,该研究还提出了 RQ-Transformer 来预测 RQ-VAE 提取的编码。对于 RQ-Transformer 的输入,该研究首先将 RQ-VAE 中的量化特征映射转换为特征向量序列;然后,RQ-Transformer 预测下一个 D 编码,以估计下一个位置的特征向量。由于 RQ-VAE 降低了特征图的分辨率,RQ-Transformer 可以显着降低计算成本并轻松学习输入的远程交互。该研究还为 RQ-Transformer 提出了两种训练技术,软标签(soft labeling)和用于 RQ-VAE 编码的随机采样。通过解决 AR 模型训练中的曝光偏差(exposure bias)进一步提高了 RQ-Transformer 的性能。
值得一提的是,该研究近日发布了在 30M 文本图像对上训练的 3.9B 参数的 RQ-Transformer 。据了解,这是公共可用模型中最大的文本到图像 (T2I) 模型。不过这一结果没有出现在该论文中。具体内容可参考 GitHub。
代码地址:https://github.com/kakaobrain/rq-vae-transformer
研究者提出了用于图像 AR 建模的 RQ-VAE 和 RQ-Transformer 两阶段框架,如下图 2 所示。RQ-VAE 使用编码簿将图像表示为 D 个离散码的堆叠图。然后,使用 RQ-Transformer 自回归预测下一个空间位置的下一个 D 码。他们还解释了使用 RQ-Transformer 解决 AR 模型训练中的曝光偏差问题。