[论文解读] OmniNet: Omnidirectional Representations from Transformers

Yi Wei Daniel Tay, Mohammad Dehghani+7 more|arXiv (Cornell University)|Mar 1, 2021

Domain Adaptation and Few-Shot Learning参考文献 26被引用 9

一句话总结

OmniNet 引入了一种新颖的 Transformer 架构，通过基于元学习器的全向注意力机制，使每个标记能够关注网络中所有层的隐藏表征。通过利用 Performer 等高效注意力机制以及低秩近似，OmniNet 在语言建模、机器翻译、长序列任务和图像识别任务上均达到最先进性能，在少样本图像分类任务上相比 ViT 最高提升 3%。

ABSTRACT

This paper proposes Omnidirectional Representations from Transformers (OmniNet). In OmniNet, instead of maintaining a strictly horizontal receptive field, each token is allowed to attend to all tokens in the entire network. This process can also be interpreted as a form of extreme or intensive attention mechanism that has the receptive field of the entire width and depth of the network. To this end, the omnidirectional attention is learned via a meta-learner, which is essentially another self-attention based model. In order to mitigate the computationally expensive costs of full receptive field attention, we leverage efficient self-attention models such as kernel-based (Choromanski et al.), low-rank attention (Wang et al.) and/or Big Bird (Zaheer et al.) as the meta-learner. Extensive experiments are conducted on autoregressive language modeling (LM1B, C4), Machine Translation, Long Range Arena (LRA), and Image Recognition. The experiments show that OmniNet achieves considerable improvements across these tasks, including achieving state-of-the-art performance on LM1B, WMT'14 En-De/En-Fr, and Long Range Arena. Moreover, using omnidirectional representation in Vision Transformers leads to significant improvements on image recognition tasks on both few-shot learning and fine-tuning setups.

研究动机与目标

• 解决标准 Transformer 限制感受野仅限于水平（序列）和局部（层内）上下文的局限性。
• 探究将注意力扩展至包含所有层中所有标记是否能提升表征学习和模型性能。
• 通过高效注意力机制和分块元学习，缓解全网络注意力带来的高计算成本。
• 评估全向表征对语言建模、机器翻译、长程推理和图像识别等多种任务的影响。
• 分析 OmniNet 的性能-计算权衡，并证明其具有帕累托最优性。

提出的方法

• OmniNet 使用一个独立的自注意力模块——元学习器——来处理网络中所有层和所有标记的隐藏表征。
• 元学习器采用高效的注意力机制，如基于核的注意力（例如 Performer）、低秩注意力或分块稀疏注意力（例如 Big Bird），以降低计算成本。
• 模型在分块的层组中应用全向注意力（例如 P=4,6,12），其中元学习器一次仅作用于 P 个连续层的组合。
• 全向注意力之后，通过池化操作从所有层中为每个标记选择最相关的值，从而增强表征的丰富性。
• 该方法同时应用于自回归语言模型和视觉 Transformer（ViT），实现跨层知识融合。
• 该架构采用标准目标端到端训练，且元学习器可微，支持梯度在所有层中反向传播。

实验结果

研究问题

• RQ1 将 Transformer 的感受野从序列和层维度扩展至包含所有层中所有标记，是否能提升模型性能？
• RQ2 如何在保留表征增益的同时，缓解全网络注意力带来的二次方计算成本？
• RQ3 整合来自多个层次化层的表征是否能提升泛化能力，尤其是在少样本学习和长序列建模中？
• RQ4 OmniNet 在准确率和效率方面与最先进模型相比，在多样化任务上的表现如何？
• RQ5 在性能与计算成本之间取得平衡时，元学习器的最佳分块大小是多少？

主要发现

• OmniNet 在 LM1B 语言建模基准上达到最先进结果，优于现有模型如 Transformer-XL。
• 在 WMT’14 英德和英法翻译任务中，OmniNet 超过 60 层的 ADMIN 模型（当时最先进 Transformer 模型）。
• 在 Long Range Arena 基准上，OmniNet 相较 Performer 提升 +8.9%，相较原始 Transformer 提升 +2.6%。
• 在 ImageNet 上，OmniNet 在少样本线性探测和微调设置下，将 ViT 性能提升最多约 3%。
• 当分块大小 P=6 时，OmniNet 在性能与计算成本之间达到最佳权衡，其准确率-浮点运算量（FLOP）比优于 P=4 和 P=12。
• 可视化结果表明，池化机制主要从早期层提取信息，表明细粒度特征有效传播至最终表征。