[Paper Review] Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMs

Angelica Chen, Ravid Schwartz-Ziv+3 more|arXiv (Cornell University)|Sep 13, 2023

Topic Modeling4 citations

TL;DR

This paper investigates abrupt training dynamics in masked language models (MLMs), identifying a sudden emergence of Syntactic Attention Structure (SAS)—where specialized attention heads focus on syntactic dependencies—coinciding with a steep drop in loss. The study demonstrates that SAS is causally necessary for grammatical capability development, and that temporarily suppressing SAS accelerates convergence and improves final performance, revealing a critical learning period and simplicity bias in pretraining.

ABSTRACT

Most interpretability research in NLP focuses on understanding the behavior and features of a fully trained model. However, certain insights into model behavior may only be accessible by observing the trajectory of the training process. We present a case study of syntax acquisition in masked language models (MLMs) that demonstrates how analyzing the evolution of interpretable artifacts throughout training deepens our understanding of emergent behavior. In particular, we study Syntactic Attention Structure (SAS), a naturally emerging property of MLMs wherein specific Transformer heads tend to focus on specific syntactic relations. We identify a brief window in pretraining when models abruptly acquire SAS, concurrent with a steep drop in loss. This breakthrough precipitates the subsequent acquisition of linguistic capabilities. We then examine the causal role of SAS by manipulating SAS during training, and demonstrate that SAS is necessary for the development of grammatical capabilities. We further find that SAS competes with other beneficial traits during training, and that briefly suppressing SAS improves model quality. These findings offer an interpretation of a real-world example of both simplicity bias and breakthrough training dynamics.

Motivation & Objective

• To understand how syntactic structure acquisition emerges during MLM pretraining, beyond final model analysis.
• To investigate whether Syntactic Attention Structure (SAS) is causally linked to the development of linguistic capabilities.
• To examine the role of SAS in competition with alternative learning strategies during training.
• To explore whether manipulating SAS during training can improve model convergence and final performance.
• To identify and characterize critical learning periods in MLM pretraining through dynamic monitoring of interpretability artifacts.

Proposed method

• Monitoring Syntactic Attention Structure (SAS) via dependency parse accuracy using attention distributions, measuring Unlabeled Attachment Score (UAS) as a proxy.
• Introducing a differentiable regularizer to suppress SAS during pretraining, enabling controlled manipulation of this structural trait.
• Tracking training dynamics through multiple metrics: MLM loss, implicit parse accuracy (UAS), GLUE, and BLiMP scores across multiple random seeds.
• Using nonparametric bootstrapping to compute 95% confidence intervals for all reported metrics.
• Comparing models with and without SAS suppression across fixed and variable training durations to isolate causal effects.
• Analyzing representational and functional similarity via Centered Kernel Alignment (CKA) and Total Variation Distance (TVD) over time.

Experimental results

Research questions

• RQ1 At what point in training does Syntactic Attention Structure (SAS) emerge, and how does it correlate with loss dynamics?
• RQ2 Is SAS causally necessary for the acquisition of grammatical capabilities in MLMs?
• RQ3 Do alternative learning strategies compete with SAS during pretraining, and if so, how does this competition affect model performance?
• RQ4 Can suppressing SAS during a specific window improve final model quality and accelerate convergence?
• RQ5 Does the critical learning period for SAS persist at scale, or do models eventually converge despite early suppression?

Key findings

• A sudden spike in SAS occurs at a consistent training point, immediately preceding a steep drop in loss, which the authors term the 'structure onset'.
• The emergence of SAS is followed by a sharp increase in grammatical capability, measured by BLiMP score, marking the 'capabilities onset'—a phase transition in functional performance.
• SAS is causally necessary for grammatical capabilities: models trained with SAS suppressed fail to develop these abilities.
• Brief suppression of SAS during the critical window improves final model quality and accelerates convergence, with the best-performing model achieving MLM loss of 1.55 ± 0.01.
• Suppressing SAS past the alternative strategy onset damages performance and blocks long-term SAS recovery, indicating a terminal point in the critical learning phase.
• Despite early divergence, long-term models (300K steps) trained with SAS suppression converge to similar performance levels as standard BERT, with no statistically significant difference in final metrics.

(a) Intrinsic dimension (TwoNN) complexity. Alternative complexity metrics in Section L.2 .