[论文解读] Infinite attention: NNGP and NTK for deep attention networks

Jiri Hron, Yasaman Bahri+2 more|arXiv (Cornell University)|Jun 18, 2020

Gaussian Processes and Bayesian Inference被引用 29

一句话总结

本文证明了多头注意力网络在无限宽度极限下收敛于高斯过程（NNGP），并表现出神经正切核（NTK）行为，解决了单头设置下注意力机制非高斯输出带来的理论挑战。论文提出了改进的注意力机制和位置编码策略，显著提升了有限宽度与无限宽度模型的性能，在不使用数据增强的情况下于CIFAR-10数据集上实现了当前最优结果，采用高斯过程推理。

ABSTRACT

There is a growing amount of literature on the relationship between wide neural networks (NNs) and Gaussian processes (GPs), identifying an equivalence between the two for a variety of NN architectures. This equivalence enables, for instance, accurate approximation of the behaviour of wide Bayesian NNs without MCMC or variational approximations, or characterisation of the distribution of randomly initialised wide NNs optimised by gradient descent without ever running an optimiser. We provide a rigorous extension of these results to NNs involving attention layers, showing that unlike single-head attention, which induces non-Gaussian behaviour, multi-head attention architectures behave as GPs as the number of heads tends to infinity. We further discuss the effects of positional encodings and layer normalisation, and propose modifications of the attention mechanism which lead to improved results for both finite and infinitely wide NNs. We evaluate attention kernels empirically, leading to a moderate improvement upon the previous state-of-the-art on CIFAR-10 for GPs without trainable kernels and advanced data preprocessing. Finally, we introduce new features to the Neural Tangents library (Novak et al., 2020) allowing applications of NNGP/NTK models, with and without attention, to variable-length sequences, with an example on the IMDb reviews dataset.

研究动机与目标

• 严格刻画具有多头注意力机制的深层神经网络在无限宽度极限下的行为。
• 解决将高斯过程与神经正切核框架扩展至注意力层时面临的理论挑战，这些挑战源于其不满足标准可交换性或条件假设。
• 提出改进的注意力机制与位置编码策略，以增强有限宽度与无限宽度网络的性能。
• 通过实验评估所提出的注意力核，并在CIFAR-10等基准数据集上展示使用高斯过程推理时性能的提升。
• 扩展Neural Tangents库以支持可变长度序列，使NNGP/NTK应用能够扩展至文本等序列数据。

提出的方法

• 在$d^{-1/2}$与$d^{-1}$缩放下，推导多头注意力网络的NNGP与NTK核，证明其在无限宽度极限下收敛于高斯过程。
• 提出一种采用$d^{-1}$缩放与基于插值的协方差更新机制的改进注意力机制，以在保持模型表达力的同时实现理论分析。
• 提出结构化位置编码，显著提升无限宽度设定下的泛化能力与性能，即使在结合NNGP/NTK推理时亦有效。
• 将理论框架适配至残差连接与层归一化，分析其对核收敛与模型行为的影响。
• 提出一种新型核更新规则，整合可学习位置嵌入的梯度贡献，通过混合插值机制修改NTK。
• 发布对Neural Tangents库的代码扩展，支持可变长度序列，使端到端NNGP/NTK推理在IMDb等数据集上成为可能。

实验结果

研究问题

• RQ1 在宽神经网络中，多头注意力是否在无限宽度极限下收敛于高斯过程？若收敛，其条件为何？
• RQ2 为何标准的NNGP与NTK推导技术在单头注意力中失效？如何将其适配至多头架构？
• RQ3 位置编码与层归一化如何影响无限宽度注意力网络中的核行为与泛化性能？
• RQ4 采用$d^{-1}$缩放的改进注意力机制能否在保持模型表达力的同时实现严谨的NNGP/NTK分析？
• RQ5 将推导出的NNGP与NTK核应用于CIFAR-10等真实世界基准数据集时，不使用数据增强，能获得多大的性能提升？

主要发现

• 多头注意力网络在无限宽度极限下收敛于高斯过程，而单头注意力则不会，原因在于共享注意力权重引发的统计依赖性。
• 本文严格推导了多头注意力在$d^{-1/2}$与$d^{-1}$缩放下的NNGP与NTK核，将理论框架成功扩展至注意力机制。
• 采用$d^{-1}$缩放与基于插值的协方差更新机制的改进注意力机制，既保持了模型表达力，又实现了稳定的核收敛，避免了先前方法中出现的池化类似行为。
• 位置编码在无限宽度极限下显著提升实证性能，即使模型未经过训练，也表明其在增强泛化能力方面具有关键作用。
• 所提方法在CIFAR-10上实现了对先前最先进方法的适度改进，采用高斯过程推理且不依赖数据增强或高级预处理。
• 作者发布了对Neural Tangents库的代码扩展，支持可变长度序列，使NNGP/NTK推理可应用于文本等序列数据，并在IMDb数据集上展示了实际应用效果。