Infinite Limits of Multi-head Transformer Dynamics

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Bordelon, Blake, Chaudhry, Hamza Tahir, Pehlevan, Cengiz
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913531386920960
author Bordelon, Blake
Chaudhry, Hamza Tahir
Pehlevan, Cengiz
author_facet Bordelon, Blake
Chaudhry, Hamza Tahir
Pehlevan, Cengiz
contents 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 that admit well-defined infinite width and depth limits, allowing 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 key/query dimension, infinite heads, 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 discuss how the parameterization qualitatively influences learned features.
format Preprint
id arxiv_https___arxiv_org_abs_2405_15712
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Infinite Limits of Multi-head Transformer Dynamics
Bordelon, Blake
Chaudhry, Hamza Tahir
Pehlevan, Cengiz
Machine Learning
Disordered Systems and Neural Networks
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 that admit well-defined infinite width and depth limits, allowing 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 key/query dimension, infinite heads, 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 discuss how the parameterization qualitatively influences learned features.
title Infinite Limits of Multi-head Transformer Dynamics
topic Machine Learning
Disordered Systems and Neural Networks
url https://arxiv.org/abs/2405.15712