[论文解读] Untangling tradeoffs between recurrence and self-attention in neural networks

Giancarlo Kerg, Bhargav Kanuparthi+4 more|arXiv (Cornell University)|Jun 16, 2020

Neural Networks and Applications参考文献 32被引用 5

一句话总结

本文对自注意力递归网络中的梯度传播进行了形式化分析，证明自注意力通过建立梯度范数的边界，缓解了梯度消失问题。提出了一种受记忆巩固启发的相关性筛选机制，实现了稀疏、可扩展的自注意力，在保持长期序列训练稳定性的同时平衡了性能与计算成本。

ABSTRACT

Attention and self-attention mechanisms, are now central to state-of-the-art deep learning on sequential tasks. However, most recent progress hinges on heuristic approaches with limited understanding of attention's role in model optimization and computation, and rely on considerable memory and computational resources that scale poorly. In this work, we present a formal analysis of how self-attention affects gradient propagation in recurrent networks, and prove that it mitigates the problem of vanishing gradients when trying to capture long-term dependencies by establishing concrete bounds for gradient norms. Building on these results, we propose a relevancy screening mechanism, inspired by the cognitive process of memory consolidation, that allows for a scalable use of sparse self-attention with recurrence. While providing guarantees to avoid vanishing gradients, we use simple numerical experiments to demonstrate the tradeoffs in performance and computational resources by efficiently balancing attention and recurrence. Based on our results, we propose a concrete direction of research to improve scalability of attentive networks.

研究动机与目标

• 正式分析自注意力如何影响递归网络中的梯度传播，特别是缓解梯度消失问题。
• 识别并量化递归与自注意力在梯度稳定性、计算成本和内存使用方面的权衡。
• 开发一种可扩展的稀疏自注意力机制，以在长期依赖关系中保持稳定的梯度流动。
• 为自注意力递归模型中的梯度范数提供理论保证，从而实现高效架构的合理设计。
• 在简单任务上通过实证验证所提出方法，展示其在可扩展性与性能权衡方面的改进。

提出的方法

• 推导出自注意力递归网络中梯度范数的理论边界，证明自注意力通过防止长序列中的梯度消失，稳定了梯度流动。
• 提出一种相关性筛选机制，动态选择用于注意力的相关过去状态，将计算图规模从序列长度的平方级降低至线性级。
• 将注意力机制建模为一种软内存访问系统，其中相关性分数通过过去隐藏状态的可学习函数预测得出。
• 引入两个超参数：ν（用于相关性预测的过去状态数量）和ρ（实际参与注意力的过去状态数量），以控制内存使用与梯度稳定性的权衡。
• 采用半参数化RNN架构，结合外部记忆和软注意力机制，实现在不使用完整递归的情况下动态访问相关历史状态。
• 通过去噪、复制、序列MNIST和PTB任务的数值实验，评估不同模型的梯度传播、训练稳定性和GPU使用情况。

实验结果

研究问题

• RQ1 自注意力如何影响递归网络中的梯度传播？我们能否形式化地界定梯度范数以防止梯度消失？
• RQ2 在结合递归与自注意力时，计算复杂度（内存与FLOPs）与梯度稳定性之间的权衡是什么？
• RQ3 相关性筛选机制能否在保持长序列稳定训练的同时降低自注意力的计算成本？
• RQ4 超参数ν（相关性预测窗口）和ρ（注意力窗口大小）如何影响梯度传播与模型性能？
• RQ5 通过记忆巩固机制实现的稀疏注意力在多大程度上能维持性能，同时减少序列长度的二次方级缩放？

主要发现

• 增大ρ（注意力窗口大小）可提升梯度稳定性，当ρ ≥ 15时可保持稳定的梯度范数，而ρ = 3会导致训练不稳定。
• 改变ν（相关性预测窗口）对梯度稳定性影响极小，尽管ν过小（如ν = 3）会因相关性估计不准确而轻微降低性能。
• 所提出的RelRNN模型结合相关性筛选机制，在序列长度T = 1000以内的复制与去噪任务中达到100%准确率，而标准LSTM在T > 500后即失效。
• 在PTB语言建模任务中，RelRNN模型在显著降低内存与FLOPs消耗的同时，保持了具有竞争力的性能，优于全注意力模型。
• 通过稀疏化注意力，将计算图规模从O(T²)降低至O(T)，实现了长序列的可扩展训练，且无梯度爆炸问题。
• 实证结果表明，性能与计算成本之间的权衡主要由ρ决定，ν的影响为次要因素，表明ρ是平衡效率与稳定性的关键超参数。