Clustering in Causal Attention Masking
Keywords: Transformers, causal attention, continuous-time interacting particle systems, clustering
TL;DR: Theoretical study of clustering phenomena in evolution of tokens representations as they propagate through layers of a causal transformer.
Abstract: This work presents a modification of the self-attention dynamics proposed in Geshkovski et al to better reflect the practically relevant, causally masked attention used in transformer architectures for generative AI. This modification translates into an interacting particle system that cannot be interpreted as a mean-field gradient flow. Despite this loss of structure, we significantly strengthen the results of Geshkovski et al in this context: While previous rigorous results focused on cases where all three matrices (key, query, and value) were scaled identities, we prove asymptotic convergence to a single cluster for arbitrary key-query matrices and value matrix equal to the identity.
Additionally, we establish a connection to the classical R\'enyi parking problem from combinatorial geometry to make initial theoretical steps towards demonstrating the existence of meta-stable states.
Primary Area: Learning theory
Submission Number: 13344
Loading