Clustering in Causal Attention Masking
This work provides incremental theoretical insights into transformer behavior for researchers in machine learning theory, focusing on causal attention mechanisms.
The paper tackles the problem of analyzing self-attention dynamics in causally masked transformers, proving asymptotic convergence to a single cluster for arbitrary key-query matrices with identity value matrices, and connects this to the Rényi parking problem to explore meta-stable states.
This work presents a modification of the self-attention dynamics proposed by Geshkovski et al. (arXiv:2312.10794) 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. (arXiv:2312.10794) 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 a value matrix equal to the identity. Additionally, we establish a connection to the classical Rényi parking problem from combinatorial geometry to make initial theoretical steps towards demonstrating the existence of meta-stable states.