PaperlayerView exact revision on arXiv

Source-linked AI summary

How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding

Yuchen Li, Yuanzhi Li, Andrej Risteski

arXiv:2303.04245v2cs.LGcs.CLstat.ML

TL;DR

Although transformers are widely used, the mechanisms by which they learn from data remain poorly understood. The paper combines theory and experiments on topic-model data and Wikipedia to study how embeddings and attention capture topic structure. It reports that both components encode topical structure, with one able to compensate when the other is not trained.

  • Problem

    Despite extensive probing of transformer capabilities, formal understanding of how gradient-trained transformers learn from data remains limited.

    Source: §1 · p. 1

  • Method

    The paper analyzes a single-layer transformer on topic-model data, characterizes embedding and attention learning in two extremal cases, and checks the patterns on synthetic data and Wikipedia.

    Source: §1 · p. 1, §2 · p. 2

  • Results

    The authors report that embeddings and attention both encode topic structure, and that either can compensate when the other is not trained.

    Source: §2 · p. 2

  • Takeaways & Limitations

    The reported findings characterize two routes by which topic structure appears in a transformer: its embeddings and its self-attention.

    Source: §2 · p. 2

    The Stage 2 analysis freezes W V, although empirically that matrix continues to fluctuate even when its norm plateaus.

    Source: §8.1 · p. 14

Abstract

from arXiv · show

While the successes of transformers across many domains are indisputable, accurate understanding of the learning mechanics is still largely lacking. Their capabilities have been probed on benchmarks which include a variety of structured and reasoning tasks -- but mathematical understanding is lagging substantially behind. Recent lines of work have begun studying representational aspects of this question: that is, the size/depth/complexity of attention-based networks to perform certain tasks. However, there is no guarantee the learning dynamics will converge to the constructions proposed. In our paper, we provide fine-grained mechanistic understanding of how transformers learn "semantic structure", understood as capturing co-occurrence structure of words. Precisely, we show, through a combination of mathematical analysis and experiments on Wikipedia data and synthetic data modeled by Latent Dirichlet Allocation (LDA), that the embedding layer and the self-attention layer encode the topical structure. In the former case, this manifests as higher average inner product of embeddings between same-topic words. In the latter, it manifests as higher average pairwise attention between same-topic words. The mathematical results involve several assumptions to make the analysis tractable, which we verify on data, and might be of independent interest as well.

1 INTRODUCTION

Despite extensive probing of transformer capabilities, the mechanisms by which gradient-trained transformers learn from data remain poorly understood. This paper studies how embeddings and attention capture topic structure, using synthetic topic-model data to analyze the learning process.

1 INTRODUCTION

  • Transformers are widely used, but their learning mechanisms remain poorly understood even for simple structured data and one-layer models.Existing work probes and visualizes trained networks, while formal understanding of gradient-based learning remains limited.Source: §1 · p. 1
  • On Wikipedia, BERT embeddings are more similar for words assigned to the same LDA topic than for words assigned to different topics.The paper measures similarity by inner product or cosine similarity.Source: §1 · p. 1
  • The authors use synthetic LDA data, theory, and experiments to study how embeddings and attention learn topic structure, and report that the structure is more pronounced in synthetic data.They also report mathematically analyzing simplified two-stage training dynamics and checking the two-stage pattern across optimizers and hyperparameters.Source: §1 · p. 1
  • Four loss-and-optimizer settings produce block-wise embedding dot products, with the ten blocks corresponding to the ten topics.The caption identifies the settings as cross-entropy or squared loss with SGD or Adam, all at learning rate 0.01.Source: Figure 1 · p. 2

2 OVERVIEW OF RESULTS

The paper characterizes topic structure in embeddings and self-attention using a single-layer transformer and a two-stage training approximation, then checks these patterns in synthetic and Wikipedia data. Its results indicate that either component can encode topic structure, and one can compensate when the other is not trained.

2 OVERVIEW OF RESULTS

  • Topic structure can be encoded in either the embedding layer or self-attention, and either component can compensate when the other is not trained.The analysis characterizes two extremal cases and checks the patterns on synthetic LDA-generated and real Wikipedia data.Source: §2 · p. 2

2.1 Topic structure is encoded in token embeddings

  • With uniform attention frozen, an optimal embedding layer has larger inner products for same-topic word pairs than for different-topic pairs.The informal result assumes topic-model data, frozen uniform attention, and other model components set to identity; the paper reports robustness across loss and optimization choices.Source: §2.1 · p. 2

2.2 Topic structure is encoded in self-attention

  • The self-attention analysis approximates training in two stages: first train W V with uniform attention, then freeze W V and analyze trained attention weights.The paper motivates the approximation by early training gradients and reports that the two-stage pattern is observed across hyperparameter settings.Source: §2.2 · pp. 2–3

2.2.1 Optimal W V in Stage 1

  • With one-hot embeddings and uniform attention, the optimal W V has block-wise structure, with entries on average larger within a topic than across topics.The result is stated under mild L2 regularization, and experiments report robustness to loss and optimizer choices.Source: §2.2.1 · p. 3

2.2.2 Optimal attention weights in Stage 2

  • With W V fixed at its first-stage block-wise optimum, optimal attention favors a convex combination of same-word and same-topic-different-word attention over different-topic attention.The paper reports that jointly trained attention weights also follow the theorem's relations.Source: §2.2.2 · pp. 3–4

2.3 Empirical results

  • The empirical section tests the theory on LDA data and Wikipedia, while Figure 4 shows the two-stage pattern in simultaneously trained models.During steps 0–400, W K and W Q norms stay near zero as W V increases; during steps 400–1000, W K and W Q increase while W V stays relatively flat.Source: §2.3 · p. 4, Figure 4 · p. 5

3 PROBLEM SETUP

The setup uses a topic-model distribution to study masked language modeling in a simplified, single-layer transformer. Its theoretical assumptions and architecture isolate how embeddings and self-attention can represent topic structure.

3.1 Topic models

  • The theoretical data model assigns words to disjoint topics and assumes infinitely long documents whose empirical token distribution matches the ground truth.The experiments instead use finite-length documents generated by LDA, allowing more variability than the theoretical setting.Source: §3.1 · pp. 5–6

3.2 Training objective

  • Training masks randomly selected tokens and predicts their original words, with the analysis using squared loss and experiments also testing cross-entropy.The cross-entropy optimum is harder to characterize because softmax logits need not be unique and may be infinite; squared loss admits more tractable finite-valued optima.Source: §3.2 · pp. 6–7

3.3 Transformer network architecture

  • The analysis uses a one-layer transformer without residual or normalization layers, with tied embedding and prediction-head weights.In selected analyses and experiments, one-hot embeddings are frozen to isolate the self-attention mechanism.Source: §3.3 · p. 7

4 TOPIC STRUCTURE CAN BE ENCODED IN TOKEN EMBEDDINGS

Under the paper’s topic-model assumptions and squared-loss setup, the optimal embedding Gram matrix has a topic-dependent block structure. This formal result predicts greater average similarity for same-topic words, a pattern the authors also report in experiments.

4 TOPIC STRUCTURE CAN BE ENCODED IN TOKEN EMBEDDINGS

  • Under the topic-model assumptions, specified bias, uniform frozen attention, and squared-loss objective, an optimal embedding matrix has topic-structured inner products.The theorem characterizes E := W E⊤W E, including its rows and columns as well as entries between vocabulary words.Source: §4 · p. 8
  • For entries of the optimal E, same-topic word pairs are larger on average and different-topic pairs are smaller on average.The authors interpret this as same-topic embeddings being more similar; when the constants are zero, the difference becomes positive versus zero.Source: §4 · p. 8
  • Experiments report the block-wise embedding pattern under both squared and cross-entropy loss, whether all layers or only embeddings are trained.Training only the embedding layer achieves loss only slightly worse than training all layers; Wikipedia examples show a similar pattern for unambiguous words.Source: §4 · p. 8
  • The theorem is proved for squared loss; comparable cross-entropy analysis is harder because the softmax can yield non-unique optima and optimal logits at infinity.The authors report empirical block patterns under cross-entropy despite this theoretical difficulty.Source: §4 · p. 8

5 TOPIC STRUCTURE CAN BE ENCODED IN SELF-ATTENTION

With embeddings fixed to one-hot vectors, the paper analyzes how self-attention can encode topic structure. It models optimization in two stages and characterizes block-structured value weights and topic-sensitive optimal attention.

5.1 The two-stage optimization process of self-attention

  • Training exhibits two stages: near-zero W_K and W_Q accompany W_V growth, then W_V plateaus before keys and queries move; the analysis approximates this with staged freezing, validated under SGD and Adam.Source: §5.1 · pp. 8–9

5.2 Optimal W_V given uniform attention

  • Under uniform attention and mild L2 regularization, optimal W_V has a blockwise pattern: each unmasked word contributes only to same-topic predictions, aggregating topic proportions.The result is for a single-layer transformer trained on the topic-model distribution with squared loss.Source: §5.2 · p. 9

5.3 Optimal attention weights

  • When the lower bound exceeds 1, attention between same-topic words is expected to exceed attention between different-topic words.The analysis characterizes optimal attention weights with W_V frozen at its first-stage optimum, under the theorem's stated assumptions.Source: §5.3 · p. 10
  • With uniform-block W_V, only a weighted average of α and β matters; for finite T and τ, the loss landscape is computed numerically rather than characterized in closed form.Different α and β can yield the same loss when that weighted average is unchanged.Source: §5.3 · pp. 10–11

6 EXPERIMENTS

Experiments test the paper’s topic-structure account on synthetic LDA data and Wikipedia. They find block-wise structure in embeddings and value matrices, and greater average attention between same-topic words.

6.1 Results on synthetic LDA-generated data

  • On synthetic LDA data, trained token embeddings develop a block-wise inner-product pattern, including when other model layers are trained.The reported pattern appears across cross-entropy or squared losses and SGD or Adam optimizers.Source: §6.1 · p. 11
  • With one-hot embeddings and uniform attention, the trained value matrix develops a block-wise pattern corresponding to topics, approximately persisting when key and query matrices are trained jointly.The joint-training result is reported in Appendix E.1.Source: §6.1 · p. 11
  • When the value matrix is trained and converges naturally to the analyzed pattern, words attend more to same-topic than different-topic words on average.Source: §6.1 · p. 11

6.2 Results on natural language data

  • On Wikipedia, same-topic words have consistently greater average embedding similarity and attention weight than different-topic words across evaluation settings.The study fits a 100-topic LDA model to Wikipedia and compares topic pairs after preprocessing; attention weights are normalized to reduce sentence-length bias.Source: §6.2 · p. 11, Table 1 · p. 12

7 RELATED WORKS

Prior work probes what trained transformers represent and studies their expressive capacity, while this paper focuses on how training dynamics learn topic structure. It also relates its assumptions and topic-model setting to work on attention dynamics and other structured data.

7 RELATED WORKS

  • Probing and attention-weight studies find information relevant to syntax and semantics in trained models; this work formalizes an intuition that embeddings capture bigram statistics [Elhage et al. (2021)].Under a topic-model distribution, the cited work’s bigram-statistics intuition corresponds to document-level co-occurrence.Source: §7 · p. 12
  • Prior theory studies transformer expressivity and representation, while this work emphasizes training dynamics because data distributions may admit multiple reasonable representational optima.The cited areas include Turing completeness, function approximation, formal languages, algebra, sample complexity, and latent representations.Source: §7 · pp. 12–13
  • Unlike Snell et al. (2021), which considers cross-attention in machine translation, this work studies self-attention on topic-model data and empirically validates its two-stage training approximation.The paper reports validation on synthetic and real data, while describing Snell et al.’s similar assumption as lacking theoretical or empirical validation.Source: §7 · p. 13
  • This analysis complements work on positional structure because topic-model attention depends on bag-of-words content, whereas the cited image-data analysis uses positional encodings [Jelassi et al. (2022)].The paper identifies interaction between position-dependent and position-independent factors as future work.Source: §7 · p. 13
  • Unlike prior topic-model analyses of loss or contextualized-representation clustering, this work analyzes transformer optimization and token-embedding parameter convergence.The cited prior work connects low loss on prediction objectives with recovering topic posteriors, or studies topic discovery by clustering contextualized representations.Source: §7 · p. 13

8 DISCUSSION

The discussion reports an approximate two-stage training pattern in multilayer Wikipedia transformers and explains why value weights can grow before key and query weights. It also cautions that the approximation is hyperparameter-sensitive and that topic-wise behavior is not fully explained by simple co-occurrence counts.

8.1 The two-stage optimization process

  • The observed two-stage pattern carries over to joint training, but depends on initialization and learning rate and is not usually visible under common default hyperparameters.Source: §8.1 · p. 14
  • Near-zero initialization makes the early key-gradient term involving W_Q small, while the value gradient uses normalized attention weights, helping W_V grow faster initially.The paper notes that columns of W_Q sum approximately to zero, whereas columns of the attention matrix sum to one.Source: §8.1 · p. 14
  • The two-stage analysis is an approximation: W_V continues changing direction at roughly constant norm even when its norm plateaus, although its final empirical convergence point matches the theory.Source: §8.1 · p. 14
  • In the first roughly 100 steps, value-weight norms increase while key and query norms stay near zero; those latter norms rise later, across all four Wikipedia-trained layers.The figure tracks weight norms and value-matrix rotation for a 4-layer, 4-head-per-layer model; the paper calls this an approximate multi-stage phenomenon.Source: §8.1 · pp. 14–15, Figure 5 · p. 15

8.2 Do topic-wise behaviors perfectly correlate with co-occurrence counts?

  • Formally disentangling topic structure from co-occurrence requires data distributions that better separate those factors, which the paper leaves for future work.Source: §8.2 · p. 16
  • For matched pair-count cutoffs, same-topic pairs have lower average attention but higher embedding cosine similarity than the most-co-occurring pairs.The comparison uses a BERT model pretrained on Wikipedia; the paper conjectures that topic-wise behavior cannot be fully explained by simple co-occurrence counts.Source: Table 2 · p. 16, §8.2 · p. 16

9 CONCLUSION

The paper initiates study of transformer training dynamics with semantic structure represented by topic models. It identifies syntactic sandboxes and disentangling semantic and syntactic learning across model components as future directions.

9 CONCLUSION

  • Extending the analysis to syntactic data such as PCFGs, and disentangling semantic and syntactic learning across attention, positional encodings, and embeddings, remain open challenges.Source: §9 · p. 16

Supplementary Material

The supplement specifies the simplified uniform-attention setup and characterizes the optimal linear transform under it. The characterization explains how the transform aggregates unmasked word counts by topic, corrects for masking, and admits multiple optima unless L2 regularization is added.

A ADDITIONAL INFORMATION ON THE SETUP

A.1 Lemma on the optimal linear transform when freezing uniform attention

  • With uniform attention, the self-attention output aggregates unmasked token counts, which the optimal linear transform uses to infer topic proportions and denoise the masking process.The transform does not predict the mask token, which was absent from the original document.Source: §A.1 · p. 24
  • The optimal transform routes each unmasked word only to predictions for words in the same topic, while assigning no contribution to words of different topics.The lemma characterizes the minimizers through row and column constraints, including free constants.Source: §A.1 · pp. 24, 26
  • Uniform attention makes the squared-loss objective non-strongly-convex, so its minimizers need not be unique; adding L2 regularization makes the objective strongly convex.The supplement attributes the rank deficiency to the rank-1 attention matrix multiplying the transform.Source: §A.1 · p. 25

B PROOF OF THEOREM 1: OPTIMAL TOKEN EMBEDDING

The proof characterizes optimal token embeddings under the topic-model distribution and shows that same-topic words have greater embedding inner products than different-topic words. It also establishes that some optimal value matrices form a family without regularization.

B PROOF OF THEOREM 1: OPTIMAL TOKEN EMBEDDING

  • Optimal token embeddings have larger inner products for same-topic word pairs than for different-topic pairs under the stated masked-language-modeling setting.The proof represents the embedding similarities through E = W_E^T W_E.Source: B PROOF OF THEOREM 1: OPTIMAL TOKEN… · p. 29
  • A subset of the characterized optimal embedding matrices is realizable as E = W_E^T W_E for an embedding matrix W_E.This factorization links the optimal similarity structure to actual word embeddings.Source: B PROOF OF THEOREM 1: OPTIMAL TOKEN… · p. 30

C.1 Optimal W V when freezing uniform attention without regularization

  • Without L2 regularization, the optimal value matrices form a family under uniform attention, rather than a unique matrix.The theorem characterizes this family through free constants.Source: §C.1 · p. 31

C.2 Proof of Theorem 2: case when adding L2 regularization

The supplement shows how mild L2 regularization selects a unique limiting value matrix from the unregularized optima. It also derives masking-probability expressions and calculates model-output terms under a topic-wise attention assumption.

C.2 Proof of Theorem 2: case when adding L2 regularization

  • As the L2-regularization strength tends to zero, the optimal value matrix converges to a unique W V* among the unregularized optima.The proof reduces selection to minimizing the Frobenius norm over the set of unregularized minimizers.Source: §C.2 · pp. 31–32

D.1 Helping lemmas on masking probabilities

  • Appendix D.1 derives the token probabilities after masking and expresses the attention normalization constant c3 in terms of α, β, and masking probabilities.The expressions distinguish tokens from selected topics and tokens outside those topics.Source: §D.1 · p. 33

D.2 Implication of topic-wise attention assumption on model output

  • Appendix D.2 calculates the expected model-output terms using z1 and z2 under the topic-wise attention assumption.The calculation distinguishes same-token, same-topic, and different-topic cases.Source: §D.2 · p. 34

D.3 Proof of Theorem 3 (optimal attention when freezing W V to uniform blocks)

This section analyzes attention conditions associated with low loss when W_V is frozen in block-wise structures, including the finite-setting landscape for uniform blocks. It also motivates a separate Stage-2 analysis for diagonal W_V optima.

D.3 Proof of Theorem 3 (optimal attention when freezing W V to uniform blocks)

  • The loss in Case 3 is strictly below Cases 1 and 2, so the minimizing γ cannot lie in either case.The proof identifies an extra positive term of scale Ω(1) in Case 2 relative to Case 3.Source: §D.3 · pp. 38–39
  • The proof gives necessary conditions on optimal γ, not an o(1)-accurate characterization, because hidden smaller-scale terms cannot reverse the Ω(1) loss separation.The analysis relies on the asymptotic setting of Assumption 3.Source: §D.3 · p. 39

D.4 Optimal attention weights (when freezing diagonal W V )

  • Without L2 regularization, a family of W_V optima depending on free constants can encode topic structure, and diagonal patterns are a special case that sometimes appears in training.The diagonal pattern is identified as a particular optimum from the family characterized in Theorem 4.Source: §D.4 · p. 40, Figure 6 · p. 40
  • For diagonal W_V optima, the paper analyzes Stage-2 optimal attention weights and states that their behavior differs from the block-wise W_V result in Theorem 3.The analysis freezes W_V to the diagonal optimum defined in Definition 2.Source: §D.4 · pp. 40–41
  • The diagonal-W V proof rules out several α, β regimes by comparing their loss lower bounds with a feasible candidate that has strictly lower loss.The comparison relies on v being a large finite constant.Source: §D.4 · p. 44

D.5 Loss landscape with respect to attention weights in the non-asymptotic setting

  • With diagonal W V, the non-asymptotic loss is lowest when β lies within an interval controlled by τ and T, while α is less than a constant multiple of β.These constraints are predicted by Theorem 5 and are visible in the landscapes for T = 100, v = 300, and τ ∈ {20, 40, 60, 80}.Source: §D.5 · p. 45, Figure 7 · p. 45

W V with uniform blocks

  • With uniform-block W V, low loss occurs in a corner-shaped region where both α and β lie within intervals whose lower bounds grow with τ.Theorem 3 predicts that a convex combination of α and β lies within an interval controlled by τ and T; this variant is visibly reflected in the non-asymptotic landscape.Source: Figure 8 · p. 45, W V with uniform blocks · p. 45

E.1 Additional results on learned value matrix W V

These appendix results examine how learned value matrices and attention weights reflect topic structure, then detail how Wikipedia attention comparisons are filtered and debiased. Attention patterns depend on the learned W_V structure, while the natural-language analysis reports mostly higher same-topic similarities and attention weights.

E.1 Additional results on learned value matrix W V

  • With uniform attention and one-hot embeddings, W_V converges to a topic-aligned block pattern under L2 regularization and without it.Without L2 regularization, Figure 9 reports block-entry means over 10 times the magnitude of sampled non-block means in its squared-loss/SGD plot.Source: §E.1 · p. 46, Figure 9 · p. 46
  • When W_K and W_Q are trained jointly with W_V, the learned W_V still approximately converges to the analyzed block-wise pattern.Figure 10 shows this pattern across cross-entropy and squared loss with SGD and Adam.Source: §E.1 · p. 46, Figure 10 · p. 46

E.2 Additional results on learned attention weights

  • With block-wise, near-uniform W_V, most settings show same-word attention above same-topic-different-word attention, which exceeds different-topic attention.The Adam 0.03 setting has larger same-topic-different-word attention; failed SGD runs showed similar averages across all three types.Source: Table 3 · p. 47, §E.2 · p. 47
  • With diagonal-pattern W_V, same-word attention is highest, followed by different-topic attention, while same-topic-different-word attention is lowest.Table 4 reports this ordering over seven selected runs.Source: §E.2 · p. 47, Table 4 · p. 47

E.3 Additional details and results on natural language data

  • The Wikipedia comparison filters ambiguous topic assignments and normalizes each sentence’s attention weights by sentence length before comparing word pairs.The ambiguity filter retains words central to a topic; the debiasing makes each sentence’s average attention weight 1/100.Source: §E.3 · p. 48
  • Across most settings, Wikipedia models show higher same-topic than different-topic embedding similarity and attention, though multi-topic assignments fall outside the stated theorems.Table 5 compares different-word pairs across ambiguity thresholds; the theoretical results do not cover multiple topics per word.Source: §E.3 · p. 48, Table 5 · p. 49
Loading 2303.04245v2…