Symmetry and rank shape attention mechanisms in large language models

On attention heads and bilinear forms

Machine Learning

Summary

Attention heads in large language models use certain mathematical forms, called bilinear forms, to decide how words relate to each other. The authors studied how these forms can be broken down into symmetric and antisymmetric parts and found that after training, they mostly look like simple rank-one forms. They also proved that the symmetric parts behave in a predictable way related to special types of mathematical forms, which holds for a large set of these models. This helps us understand the math behind how language models focus on information.

What this means in practice

  • For language model developers: Design or analyze attention modules with an understanding of how symmetric parts decompose mathematically in trained models.
  • For neural network theorists: Use the proved structural properties to explore model behavior and generalization related to attention bilinear forms.

A theory result. No direct application yet.

Authors

Andrew O'Desky

Abstract

We study the symmetric and antisymmetric parts of bilinear forms in the attention heads of trained large language models. We introduce an orthogonally invariant profile map from real bilinear forms to a three-dimensional simplex and observe that profiles of trained bilinear forms accumulate near profiles of rank-one bilinear forms. We prove that the symmetric part of a bilinear form in an attention head is the sum of a hyperbolic form and a zero form for a Zariski-dense subset of query-key matrices.