How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding

Yuchen LiYuanzhi LiAndrej Risteski

article2023ICML87 citations

Proves mathematically and verifies empirically on topic-modeled data how transformer embeddings and self-attention mechanisms independently learn word co-occurrence structures through a distinct two-stage training dynamic.

Listen

Transformer neural networks serve as the backbone for modern artificial intelligence and natural language processing, yet a formal mathematical understanding of how they learn semantic structure from data remains limited. The article addresses this gap by investigating the precise learning mechanics through which transformers capture topic and word co-occurrence structure. Its primary objective is to demonstrate how individual model components—specifically the token embedding layer and the self-attention mechanism—encode topical relationships during standard masked language modeling training.

To establish these mechanisms, the authors combine mathematical analyses of optimization dynamics with empirical experiments on both synthetic data generated by Latent Dirichlet Allocation topic models and natural text from Wikipedia. The evaluation examines two complementary extremes: one where attention is held uniform while only token embeddings are trained, and another where token embeddings are fixed to one-hot representations while attention parameters are trained. The study also analyzes pretrained production models, including BERT, RoBERTa, ALBERT, BART, and ELECTRA, across various optimizers and loss formulations.

Key findings show that topic structure can be captured independently by either the embedding layer or the self-attention mechanism, meaning each component can compensate if the other is restricted. When training token embeddings alone, word representations converge such that word pairs within the same topic exhibit significantly higher inner products and similarity scores than pairs from different topics. When training the self-attention layer, learning naturally exhibits a two-stage dynamic where the value matrix first develops a block-wise topic structure before key and query matrices begin adjusting. Consequently, optimal attention heads assign substantially higher average pairwise attention weights to words from the same topic compared to different topics, a pattern verified across multiple real-world transformer models.

These results show that transformers do not merely memorize surface co-occurrence statistics but systematically organize internal representations to mirror underlying latent topic distributions. This structural insight provides a mechanistic foundation for understanding representation learning, aiding model interpretability, diagnostic probing, and future architectural design. Organizations developing or fine-tuning transformer architectures can use these insights to monitor representation convergence and evaluate whether attention heads are correctly capturing semantic domains.

Confidence in the core conclusions is high, supported by analytical proofs and consistent experimental validation across synthetic benchmarks, multiple optimization algorithms, and several pretrained language models. Nevertheless, certain limitations remain: the theoretical derivations assume a simplified single-layer architecture without residual connections or normalization, infinitely long documents, and strictly disjoint topic assignments where words belong to only one topic. Future work should extend this mechanistic analysis to deeper multi-layer models and datasets with complex syntactic structures.

Carnegie Mellon University Microsoft attention mechanism encoder-decoder reasoning topic structure mechanistic understanding Latent Dirichlet Allocation

YuchenLi01/transformer_topic_model_LDA Li et al (2023).pdf

Paper: Attention Is All You Need, Ashish Vaswani et al. (2017). Vaswani et al. introduced the foundational multi-head self-attention and transformer architecture whose internal representational dynamics and semantic encoding mechanisms this paper rigorously investigates.
Paper: Neural Word Embedding as Implicit Matrix Factorization, Omer Levy et al. (2014). Levy and Goldberg established the theoretical connection between word co-occurrence statistics and inner-product representations in neural embeddings, providing the conceptual groundwork for analyzing how semantic topic structures emerge in embedding layers.
Paper: What Does BERT Look at? An Analysis of BERT’s Attention, Kevin Clark et al. (2019). Clark et al. pioneered empirical investigations into how transformer attention heads capture linguistic and semantic structure, which this paper formalizes through theoretical analysis and topic-modeling dynamics.
Paper: How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings, Kawin Ethayarajh (2019). Ethayarajh analyzed the geometric structure and anisotropy of contextualized word representations across transformer layers, establishing key context for understanding how word inner products reflect semantic relationships.

Paper: How Transformers Learn Causal Structure with Gradient Descent, Eshaan Nichani et al. (2024). Nichani et al. extend the theoretical study of transformer training dynamics from topic and co-occurrence structures to provably uncovering latent causal graphs and in-context learning mechanisms via gradient descent.
Paper: One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context, Skanda Athreya et al. (2026). Athreya and Wang advance the mechanistic understanding of simplified transformer layers by theoretically proving how softmax attention dynamics learn non-parametric classification rules.

Cover for How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding

Abstract

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.

Topics

attention mechanism encoder-decoder reasoning synthetic data word embedding latent factor latent variable part of speech

Keywords

topic structure mechanistic understanding Latent Dirichlet Allocation attention pattern training dynamics attention mechanism self-attention layer BERT model word embeddings

1. INTRODUCTION
2. OVERVIEW OF RESULTS
2.1. Topic structure is encoded in token embeddings
2.2. Topic structure is encoded in self-attention
2.2.1. OPTIMAL W V IN STAGE 1
2.2.2. OPTIMAL ATTENTION WEIGHTS IN STAGE 2
2.3. Empirical results
3. PROBLEM SETUP
3.1. Topic models
3.2. Training objective
3.3. Transformer network architecture
4. TOPIC STRUCTURE CAN BE ENCODED IN TOKEN EMBEDDINGS
5. TOPIC STRUCTURE CAN BE ENCODED IN SELF-ATTENTION
5.1. The two-stage optimization process of self-attention
5.2. Optimal W V given uniform attention
5.3. Optimal attention weights
6. EXPERIMENTS
6.1. Results on synthetic LDA-generated data
6.2. Results on natural language data
7. RELATED WORKS
8. DISCUSSION: TWO-STAGE TRAINING
9. CONCLUSION
References
A. ADDITIONAL DISCUSSIONS
A.1. The two-stage optimization process (continued)
A.2. Do topic-wise behaviors perfectly correlate with co-occurrence counts?
B. ADDITIONAL DISCUSSIONS ON RELATED WORKS
C. ADDITIONAL INFORMATION ON THE SETUP
C.1. Lemma on the optimal linear transform when freezing uniform attention
D. PROOF OF THEOREM 1: OPTIMAL TOKEN EMBEDDING
E. PROVING OPTIMAL W V IN SELF-ATTENTION
E.1. Optimal W V when freezing uniform attention without regularization
E.2. Proof of Theorem 2: case when adding L 2 regularization
Step 1: the optima converges to one outlined in Lemma 1
Step 2: solve for the coefficients that minimize the L 2 penalty
F. ADDITIONAL RESULTS ON ATTENTION WEIGHTS
F.1. Helping lemmas on masking probabilities
F.2. Implication of topic-wise attention assumption on model output
F.3. Proof of Theorem 3 (optimal attention when freezing W V to uniform blocks)
Comparing the above cases
F.4. Optimal attention weights (when freezing diagonal W V )
F.5. Loss landscape with respect to attention weights in the non-asymptotic setting
G. ADDITIONAL EMPIRICAL RESULTS
G.1. Additional results on learned value matrix W V
G.2. Additional results on learned attention weights
G.3. Additional details and results on natural language data

Knowls

Knowl 1 — Optimal Word Embeddings Encode Topic Structure
theoretical result
Consider a single-layer transformer trained on data generated from a topic model with $T$ disjoint topics and $v$ words per topic (total vocabulary $V = Tv + 1$ including token $[\text{MASK}] = 0$ ). In the infinite document length limit where empirical topic frequencies match ground truth, let each document contain $\tau$ distinct topics with probability $p_w(i) = \frac{1}{\tau v}$ for words in active topics and $0$ otherwise. The model architecture uses tied prediction head and embedding weights $W^{pred} = W^E \in \mathbb{R}^{d \times (Tv+1)}$ , frozen uniform attention $W^K = 0, W^Q = 0$ , frozen identity value matrix $W^V = I$ , and fixed output bias $b_i^{pred} = -\frac{p_m p_r}{(1-(1-p_c)p_m)Tv}$ for $i \in \{1, \dots, Tv\}$ . Masking parameters are masking rate $p_m$ , correct word probability $p_c$ , and random word probability $p_r$ .

Under the masked language modeling objective with squared loss, the optimal word embedding Gram matrix $E := (W^E)^\top W^E \in \mathbb{R}^{(Tv+1) \times (Tv+1)}$ satisfies the following conditions for some constants $u_0, u_1, \dots, u_{Tv} \in \mathbb{R}$ :
1. 0-th Row ( $[\text{MASK}]$ interactions): $E_{00} = -\left(\frac{1}{p_m(1-p_c-p_r)} - 1\right) u_0$ $\forall t \in [T], \quad \sum_{l \in t} E_{0l} = u_0 v$
2. 0-th Column: $\forall i \in \{1, \dots, Tv\}, \quad E_{i0} = -\left(\frac{1}{(1-p_c-p_r)p_m} - 1\right) u_i$
3. Word-to-Word Blocks ( $i, j \in \{1, \dots, Tv\}$ ): $\sum_{l \in \text{topic}(i)} E_{il} = u_i v + \frac{1}{1-(1-p_c)p_m}$ $\forall t \in [T] \text{ such that } \text{topic}(i) \neq t, \quad \sum_{l \in t} E_{il} = u_i v$
This guarantees that the average inner product of embeddings between tokens belonging to the same topic is strictly larger (by $\frac{1}{v(1-(1-p_c)p_m)}$ ) than between tokens belonging to different topics.
Knowl 2 — Optimal Value Matrix under Uniform Attention and Regularization
theoretical result
Consider a single-layer transformer with frozen one-hot word embeddings ( $W^E = I \in \mathbb{R}^{(Tv+1) \times (Tv+1)}$ ), uniform attention ( $W^K = 0, W^Q = 0$ ), and zero prediction bias ( $b^{pred} = 0$ ), trained on topic model data under the $L_2$ -regularized masked language modeling squared loss $L_{\text{reg}}(W^V) = L(W^V) + \lambda \|W^V\|_F^2$ .

In the vanishing regularization limit $\lim_{\lambda \to 0} \arg\min W^V$ , the optimal value matrix $W^{V*} \in \mathbb{R}^{(Tv+1) \times (Tv+1)}$ converges uniquely to a block-wise structure:
1. 0-th Row: $\forall j \in \{0, \dots, Tv\}, \quad W^{V*}_{0j} = 0$
2. 0-th Column: $\forall i \in \{1, \dots, Tv\}, \quad W^{V*}_{i0} = \frac{c_2 c_3 - c_1 Tv}{c_2^2 + Tv}$
3. Word-to-Word Entries ( $i, j \in \{1, \dots, Tv\}$ ): $W^{V*}_{il} = W^{V*}_{\text{diff-topic}} := -\frac{c_1 c_2 + c_3}{c_2^2 + Tv} \quad \text{for } l \notin \text{topic}(i)$ $W^{V*}_{il} = W^{V*}_{\text{same-topic}} := W^{V*}_{\text{diff-topic}} + \frac{c_3}{v} \quad \text{for } l \in \text{topic}(i)$
where the positive constants are: $c_1 = \frac{p_r}{(1-p_c-p_r)(1-(1-p_c)p_m)Tv} \in (0, 1)$ $c_2 = \frac{1}{(1-p_c-p_r)p_m} - 1 \in (0, \infty)$ $c_3 = \frac{1}{1-(1-p_c)p_m} \in (1, \infty)$

Thus, each unmasked token provides a positive contribution to predicting tokens of its own topic and negative contribution to different-topic tokens.
Knowl 3 — Optimal Attention Weights with Block-Wise Value Matrix
theoretical result
Let a single-layer transformer with frozen one-hot embeddings have its value matrix $W^V$ frozen to the uniform block-wise optimum $W^{V*}$ . Assume the attention matrix $A(\tilde{X}) \in (0, 1)^{N \times N}$ satisfies symmetric topic-wise weights: $A(\tilde{X})_{ij} = \begin{cases} \beta c_3 & \text{if } \tilde{w}_i = \tilde{w}_j \\ \alpha c_3 & \text{if } \tilde{w}_i \neq \tilde{w}_j \land \text{topic}(\tilde{w}_i) = \text{topic}(\tilde{w}_j) \\ c_3 & \text{if } \text{topic}(\tilde{w}_i) \neq \text{topic}(\tilde{w}_j) \end{cases}$ where $c_3$ normalizes each column of attention to 1.

Under the asymptotic conditions $T \to \infty$ , $\tau \to \infty$ with $\tau = o(T)$ (sparse documents), $v > (\frac{1}{1-(1-p_c)p_m} + 1)^2 + 1$ , $p_m < 1/2$ , and $p_c = p_r \in (0, 1/2)$ , the optimal parameters $(\alpha, \beta)$ that minimize the masked language modeling squared loss satisfy: $\frac{v-1}{v}\alpha + \frac{1}{v}\beta \in \left(\lambda_1(\tau - 1), \lambda_2 T\right)$ where: $\lambda_1 := \frac{(1-(1-p_c)p_m + p_m p_r)(1 + (1-p_c)p_m)}{2(1-(1-p_c)p_m)}$ $\lambda_2 := 100 \left(\frac{1-(1-p_c)p_m}{p_m p_r} + 1\right)$

Because $\tau \to \infty$ , the weighted average $\frac{v-1}{v}\alpha + \frac{1}{v}\beta > 1$ , showing that same-topic attention is on average strictly larger than different-topic attention.
Knowl 4 — Optimal Attention Weights with Diagonal Value Matrix
theoretical result
When a single-layer transformer is trained without $L_2$ regularization, $W^V$ can converge to a diagonal optimum $D^V \in \mathbb{R}^{(Tv+1) \times (Tv+1)}$ where $D^V_{ij} = 0$ for all $i \neq j$ ( $i, j \in \{1, \dots, Tv\}$ ). Freezing $W^V = D^V$ and training the attention parameters $(\alpha, \beta)$ under the topic-wise attention parameterization yields optimal values satisfying: $\lambda_3 \tau < \beta < \lambda_4 T$ $\alpha < \lambda_5 \beta$

where: $\lambda_3 := \frac{1-(1-p_c)p_m + p_m p_r}{100} v$ $\lambda_4 := \frac{1-(1-p_c)p_m}{\sqrt{v-1} - 2 + (1-p_c)p_m} \cdot \frac{1-(1-p_c-p_r)p_m}{p_m p_r} v$ $\lambda_5 := \frac{1}{(v-1)(1-(1-p_c)p_m)}$

Under this regime, same-word attention $\beta$ is strongly amplified compared to cross-topic attention ( $1$ ), while same-topic-different-word attention $\alpha$ is suppressed relative to $\beta$ .
Knowl 5 — Two-Stage Learning Dynamics in Self-Attention Optimization
theoretical result
When a single-layer transformer's key matrix $W^K$ , query matrix $W^Q$ , and value matrix $W^V$ are initialized to random matrices near zero and trained jointly with gradient descent on masked language modeling, the learning dynamics naturally separate into two distinct stages:
1. Stage 1 (Norm Growth of $W^V$ ): The gradient $\nabla_{W^K} L$ is proportional to $W^Q \approx 0$ , where each column of $W^Q$ sums to approximately $0$ . Conversely, $\nabla_{W^V} L$ scales with the softmax-normalized attention matrix $A(\tilde{X})$ , whose columns sum to exactly $1$ . Consequently, $\|\nabla_{W^V} L\| \gg \|\nabla_{W^K} L\|$ , causing $\|W^V\|_F$ and the orientation of $W^V$ to change rapidly while $\|W^K\|_F$ and $\|W^Q\|_F$ stay near zero.
2. Stage 2 (Norm Growth of $W^K, W^Q$ ): After $\|W^V\|_F$ plateaus and exhibits a block-wise or diagonal topic structure, $\|W^K\|_F$ and $\|W^Q\|_F$ grow significantly to optimize the attention distribution conditioned on the learned $W^V$ .
Knowl 6 — Topic Model Data Distribution and Masked Language Modeling Objective
definition
The theoretical data sandbox consists of $T$ topics $\{1, \dots, T\}$ and a vocabulary of $Tv$ words $\{1, \dots, Tv\}$ . Each word $i$ belongs to a unique topic $\text{topic}(i) = \lceil i/v \rceil$ .

Each document $w = w_1 \dots w_N$ is generated by:
1. Choosing $\tau$ distinct topics $\{t_1, \dots, t_\tau\}$ uniformly from $[T]$ .
2. For each position $n \in [N]$ , drawing topic $t \in \{t_1, \dots, t_\tau\}$ uniformly, then drawing word $w_n$ uniformly from topic $t$ .
The masked document $\tilde{w}$ applies masking probabilities $p_m, p_c, p_r \in (0, 1)$ :
- Each position $i \in [N]$ is selected for masking independently with probability $p_m$ .
- If selected, $\tilde{w}_i = w_i$ with probability $p_c$ , $\tilde{w}_i = \text{uniform}(1, \dots, Tv)$ with probability $p_r$ , and $\tilde{w}_i = [\text{MASK}] := 0$ with probability $1 - p_c - p_r$ .
Given one-hot encoding $\tilde{X} \in \{0, 1\}^{(Tv+1) \times N}$ , the network prediction $\hat{X} = f_\theta(\tilde{X}) \in \mathbb{R}^{(Tv+1) \times N}$ is trained by minimizing: $L(\theta) = \mathbb{E}_{X \sim D_X, M} \left[\frac{1}{|M|} \sum_{j \in M} l(f_\theta(\tilde{X})_{:j}, X_{:j})\right]$ with squared loss $l(x, y) = \|x - y\|_2^2$ or cross-entropy loss $l(x, y) = -\log \frac{\exp(x_i)}{\sum_j \exp(x_j)}$ .

Knowl 7 — Topic Structure Alignment in Pretrained Language Models

data/table

Evaluation of pretrained transformers on Wikipedia with an LDA topic model ( $T = 100$ topics) shows that token embeddings and attention weights systematically encode topic structure. Embeddings between same-topic words have higher cosine similarity and dot products than different-topic words, and debiased attention between same-topic words is higher than cross-topic words.

Model	Ambiguity	Avg Embedding	Avg Embedding	Avg Attn Weight
	Threshold	Cosine Sim (S/D)	Dot Product (S/D)	Ratio (S/D)
BERT	0.0005	1.21	1.19	1.32
	0.001	1.13	1.15	1.28
	0.002	1.11	1.13	1.22
ALBERT	0.0005	5.64	6.29	1.33
	0.001	4.18	3.74	1.28
	0.002	3.24	2.93	1.22
BART	0.0005	2.80	2.67	1.35
	0.001	1.95	1.92	1.31
	0.002	1.63	1.62	1.23
ELECTRA	0.0005	5.98	5.37	2.14
	0.001	7.70	7.35	2.09
	0.002	7.46	8.08	1.95
RoBERTa	0.0005	6.44	6.81	1.40
	0.001	5.73	6.31	1.31
	0.002	5.24	5.30	1.22
BERT (Random)	0.0005	1.00080	1.00063	0.99943
	0.001	0.99974	1.00036	0.99996
	0.002	1.00016	1.00027	1.00007

Ambiguity threshold $\theta$ denotes the fraction of highest-likelihood words kept per topic. In contrast to trained models where same-to-different topic ratios strictly exceed 1, a randomly initialized BERT model yields ratios near 1.0.

Knowl 8 — Sentence Length Debiasing and Topic Filtering Protocol for Attention Probing
algorithm
To evaluate topic-structure encoding in pretrained language models on natural corpora without confounding length bias or polysemy, the following filtering and normalization algorithm is used:
Input: Pretrained transformer model $f_\theta$ , corpus documents $\mathcal{D}$ , LDA model with $T$ topics, ambiguity threshold $\theta \in (0, 1)$
Output: Debiased same-topic and different-topic attention weights
for each topic $t \in [T]$ do
    Sort vocabulary words by LDA likelihood: $(w_1, p_1), \dots, (w_n, p_n)$ in descending order
    Assign words $\{w_i : i \le \theta n\}$ to topic $t$ , excluding stop words
end for
Filter vocabulary so that each word is assigned to at most one topic
for each sentence $s = (w_1, \dots, w_L) \in \mathcal{D}$ do
    Compute raw pairwise attention matrix $A \in \mathbb{R}^{L \times L}$
     $A^{\text{debiased}} \leftarrow A \cdot \frac{L}{100}$
    for each pair $(i, j)$ with $i \neq j$ do
        if $\text{topic}(w_i) == \text{topic}(w_j)$ then
            Accumulate $A^{\text{debiased}}_{ij}$ into same-topic attention pool
        else if $\text{topic}(w_i) \neq \text{topic}(w_j)$ then
            Accumulate $A^{\text{debiased}}_{ij}$ into different-topic attention pool
        end if
    end for
end for
return Average of same-topic attention pool, Average of different-topic attention pool
Knowl 9 — Attention Weight Ordering Under Uniform-Block versus Diagonal Value Matrix Optima
data/table
Empirical attention measurements on synthetic LDA data under cross-entropy loss verify that the structure of $W^V$ determines the relative ordering of attention types:
1. Uniform-Block $W^V$ (Adam with LR 0.01): Attention between same-topic words exceeds attention between different-topic words:
- Average Same-Word Attention: $0.00811 \pm 0.000705$
- Average Same-Topic-Different-Word Attention: $0.0100 \pm 0.000392$
- Average Different-Topic Attention: $0.00707 \pm 0.000178$
1. Diagonal $W^V$ (Adam with LR 0.01): Same-word attention dominates, while same-topic-different-word attention is suppressed below different-topic attention:
- Average Same-Word Attention: $0.0918 \pm 0.00244$
- Average Same-Topic-Different-Word Attention: $0.00182 \pm 0.000474$
- Average Different-Topic Attention: $0.00256 \pm 0.000109$

Coverage note — Omitted intermediate algebraic lemmas (such as Lemma 1 and Propositions 1-3), non-asymptotic numeric loss landscape grid search plots (Figures 7-8), and the non-exclusive multi-topic empirical extension table (Table 5) as their main qualitative insights are fully captured by the primary theorems and tables.

References

1.Arora, S., Ge, R., Koehler, F., Ma, T., and Moitra, A. Provable algorithms for inference in topic models. In Balcan, M. F. and Weinberger, K. Q. (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2859–2867, New York, New York, USA, 20–22 Jun 2016. PMLR. URL https://proceedings.mlr.press/v48/arorab16.html.
2.Awasthi, P. and Risteski, A. On some provably correct cases of variational inference for topic models. In Cortes, C., Lawrence, N., Lee, D., Sugiyama, M., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL https://proceedings.neurips.cc/paper/2015/file/68a83eeb494a308fe5295da69428a507-Paper.pdf.
3.Belinkov, Y. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207–219, March 2022. doi: 10.1162/coli_a_00422. URL https://aclanthology.org/2022.cl-1.7.
4.Bhattamishra, S., Ahuja, K., and Goyal, N. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096–7116, Online, November 2020a. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.576. URL https://aclanthology.org/2020.emnlp-main.576.
5.Bhattamishra, S., Patel, A., and Goyal, N. On the computational power of transformers and its implications in sequence modeling. In Proceedings of the 24th Conference on Computational Natural Language Learning, pp. 455–475, Online, November 2020b. Association for Computational Linguistics. doi: 10.18653/v1/2020.conll-1.37. URL https://aclanthology.org/2020.conll-1.37.
6.Blei, D. M., Ng, A. Y., and Jordan, M. I. Latent dirichlet allocation. J. Mach. Learn. Res., 3(null):993–1022, mar 2003. ISSN 1532-4435.
7.Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877–1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
8.Clark, K., Khandelwal, U., Levy, O., and Manning, C. D. What does BERT look at? an analysis of BERT’s attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276–286, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4828. URL https://aclanthology.org/W19-4828.
9.Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423.
10.Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=YicbFdNTTy.
11.Ebrahimi, J., Gelda, D., and Zhang, W. How can self-attention networks recognize Dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4301–4306, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.findings-emnlp.384. URL https://aclanthology.org/2020.findings-emnlp.384.
12.Edelman, B. L., Goel, S., Kakade, S., and Zhang, C. Inductive biases and variable creation in self-attention mechanisms. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 5793–5831. PMLR, 17–23 Jul 2022. URL https://proceedings.mlr.press/v162/edelman22a.html.
13.Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://transformer-circuits.pub/2021/framework/index.html.
14.Gers, F. and Schmidhuber, J. Lstm recurrent networks learn simple context-free and context-sensitive languages. IEEE transactions on neural networks, 12 6:1333–40, 2001.
15.Hewitt, J. and Liang, P. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 2733–2743, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1275. URL https://aclanthology.org/D19-1275.
16.Hewitt, J. and Manning, C. D. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4129–4138, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1419. URL https://www.aclweb.org/anthology/N19-1419.
17.Hewitt, J., Hahn, M., Ganguli, S., Liang, P., and Manning, C. D. RNNs can generate bounded hierarchical languages with optimal memory. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1978–2010, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.156. URL https://www.aclweb.org/anthology/2020.emnlp-main.156.
18.Htut, P. M., Phang, J., Bordia, S., and Bowman, S. R. Do attention heads in bert track syntactic dependencies? ArXiv, abs/1911.12246, 2019.
19.Jelassi, S., Sander, M. E., and Li, Y. Vision transformers provably learn spatial structure. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=eMW9AkXaREI.
20.Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583–589, 2021.
21.Kovaleva, O., Romanov, A., Rogers, A., and Rumshisky, A. Revealing the dark secrets of BERT. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4365–4374, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1445. URL https://aclanthology.org/D19-1445.
22.Li, X. and Gong, H. Robust optimization for multilingual translation with imbalanced data. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 25086–25099. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/file/d324a0cc02881779dcda44a675fdcaaa-Paper.pdf.
23.Li, Y. and Risteski, A. The limitations of limited context for constituency parsing. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 2675–2687, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.208. URL https://aclanthology.org/2021.acl-long.208.
24.Liu, B., Hsu, D., Ravikumar, P. K., and Risteski, A. Masked prediction: A parameter identifiability view. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=Hbvlb4D1aFC.
25.Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=De4FYqjFueZ.
26.Liu, L., Liu, X., Gao, J., Chen, W., and Han, J. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747–5763, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.463. URL https://aclanthology.org/2020.emnlp-main.463.
27.Luo, Z., Weng, C., Wu, S., Zhou, M., and Ge, R. One objective for all models–self-supervised learning for topic models. arXiv preprint arXiv:2203.03539, 2022.
28.Meng, Y., Zhang, Y., Huang, J., Zhang, Y., and Han, J. Topic discovery via latent space clustering of pretrained language model representations. In Proceedings of the ACM Web Conference 2022, WWW ’22, pp. 3143–3152, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450390965. doi: 10.1145/3485447.3512034. URL https://doi.org/10.1145/3485447.3512034.
29.Merrill, W. Sequential neural networks as automata. In Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges, pp. 1–13, Florence, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-3901. URL https://www.aclweb.org/anthology/W19-3901.
30.Nguyen, T. Q. and Salazar, J. Transformers without tears: Improving the normalization of self-attention. In Proceedings of the 16th International Conference on Spoken Language Translation, Hong Kong, November 2-3 2019. Association for Computational Linguistics. URL https://aclanthology.org/2019.iwslt-1.17.
31.Press, O. and Wolf, L. Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, pp. 157–163, Valencia, Spain, April 2017. Association for Computational Linguistics. URL https://aclanthology.org/E17-2025.
32.Sia, S., Dalmia, A., and Mielke, S. J. Tired of topic models? clusters of pretrained word embeddings make for fast and good topics too! In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1728–1736, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.135. URL https://aclanthology.org/2020.emnlp-main.135.
33.Siegelmann, H. T. and Sontag, E. D. On the computational power of neural nets. In Proceedings of the Fifth Annual Workshop on Computational Learning Theory, COLT ’92, pp. 440–449, New York, NY, USA, 1992. Association for Computing Machinery. ISBN 089791497X. doi: 10.1145/130385.130432. URL https://doi.org/10.1145/130385.130432.
34.Snell, C., Zhong, R., Klein, D., and Steinhardt, J. Approximating how single head attention learns, 2021.
35.Sontag, D. and Roy, D. Complexity of inference in latent dirichlet allocation. In Shawe-Taylor, J., Zemel, R., Bartlett, P., Pereira, F., and Weinberger, K. (eds.), Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011. URL https://proceedings.neurips.cc/paper/2011/file/3871bd64012152bfb53fdf04b401193f-Paper.pdf.
36.Sun, K. and Marasović, A. Effective attention sheds light on interpretability. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pp. 4126–4135, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.findings-acl.361. URL https://aclanthology.org/2021.findings-acl.361.
37.Sun, X. and Lu, W. Understanding attention for text classification. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 3418–3428, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-main.312. URL https://aclanthology.org/2020.acl-main.312.
38.Suzgun, M., Belinkov, Y., Shieber, S., and Gehrmann, S. LSTM networks can perform dynamic counting. In Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges, pp. 44–54, Florence, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-3905. URL https://www.aclweb.org/anthology/W19-3905.
39.Talebpour, M., García Seco de Herrera, A., and Jameel, S. Topics in contextualised attention embeddings. In Kamps, J., Goeuriot, L., Crestani, F., Maistro, M., Joho, H., Davis, B., Gurrin, C., Kruschwitz, U., and Caputo, A. (eds.), Advances in Information Retrieval, pp. 221–238, Cham, 2023. Springer Nature Switzerland. ISBN 978-3-031-28238-6.
40.Tenney, I., Das, D., and Pavlick, E. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4593–4601, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1452. URL https://aclanthology.org/P19-1452.
41.Thompson, L. and Mimno, D. Topic modeling with contextualized word representation clusters, 2020.
42.Tosh, C., Krishnamurthy, A., and Hsu, D. Contrastive estimation reveals topic posterior information to linear models. J. Mach. Learn. Res., 22(1), jan 2021. ISSN 1532-4435.
43.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. u., and Polosukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
44.Vig, J. and Belinkov, Y. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 63–76, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4808. URL https://aclanthology.org/W19-4808.
45.Wei, C., Chen, Y., and Ma, T. Statistically meaningful approximation: a case study on approximating turing machines with transformers, 2021. URL https://arxiv.org/abs/2107.13163.
46.Weiss, G., Goldberg, Y., and Yahav, E. On the practical computational power of finite precision RNNs for language recognition. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 740–745, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10.18653/v1/P18-2117. URL https://www.aclweb.org/anthology/P18-2117.
47.WikimediaFoundation. Wikimedia downloads. Wikimedia Downloads, 2023. URL https://dumps.wikimedia.org.
48.Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38–45, Online, October 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/2020.emnlp-demos.6.
49.Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T.-Y. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020.
50.Yao, S., Peng, B., Papadimitriou, C., and Narasimhan, K. Self-attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3770–3785, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.292. URL https://aclanthology.org/2021.acl-long.292.
51.Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S., and Kumar, S. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=ByxRM0Ntvr.
52.Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. Why are adaptive methods good for attention models? In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 15383–15393. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/b05b57f6add810d3b7490866d74c0053-Paper.pdf.
53.Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., and Wagner, T. Unveiling transformers with lego: a synthetic reasoning task, 2022a. URL https://arxiv.org/abs/2206.04301.
54.Zhang, Z., Fang, M., Chen, L., and Namazi Rad, M. R. Is neural topic modelling better than clustering? an empirical study on clustering with contextual embeddings for topics. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 3886–3893, Seattle, United States, July 2022b. Association for Computational Linguistics. doi: 10.18653/v1/2022.naacl-main.285. URL https://aclanthology.org/2022.naacl-main.285.

Citation

MLA

Li, Y., et al. “How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding”. International Conference on Machine Learning, vol. 202, 2023, pp. 19689–729, https://proceedings.mlr.press/v202/li23p.html.

APA

Li, Y., Li, Y., & Risteski, A. (2023). How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding. International Conference on Machine Learning, 202, 19689–19729. https://proceedings.mlr.press/v202/li23p.html

Chicago

Li, Y., Y. Li, and A. Risteski. 2023. “How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding”. International Conference on Machine Learning 202: 19689–729. https://proceedings.mlr.press/v202/li23p.html.

Harvard

Li, Y., Li, Y. and Risteski, A. (2023) “How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding”, International Conference on Machine Learning. PMLR, pp. 19689–19729. Available at: https://proceedings.mlr.press/v202/li23p.html.

Vancouver

1. Li Y, Li Y, Risteski A (2023) How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding. In: International Conference on Machine Learning. PMLR, pp 19689–19729

BibTeX

@InProceedings{pmlr-v202-li23p,
  title = 	 {How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding},
  author =       {Li, Yuchen and Li, Yuanzhi and Risteski, Andrej},
  booktitle = 	 {Proceedings of the 40th International Conference on Machine Learning},
  pages = 	 {19689--19729},
  year = 	 {2023},
  editor = 	 {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan},
  volume = 	 {202},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {23--29 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v202/li23p/li23p.pdf},
  url = 	 {https://proceedings.mlr.press/v202/li23p.html},
  abstract = 	 {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.}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Optimal Word Embeddings Encode Topic Structure

Knowl 2 — Optimal Value Matrix under Uniform Attention and Regularization

Knowl 3 — Optimal Attention Weights with Block-Wise Value Matrix

Knowl 4 — Optimal Attention Weights with Diagonal Value Matrix

Knowl 5 — Two-Stage Learning Dynamics in Self-Attention Optimization

Knowl 6 — Topic Model Data Distribution and Masked Language Modeling Objective