Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Gong, Zixuan, Li, Shijia, Liu, Yong, Teng, Jiaye
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917003910971392
author Gong, Zixuan
Li, Shijia
Liu, Yong
Teng, Jiaye
author_facet Gong, Zixuan
Li, Shijia
Liu, Yong
Teng, Jiaye
contents Transformers may exhibit two-stage training dynamics during the real-world training process. For instance, when training GPT-2 on the Counterfact dataset, the answers progress from syntactically incorrect to syntactically correct to semantically correct. However, existing theoretical analyses hardly account for this feature-level two-stage phenomenon, which originates from the disentangled two-type features like syntax and semantics. In this paper, we theoretically demonstrate how the two-stage training dynamics potentially occur in transformers. Specifically, we analyze the feature learning dynamics induced by the aforementioned disentangled two-type feature structure, grounding our analysis in a simplified yet illustrative setting that comprises a normalized ReLU self-attention layer and structured data. Such disentanglement of feature structure is general in practice, e.g., natural languages contain syntax and semantics, and proteins contain primary and secondary structures. To our best knowledge, this is the first rigorous result regarding a feature-level two-stage optimization process in transformers. Additionally, a corollary indicates that such a two-stage process is closely related to the spectral properties of the attention weights, which accords well with our empirical findings.
format Preprint
id arxiv_https___arxiv_org_abs_2502_20681
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers
Gong, Zixuan
Li, Shijia
Liu, Yong
Teng, Jiaye
Computation and Language
Artificial Intelligence
Machine Learning
Transformers may exhibit two-stage training dynamics during the real-world training process. For instance, when training GPT-2 on the Counterfact dataset, the answers progress from syntactically incorrect to syntactically correct to semantically correct. However, existing theoretical analyses hardly account for this feature-level two-stage phenomenon, which originates from the disentangled two-type features like syntax and semantics. In this paper, we theoretically demonstrate how the two-stage training dynamics potentially occur in transformers. Specifically, we analyze the feature learning dynamics induced by the aforementioned disentangled two-type feature structure, grounding our analysis in a simplified yet illustrative setting that comprises a normalized ReLU self-attention layer and structured data. Such disentanglement of feature structure is general in practice, e.g., natural languages contain syntax and semantics, and proteins contain primary and secondary structures. To our best knowledge, this is the first rigorous result regarding a feature-level two-stage optimization process in transformers. Additionally, a corollary indicates that such a two-stage process is closely related to the spectral properties of the attention weights, which accords well with our empirical findings.
title Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2502.20681