Title: A Two-Character Change in Transformer Architecture Promotes Ideal Token Geometry
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Transformers, Neural Collapse, Attention

Abstract: We hypothesize that in the optimal geometric configuration of token embeddings for transformers, tokens should collapse to single points according to their classes, and these points themselves should exhibit Neural Collapse. We study whether current transformers achieve this configuration through principal component projections, cosine similarity measurements, analysis of variance on token embeddings, and Neural Collapse measurements, and find that they fall far short of the conjectured ideal. To address this, we introduce a simple modification to attention that brings token embeddings markedly closer to the conjectured configuration and yields consistent performance improvements across benchmarks.

TL;DR: We conjecture that ideal transformers collapse tokens, but we find that current transformers don't collapse tokens well, so we modify attention to achieve more collapse and better performance.