Context Staircase: Signature-Aligned Dynamics of Token Embeddings under Small Initialization
2026-08-31 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors studied how language models learn token embeddings, which connect words to computations. They found that embeddings develop progressively, first capturing simple relationships between tokens and labels without context, then gradually including more complex context-based information. This process, called the Context Staircase, was explained using math for different model architectures and observed in real training. The authors showed these evolving embeddings are important for learning tasks and representing meanings. Overall, their work explains how data patterns and model design shape embeddings step-by-step.
token embeddingslanguage modelsgradient-based trainingcontext Staircaseprobability signaturesfeed-forward architectureself-attentionembedding evolutionlow-order statisticscontextual distributions
Authors
Junjie Yao, Liangkai Hang, Zhi-Qin John Xu
Abstract
Token embeddings are the basic representational units that connect discrete tokens with continuous computation in language models. Although modern language models learn embeddings from random initialization through gradient-based training, the dynamical mechanism by which meaningful embedding structures emerge remains unclear. In this work, we identify that the evolving embedding structures are closely related to token-conditioned label and contextual distributions, which we formalize as probability signatures. We observe a progressive learning process, which we term Context Staircase: embeddings learn the low-order statistic signatures of the data before the high-order ones. More specifically, we observe that early in training they align with the simplest, context-free signature linking a token to its label, and as training proceeds, they progressively reflect signatures involving more and more context tokens. We then analyze the gradient flow of embeddings under small initialization to explain this phenomenon, deriving embedding evolution equations for feed-forward and self-attention architectures. We further extend these observations to real language-model training. Finally, we show that these embedding structures play an important role in both task learning and the incorporation of semantic structure into the embedding space. Overall, our results provide a dynamic explanation of how data statistics and architecture jointly shape token embeddings in language models, and reveal an implicit bias in the space of data statistics: training proceeds from simpler, low-order statistical relations toward increasingly complex, context-dependent ones.