# Towards understanding how attention mechanism works in deep learning

arXiv: https://arxiv.org/abs/2412.18288 · HTML: https://mlsanity.com/paper/2412.18288/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 57 references that are in the map.
- First posted on arXiv 24 Dec 2024.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Multi30K: Multilingual English-German Image Descriptions](https://mlsanity.com/paper/1605.00459/) (2016)
- [The Lipschitz Constant of Self-Attention](https://mlsanity.com/paper/2006.04710/) (2020)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [White-Box Transformers via Sparse Rate Reduction](https://mlsanity.com/paper/2306.01129/) (2023)
- [A Mathematical Theory of Attention](https://mlsanity.com/paper/2007.02876/) (2020)
- [Manifold Fitting under Unbounded Noise](https://mlsanity.com/paper/1909.10228/) (2019)
- [Deep Residual Networks Learn the Geodesic Curve in the Wasserstein Space](https://mlsanity.com/paper/2102.09235/) (2021)

## Built on by (most cited first, 3 of 3)

- [Attention Mechanisms Perspective: Exploring LLM Processing of Graph-Structured Data](https://mlsanity.com/paper/2505.02130/) (2025)
- [MomentKV: Closing the Directional Gap in KV Cache Eviction for Long-Context Inference](https://mlsanity.com/paper/2606.01563/) (2026)
- [Data-Adaptive Mahalanobis Metric Learning for Cross-Head Attention in Transformers](https://mlsanity.com/paper/2610.13261/) (2026)

## Abstract (arXiv)

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.