Skip to yearly menu bar Skip to main content


Poster

Infinite Limits of Multi-head Transformer Dynamics

Blake Bordelon · Hamza Chaudhry · Cengiz Pehlevan

East Exhibit Hall A-C #4804
[ ]
Fri 13 Dec 11 a.m. PST — 2 p.m. PST

Abstract:

In this work we analyze various scaling limits of the training dynamics of transformer models in the feature learning regime. We identify the set of parameterizations which admit well defined infinite width and depth limits that allow the attention layers to update throughout training, a relevant notion of feature learning in these models. We then use tools from dynamical mean field theory (DMFT) to analyze various infinite limits (infinite heads, infinite key/query dimension, and infinite depth) which have different statistical descriptions depending on which infinite limit is taken and how attention layers are scaled. We provide numerical evidence of convergence to the limits and show they maintain the correct scale of updates for both SGD and Adam.

Live content is unavailable. Log in and register to view live content