[论文解读] Towards understanding how attention mechanism works in deep learning

Tianyu Ruan, Shihua Zhang|arXiv (Cornell University)|Dec 24, 2024

Neural Networks and Applications被引用 4

一句话总结

本文揭示了深度学习中的自注意力机制作为基于流形学习和扩散原理的可学习、自适应相似性计算过程，其收敛于漂移-扩散过程，在特定条件下可进一步收敛为热方程。本文提出了一种新型机制——度量注意力（metric-attention），通过将度量学习整合到注意力机制中，利用可学习的伪度量，显著提升了训练效率、准确率与鲁棒性。

ABSTRACT

Attention mechanism has been extensively integrated within mainstream neural network architectures, such as Transformers and graph attention networks. Yet, its underlying working principles remain somewhat elusive. What is its essence? Are there any connections between it and traditional machine learning algorithms? In this study, we inspect the process of computing similarity using classic metrics and vector space properties in manifold learning, clustering, and supervised learning. We identify the key characteristics of similarity computation and information propagation in these methods and demonstrate that the self-attention mechanism in deep learning adheres to the same principles but operates more flexibly and adaptively. We decompose the self-attention mechanism into a learnable pseudo-metric function and an information propagation process based on similarity computation. We prove that the self-attention mechanism converges to a drift-diffusion process through continuous modeling provided the pseudo-metric is a transformation of a metric and certain reasonable assumptions hold. This equation could be transformed into a heat equation under a new metric. In addition, we give a first-order analysis of attention mechanism with a general pseudo-metric function. This study aids in understanding the effects and principle of attention mechanism through physical intuition. Finally, we propose a modified attention mechanism called metric-attention by leveraging the concept of metric learning to facilitate the ability to learn desired metrics more effectively. Experimental results demonstrate that it outperforms self-attention regarding training efficiency, accuracy, and robustness.

研究动机与目标

• 揭示深度学习中注意力机制背后的数学与物理原理。
• 建立注意力机制与经典机器学习算法（如流形学习、聚类和k-NN）之间的联系。
• 将注意力机制形式化为由偏微分方程（PDE）控制的连续动力系统。
• 通过整合度量学习原理，提出一种新型注意力变体——度量注意力，以提升训练效率与鲁棒性。
• 对一般伪度量下的注意力机制进行一阶分析，阐明其在学习到的伪度量下可被解释为最近邻更新。

提出的方法

• 将注意力机制分解为可学习的伪度量函数与基于相似性的信息传播过程。
• 采用连续建模方法，证明当伪度量为度量的变换且在合理假设下，自注意力机制收敛于漂移-扩散过程。
• 在新度量下，将漂移-扩散过程转化为热方程，从而赋予注意力机制以物理直觉与可解释性。
• 对一般伪度量函数进行一阶分析，表明注意力更新等价于在学习到的伪度量定义的空间中的最近邻更新。
• 通过将度量学习嵌入注意力框架，提出度量注意力机制，实现从任务标签端到端直接学习伪度量。
• 理论分析基于PDE与测度论动力系统，与以往聚焦于常微分方程（ODE）或流映射的研究形成鲜明区别。

实验结果

研究问题

• RQ1 自注意力机制如何与涉及相似性计算的经典机器学习算法相关联？
• RQ2 注意力机制的连续极限行为是什么？能否由偏微分方程（PDE）描述？
• RQ3 注意力机制能否被解释为扩散过程？在何种条件下其会收敛为热方程？
• RQ4 使用可学习伪度量相较于标准自注意力机制，如何提升性能？
• RQ5 在学习到的伪度量下，注意力更新的物理与几何解释是什么？

主要发现

• 在伪度量为度量的变换这一假设下，自注意力机制被形式化证明收敛于漂移-扩散过程。
• 在相同假设下，漂移-扩散过程可通过新度量变换为热方程，从而为注意力机制提供作为扩散过程的物理解释。
• 一阶分析表明，注意力更新等价于在学习到的伪度量定义的空间中的最近邻更新。
• 所提出的度量注意力机制在各项实验中均优于标准自注意力机制，在训练效率、准确率与鲁棒性方面表现更优。
• 本研究通过共享的相似性计算与信息传播原理，建立了注意力机制与经典算法（如k-NN、k-均值与扩散图）之间牢固的概念与数学联系。
• 研究结果表明，训练注意力模块等价于学习一个有用的伪度量，其目标与度量学习一致，但因隐式、任务驱动的优化而具备更高的灵活性。