Clustering Head: A Visual Case Study of the Training Dynamics in Transformers

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Odonnat, Ambroise, Bouaziz, Wassim, Cabannes, Vivien
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909472504414208
author Odonnat, Ambroise
Bouaziz, Wassim
Cabannes, Vivien
author_facet Odonnat, Ambroise
Bouaziz, Wassim
Cabannes, Vivien
contents This paper introduces the sparse modular addition task and examines how transformers learn it. We focus on transformers with embeddings in $\R^2$ and introduce a visual sandbox that provides comprehensive visualizations of each layer throughout the training process. We reveal a type of circuit, called "clustering heads," which learns the problem's invariants. We analyze the training dynamics of these circuits, highlighting two-stage learning, loss spikes due to high curvature or normalization layers, and the effects of initialization and curriculum learning.
format Preprint
id arxiv_https___arxiv_org_abs_2410_24050
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Clustering Head: A Visual Case Study of the Training Dynamics in Transformers
Odonnat, Ambroise
Bouaziz, Wassim
Cabannes, Vivien
Machine Learning
This paper introduces the sparse modular addition task and examines how transformers learn it. We focus on transformers with embeddings in $\R^2$ and introduce a visual sandbox that provides comprehensive visualizations of each layer throughout the training process. We reveal a type of circuit, called "clustering heads," which learns the problem's invariants. We analyze the training dynamics of these circuits, highlighting two-stage learning, loss spikes due to high curvature or normalization layers, and the effects of initialization and curriculum learning.
title Clustering Head: A Visual Case Study of the Training Dynamics in Transformers
topic Machine Learning
url https://arxiv.org/abs/2410.24050