Reproduce “Implicit Regularization of Gradient Flow on One-Layer Softmax Attention” on OpenResearch

About this paper

Heejune Sheen, Siyu Chen, Tianhao Wang, Harrison H. Zhou

This paper from Yale University's Department of Statistics and Data Science provides a theoretical characterization of the implicit regularization in a one-layer softmax attention model. It demonstrates that when key and query matrices are trained separately, gradient flow implicitly minimizes the nuclear norm of the combined attention weights, which contrasts with prior findings of Frobenius norm minimization when a single combined weight matrix is trained.

Key insights

• The specific parameterization of an attention model (training key and query matrices separately vs. a combined matrix) fundamentally alters the implicit bias of gradient flow.
• Separately training key (K) and query (Q) matrices leads to an implicit bias towards minimizing the nuclear norm of the combined attention weight matrix W = KQ^T.
• This nuclear norm minimization provides a theoretical basis for empirically observed low-rank structures in trained transformer attention weights.

Summary by alphaXiv.