Rigoni, Nathan
We propose the hypothesis that token embeddings serve not merely as lookup tables but as the query for attention, and introduce No-Q attention to improve performance and reduce parameters.
In transformers, the embedding layer is typically regarded as a simple input transformation, but the authors argue that embeddings critically influence the geometric structure of attention. Prior work either overlooks the role of embeddings or faces issues where Fourier-domain manipulations are unsafe for language models.
We proceed in four stages: (1) Prescribed Fourier Frequency Training (PFFT): Imposing near-Nyquist frequency modes in the embedding gradient to accelerate grokking. (2) Sounding Hammer diagnostic: Discovering that BPE vocabulary gradients are spectrally flat (ρ=0.42), making Fourier steering unsafe, and proposing Natural Ordering Conditions (NOC). (3) Fourier Gradient Projection (FGP): A general gradient-domain tool that dynamically tracks important frequency modes during training. (4) Behavioral weight trajectory analysis: Revealing that all weight matrices (Q, K, V, MLP) inherit the same two-arm trajectory shape from the embedding, motivating No-Q attention (removing the query projection).
PFFT reduces grokking epochs by 92.7% on modular arithmetic. No-Q attention improves validation BPC by 3.18% on TinyStories and 2.24% on FineWeb, reduces parameters by 8%, and accelerates grokking by 51.0% on modular arithmetic. This work provides a new perspective that embeddings act as the query for attention, offering important insights for future transformer design.