Attention (as Discrete-Time Markov) Chains

TL;DR

This paper presents a novel interpretation of attention mechanisms in transformers as discrete-time Markov chains, enabling new insights, improved segmentation, and enhanced image generation through the analysis of metastable states and TokenRank.

Contribution

It introduces a Markov chain perspective on attention, allowing analysis of token importance and attention dynamics, leading to state-of-the-art zero-shot segmentation and improved image generation.

Findings

Metastable states correspond to semantically similar regions.

TokenRank improves image generation quality and diversity.

The framework enhances segmentation performance on benchmarks.

Abstract

We introduce a new interpretation of the attention matrix as a discrete-time Markov chain. Our interpretation sheds light on common operations involving attention scores such as selection, summation, and averaging in a unified framework. It further extends them by considering indirect attention, propagated through the Markov chain, as opposed to previous studies that only model immediate effects. Our key observation is that tokens linked to semantically similar regions form metastable states, i.e., regions where attention tends to concentrate, while noisy attention scores dissipate. Metastable states and their prevalence can be easily computed through simple matrix multiplication and eigenanalysis, respectively. Using these lightweight tools, we demonstrate state-of-the-art zero-shot segmentation. Lastly, we define TokenRank -- the steady state vector of the Markov chain, which measures global token importance. We show that TokenRank enhances unconditional image generation, improving both quality (IS) and diversity (FID), and can also be incorporated into existing segmentation techniques to improve their performance over existing benchmarks. We believe our framework offers a fresh view of how tokens are being attended in modern visual transformers.