The Brownian motion in the transformer model

TL;DR

This paper analyzes the transformer model's multi-head self-attention as a Brownian motion process, revealing new insights into token dynamics and proposing a second-order optimizer for improved training.

Contribution

It introduces a novel perspective of MHSA as a Brownian motion and develops a second-order optimizer based on this insight, advancing understanding of transformer training dynamics.

Findings

Tokens mapped to high-dimensional hyper-sphere

Attention acts as Markov transition matrix on the sphere

Training process corresponds to Brownian motion

Abstract

Transformer is the state of the art model for many language and visual tasks. In this paper, we give a deep analysis of its multi-head self-attention (MHSA) module and find that: 1) Each token is a random variable in high dimensional feature space. 2) After layer normalization, these variables are mapped to points on the hyper-sphere. 3) The update of these tokens is a Brownian motion. The Brownian motion has special properties, its second order item should not be ignored. So we present a new second-order optimizer(an iterative K-FAC algorithm) for the MHSA module. In some short words: All tokens are mapped to high dimension hyper-sphere. The Scaled Dot-Product Attention $softmax(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d}})$ is just the Markov transition matrix for the random walking on the sphere. And the deep learning process would learn proper kernel function to get proper positions of these tokens. The training process in the MHSA module corresponds to a Brownian motion worthy of further study.