Abrupt Learning in Transformers: A Case Study on Matrix Completion

TL;DR

This paper investigates the abrupt loss drops during Transformer training by analyzing a matrix completion task, revealing a transition from copying to accurate prediction, with interpretable attention and embeddings.

Contribution

It demonstrates the phenomenon of sudden loss drops in Transformers and provides interpretability insights into the model's internal transition during training.

Findings

Loss curves show a plateau followed by a sharp drop.

Attention heads develop interpretable patterns after the drop.

Embeddings encode task-relevant information.

Abstract

Recent analysis on the training dynamics of Transformers has unveiled an interesting characteristic: the training loss plateaus for a significant number of training steps, and then suddenly (and sharply) drops to near--optimal values. To understand this phenomenon in depth, we formulate the low-rank matrix completion problem as a masked language modeling (MLM) task, and show that it is possible to train a BERT model to solve this task to low error. Furthermore, the loss curve shows a plateau early in training followed by a sudden drop to near-optimal values, despite no changes in the training procedure or hyper-parameters. To gain interpretability insights into this sudden drop, we examine the model's predictions, attention heads, and hidden states before and after this transition. Concretely, we observe that (a) the model transitions from simply copying the masked input to accurately predicting the masked entries; (b) the attention heads transition to interpretable patterns relevant to the task; and (c) the embeddings and hidden states encode information relevant to the problem. We also analyze the training dynamics of individual model components to understand the sudden drop in loss.