🤖 AI Summary
Self-attention lacks an interpretable dynamical systems characterization, hindering theoretical understanding of its relationship with classical machine learning paradigms such as similarity computation and information propagation.
Method: We establish, for the first time, that self-attention converges to a drift-diffusion process in the continuous limit, yielding an equivalent heat equation; based on this, we propose metric-attention—a novel attention mechanism grounded in learnable pseudo-metrics that unifies metric learning and attention modeling. Our approach integrates continuous-time modeling, manifold analysis, and diffusion PDE derivation.
Contribution/Results: Experiments demonstrate that metric-attention significantly outperforms standard self-attention in training efficiency, accuracy, and adversarial robustness. This work provides a theoretically rigorous yet practically superior paradigm for attention, bridging foundational principles of dynamical systems, geometry, and statistical learning.
📝 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.