[Paper Review] Towards understanding how attention mechanism works in deep learning

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

Neural Networks and Applications4 citations

TL;DR

This paper reveals that the self-attention mechanism in deep learning operates as a learnable, adaptive similarity computation process rooted in manifold learning and diffusion principles, converging to a drift-diffusion process and, under certain conditions, a heat equation. It introduces metric-attention, a novel mechanism that enhances training efficiency, accuracy, and robustness by integrating metric learning into attention via a learnable pseudo-metric.

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.

Motivation & Objective

• To uncover the mathematical and physical principles underlying attention mechanisms in deep learning.
• To establish connections between attention mechanisms and classical machine learning algorithms such as manifold learning, clustering, and k-NN.
• To formalize the attention mechanism as a continuous dynamical system governed by partial differential equations (PDEs).
• To propose a new attention variant, metric-attention, by integrating metric learning principles to improve training efficiency and robustness.
• To provide a first-order analysis of attention with general pseudo-metrics and clarify its interpretation as nearest-neighbor updates under a learned metric.

Proposed method

• The attention mechanism is decomposed into a learnable pseudo-metric function and a similarity-based information propagation process.
• Continuous modeling is used to show that self-attention converges to a drift-diffusion process when the pseudo-metric is a transformation of a metric and under reasonable assumptions.
• The drift-diffusion process is transformed into a heat equation under a new metric, enabling physical intuition and interpretability.
• A first-order analysis is conducted for general pseudo-metric functions, showing that attention updates are equivalent to nearest-neighbor updates under the learned pseudo-metric.
• The metric-attention mechanism is proposed by embedding metric learning into the attention framework, allowing direct end-to-end learning of the pseudo-metric from task labels.
• Theoretical analysis is grounded in PDEs and measure-theoretic dynamics, distinguishing it from prior work that focused on ODEs or flow maps.

Experimental results

Research questions

• RQ1 How does the self-attention mechanism relate to classical machine learning algorithms involving similarity computation?
• RQ2 What is the continuous limit behavior of the attention mechanism, and can it be described by a PDE?
• RQ3 Can the attention mechanism be interpreted as a diffusion process, and under what conditions does it converge to a heat equation?
• RQ4 How does the use of a learnable pseudo-metric improve performance compared to standard self-attention?
• RQ5 What is the physical and geometric interpretation of attention updates under a learned pseudo-metric?

Key findings

• The self-attention mechanism is formally shown to converge to a drift-diffusion process under the assumption that the pseudo-metric is a transformation of a metric.
• Under the same assumptions, the drift-diffusion process can be transformed into a heat equation under a new metric, providing a physical interpretation of attention as a diffusion process.
• The first-order analysis reveals that attention updates are equivalent to nearest-neighbor updates in the space defined by the learned pseudo-metric.
• The proposed metric-attention mechanism outperforms standard self-attention in training efficiency, accuracy, and robustness across experiments.
• The study establishes a strong conceptual and mathematical link between attention mechanisms and classical algorithms such as k-NN, k-means, and diffusion maps through shared principles of similarity computation and information propagation.
• The findings suggest that training attention blocks is equivalent to learning a useful pseudo-metric, aligning the objective with metric learning but with greater flexibility due to implicit, task-driven optimization.