PaperlayerView exact revision on arXiv

Source-linked AI summary

Context Staircase: Signature-Aligned Dynamics of Token Embeddings under Small Initialization

Junjie Yao, Liangkai Hang, Zhi-Qin John Xu

arXiv:2608.30315v1cs.LGcs.CL

TL;DR

The paper asks how meaningful token-embedding structures emerge during gradient-based language-model training. It introduces probability signatures and analyzes gradient flow under small initialization across feed-forward, self-attention, and real language-model settings. The results identify Context Staircase, in which low-order structures appear first and richer contextual signatures become visible later, while their usefulness depends on downstream computation.

  • Problem

    It remains unclear how evolving token-embedding geometry relates to identifiable data statistics and model architecture during training.

  • Method

    The paper defines token-conditioned probability signatures and uses gradient-flow analysis under small initialization for feed-forward and self-attention architectures, extending the analysis to real language-model training.

  • Results

    The study identifies Context Staircase: early embedding geometry aligns with simple label signatures, while later geometry jointly reflects higher-order contextual signatures and architecture-specific paths.

  • Takeaways & Limitations

    The signature perspective provides a unified account of how statistical structures enter embeddings, change during training, and become useful for task solving and semantic organization.

  • Takeaways & Limitations

    Signature-aligned geometry alone does not ensure task success; usefulness depends on whether subsequent computations can exploit the encoded structure.

Abstract

from arXiv · show

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.

1 Introduction

The paper investigates how token-embedding geometry emerges during training and identifies Context Staircase: embeddings progress from low-order statistical signatures to richer context-dependent structures. Gradient-flow analysis under small initialization explains this ordering, which extends across architectures and relates to task learning and semantic organization.

  • Probability signatures formalize token-conditioned relationships among a token, its label, and contextual tokens.
  • Small-initialization gradient-flow analysis explains the progression by suppressing higher-order contributions initially and ordering their later influence.
  • Context Staircase describes embeddings aligning with low-order signatures early and higher-order signature structures later in training.
  • The analysis extends from feed-forward networks to self-attention architectures and real language-model training.
  • Signature-aligned embedding structures affect task learning and the incorporation of semantic structure into the embedding space.

2 Related Work

Prior work established that learned representations encode semantic, syntactic, and distributional structure, but mostly examined representations after training. This paper addresses the unresolved dynamics connecting evolving embedding geometry to data statistics and initialization.

  • Earlier embedding methods represented distributional properties through co-occurrence statistics, matrix factorization, or context-prediction objectives.
  • Research on contextualized and large language models largely analyzed the geometric and linguistic information present in learned representations.
  • Initialization scale can produce qualitatively different representation structures and learning behaviors, with small initialization yielding more structured task-reflective spaces.
  • It remained unclear which data-dependent statistical structures evolving embeddings correspond to and how those structures emerge through gradient-based training.
  • This work relates embedding training dynamics under small initialization to a hierarchy of probability signatures derived from the data.

3 Preliminary

The paper formulates sequence prediction, token-conditioned probability signatures, and small-initialization dynamics as the foundation for analyzing Context Staircase. The preliminary result is that low-order structure dominates initially while higher-order contextual structure emerges as training progresses.

  • 3.1 Prediction Problem: The prediction problem maps an input sequence X ∈ V^L to a target token y ∈ V using input embeddings, an output unembedding, and cross-entropy training.
  • 3.2 Token-conditioned Signatures: Probability signatures are normalized token-conditioned distributions over labels and sampled context tokens, with marginalization linking order-m signatures to order-(m−1) signatures.
  • 3.2 Token-conditioned Signatures: The initial analysis assumes exchangeable context tokens, while self-attention later relaxes this by distinguishing token positions.
  • 3.3 Parameter Initialization: Small initialization is controlled by γ, with larger γ corresponding to smaller parameter scales; the analysis uses γ > 1.
  • 3.4 Main Result: Context Staircase: Context Staircase begins with low-order, especially zeroth-order label signatures and later reveals richer higher-order contextual structures.
  • 3.4 Main Result: Context Staircase: Higher-order signatures are initially suppressed by additional embedding factors but gain influence as parameter scale grows.
  • 3.4 Main Result: Context Staircase: Late-stage embedding geometry can reflect multiple signature orders jointly, with the highest accessible order strongest under the analyzed conditions.

4 Signature-Aligned Dynamics in Feed-Forward Models

Learned token embeddings follow a Context Staircase: they align first with low-order label signatures, then progressively reflect higher-order, context-dependent structure. Small-initialization gradient flow explains this ordering by suppressing higher-order contributions initially, while later training allows multiple signature orders to shape embedding geometry.

  • Empirical signature alignment: Early embedding geometry closely matches the zero-order label signature across controlled arithmetic tasks.Embedding-logit distributions show high correspondence with token-conditioned label distributions, and similarity matrices or PCA projections reproduce their task-specific structures.
  • Empirical signature alignment: Addition and modular addition diverge early because only addition has distinguishable zero-order label distributions.Addition forms ordered embeddings, whereas modular-addition embeddings nearly align and fail to learn during the early stage.
  • Progressive signature learning: As training continues, zero-order alignment weakens while first-order structure becomes a better predictor of embedding geometry.For modular addition, the zero-order signature is constant; the optimized first-order signature later achieves higher correlation and lower prediction error.
  • Gradient-flow explanation: Theorem 1 shows that embedding gradients explicitly contain label-signature contributions from order 0 through the largest accessible order.This identifies the statistical components available to drive embedding updates and connects their empirical geometric alignment to gradient dynamics.
  • Gradient-flow explanation: At initialization, higher-order signature terms are suppressed by additional powers of the embedding scale, making the zero-order term dominant.The resulting magnitude hierarchy is m = 0 > m = 1 > m = 2 > · · ·.
  • Relations across signature orders: Later, signature orders share the same leading parameter scale, with the highest-order contribution potentially strongest and spectral refinement adding task-specific structure.Lower-order signatures are contained in higher-order zero-frequency projections, while nonzero-frequency components encode additional higher-order structure; random seeds can select different frequency components.
  • Task-learning boundary: The ordering does not guarantee faster learning: when low-order signatures lack task information, small initialization can delay progress.In modular addition, substantial learning begins only after higher-order signatures become influential.

5 Attention-Modulated Signature Dynamics

Self-attention preserves signature-based embedding dynamics while making them position-sensitive and weighting each token’s contribution by attention. Experiments show early zero-order alignment, later first-order structure, and faster evolution for highly attended tokens.

  • Gradient-flow analysis: Self-attention expresses embedding dynamics as position-specific label signatures modulated by attention-dependent weights.The gradient decomposes into value, key, and query contributions, with each signature weighted by how strongly corresponding tokens are attended.
  • Gradient-flow analysis: Attention scores route statistical information non-uniformly, so tokens with greater attention mass exert stronger influence on their embedding updates.The attention mechanism distinguishes token positions and scales signature contributions according to routed attention.
  • Empirical consequences: In length-2 addition and modular-addition tasks, zero-order label signatures reshape embeddings early, while first-order effects become visible later.Attention scores are uniform in these tasks, so signature structure primarily governs the dynamics.
  • Empirical consequences: In the parity task, anchor-token attention remains larger than noise-token attention, producing faster anchor embedding changes and slower noise-token evolution.The noise-token attention scores decay nearly to zero, making their associated gradient magnitudes negligible.

6 Path-Specific Signatures in Transformer Training

A full decoder-only Transformer combines direct, attention, and FFN paths, each contributing distinct signature information to embedding updates. Training and probing experiments support an early next-token alignment followed by richer, position- and context-dependent structure.

  • Path decomposition: The full Transformer embedding gradient is an additive sum of direct, attention, and FFN path contributions.This decomposition connects the modular analyses of the individual components to full language-model embedding dynamics.
  • Direct path: The direct readout path updates a token embedding according to the distribution of tokens that follow it, namely its next-token signature.This is the most immediate statistical signal before indirect attention and FFN effects are considered.
  • FFN path: Early in training, the FFN path induces only zero-order next-token signatures because its effective input at each position has length one.As attention outputs and the residual stream grow comparable, weighted prior-token inputs allow higher-order signatures to enter the FFN gradient.
  • Attention path: The early attention path contains an attention-weighted zero-order signature, while later key- and query-dependent terms introduce position-specific and higher-order signatures.Tokens receiving larger attention mass exert stronger influence on their corresponding embedding updates.
  • Real language-model training: In a 0.7B-parameter Transformer trained on 36B tokens, next-token similarity peaks near the beginning, then declines to a stable nonzero level as additional structures emerge.A first-order signature probe achieves consistently low held-out similarity MSE and generalizes across unseen anchor tokens.
  • Real language-model training: The attention-path representations show clear early-stage correlation with the corresponding attention-weighted signature.This supports the path-specific prediction from the theoretical analysis.

7 The Role of Signatures in Model Learning

Signature-aligned embedding geometry can support task learning, but its usefulness depends on whether the downstream architecture can exploit the relevant information and when task distinctions emerge.

  • Signature structure alone is not sufficient: A structured embedding geometry alone does not ensure task learning: the FFN succeeds, whereas the attention-only model fails despite learning first-order signature structure.The attention-only model is not expressive enough to transform the learned representation into the target operation.
  • Task-relevant signatures can accelerate learning: Fixed first-order signature embeddings converge substantially faster than the standard trainable-embedding baseline, showing that these signatures contain task-relevant information.Fixed random embeddings also outperform trainable embeddings early when zeroth-order signatures cannot distinguish number tokens.
  • Signature order controls learning onset: When M divides N, zeroth-order signatures are identical and modular addition requires higher-order information; for M = 4, 5, 10, and 20, learning takes thousands of epochs.For other tested moduli, zeroth-order signatures distinguish residue classes and the task typically learns within a few hundred epochs.
  • Implications: The task-learning benefit of signatures depends jointly on the processing architecture and the order at which discriminative information appears.Signature-based representations provide useful information only when the architecture can access and transform it.
  • An example of task solving via signatures: Additive superposition of zero-order signatures is sufficient to support the min–max prediction rule when the architecture can access that information.The construction demonstrates sufficiency of the statistical information, without claiming that every trained model implements the exact score.
  • Semantic structure: Corpus-derived signatures already display semantic organization: numerical tokens follow an ordered trajectory, weekdays form a cycle, and token categories exhibit structure.These geometries provide a possible statistical origin for semantic organization in learned embeddings.

8 Conclusion

The paper characterizes token-embedding evolution as progressive alignment with increasingly complex signatures and explains it through gradient-flow dynamics. Controlled and real-language-model analyses connect these dynamics to architecture, semantic organization, and task usefulness, while showing that geometry alone does not guarantee task success.

  • Main findings: Across controlled tasks, early embedding geometry aligns with simple token-conditioned label distributions, while later training reveals higher-order contextual signatures.Higher-order signatures become jointly visible rather than replacing lower-order structure one at a time.
  • Semantic organization: Signatures from a 36B-token corpus already contain semantic organization, including ordered numerical structure and cyclic weekday structure.The figure presents PCA projections of numerical tokens, weekdays, and multiple semantic categories.
  • Mechanism: Small initialization suppresses higher-order contributions initially, while self-attention modulates position-specific signature effects through attention scores.The resulting analysis provides a gradient-flow explanation for progressive signature alignment.
  • Real language-model training: In a real language-model run, early embedding geometry tracks next-token signatures from direct and FFN paths, while attention-path representations track attention-weighted signatures across layers.These correlations later decay as additional structures enter the model.
  • Scope: Signature-aligned geometry can accelerate optimization and support prediction only when the architecture can exploit the encoded information.A structured embedding space may emerge without successful task learning when the available computations cannot use it.

A.1.2 Prescribed Label-Distribution Experiment

The appendix specifies controlled experiments for prescribed label distributions, arithmetic tasks, signature measurements, attention analyses, fixed embeddings, and real-language-model training, together with the gradient-flow derivation setup.

  • Experimental data: The prescribed label-distribution experiment uses an anchor token, context-token sets, sampled target distributions, and fixed labels to construct controlled sequence-prediction data.The appendix considers uniform, normal, Poisson, and exponential distributions, with explicit input-set and checkpoint settings.
  • Signature measurements: Embedding and signature geometries are compared through pairwise similarities, Pearson correlations, projected signatures, and frequency-resolved cosine-similarity matrices.Probe vectors are optimized with five-fold cross-validation, while heatmaps sample epochs every 100 epochs.
  • Attention analysis: The attention experiment uses causal masking, learned positional embeddings, and no residual, normalization, or FFN modules, while tracking attention mass assigned to anchor and noise tokens.Figure 5K reports these quantities averaged over the full dataset.
  • Embedding settings: Fixed signature embeddings use empirical first-order label signatures followed by a fixed random linear projection, with trainable scale parameters compared against trainable and fixed-random embeddings.The fixed-signature and fixed-random vectors remain fixed during training while their scale parameter is trainable.
  • Language-model experiment: The real-language-model experiment trains a 0.7B pre-norm decoder-only Transformer on a 36B-token bilingual corpus using next-token prediction and AdamW.The model uses 24 layers, hidden size 1024, 16 attention heads, and a 4096-dimensional MLP intermediate layer.
  • Gradient-flow derivation: The gradient-flow derivation expands early softmax behavior and polynomial terms to identify embedding-gradient contributions associated with different signature orders.The analysis uses small-logit first-order softmax expansion, multinomial expansion, symmetrization, and tensor representations.

B.3 Proof of Proposition 1

The proof establishes normalization and marginal consistency for signatures, then bounds higher-order label-signature contractions relative to the zeroth-order contraction at initialization.

  • Co-occurrence signatures: The co-occurrence signatures satisfy the same normalization, marginal-consistency, and norm-monotonicity properties after omitting the label index.The co-occurrence proof follows the same argument as the label-signature case.
  • Signature properties: Label signatures are joint distributions over the label and tokens at uniformly sampled distinct context positions, so their entries are nonnegative and sum to one.Deleting one context coordinate and summing over its token value yields the lower-order marginal.
  • Initialization scaling: For fixed order m, the lowest-degree contribution is k = m + 1, and higher-degree terms are smaller by additional powers of d^-γ.This identifies the leading initialization scale of each contracted order-m operator.
  • Proposition 1: Every fixed higher-order label-signature contraction is asymptotically smaller than the zero-order contraction at initialization.The order-zero case is dominated by its degree-one term, while higher-degree terms are suppressed.

B.4 Proof of Proposition 2

Proposition 2 establishes that, after initialization-induced suppression is lifted, the highest-degree terms dominate within each signature order, all orders share a common late-stage scale, and the highest accessible order becomes largest.

  • For each fixed signature order m, the highest-degree term dominates all lower-degree terms once the embedding scale is sufficiently large.The proof uses the vanishing ratio of lower-degree to highest-degree terms and a common threshold S∗ across finitely many orders.
  • All signature orders have the same late-stage parameter scaling, so their hierarchy is no longer determined by different powers of s_E(t).Their relative magnitudes instead depend on order-specific prefactors, combinatorial coefficients, normalized contractions, and signature factors.
  • When s_E(t) ≥ S∗ and K is sufficiently large, the highest-order prefactor is largest among the same-scale signature contributions.The comparison is controlled by the combinatorial dependence on K together with non-degenerate normalized contractions and signatures.
  • Consequently, the order-(L − 1) contribution is the largest among all signature contributions in the late-stage regime.This conclusion follows after combining degree dominance with the comparison of order-dependent prefactors.

B.5 Proof of Theorem 3

Theorem 3 rewrites self-attention embedding gradients as signature terms whose statistical contributions are weighted by attention, including value-path and score-path effects.

  • The value-path contribution can be expressed using label signatures and an attention-weighting vector.The construction separates the statistical signature from the attention mass assigned to the corresponding token.
  • The first signature term has the same form as a label-signature update, except that the label signature is reweighted by the attention mass assigned to token α.This identifies the direct attention-modulated contribution to the embedding gradient.
  • The score-path contribution arises from differentiating attention scores and includes distinct key-score and query-score contributions.The query-score term depends on joint distributions involving labels and contextual tokens, while attention-score derivatives determine the weighting.
  • Conditioning on the selected token introduces label–context signatures and attention-product statistics into the score-path decomposition.These include joint label-last-token and multi-token distributions together with conditional expectations of attention scores and their products.

Appendix C. Multi-stage embedding evolution in length-three sequences.

In a length-three modular-addition task, identical zeroth-order signatures initially produce collapsed embeddings, while contextual signatures later generate increasingly structured geometry shaped jointly by first- and second-order statistics.

  • The length-three task uses inputs from A × A × X with target y = F(a1, a2, x) = (a1 + a2 + x) mod |X|.The anchor tokens and contextual token occupy disjoint ranges.
  • All anchor tokens share identical zeroth-order label signatures, whereas first- and second-order signatures distinguish them through joint token–label statistics.This construction isolates higher-order contextual information as the source of anchor-token differentiation.
  • Early embeddings collapse into a compact cluster with nearly uniform cosine similarities, consistent with identical zeroth-order signatures.At the intermediate stage, a banded similarity pattern and smooth one-dimensional PCA curve indicate emerging contextual-signature effects.
  • Late embeddings develop more complex geometry consistent with the joint influence of first- and second-order signatures.The late-stage structure is not a replacement of one order by another but a coexistence of accessible higher-order effects.
  • Correlations with S(1) and S(2) rapidly approach one, while S(2) provides a consistently stronger later-stage fit than S(1).The zeroth-order correlation is transient, briefly approaching one before becoming strongly negative and returning toward zero.

Appendix D. Different Seeds, Different Embedding

Final embedding geometries vary across random seeds even with fixed training conditions, but the differences are explained by alignment with different Fourier components of the same higher-order signature.

  • Different random seeds produce visibly different final embedding geometries despite identical datasets, architectures, and training configurations.The phenomenon appears in both addition and modular-addition tasks.
  • For each seed, the analysis compares embedding similarity with the similarity structure induced by each Fourier frequency k of the corresponding signature.The selected frequency is the one associated with maximum Pearson correlation.
  • The selected frequency k∗ varies across seeds, and the embedding PCA geometry closely matches the projection of the signature component at that frequency.Figures 14 and 15 show the correlation profiles and paired two-dimensional projections for ten seeds at epoch 1,000.
  • Seed-dependent geometries are therefore distinct but governed by different spectral components of the same higher-order signature.The result links initialization-dependent variation to spectral selection rather than arbitrary embedding organization.
Loading 2608.30315v1…