The Flow of Attention

2026-06-17 · Source: Towards AI - Medium · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Mathematics & Computational Sciences · Depth: Expert, extended

Summary

An input prompt to a large language model can be visualized as a cloud of N token embeddings in a d-dimensional vector space E. As a Transformer processes this prompt through L layers, the cloud reconfigures itself, with each token's position shifting to reflect its contextual relationships. This article introduces the EMF model, which reduces attention to two operators per layer: a bilinear form M for relevance scoring and a linear operator F for content extraction, both acting in space E. The layer-wise updates are additive, forming a "residual stream" where each token traces a trajectory. This collective motion is characterized as a transport flow in Wasserstein space, bearing structural resemblances to a Wasserstein Gradient Flow, even though it's not a strict WGF due to discrete steps, changing landscapes, and asymmetric attention. The flow organizes tokens into semantically related clusters, and the final position of the last token in E determines the next-token prediction.

Key takeaway

For AI Scientists analyzing Transformer behavior, understanding attention as a Wasserstein transport flow reveals how contextual meaning emerges. You should investigate the layer-wise dynamics of token embeddings, recognizing that additive updates preserve semantic geometry. This perspective helps interpret why tokens cluster and how the final token's position drives prediction, guiding efforts in model interpretability and design.

Key insights

The Transformer's attention mechanism reconfigures token embeddings as a coupled particle flow in Wasserstein space, organizing contextual meaning.

Principles

• Additive updates preserve semantic geometry across Transformer layers.
• Attention's core mechanism is a two-operator (M, F) system in embedding space.
• Softmax normalization prevents attention weights from collapsing.

Method

The EMF model reduces attention to a bilinear form M (relevance scoring) and a linear operator F (content extraction) acting in a single embedding space E, with contextual updates as additive displacements.

In practice

• Visualize token embeddings as a dynamic cloud in high-dimensional space.
• Analyze residual stream trajectories to understand contextualization.
• Consider the impact of additive updates on semantic preservation.

Topics

• Transformer Architecture
• Attention Mechanism
• Embedding Space Geometry
• Wasserstein Distance
• Residual Stream
• Large Language Models

Best for: Research Scientist, AI Scientist, Machine Learning Engineer

Related on AIssential

• Transformer Architecture: Embedding se Output tak AI ke andar actually kya hota hai Token… — The Transformer is a multi-stage pipeline that processes text into context-aware representations for next token prediction.
• The Secret Sauce of Modern AI: Self-Attention Explained Like You’re Hearing It for the First Time — Self-Attention creates contextualized word embeddings by learning relationships between all words in a sequence.
• Attention Is All You Need: But What Does That Actually Mean? — Transformers build meaning by calculating weighted relationships between all words in a sequence, not by human-like understanding.
• Contribution Weights: A Geometrical Analysis of Self-Attention Transformers — Contribution Weights offer a geometrically-informed, more faithful measure of token importance and reveal active roles for attention sinks in LLMs.
• How Aara Reads: The Secret Language Beneath the Words — LLMs translate human language into numerical vectors via tokenization, embedding, and attention to process meaning.
• The Learned Kernel: How Transformers Actually Do In-Context Learning — Transformers excel at in-context learning by dynamically learning optimal similarity functions through their attention mechanisms.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Towards AI - Medium.