Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers

## Abstract

Abstract In addition to the ability to generate fluent text in various languages, large language models have been successful at tasks that involve basic forms of logical “reasoning” over their context. Recent work found that selectively removing certain components from weight matrices in pre-trained models can improve such reasoning capabilities. We investigate this phenomenon further by carefully studying how certain global associations tend to be stored in specific weight components or Transformer blocks, in particular feed-forward layers. Such associations may hurt predictions in reasoning tasks, and removing the corresponding components may then improve performance.
We analyze how this arises during training, both empirically and theoretically, on a two-layer Transformer trained on a basic reasoning task with noise, a toy associative memory model, and on the Pythia family of pre-trained models tested on simple reasoning tasks.

## 1 Introduction

Large language models (LLMs) have shown impressive capabilities on a variety of tasks, from generating coherent and grammatically correct text, to language understanding and basic mathematical reasoning (Brown et al., ; Touvron et al., ).
At the heart of this success is the Transformer architecture (Vaswani et al., ), which relies on a sequence of self-attention and feed-forward layers to efficiently combine information from the input context and patterns learned from training data.
Despite recent progress on interpreting the mechanisms learned by different layers (Meng et al., ; Wang et al., ), these models remain largely black boxes.
A better understanding of the role of Transformer layers and how they are affected by the training process could enable new monitoring and editing techniques, better training data, and ultimately more reliable LLMs.

The task of next-token prediction in language modeling inherently involves different subtasks that may be at odds with each other. For instance, given the context “John gave a book to”, the word “the” is a natural and grammatically correct next word to predict, and relying on global bigram statistics might be enough to predict it given the last word “to”. Nonetheless, if another character is present in the context, say Mary, then the name “Mary” may be a better prediction, and this would require a more involved form of “reasoning” over the context to retrieve this name. Previous work on interpretability has found that “circuits” of attention heads seem responsible for such in-context predictions (Wang et al., ), while feed-forward layers may be storing more global statistics such as the bigram “to the” or general factual knowledge (Geva et al., ; Meng et al., ).
The recent work (Sharma et al., ) found that selectively replacing certain layer weights to their low-rank approximation may improve performance on various reasoning benchmarks, and observed that the truncated components were often responsible for predicting “generic” words such as “the”.

In this paper, we provide a finer understanding of these phenomena by studying how such mechanisms arise during training, in particular how global associations, such as the bigram “to the”, can be localized to specific components or layers of the model weights. We first investigate this on pre-trained language models, namely the Pythia family, which has checkpoints available at different training steps (Biderman et al., ). We then provide a fine-grained study of dynamics on simple data models and architectures exhibiting similar properties:

- •
In a two layer transformer architecture trained on an in-context recall task similar to Bietti et al. (), but with additional noise on in-context tokens, we show that the noise is mainly learned in feed-forward layers, even for large noise levels. Removing those layers then leads to clean in-context predictions. We provide some theoretical justification through the first gradient step.
- •
In a linear associative memory model trained on data involving a common noise token, we show that the noise can be identified in a rank-one subspace of the weights. When the noise level is small, low-rank truncation can filter it out and predict clean outputs.

Overall, we provide a useful description of how global associations and in-context reasoning mechanisms are learned during training, and tend to be disentangled in different parts of the model, such that selectively removing certain components may lead to better predictions in reasoning tasks.

#### Related work.

Sharma et al. () recently empirically observed that a low-rank approximation of some weights in some pre-trained LLMs can improve reasoning capabilities.
Several interpretability works have looked at the role of attention versus feed-forward layers for different tasks. The prominence of feed-forward/MLP layers for storing “global” or “persistent” associations or facts has been observed in (Sukhbaatar et al., ; Geva et al., ; Meng et al., ; Geva et al., ). In contrast, several works have investigated the role of attention heads for “reasoning” or computation over the context, *e.g.*, for simple copying mechanisms with so-called induction heads (Elhage et al., ; Olsson et al., ; Bietti et al., ), or for more complex tasks (Merrill et al., ; Wang et al., ; Zhang et al., ; Liu et al., ; Sanford et al., ).

Training dynamics of transformers and attention have been studied in various works (Snell et al., ; Jelassi et al., ; Li et al., ; Oymak et al., ; Tian et al., ; Bietti et al., ; Reddy, ; Tian et al., ; Zhang et al., ; Nichani et al., ; Edelman et al., ). In particular, the two-layer model and copy task we consider are similar to Bietti et al. (), yet their data model does not involve noise on in-context predictions, and they do not study learning of global associations. Reddy () study in-context vs. in-weights learning empirically, on a different task than ours. Cabannes et al. () study training dynamics of linear associative memories, but focuses on deterministic data while our setup has noise.
Training dynamics were also studied empirically for interpretability (Olsson et al., ; Nanda et al., ; Quirke et al., ; Chen et al., ).
Edelman et al. (); Bai et al. (); Abernethy et al. () studied sample complexity of self-attention and in-context learning operations, but did not consider training dynamics.

## 2 Background and Motivation

In this section, we provide some background and motivation on reasoning tasks and rank reduction, and conduct initial investigations on pre-trained language models.

### 2.1 Reasoning from Context

Recent LLMs have shown promising results in more complex “reasoning” tasks which may involve multiple steps of logical or computational processing from context or prompt (Srivastava et al., ; Wei et al., ; Bubeck et al., ; Dziri et al., ), as opposed to simple pattern matching or memorization of training data, for instance using learned n-gram predictions.

While it is difficult to clearly separate reasoning from memorization, in this work we will make the simplifying distinction that reasoning involves dependencies between *multiple tokens* potentially far away in the context, while we consider global associations as simpler predictions that only depend on the *last token*, e.g., through a global bigram model. Thus, reasoning will typically require using attention operations in Transformers over context, while feed-forward layers should suffice for learning global associations.

Under this definition, we list a few simple examples of reasoning that we will consider in the sequel:

- •
*In-context recall*: when the last token is a, we’d like to copy the token that follows previous occurrences of a in the context. This [.. a b .. a] $\to$ b pattern typically requires a two-layer attention mechanism known as an *induction head* (Elhage et al., ; Bietti et al., );
- •
*Indirect object identification (IOI)*: we consider contexts of the form “When Mary and John went to the store, John gave the ice cream to” where the prediction should be “Mary” (IO, the indirect object), instead of “John” (S, the subject). Wang et al. () found a circuit of several attention heads that perform this task by copying the name which only occurs once in the context;
- •
*Factual recall*: sentences of the form “Paul Citroen is a native speaker of” with target “Dutch” as in (Sharma et al., ). While this may be seen as retrieving a global association, we will treat it here as reasoning since it involves combining the subject and relation from the context, while a global bigram that only depends on the last token “of” might instead predict the word “the.”

We note that our assumption of global associations depending only on the last token is mainly for convenience of our analysis. In practice, the last token’s representation at intermediate layers of the Transformer may contain additional information from the context, and our arguments can easily extend to global associations that only depend on that representation. For instance, this could include previous tokens thanks to position-based attention heads (Voita et al., ; Elhage et al., ; Akyürek et al., ), which allows global n-grams instead of just bigrams.

### 2.2 LASER: Layer-Selective Rank Reduction

Sharma et al. () observed that reducing the rank of MLP matrices in certain layers of LLMs effectively brings better performance on several reasoning benchmarks. Their proposed method, Layer-Selective Rank Reduction (LASER), replaces any matrix in the full model by its low-rank approximation with fraction $\rho$, i.e., a matrix $\mathbf{W}\in\mathbb{R}^{d_{\text{in}},d_{\text{out}}}$ would be replaced by its rank-$\lfloor\rho\cdot\min\{d_{\text{in}},d_{\text{out}}\}\rfloor$ approximation via Singular Value Decomposition (SVD). After searching for the best parameters of different models on different datasets, Sharma et al. () concludes that the best practice for LLMs is to conduct LASER on weight matrices of MLPs on relatively deep layers.
The optimal $\rho$ is smaller than $0.2$ for many datasets. We refer to their Table 3 for more results of the parameters after searching.

Another observation from Sharma et al. () is that, when LASER improves the model’s prediction on some samples, the full model often predicts “generic” words while the improved model is able to predict the ground-truth answer.
For instance, given an input “Madrid is located in”, the full model predicts “the” while the truncated model predicts the target “Spain” in Table 1.
Here, the generic word is consistent with our definition of global associations in Section 2.1, as it may naturally follow from a bigram distribution conditioned on “in”, while the factual answer is more akin to reasoning from context. Thus, we would like to better understand how LASER improves the model from predicting generic words to inferring the answer from context, and how such a gap appears during training.

**Table 1: Probabilities of the top-5 next-tokens in Pythia-1B before and after LASER. The input prompt is “Madrid is located in”. Probabilities of two generic words, *i.e.*, “the” and “a”, drop sharply after LASER, while probabilities of meaningful words increase, especially the target “Spain”.**
|  | “the” | “Spain” | “a” | “southern” | “northern” |
| --- | --- | --- | --- | --- | --- |
| Full | 0.499 | 0.079 | 0.069 | 0.023 | 0.021 |
| LASER | 0.027 | 0.300 | 0.002 | 0.044 | 0.046 |

### 2.3 An Investigation on GPT-2 Small and Pythia Models

In this section, we empirically investigate how LLMs process in-context vs global associations, and how this evolves during training. We consider GPT-2 small and Pythia models on the indirect object identification (IOI) and factual recall tasks described in Section 2.1.

IOI on GPT2 Small. Different from Wang et al. (), we would like to consider whether a model proposes an output beyond the input $x$. A quick demonstration is to consider the IOI task with input $x=$“When Mary and John went to a store, John gave a drink to”(^2^22Note that here we use “a” store instead of “the” store in the original example of Wang et al. (). The reason is to rule out the word “the“ from the input context.). The top 4 predicted tokens for GPT-2 Small (Radford et al., ) on $x$ are [“Mary”, “them”, “the”, “John”]. Although GPT-2 Small successfully predicts Mary (the IO target) instead of John (S), the other two top candidate tokens, i.e., “them” and “the”, do not even appear in the context.
This prominence of such “generic” words is similar to the factual recall example from Section 2.2, and plausibly follows from a global associative mechanism conditioned on the preposition “to”.

Therefore, for the above input $x$, we naturally extend the candidate set as $\mathcal{C}=\{$“Mary”, “them”, “the”, “John”$\}$. To verify whether or not the emergence of “the” is connected to the mechanism of LASER, we examine how the probability of each $c\in\mathcal{C}$ change after running LASER on different layers on GPT-2 Small in Figure 1. LASER on Layer 9, 10 and 11 turns out to significantly decrease the probability of predicting “the” and “them” compared with the full model.

Figure: Figure 1: Predicted probability for $c\in\{$“Mary”, “them”, “the”, “John”$\}$. LASER is conducted on input matrices of MLP layers on the layer $l=9,10,11,12$ of GPT-2 Small. The input is “When Mary and John went to a store, John gave a drink to”. The horizontal is the fraction of perserved rank, $\rho\in[0,1]$, where $\rho=1$ stands for the full model. It turns out LASER clearly decreases probability of “the” and “them” when $\rho\in[0.1,0.8]$ for layer $l=9,10,11$, compared with the full model.
Refer to caption: x1.png

The above demonstration on GPT-2 Small implies that, when a model introduces extra candidates beyond the input $x$, LASER may decrease the probability of predicting these extra candidates, which means LASER may enhance the model’s performance on contextual tasks.

IOI on Pythia-1B. Now we would like to verify this observation on more models and, more comprehensively, track the behavior of these models along training. We choose to conduct the IOI experiments on Pythia (Biderman et al., ), a family of models ranging in sizes from 14M to 12B trained on web data, with hundreds of training checkpoints for each size. We generate an IOI dataset of 100 sentences with random names for [IO] and [S] in each sample. Figure 2 reports the test results of Pythia-1B along training. Here LASER is conducted on MLP weights, with parameters given in Appendix B.2. LASER boosts the probability ratio of [IO] over “the”
from 2.3$\times$ to 12.3$\times$ at 14K steps.

Factual recall on Pythia-1B. As in Table 1, we verify factual recall with input as “Madrid is located in”. The full model of Pythia-1B generates “Madrid is located in the north of Spain”, while the model after LASER generates “Madrid is located in Spain”. We track the probability of predicting “Spain” and “the” along training in Figure 2. LASER turns out to boost the probability ratio of “Spain” over “the” from 0.16$\times$ to 11.3$\times$ at 14K steps.
We note that better prompting could avoid the need for LASER in this case (e.g., “Madrid is located in the country of” predicts “Spain”), but increases the context length and thus the inference cost, though this is outside the scope of this paper.

Figure: Figure 2: Left: average probability of tokens [IO], [S] and “the” in IOI task in the prediction by Pythia-1B along training. Right: average probability of tokens “Spain” and “the” in a factual task predicted by Pythia-1B along training, with input as “Madrid is located in”. In both tasks, the full model learns to predict “the” with high probability starting from $\sim$10 steps, and then learns to solve the tasks. LASER boosts the probability of correct answers against “the” in both tasks: the average probability ratio of correct answers against “the” improves from 2.3$\times$ to 12.3$\times$ (in IOI) and from 0.16$\times$ to 11.3$\times$ (in factual) at 14K steps.
Refer to caption: x2.png

Training dynamics on Pythia.
The behavior of the Pythia models on the IOI and factual recall tasks during their pre-training process displays several phases, as shown in Figure 2.
For IOI, we observe:

- i.
Initialization: all tokens have similar logits since the weights are random initialized.
- ii.
Between 10 and 1000 steps: the models consistently output “the”. They cannot solve IOI task at all, as long as they have almost the same output for [IO] and [S]. After 500 steps, [IO] starts the growth towards one of the top predictions.
- iii.
After 2000 steps: Pythia starts to be able to solve IOI task by preferring [IO] than [S] and “the”. Meanwhile, the benefit of LASER appears as enhancing the leading position of [IO].

Therefore, the training process reveals the capacity of predicting “the” is learnt much earlier than predicting [IO]. The reason might be that predicting “the” requires a simpler grammar structure, while predicting [IO] requires a complicated architecture of attention heads of different roles across layer (Wang et al., ). Then we note that the IOI task always has “to” before the masked [IO], which means “to” may be an indicator for the model to predict “the” with non-negligible probability.
Similarly, for factual recall we see early learning of the “generic” answer, while the factual answer is learned later.
Conceptually, if LLMs are able to write natural text or have been trained sufficiently with natural texts, it is not surprising for the model to predict “the” with high probability after seeing “to”. This is verified in Appendix B.1.

Implications from experiments. We summarize our main experimental observations of this section.

###### Observation 1 .

Global associations may “distract” LLMs away from in-context predictions, hurting performance on reasoning tasks.

###### Observation 2 .

LASER on MLP weights in LLMs helps inhibit predictions of global associations, thus improving in-context predictions.

###### Observation 3 .

During pre-training, global associations are learned earlier than complex reasoning.

These observations raise the following questions, which we investigate in the next sections.

Q1: *Why are global associations learned before than complex reasoning?*

Q2: *Are feed-forward layers responsible of learning global associations?*

###### Observation 1 .

###### Observation 2 .

###### Observation 3 .

## 3 Two-layer Transformer on Noisy In-context Recall

In this section, we consider two-layer transformers on an in-context recall task with added global noise, which allows us to study some key properties observed in Section 2 in a controlled setting. We empirically show how transformers solve this task by storing the noise in feed-forward layers, while attention implements the in-context mechanism. We then provide theory showing why feed-forward layers are more likely to store the global noise association, by studying gradients at initialization.

Data and task. The data model we consider is similar to Bietti et al. (), with additional noise. Consider a vocabulary $\mathcal{V}=\{1,2,\dots,N,N+1\}$. The token $N+1$ is the noise token. We fix a *trigger* token $q\in[N]$, which governs in-context recall, and a context length $T$. Each sequence of tokens $z_{1:T}=[z_{1},z_{2},\dots,z_{T}]$ is generated as follows:

- i.
Sample a correct output token $\bar{y}$ uniformly in $[N]$.
- ii.
Sample $z_{1:T-1}$ according to the following Markov process ($\pi_{u},\pi_{b}$ are distributions on $[N]$ defined later): $z_{1}\sim\pi_{u}(\cdot)$, and
$\displaystyle z_{t+1}|z_{t}\sim\begin{cases}\pi_{b}(\cdot|z_{t}),&\text{if }z_
{t}\neq q,\\
p_{\alpha,\bar{y}}(\cdot),&\text{otherwise,}\end{cases}~{}~{}~{}~{}~{}~{}~{}~{
}~{}~{}p_{\alpha,\bar{y}}(x)=\begin{cases}1-\alpha,&\text{if }x=\bar{y},\\
\alpha,&\text{if }x=N+1,\\
0,&\text{otherwise}.\end{cases}$
- iii.
Set $z_{T}=q$, and sample the final output $y=z_{T+1}\sim p_{\alpha,\bar{y}}(\cdot)$.

Note that the true $\bar{y}$ varies across sequences, so that the model needs to infer it from context, e.g., using an induction head as in (Bietti et al., ).
Predicting $\bar{y}$ may thus be seen as a basic “reasoning” task, yet when training with $\alpha>0$, the noisy output also requires the model to learn a global trigger-noise association, similar to the “to the” bigram discussed in Section 2. We also consider using multiple trigger tokens in Appendix A.3 and Figure 7.

Two-layer transformer. We consider a simplified two-layer transformer formulated below. The input is a sequence of tokens $z_{1:T}=[z_{1},\dots,z_{T}]\in[N+1]^{T}$, and the output is $\xi$. The embedding matrix $\mathbf{W}_{E}\in\mathbb{R}^{(N+1)\times d}$ and un-embedding matrix $\mathbf{W}_{E}\in\mathbb{R}^{(N+1)\times d}$ are fixed at random initialization. The two attention layers have learnable weights $\mathbf{W}_{KQ}^{1},\mathbf{W}_{V}^{1},\mathbf{W}_{KQ}^{2},\mathbf{W}_{V}^{2}
\in\mathbb{R}^{d\times d}$ with $\sigma(\cdot)$ the softmax on a vector. The two feed-forward layers $F_{1},F_{2}$ are also learnable, and typically we set them as two-layer MLPs with ReLU activation. We will discuss different architectural choices of $F_{1},F_{2}$ in Appendix A.4.
We use the cross-entropy loss to predict $y=z_{T+1}$ from the logits $\xi_{T}\in\mathbb{R}^{N+1}$.

$$ $\displaystyle x_{t}$ $\displaystyle\triangleq\mathbf{W}_{E}(z_{t})+p_{t},$ $\displaystyle h_{t}^{1}$ $\displaystyle\triangleq\sum_{s\leq t}\left[\sigma(x_{t}^{\top}\mathbf{W}_{KQ}^ {1}x_{1:t})\right]_{s}\cdot\mathbf{W}_{V}^{1}x_{s},$ $\displaystyle x_{t}^{1}$ $\displaystyle\triangleq x_{t}+h_{t}^{1}+F_{1}(x_{t}+h_{t}^{1}),$ $\displaystyle h_{t}^{2}$ $\displaystyle\triangleq\sum_{s\leq t}\left[\sigma({x_{t}^{1}}^{\top}\mathbf{W} _{KQ}^{2}x_{1:t}^{1})\right]_{s}\cdot\mathbf{W}_{V}^{2}x_{s}^{1},$ $\displaystyle x_{t}^{2}$ $\displaystyle\triangleq x_{t}^{1}+h_{t}^{2}+F_{2}(x_{t}^{1}+h_{t}^{2}),$ $\displaystyle\xi_{t}$ $\displaystyle\triangleq\mathbf{W}_{U}x_{t}^{2}.$ $$

Experimental observations. Following Bietti et al. (), we take $\pi_{u}$ and $\pi_{b}$ to be the unigram and brigram character-level distributions estimated from the tiny Shakespeare dataset with $N=65.$ The model setup includes $d=256$ and two-layer MLPs with ReLU for both $F_{1},F_{2}$. The training setup includes batch size as $512$ and the context length $T=256$.
When evaluating trained models, we consider LASER on the input weight $U_{in}$ of $F_{2}$.
We consider a noise level $\alpha=0.5$ for training data (though any other constant value would lead to similar observations).
During test time, we set $\alpha=0$ to compute the test loss, aiming to measure how likely the (full or after-LASER) model predicts the ground-truth $\bar{y}$.

Experimental results are reported in Figure 3 and 4. The full model predicts noise with probability close to $\alpha$, which is expected since it is trained to predict the noise token w.p. $\alpha$. However, when dropping the second-layer MLP $F_{2}$, the truncated model predicts the ground-truth $\bar{y}$ with an almost perfect probability $\approx 0.98$. This suggests that $F_{2}$ is responsible for storing the global association “[trigger] + [noise]”. Another observation is that the full model first learns to predict the noise with high probability in very early steps, after which it starts learning to predict the correct $\bar{y}$, which resembles the dynamics observed for learning the “to the” bigram in Pythia models in Figure 2. This suggests that learning the (global) trigger-noise association is easier than predicting $\bar{y}$, and we will study this theoretically in Section 3.1.

After the global noise association is learned, we observe a slower learning of an induction head mechanism, with similar dynamics to Bietti et al. ().
Compared to Bietti et al. (), we notice that the induction head (i.e., the second layer attention head) filters out the noise tokens and only attends to non-noisy output tokens following the trigger, corresponding to the correct $\bar{y}$, as shown in Figure 5. We present primitive exploration into this mechanism in Section 3.2.
Appendix A.1 summarizes roles of all components in the two-layer transformer in this task.

Figure: Figure 3: Average probability of predicting correct and noise tokens, and test loss on clean data ($\alpha=0$), with different fractions $\rho$ of preserved rank in $U_{in}$ of the second-layer MLP $F_{2}$. The full model learns to predict noise with probability around $\alpha=0.5$, as expected from training data. When $F_{2}$ is dropped ($\rho=0$), the model predicts the correct token $\bar{y}$ with probability $\approx 0.98$.
Refer to caption: x3.png

Figure: Figure 4: Fractions of predicting the noise token and the other non-noise tokens with $\alpha=0.5$. (Left) pretraining steps on noisy data; (right) finetuning steps on noisy data, after pretraining on clean data with $\alpha=1$. In both cases, the models learn to predict noise with probability nearly 0.5. In the first few ($\sim 5$) steps, the models quickly learn to predict noise with probability close to 1.
Refer to caption: x4.png

Figure: Figure 5: The second-layer attention scores of models trained with noise (left), fine-tuned with noise (right, initialized as a model pre-trained without noise), given the same input. It turns out both models learn to attend to the informative structure “[trigger]+$\bar{y}$” instead of “[trigger]+noise”. This implies that the attention in these models is only responsible to predict $\bar{y}$, although the training input and output have noise with probability $\alpha=\Theta(1)$.
Refer to caption: x5.png

### 3.1 Theoretical analysis: how and why do feed-forward layers store the noise?

As we saw in Figure 3 and 4, the model very quickly learns to predict the noise token after a few steps. Then the gap between $\rho=0$ and $1$ in Figure 3 suggests that the feed-forward layer $F_{2}$ is responsible for storing the global association about noise, which is verified in Figure 6 (middle).
We now provide theoretical justification for this behavior.
Understanding the full dynamics of the model used in our experiments is out of the scope of the present paper,
due to the many moving parts and the complexity of non-linear MLPs.
Instead, we focus on a simpler model involving one linear feed-forward layer and one attention layer, and look at the gradient dynamics near initialization. In particular, we will show that the gradients over the feed-forward parameters are much more informative than the attention gradient, which is dominated by noise unless the sample size is very large. This shows that the feed-forward layer is much more likely to capture the global association.

Simplified architecture and data.
Consider the input $x_{t}\in\mathbb{R}^{d}$ at position $t$ defined as
$x_{t}\triangleq\mathbf{W}_{E}(z_{t}),$
where $z_{t}\in[N+1]$ is the token at position $t$ and $\mathbf{W}_{E}(\cdot)$ returns its (untrained) embedding.
Here we ignore positional encoding for simplicity as it carries little signal at initialization, noting it could be easily incorporated.
For data generation, $\pi_{u}$ and $\pi_{b}$ are uniform distributions on $[N].$

Given a sequence of inputs, $x_{1:T}\in\mathbb{R}^{T\times d}$, the output of model is $\xi\triangleq\xi_{\text{attn}}+\xi_{\text{ff}}$ as

$$ $\displaystyle\xi_{\text{attn}}(x_{1:T})$ $\displaystyle\triangleq\mathbf{W}_{U}\phi(x_{T},x_{1:T})\in\mathbb{R}^{N+1},$ (1) $\displaystyle\xi_{\text{ff}}(x_{1:T})$ $\displaystyle\triangleq\mathbf{W}_{U}F(x_{T})=\mathbf{W}_{U}\mathbf{W}_{F}x_{T }\in\mathbb{R}^{N+1},$ $\displaystyle\phi(x_{T},x_{1:T})$ $\displaystyle\triangleq\sum_{t\leq T}\left[\sigma\left(x_{T}^{\top}\mathbf{W}_ {KQ}x_{1:T}\right)\right]_{t}\cdot\mathbf{W}_{V}x_{t}\in\mathbb{R}^{d},$ $$

where $\mathbf{W}_{U}\in\mathbb{R}^{(N+1)\times d}$ is the unembedding matrix, $\phi(s,t)$ is the attention module with query $s$ and context $t$, and $F(\cdot)$ is a linear feed-forward layer.
This architecture is similar to a one-layer transformer, but already highlights the difference between feed-forward and attention layers in a way that we expect to still hold for more layers.
In the above parametrization, the learnable matrices are $\mathbf{W}_{KQ},\mathbf{W}_{F},\mathbf{W}_{V}\in\mathbb{R}^{d\times d}$. At initialization, we set $\mathbf{W}_{KQ},\mathbf{W}_{F},\mathbf{W}_{V}=0$, noting that random initialization in high dimension would lead to similar behaviors thanks to near-orthogonality. Hence we assume all embeddings follow Assumption D.1.
We now look at the first gradient step from initialization, which has commonly been used to understand feature learning and sample complexity in neural networks (Damian et al., ; Ba et al., ; Dandi et al., ; Oymak et al., ; Bietti et al., ).
Note that $\mathbf{W}_{KQ}$ has no gradient at initialization, so that the gradient of $W_{V}$ is most relevant initially (see also Snell et al., ; Li et al., ; Oymak et al., ; Bietti et al., ).

###### Theorem 1 (Logits after one gradient step) .

Assume $N,T\gg 1,\alpha=\Theta(1)$. Consider a one gradient step update from zero-initialization on $m$ i.i.d. samples of $z_{1:T}$ with separate learning rates $\eta_{f}$ for $\mathbf{W}_{F}$ and $\eta_{v}$ for $\mathbf{W}_{V}$ (note that the gradient on $\mathbf{W}_{KQ}$ is zero).
With probability $1-\delta$, the resulting logits for the feed-forward and attention blocks satisfy,
for any test sequence $z_{1:T}$,

$$ $\displaystyle\left|\Delta(\xi_{\text{ff}}(x_{1:T}))-\eta_{f}\cdot\alpha\right|$ $\displaystyle\leq\eta_{f}\cdot O\left(\sqrt{\frac{\ln\frac{2(N+1)}{\delta}}{m} }\right),$ $\displaystyle\left|\Delta(\xi_{\text{attn}}(x_{1:T}))-\frac{\eta_{v}}{N}\cdot \hat{\alpha}\right|$ $\displaystyle\leq\eta_{v}\cdot O\left(\sqrt{\frac{(\frac{1}{TN}+\frac{1}{N^{2} })\ln\frac{2(N+1)}{\delta}}{m}}+\frac{\ln\frac{2(N+1)}{\delta}}{m}\right),$ $$

where $\Delta(\xi)=\xi_{N+1}-\max_{j\in[N]}\xi_{j}$ is the margin of predicting the noise token and $\hat{\alpha}=(\alpha^{2}\hat{q}+\alpha(1-\hat{q}))$, where $\hat{q}=\frac{1}{T}\sum_{t\leq T}\mathbbm{1}\{z_{t}=N+1\}$ is the fraction of noise tokens in $z_{1:T}$.

The margin $\Delta(\xi)$ reflects how much signal there is in the logits for predicting the noise token, and the theorem provides concentration bounds on the contributions of the updates on $\mathbf{W}_{F}$ and $\mathbf{W}_{V}$ to the margin. Note that $\hat{q}\ll 1$ w.h.p. for large $N,T$, so $\hat{\alpha}\approx\alpha$. We make the following observations:

- i.
When $m=\tilde{\Omega}(1)$, there is enough signal in $\mathbf{W}_{F}$ to predict the noise, say with $\eta_{f}=1$, and a choice of $\eta_{v}=O(1)$ will lead to
a small but controlled contribution to the prediction from $\mathbf{W}_{V}$.
- ii.
When $m=\tilde{\Omega}(N)$, $\mathbf{W}_{V}$ can also reliably predict the noise by setting $\eta_{v}=\Theta(N)$ (i.e., with small deviation on the r.h.s.), at the cost of many more samples.

Our result thus shows that in the initial phase of training, feed-forward layers are more likely to pick up the noise token, while attention will be slower due to additional noise and possibly smaller step-sizes. We may then expect the attention layers to focus instead on learning the induction head mechanism, as we observe empirically. Understanding this trade-off requires studying the dynamics of other attention parameters including key-query matrices, a much more involved endeavor which we leave to future work.

###### Theorem 1 (Logits after one gradient step) .

### 3.2 Theoretical insight: attention avoids attending to noise tokens

When the feed-forward weight learns to predict the noise as shown in Theorem 1, Figure 5 reveals that the second-layer attention in the two-layer model attends only towards the correct tokens. In contrast, a model pre-trained without noise has second-layer attention attend towards all tokens just after the triggers (Bietti et al., ), as observed in the attention pattern at the first step in Figure 5(right). Then, after being fine-tuned on noise data, the attention becomes only focused on the correct tokens. Understanding this mechanism requires the analysis of the dynamics of $\mathbf{W}_{KQ}$.

Following the simplified model and data distribution in Section 3.1, we take a step towards understanding how attention “avoids” the noise tokens, detailed in Appendix A.2. Concretely, this mechanism appears because, after the initial training phase, $\mathbf{W}_{V}$ has a minor structure that has a smaller projection onto $\mathbf{W}_{U}(N+1)\mathbf{W}_{E}(N+1)^{\top}$ as in Table 2, which makes $\mathbf{W}_{KQ}$ move negative in the direction of $\mathbf{W}_{E}(N+1)\mathbf{W}_{E}(q)^{\top}$.
A more detailed analysis of the dynamics of $\mathbf{W}_{KQ}$ throughout the training process would be an interesting avenue for future work.

## 4 Linear Associative Memory

In Section 3, we showed that *fully* truncating a feed-forward layer can be helpful for reasoning. We now present a setting where noisy associations are stored in a rank-one subspace of a layer, so that *intermediate* levels of truncation are more useful to remove noise.

Model and data. We consider a simple associative memory setting where the goal is learn an fixed permutation from input tokens to output tokens (w.l.o.g. taken to be the identity), with a linear model similar to Cabannes et al. (). Consider a learnable weight matrix $\mathbf{W}\in\mathbb{R}^{d\times d}$. Consider embeddings for $n$ input tokens as $\{e_{i}\}_{i=1}^{n}\subset\mathbb{R}^{d}$ and embeddings for $c$ output tokens as $\{u_{i}\}_{i=1}^{c}\subset\mathbb{R}^{d}$. In contrast to Cabannes et al. (), we consider an additional “common noise” output token $c=n+1$, which is chosen for any input with probability $\alpha\in(0,1)$.
For any input $x\in[n]$, the target distribution $p_{\alpha}(\cdot|x)$ is defined by
$p_{\alpha}(y|x)=(1-\alpha)\cdot\mathbbm{1}\{y=x\}+\alpha\cdot\mathbbm{1}\{y=c\}.$

In other words, the last channel ($c$) for output is the common noise with probability $\alpha$ for any input. The training dataset $\mathcal{D}_{\alpha}$ consists of uniformly distributed inputs $x\in[n]$, and outputs conditionally sampled as $y|x\sim p_{\alpha}(\cdot|x)$.
Given any pair of input and output tokens, the associative memory model takes the form
$f(i,j;\mathbf{W})\triangleq\langle u_{j},\mathbf{W}e_{i}\rangle,~{}\forall~{}i
,j\in[n]\times[c].$
When $k\leq d$, we denote the rank-$k$ approximation of $f$ as $f^{(k)}$ by replacing $\mathbf{W}$ with $\mathbf{W}^{(k)}$, where $\mathbf{W}^{(k)}$ is its rank-$k$ approximation.

Experiments. During training, the dataset $\mathcal{D}_{\alpha}$ is generated with non-zero noise probability $\alpha>0$. At test time, the dataset $\mathcal{D}_{0}$ is without noise as $\alpha=0$, so the computed loss is called pure-label loss.
The full model is trained with Gradient Descent (GD) subjected to cross-entropy loss. The results are reported in Figure 13, with discussions in Appendix C.1.

Low-rank subspace stores noise. In Figure 13, the rank-1 subspace corresponding to the smallest non-zero singular value is responsible to store the noise. We prove this mechanism as follows.

###### Theorem 2 .

Assume Assumptions C.1 and C.2 hold, considering $n=2,c=3$ and $\alpha\in(0.2,0.4)$, we train the full model $f(\cdot,\cdot;\mathbf{W})$ with gradient flow. Denote $P(i,j;\mathbf{W})$ as the model’s predicted probability for output $j$ conditioned on input $i$. Then, for $t\rightarrow\infty$ and $i\in\{1,2\}$, we have

$$ $\displaystyle P(i,j;\mathbf{W})$ $\displaystyle=(1-\alpha)\cdot\mathbbm{1}\{j=i\}+\alpha\cdot\mathbbm{1}\{j=c\},$ $\displaystyle P(i,j;\mathbf{W}^{(1)})$ $\displaystyle=(1-\Theta(t^{-\nicefrac{{1}}{{2}}}))\cdot\mathbbm{1}\{j=i\}+ \Theta(t^{-\nicefrac{{1}}{{2}}})\cdot\mathbbm{1}\{j=c\}.$ $$

The above theorem implies, the full model always predicts noise w.p. $\alpha$, while the rank-1 model eventually predicts correctly without noise, although training is only on the full model with noise.

###### Theorem 2 .

## 5 Discussion and Limitations

In this paper, we studied the questions of how transformer language models learn to process global associations differently than in-context inputs, and how truncating specific weights or layers, particularly feed-forward layers, can help reasoning tasks.
While our work provides some initial theoretical understanding of how this may arise on simple controlled settings, our analysis is heavily simplified, and many questions remain open: (i) what are the training dynamics and truncation behaviors in richer data models where there are many more places and ways to choose between in-context and global associations? (ii) in some architectures, with an example reported in Appendix A.4, it appears that global associations are not stored in MLPs, but rather in attention – does this happen more broadly, for instance in attention sinks or registers (Darcet et al., ; Xiao et al., )? (iii) can we provide a more granular study of the training dynamics of SGD, jointly over feed-forward, value, and key-query matrices, and throughout the different phases?
We believe these are all interesting directions for future work.

## Acknowledgements

We are grateful to Yifang Chen, Ekin Akyürek and Denny Wu for helpful discussions.

## Appendix A How Does the Two-layer Model Solve Noisy In-context Recall?

### A.1 Summarizing: roles of key components in the two-layer transformer

Recall the architecture of two-layer transformers in Section 3 as

$$ $\displaystyle x_{t}$ $\displaystyle\triangleq\mathbf{W}_{E}(z_{t})+p_{t},$ $\displaystyle h_{t}^{1}$ $\displaystyle\triangleq\sum_{s\leq t}\left[\sigma(x_{t}^{\top}\mathbf{W}_{KQ}^ {1}x_{1:t})\right]_{s}\cdot\mathbf{W}_{V}^{1}x_{s},$ $\displaystyle x_{t}^{1}$ $\displaystyle\triangleq x_{t}+h_{t}^{1}+F_{1}(x_{t}+h_{t}^{1}),$ $\displaystyle h_{t}^{2}$ $\displaystyle\triangleq\sum_{s\leq t}\left[\sigma({x_{t}^{1}}^{\top}\mathbf{W} _{KQ}^{2}x_{1:t}^{1})\right]_{s}\cdot\mathbf{W}_{V}^{2}x_{s}^{1},$ $\displaystyle x_{t}^{2}$ $\displaystyle\triangleq x_{t}^{1}+h_{t}^{2}+F_{2}(x_{t}^{1}+h_{t}^{2}),$ $\displaystyle\xi_{t}$ $\displaystyle\triangleq\mathbf{W}_{U}x_{t}^{2}.$ $$

When the task is without noise, i.e., $\alpha=0$, Bietti et al. [] point out the first-layer attention attends to the previous token through $\mathbf{W}_{KQ}^{1}=\sum_{t=2}^{T}p_{t-1}p_{t}^{\top}$. Therefore, when $z_{t}=\bar{y}$ with $z_{t-1}=q$, the output of the first layer is $x_{t}^{1}\approx\mathbf{W}_{E}(\bar{y})+\mathbf{W}_{V}^{1}\mathbf{W}_{E}(q)$. Then they show that the second-layer attention matches such $x_{t}^{1}$ with $z_{T}=q$ by $\mathbf{W}_{KQ}^{2}=(\mathbf{W}_{V}\mathbf{W}_{E}(q))\mathbf{W}_{E}(q)^{\top}$, through which the information of $\bar{y}$ in $x_{t}^{1}$ is copied to last token as $h_{T}^{2}\approx\mathbf{W}_{V}^{2}\mathbf{W}_{E}(\bar{y})$. Finally $\mathbf{W}_{V}^{2}=\sum_{z\in[N]}\mathbf{W}_{U}(z)\mathbf{W}_{E}(z)^{\top}$ helps output the correct label of $\bar{y}$.

In our work with noise $\alpha>0$, the key difference is that there is a fixed probability $\alpha$ for a noise token $N+1$ to appear after each trigger $q$. This requires $\mathbf{W}_{KQ}^{2}$ to not only match the trigger but also avoid the noise token after trigger. Let’s first summarize the whole pipeline of this model for our task.

Roles of key components.
The first layer will be basically the same as Bietti et al. [], where $\mathbf{W}_{KQ}^{1}=\sum_{t=2}^{T}p_{t-1}p_{t}^{\top}$ attends to the previous token. Consider two positions $t_{1},t_{2}$ with $z_{t_{1}-1}=z_{t_{2}-1}=q,z_{t_{1}}=\bar{y},z_{t_{2}}=N+1$, then outputs of the first layer at these two positions are $x_{t_{1}}^{1}\approx\mathbf{W}_{E}(\bar{y})+\mathbf{W}_{V}^{1}\mathbf{W}_{E}(q)$, $x_{t_{2}}^{1}\approx\mathbf{W}_{E}(N+1)+\mathbf{W}_{V}^{1}\mathbf{W}_{E}(q)$. Then the second-layer attention $\mathbf{W}_{KQ}=(\mathbf{W}_{V}\mathbf{W}_{E}(q)-c\cdot\mathbf{W}_{E}(N+1))
\mathbf{W}_{E}(q)^{\top}$ with some positive $c$ makes the attention attend to $t_{1}$ and avoid $t_{2}$ simultaneously, matching with the last token $z_{T}=q$. Therefore, the output of the second-layer attention at $T$ is basically $h_{T}^{2}\approx\mathbf{W}_{V}^{2}\mathbf{W}_{E}(\bar{y})$. Similar to the noiseless case, $\mathbf{W}_{V}^{2}=\sum_{z\in[N]}\mathbf{W}_{U}(z)\mathbf{W}_{E}(z)^{\top}$ helps output the correct label of $\bar{y}$.
Meanwhile, note that $x_{T}^{1}$ actually contains $\mathbf{W}_{E}(q)$ through $x_{T}$, so $F_{2}$ is able to predict the noise $N+1$ when seeing a fixed $\mathbf{W}_{E}(q)$. As a result, combining the two streams from $h_{T}^{2}$ and $F_{2}(x_{T}^{1})$, the full model is able to predict any $\bar{y}$ w.p. $1-\alpha$ and predict the noise $N+1$ w.p. $\alpha$.

Evidence.
Figure 5 illustrates that the second-layer attention learns to attend to $z_{t_{1}}=\bar{y}$ and avoid $z_{t_{2}}=N+1$, with Appendix A.2 presenting a primitive exploration on how the avoidance is learnt in a simplified setting. Figure 6 (left) shows the attention pattern from $\mathbf{W}_{KQ}^{1}$ of attending to the previous token. Figure 6 (middle) shows the memory recall of $\mathbf{W}_{U}(N+1)^{\top}F_{2}(\mathbf{W}_{E}(q))$ to predict the noise. Figure 6 (right) illustrates the memory recall of $\mathbf{W}_{U}(i)^{\top}\mathbf{W}_{V}^{2}\mathbf{W}_{E}(i)$ to predict the correct token.

Figure: Figure 6: Left: first-layer attention attending to the previous token from the current token. Middle: logits to predict noise from $\langle F_{2}(\mathbf{W}_{E}(i)),\mathbf{W}_{U}(j)\rangle$ with input $i\in[N+1]$ and output $j\in[N+1]$, where the output channel 2 is set as the noise channel. It turns out, for all input $i$, the logits on output 2 are large, which matches our construction that, at least for trigger $q$ as input, the output 2 has large logits. Right: logits to predict singal from $\langle\mathbf{W}_{V}^{2}\mathbf{W}_{E}(i),\mathbf{W}_{U}(j)\rangle$ for input $i\in[N+1]$ and output $j\in[N+1]$. It matches our construction that $i=j$ has large logits. Meanwhile, $i=j=2$ does not have large logits since 2 is the noise channel.
Refer to caption: x6.png

### A.2 How does attention attend less towards the noise token?

We use the same simplified model as in Section 3.1 to understand how the second-layer attention learns to avoid the noise.
When using the same learning rate $\eta=\eta_{v}=\eta_{f}$, Theorem 1 implies that the feed-forward $\mathbf{W}_{F}$ makes the most contribution for predicting the noise after the first-step update. Denote the logits for the noise of the model at time $t$ as $\xi_{t}$. The arguments in this section make the following assumptions, which hold at least after the first-step update:

- i.
$\mathbf{W}_{F}$ dominates the logits $\xi_{t}$ of predicting the noise token, compared with $\mathbf{W}_{V}$.
- ii.
Logits for predicting any $k\leq N$ is close to 0, which means the predicted probability $p_{t}$ is approximately $p_{t}\approx\frac{\exp(\xi_{t})}{N+\exp(\xi_{t})}$.
- iii.
The predicted probability $p_{t}<\alpha$.
- iv.
The attention matrix $\mathbf{W}_{KQ}$ is approximately 0, inducing a uniform attention.
- v.
The dataset has $T,N\gg 1$ and $m\rightarrow\infty$, so the gradient is from population loss.

The first assumption holds after the first step from Theorem 1 with $\eta_{f}=\eta_{v}$.

Then, since $|\mathbf{W}_{U}(k)^{\top}(\nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{E}(q)|=O(\frac{
1}{N})\cdot|\mathbf{W}_{U}(N+1)^{\top}(\nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{E}
(q)|$ for any $k\leq N$ in Lemma D.1, the second assumption holds. Meanwhile, the projection of $\nabla_{\mathbf{W}_{V}}L$ onto any direction in Lemma D.2 is also smaller than $\mathbf{W}_{U}(N+1)^{\top}(\nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{E}(q)$ by a factor of $O(\nicefrac{{1}}{{N}})$.

Let’s check the condition of the third assumption.
In the proof of Lemma D.1, the gradient of $\mathbf{W}_{F}$ has the form of

$$ $\displaystyle\mathbf{W}_{U}(N+1)^{\top}(-\nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{ E}(q)=\alpha-p_{t}.$ $$

This update induces $\xi_{t}$ to increase by $\eta(\alpha-p_{t})$. This implies

$$ $\displaystyle\xi_{t}\approx\xi_{t-1}+\eta\bigg{(}\alpha-\frac{\exp(\xi_{t})}{N +\exp(\xi_{t})}\bigg{)},~{}~{}\forall~{}t\geq 1.$ $$

This sequence $\{\xi_{t}\}_{t\geq 1}$ has stationary point $\xi^{*}=\log N+\log(\frac{\alpha}{1-\alpha})$. Denoting $\hat{\xi}_{t}\triangleq\xi_{t}-\xi^{*}$ with $\hat{\xi}_{1}=-\xi^{*}<0$, the iteration becomes

$$ $\displaystyle\hat{\xi}_{t+1}\approx\hat{\xi}_{t}+\eta\bigg{(}\alpha-\frac{\exp (\hat{\xi}_{t})}{\frac{1-\alpha}{\alpha}+\exp(\hat{\xi}_{t})}\bigg{)}.$ $$

If we would like to have $\hat{\xi}_{t}$ not hit the positive region by controlling $\eta$, it suffices to bound $\eta$ with any $\hat{\xi}<0$,

$$ $\displaystyle\eta\leq\frac{\hat{\xi}}{\frac{\exp(\hat{\xi})}{\frac{1-\alpha}{ \alpha}+\exp(\hat{\xi})}-\alpha},$ $$

where RHS is continuous and decreasing on $\xi<0$ when $\alpha<0.5$. Hence, we have $\eta\leq\frac{1}{\alpha(1-\alpha)}$ evaluated at $\hat{\xi}=0$ by L’Hospital rule. This bound of $\eta$ is very strong, since $\eta=O(\log N)$ can still have $\hat{\xi}<0$ after one step.

The fourth assumption is basically from what we will show at the end of this section, as the second observation.

Then consider the dynamics of $\mathbf{W}_{V}$, which is much slower than $\mathbf{W}_{F}$. From the proof of Lemma D.2, the gradient of $\mathbf{W}_{V}$ satisfies

$$ $\displaystyle\nabla_{\mathbf{W}_{V}}L$ $\displaystyle=\mathbb{E}_{x}\left[\sum_{k=1}^{N+1}(p_{\mathbf{W}}(k|x)- \mathbbm{1}\{y=k\})\mathbf{W}_{U}(k)\bigg{(}\frac{1}{T}\sum_{t=1}^{t}x_{t} \bigg{)}^{\top}\right],$ (2) $\displaystyle\mathbf{W}_{U}(N+1)^{\top}(-\nabla_{\mathbf{W}_{V}}L)\mathbf{W}_{ E}(k)$ $\displaystyle\approx\frac{1}{N}\sum_{t\geq 1}(\alpha-p_{t})(\mathbbm{1}\{k\leq N \}+\alpha\cdot\mathbbm{1}\{k=N+1\})$ $\displaystyle\triangleq c\cdot\mathbbm{1}\{k\leq N\}+c\cdot\alpha\cdot\mathbbm {1}\{k=N+1\}=\Theta(\frac{1}{N}),$ $$

where the projection on $W_{E}(N+1)$ is always positive and smaller than that on other directions when $p_{t}<\alpha$. Projections onto other directions $\mathbf{W}_{U}(j)\mathbf{W}_{E}(k)^{\top}$, $\forall~{}j\leq N$, are smaller as $\Theta(\frac{1}{N^{2}})$.

Finally, let’s consider the dynamics of $\mathbf{W}_{KQ}$. At initialization, $\mathbf{W}_{KQ}=0$ and $\nabla_{\mathbf{W}_{KQ}}L=0$ due to zero initialization of $\mathbf{W}_{V}$. After one-step, $\mathbf{W}_{V}$ has such a structure in Eq.(2). Then, with $\bar{x}_{1:T}\triangleq\frac{1}{T}\sum_{1\leq t\leq T}x_{t}$ from uniform attention, the gradient of $\mathbf{W}_{KQ}$ satisfies

$$ $\displaystyle-\nabla_{\mathbf{W}_{KQ}}L$ $\displaystyle=\mathbb{E}_{x}\left[\sum_{k=1}^{N}(\mathbbm{1}\{y=k\}-p_{\mathbf {W}}(k|x))\frac{1}{T}\sum_{t=1}^{T}(\mathbf{W}_{U}(k)^{\top}\mathbf{W}_{V}x_{t })\cdot(x_{t}-\bar{x}_{1:T})\mathbf{W}_{E}(q)^{\top}\right]$ (3) $\displaystyle\approx\sum_{k=1}^{N}\bigg{(}\frac{1-\alpha}{N}-\frac{1-p_{t}}{N} \bigg{)}\underbrace{\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}\mathbf{W}_{U}(k) ^{\top}\mathbf{W}_{V}x_{t}\cdot(x_{t}-\bar{x}_{1:T})\mathbf{W}_{E}(q)^{\top} \right]}_{\triangleq A}$ $\displaystyle~{}~{}~{}~{}+(\alpha-p_{t})\underbrace{\mathbb{E}\left[\frac{1}{T }\sum_{t=1}^{T}(\mathbf{W}_{U}(N+1)^{\top}\mathbf{W}_{V}x_{t})\cdot(x_{t}-\bar {x}_{1:T})\mathbf{W}_{E}(q)^{\top}\right]}_{\triangleq B}.$ $$

Then, we have

$$ $\displaystyle~{}~{}~{}~{}~{}~{}\mathbf{W}_{E}(N+1)^{\top}B\mathbf{W}_{E}(q)= \mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}(\mathbf{W}_{U}(N+1)^{\top}\mathbf{W} _{V}x_{t})\cdot\mathbf{W}_{E}(N+1)^{\top}(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\mathbb{E}\left[\frac{1}{T}\sum_ {t=1}^{T}(c+c(\alpha-1)\cdot\mathbbm{1}\{z_{t}=N+1\})\cdot\mathbf{W}_{E}(N+1)^ {\top}(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{=}}\mathbb{E}\left[\frac{1}{T}\sum_ {t=1}^{T}(c(\alpha-1)\cdot\mathbbm{1}\{z_{t}=N+1\})\cdot\mathbf{W}_{E}(N+1)^{ \top}(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle=\frac{\alpha}{N}\cdot c(\alpha-1)(1-\frac{\alpha}{N})=\Theta( \frac{1}{N^{2}})<0.$ $$

where (a) is from Eq.(2), (b) is due to $\bar{x}_{1:T}=\frac{1}{T}\sum_{t}x_{t}$ and note that $c=\Theta(\frac{1}{N})$.

Similarly, we also have

$$ $\displaystyle\mathbf{W}_{E}(N+1)^{\top}A\mathbf{W}_{E}(q)$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}(\mathbf{W}_{U}(k)^{ \top}\mathbf{W}_{V}x_{t})\mathbf{W}_{E}(N+1)^{\top}\cdot(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}\Theta(\frac{1}{N^{2}}) \cdot\mathbbm{1}\{z_{t}=N+1\}\mathbf{W}_{E}(N+1)^{\top}\cdot(x_{t}-\bar{x}_{1: T})\right]=\Theta(\frac{1}{N^{3}}).$ $$

For any $k\leq N$, we have

$$ $\displaystyle~{}~{}~{}~{}~{}~{}\mathbf{W}_{E}(k)^{\top}B\mathbf{W}_{E}(q)= \mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}(\mathbf{W}_{U}(N+1)^{\top}\mathbf{W} _{V}x_{t})\cdot\mathbf{W}_{E}(k)^{\top}(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}(c(\alpha-1)\cdot \mathbbm{1}\{z_{t}=k\})\cdot\mathbf{W}_{E}(N+1)^{\top}(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle=\frac{\alpha}{N}\cdot c(\alpha-1)(-\frac{1}{N})=\Theta(\frac{1}{ N^{3}})>0,$ $$

and

$$ $\displaystyle\mathbf{W}_{E}(k)^{\top}A\mathbf{W}_{E}(q)$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}(\mathbf{W}_{U}(k)^{ \top}\mathbf{W}_{V}x_{t})\mathbf{W}_{E}(k)^{\top}\cdot(x_{t}-\bar{x}_{1:T})\right]$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}\Theta(\frac{1}{N^{2}}) \cdot\mathbbm{1}\{z_{t}=N+1\}\mathbf{W}_{E}(k)^{\top}\cdot(x_{t}-\bar{x}_{1:T} )\right]=\Theta(\frac{1}{N^{4}}).$ $$

Combining the above four esimation of projections of $A$ and $B$ with Eq.(3), we have

$$ $\displaystyle\mathbf{W}_{E}(N+1)^{\top}(-\nabla_{\mathbf{W}_{KQ}}L)\mathbf{W}_ {E}(q)$ $\displaystyle=\Theta(\frac{1}{N^{2}})<0,$ $\displaystyle\forall~{}k\leq N,~{}~{}\mathbf{W}_{E}(k)^{\top}(-\nabla_{\mathbf {W}_{KQ}}L)\mathbf{W}_{E}(q)$ $\displaystyle=\Theta(\frac{1}{N^{3}})>0.$ $$

Then we have three observations

- i.
$\mathbf{W}_{KQ}$ in this phase avoids the noise token $N+1$ and uniformly attends to all tokens $k\leq N$.
- ii.
The update of $\mathbf{W}_{KQ}$ is in $\Theta(\frac{1}{N^{2}})$, while the update of $\mathbf{W}_{F}$ is $\Theta(1)$ in Lemma D.1 and that of $\mathbf{W}_{V}$ is $\Theta(\frac{1}{N})$ in Lemma D.2. These three levels of updating speed also coincide with the assumptions that $\mathbf{W}_{F}$ dominates first and then $\mathbf{W}_{V}$ has a micro structure that induces the evolving of $\mathbf{W}_{KQ}$.
- iii.
The current proof for $\mathbf{W}_{KQ}$ strongly depends on the fact that the noise token appears less than other token by a factor $\alpha$ in expectation. The proof will have the opposite result if the noise token is made to appear more by manipulating the data distribution. Therefore, we leave a new proof that is robust to such an assumption in data distribution as future work.

### A.3 Multiple Triggers

In Section 3, we assume there is only one fixed trigger $q\in[N]$ for simplicity. Actually the case of multiple triggers has the same mechanism. As discussed by Bietti et al. [] and Appendix A.1, for one trigger, the second-layer attention has large logits in $\langle\mathbf{W}_{V}^{1}\mathbf{W}_{E}(i)^{\top},\mathbf{W}_{KQ}^{2}\mathbf{W
}_{E}(j)\rangle$ only for $i=j=q$. For multiple triggers, basically $\langle\mathbf{W}_{V}^{1}\mathbf{W}_{E}(i)^{\top},\mathbf{W}_{KQ}^{2}\mathbf{W
}_{E}(j)\rangle$ only have large values when $q\in Q$. This is verified in Figure 7.

Figure: Figure 7: Logits of $\langle\mathbf{W}_{V}^{1}\mathbf{W}_{E}(i)^{\top},\mathbf{W}_{KQ}^{2}\mathbf{W }_{E}(j)\rangle$ for input $i$ and output $j$ when there is one trigger (left, $q=1$) and five triggers (right, $q\in Q=\{1,39,43,53,58\}$). In both cases, the logits only have large values when $i=j=q$, verifies the matching mechanism in Appendix A.1.
Refer to caption: x7.png

### A.4 Architectural Choices

In Section 3 and Appendix A.1, we were focused on experiments with both $F_{1},F_{2}$ being two-layer ReLU MLPs. Meanwhile, we have also tried other choices of $F_{1},F_{2}$ and then search for the best truncation method for each architecture. In this section, we would like to summarize our experimental results for better understanding of all modules in the two-layer transformer.

Generally, the feed-forward layer can be two-layer ReLU MLPs, one-layer Linear or “None”, where None stands for there is no feed-forward layer so that the value matrices in attention layers are the only weight matrices that transform features.

Both $F_{1},F_{2}$ are two-layer MLPs. This is our main setting. The best truncation method is to fully drop $F_{2}$. We also try to fully drop $F_{1}$, as reported in Figure 8. It turns out fully dropping $F_{1}$ makes the model predict the noise with high probability.

Figure: Figure 8: Test performance of fully dropping $F_{1},F_{2}$ when both $F_{1},F_{2}$ are two-layer MLPs. It turns out, while dropping $F_{2}$ makes the model predict correctly w.p. near 1, dropping $F_{1}$ has the model predict noise with high probability.
Refer to caption: x8.png

$F_{1}$ is MLPs and $F_{2}$ is Linear. Figure 9 reports the results. Dropping $F_{1}$ and $F_{2}$ both improve the correct prediction, and dropping $F_{1}$ is better with lower test loss. Note that, when test accuracies are near $100\%$, lower test loss is a better measurement of the prediction quality, because accuracies are taken by argmax over the output logits while test loss are about the exactly predicted probability.

Figure: Figure 9: Test performance of fully dropping $F_{1},F_{2}$ when both $F_{1}$ is MLPs and $F_{2}$ Linear. Both dropping methods turn out to help predict more correctly than the full model. Meanwhile, dropping the MLP $F_{1}$ is better with lower test loss.
Refer to caption: x9.png

$F_{1}$ is Linear and $F_{2}$ is MLPs. Figure 10 reports the results. Dropping $F_{2}$ improves the correct prediction while dropping $F_{1}$ makes the model predict noise more.

Figure: Figure 10: Test performance of fully dropping $F_{1},F_{2}$ when both $F_{1}$ is Linear and $F_{2}$ MLPs. Only dropping $F_{2}$ helps predict more correctly. Dropping $F_{1}$ makes the model predicting noise more.
Refer to caption: x10.png

Both $F_{1}$ and $F_{2}$ are None. Figure 11 reports the results. While there is no feed-forward layer any more, low-rank truncating a part $\mathbf{W}_{O}^{1}$ of the first-layer matrix improves the model’s prediction a little. This implies that, when there is not feed-forward layers, the noise association is possible stored in the first-layer value matrix of attention. Note that the improvement of such low-rank truncation is clearly smaller than fully dropping one of feed-forward layers in the previous cases. Meanwhile, a smaller $\rho=0.01$ destroys the model’s performance. This implies fully dropping is not the optimal choice for low-rank truncation of the value matrix, and there is low-rank subspace in it that is useful for predicting the correct tokens. Our discussion of the role of $\mathbf{W}_{V}^{1}$ in Appendix A.1 is a possible answer to this phenomena.

Figure: Figure 11: Test performance of low-rank truncating of $\mathbf{W}_{O}^{1}$ when there is no $F_{1},F_{2}$. Here $\rho$ is the fraction of preserved rank of $\mathbf{W}_{O}^{1}$, where actually we re-parametrize the first-layer value matrix in attention as $\mathbf{W}_{O}^{1}\mathbf{W}_{V}^{1}\in\mathbb{R}^{d\times d}$. It turns out the best $\rho=0.05$ improves the model’s prediction a little. Meanwhile, a smaller $\rho$ destroys the model’s performance.
Refer to caption: x11.png

### A.5 Training Details about Experiments

All of the training is with SGD optimization with learning rate in $\{0.001,0.03\}$. The batch size is 512. The dimension is 256. The context length is 256. All results in the experiments are stable for any learning rate between 0.001 and 0.03. Each run of experiments is on a single Nvidia Tesla V100 GPU. It takes 3 hours to finish each run for 2K steps, which probably can be optimized a lot since we are tracking a lot of measurement along training, not limited to hundreds of possible truncations at each test time.

## Appendix B More Experiments on Pythia

### B.1 Learning Association with Prepositions

We would like to verify our guess about the structure of “to + the” in Pythia in Section 2.3. To make the argument generalizable than IOI dataset, we consider a structure of “[preposition] + the”, where [preposition] has a pool of 30 prepositions in English, including “to”. The input is a raw “[preposition]” or a random sentence ending with “[preposition]”, with some examples in Appendix G.1. For both kinds of inputs, Pythia-160M/410M/1B turns out to learn the structure of “[preposition] + the” around 10 steps, as shown in Figure 12.

Figure: Figure 12: Average ranking of tokens “the” in the prediction by Pythia-160M/410M/1B along training. The inputs are 30 preposition words (left) and 40 sentences ending with prepositions. It turns out “the” becomes one of top predictions around 10 steps.
Refer to caption: x12.png

### B.2 LASER Parameters for Evaluated LLMs

Following the definition of LASER in Section 2.2, we search for the optimal layer, $\rho$ and target weights in Pythia models and GPT-2 Small for each dataset.

IOI on Pythia-1B. The model has 16 layers. The truncation is on the input matrix of MLPs on the 11-th layer with $\rho=0.008$.

Factual recall on Pythia-1B. The truncation is on the input matrix of MLPs on the 16-th layer with $\rho=0.0125$.

IOI on GPT2 Small. Related parameters have been contained in Section 2.3.

## Appendix C Linear Associative Memory

### C.1 Experiments and Discussions

In Section 3, we showed that *fully* truncating a feed-forward layer can be helpful for reasoning. We now present a setting where noisy associations are stored in a rank-one subspace of a layer, so that *intermediate* levels of truncation are more useful to remove noise.

Model and data. We consider a simple associative memory setting where the goal is learn an fixed permutation from input tokens to output tokens (w.l.o.g. taken to be the identity), with a linear model similar to Cabannes et al. []. Consider a learnable weight matrix $\mathbf{W}\in\mathbb{R}^{d\times d}$. Consider embeddings for $n$ input tokens as $\{e_{i}\}_{i=1}^{n}\subset\mathbb{R}^{d}$ and embeddings for $c$ output tokens as $\{u_{i}\}_{i=1}^{c}\subset\mathbb{R}^{d}$. In contrast to Cabannes et al. [], we consider an additional “common noise” output token $c=n+1$, which is chosen for any input with probability $\alpha\in(0,1)$.
For any input $x\in[n]$, the target distribution $p_{\alpha}(\cdot|x)$ is defined by

$$ $\displaystyle p_{\alpha}(y|x)=(1-\alpha)\cdot\mathbbm{1}\{y=x\}+\alpha\cdot \mathbbm{1}\{y=c\}.$ (4) $$

In other words, the last channel ($c$) for output is the common noise with probability $\alpha$ for any input. The training dataset $\mathcal{D}_{\alpha}$ consists of uniformly distributed inputs $x\in[n]$, and outputs conditionally sampled as $y|x\sim p_{\alpha}(\cdot|x)$.

Given any pair of input and output tokens, the associative memory model takes the form

$$ $\displaystyle f(i,j;\mathbf{W})$ $\displaystyle\triangleq\langle u_{j},\mathbf{W}e_{i}\rangle,~{}~{}~{}~{} \forall~{}i,j\in[n]\times[c],$ (5) $$

When $k\leq d$, we denote the rank-$k$ approximation of $f$ as $f^{(k)}$ by replacing $\mathbf{W}$ with $\mathbf{W}^{(k)}$, where $\mathbf{W}^{(k)}$ is the rank-$k$ approximation of $\mathbf{W}$.

Training. During training, the dataset $\mathcal{D}_{\alpha}$ is generated with non-zero noise probability $\alpha>0$. At test time, the dataset $\mathcal{D}_{0}$ is without noise as $\alpha=0$, so the computed loss is called pure-label loss.
The model is trained with Gradient Descent (GD) subjected to cross-entropy loss.

Experiments with randomness. Assume both $\{e_{i}\}_{i=1}^{n}$ and $\{u_{i}\}_{i=1}^{c}$ are i.i.d. uniformly drawn from sphere $\mathbb{S}^{d-1}$. Also assume the model is initialized as $\mathbf{W}_{i,j}\sim\mathcal{N}(0,\frac{1}{d})$. Due to randomness from embeddings and model initialization, let’s first conduct 20 runs of experiments to obtain significant factors before moving the theoretical argument.

Note that only full models are trained, and we track loss for low-rank models by conducting SVD in each step without manipulating training.
In Figure 5, we illustrate the pure-label loss v.s. training steps for models of different ranks, where $n=3$, $\alpha=0.03$ and $d=8$ or $12$. It turns out, while the full model (rank$\geq 3$) has a constant pure-label loss ($\sim 0.03$, dependent on $\alpha$), the rank-2 model is very likely to have a significant loss than the full model. Meanwhile, the larger $d$ has more stable results than small $d$.

Figure: Figure 13: Pure-label loss for rank-1,2,3,4 models with $n=3,\alpha=0.03$ and $d=12$ (left) or $8$ (right). Only full models are trained, and we report low-rank results by conducting SVD in each step without manipulating the training. In both figures, the experiments are run for 20 times to examine the randomness. For each rank, we plot curves of the median, $25\%$ and $75\%$ out of 20 runs. It turns out: i) rank-2 models are very likely to have significantly lower pure-label loss thant full models (rank$\geq 3$), and ii) the larger dimension $d$ has more stable results.
Refer to caption: x13.png

Therefore, we can qualify the following important factors for this model:

- i.
$d$ v.s. $n,c$: when $d\gg n,c$, random drawn embeddings tend to be orthogonal to each other, with inner product in $O(\nicefrac{{1}}{{\sqrt{d}}})$. If $n,c=\Omega(d)$, embeddings will be in strong correlations, making the problem extremely difficult to understand. Cabannes et al. [] also discussed about such particle interaction in associative memory.
- ii.
Low-rank subspace storing the noise. In Figure 13, the rank-1 subspace between the full and rank-2 models is responsible to store the noise, removing which will induce a model ideally predicting the ground-truth without noise. This is understandable if the embeddings are orthogonal, as shown in Theorem 2.
- iii.
$\alpha$ v.s. $n$. When $n$ is large, orthogonal embeddings still induces a low-rank subspace storing the noise, but $\alpha$ decides whether the low-rank subspace corresponds to the smallest singular values of $\mathbf{W}$. If not, it requires more careful manipulation of the spectrum instead of low-rank approximation of $\mathbf{W}$.

Now we present a theoretical analysis of this problem with some assumptions.

###### Assumption C.1 (Orthonormality) .

Embeddings of input and output tokens are orthonormal, i.e., $e^{\top}_{i}e_{j}=\mathbbm{1}\{i=j\},\forall~{}i,j$ and $u^{\top}_{i}u_{j}=\mathbbm{1}\{i=j\},\forall~{}i,j.$

###### Assumption C.2 (Initialization) .

The learnable matrix $\mathbf{W}$ is initialized from $\mathbf{0}$ when $t=0$.

###### Theorem 3 (Restatement of Theorem 2 ) .

Assume Assumptions C.1 and C.2 hold, considering $n=2,c=3$ and $\alpha\in(0.2,0.4)$, we train the full model $f(\cdot,\cdot;\mathbf{W})$ with gradient flow. Denote $P(i,j;\mathbf{W})$ as the model’s predicted probability for output $j$ conditioned on input $i$. Then, for $t\rightarrow\infty$ and $i\in\{1,2\}$, we have

$$ $\displaystyle P(i,j;\mathbf{W})$ $\displaystyle=(1-\alpha)\cdot\mathbbm{1}\{j=i\}+\alpha\cdot\mathbbm{1}\{j=c\},$ $\displaystyle P(i,j;\mathbf{W}^{(1)})$ $\displaystyle=(1-\Theta(t^{-\nicefrac{{1}}{{2}}}))\cdot\mathbbm{1}\{j=i\}+ \Theta(t^{-\nicefrac{{1}}{{2}}})\cdot\mathbbm{1}\{j=c\}.$ $$

###### Remark 1 .

Note that here the assumption $\alpha\in(0.2,0.4)$ is a technical choice. In experiments, any value $\alpha\in(0,0.4)$ still has the same result.

###### Proof.

W.l.o.g., we assume the embeddings are standard basis in $\mathbb{R}^{d}$.
For any $\mathbf{W}$, the gradient $\nabla_{\mathbf{W}}L$ can be decomposed as

$$ $\displaystyle\nabla_{\mathbf{W}}L=\gamma_{1}\begin{bmatrix}1\\ -1\\ 0\end{bmatrix}\begin{bmatrix}1&-1&0\end{bmatrix}+\gamma_{2}\begin{bmatrix}1\\ 1\\ -2\end{bmatrix}\begin{bmatrix}1&1&0\end{bmatrix}.$ (6) $$

Since $\mathbf{W}$ initializes from zero, this implies $\mathbf{W}$ can always be decomposed with the same basis

$$ $\displaystyle\mathbf{W}=\beta_{1}\begin{bmatrix}1\\ -1\\ 0\end{bmatrix}\begin{bmatrix}1&-1&0\end{bmatrix}+\beta_{2}\begin{bmatrix}1\\ 1\\ -2\end{bmatrix}\begin{bmatrix}1&1&0\end{bmatrix}.$ (7) $$

Then gradient flow gives the following ODE

$$ $\displaystyle\dot{\beta}_{1}=-\gamma_{1}$ $\displaystyle=\frac{\exp(-\beta_{1}+\beta_{2})-\exp(\beta_{1}+\beta_{2})}{\exp (-\beta_{1}+\beta_{2})+\exp(\beta_{1}+\beta_{2})+\exp(-2\beta_{2})}+1-\alpha$ (8) $\displaystyle=\frac{\exp(-2\beta_{1})-1}{\exp(-2\beta_{1})+\exp(-\beta_{1}-3 \beta_{2})+1}+1-\alpha,$ $\displaystyle\dot{\beta}_{2}=-\gamma_{2}$ $\displaystyle=\frac{3\exp(-2\beta_{2})}{\exp(-\beta_{1}+\beta_{2})+\exp(\beta_ {1}+\beta_{2})+\exp(-2\beta_{2})}-3\alpha$ $\displaystyle=\frac{3\exp(-\beta_{1}-3\beta_{2})}{\exp(-2\beta_{1})+\exp(- \beta_{1}-3\beta_{2})+1}-3\alpha.$ $$

Denoting $a=-2\beta_{1},b=-\beta_{1}-3\beta_{2}$, the ODE becomes

$$ $\displaystyle\dot{a}$ $\displaystyle=\frac{2-2\exp(a)}{\exp(a)+\exp(b)+1}-2+2\alpha,$ (9) $\displaystyle\dot{b}$ $\displaystyle=\frac{2-8\exp(b)}{\exp(a)+\exp(b)+1}-2+10\alpha.$ $$

Lemma F.3 gives the solution as, when $t\rightarrow\infty$,

$$ $\displaystyle a\rightarrow-\log(t)-\log(1-\alpha)(4-2\alpha),~{}~{}~{}~{}b \rightarrow\log\frac{\alpha}{1-\alpha}.$ $$

For the full model, taking the scores $\mathbf{W}_{1,:}$ of the first input token as an example, we have $\mathbf{W}_{11}=\beta_{1}+\beta_{2},\mathbf{W}_{12}=-\beta_{1}+\beta_{2},
\mathbf{W}_{13}=-2\beta_{2}$, so the margins are

$$ $\mathbf{W}_{11}-\mathbf{W}_{12}=2\beta_{1}=-a,\mathbf{W}_{11}-\mathbf{W}_{13}= \beta_{1}+3\beta_{2}=-b.$ $$

For the rank-1 model (assuming $\beta_{1}>\beta_{2}$), the margins are

$$ $\mathbf{W}^{(1)}_{11}-\mathbf{W}^{(1)}_{12}=2\beta_{1},\mathbf{W}^{(1)}_{11}- \mathbf{W}^{(1)}_{13}=\beta_{1}.$ $$

The proof finishes by computing softmax on the margins.
∎

###### Assumption C.1 (Orthonormality) .

###### Assumption C.2 (Initialization) .

###### Theorem 3 (Restatement of Theorem 2 ) .

###### Remark 1 .

###### Proof.

## Appendix D Proof for Theorem 1

###### Assumption D.1 (Orthonormal embeddings) .

The embeddings $u_{k}\in\mathbb{R}^{d}$ are assumed to be orthonormal, i.e., $u_{i}^{\top}u_{j}=\mathbbm{1}\{i=j\}$.

###### Theorem 4 (Restatement of Theorem 1 ) .

Assume $N,T\gg 1,\alpha=\Theta(1)$. Consider a one gradient step update from zero-initialization on $m$ i.i.d. samples of $z_{1:T}$ with separate learning rates $\eta_{f}$ for $\mathbf{W}_{F}$ and $\eta_{v}$ for $\mathbf{W}_{V}$ (note that the gradient on $\mathbf{W}_{KQ}$ is zero). For a test sequence $z_{1:T}$, the resulting logits for the feed-forward and attention blocks satisfy, with probability $1-\delta$

$$ $\displaystyle\left|\Delta(\xi_{\text{ff}}(x_{1:T}))-\eta_{f}\cdot\alpha\right|$ $\displaystyle\leq\eta_{f}\cdot O\left(\sqrt{\frac{\ln\frac{2(N+1)}{\delta}}{m} }\right),$ $\displaystyle\left|\Delta(\xi_{\text{attn}}(x_{1:T}))-\frac{\eta_{v}}{N}\cdot( \alpha^{2}\hat{q}+\alpha(1-\hat{q}))\right|$ $\displaystyle\leq\eta_{v}\cdot O\left(\sqrt{\frac{(\frac{1}{TN}+\frac{1}{N^{2} })\ln\frac{2(N+1)}{\delta}}{m}}+\frac{\ln\frac{2(N+1)}{\delta}}{m}\right),$ $$

where $\Delta(\xi)=\xi_{N+1}-\max_{j\in[N]}\xi_{j}$ is the margin of predicting the noise token and $\hat{q}=\frac{1}{T}\sum_{t\leq T}\mathbbm{1}\{z_{t}=N+1\}$.

###### Proof.

For $\mathbf{W}_{F}$, since the input is always $z_{T}=q$, the logits will be $[\xi_{\text{ff}}]_{k}=\mathbf{W}_{U}(k)^{\top}\mathbf{W}_{F}\mathbf{W}_{E}(q)$, $\forall~{}k\in[N+1]$. As $\mathbf{W}_{F}$ is initialized from 0 and updated by GD with learning rate $\eta_{f}$, after one-step update, we have

$$ $\displaystyle\xi_{\text{ff}}=\mathbf{W}_{U}(k)^{\top}\bigg{(}-\eta_{f}\nabla_{ \mathbf{W}_{F}}\hat{L}\bigg{|}_{\mathbf{W}_{F}=0}\bigg{)}\mathbf{W}_{E}(q)\in \mathbb{R}^{N+1}.$ $$

By Lemma D.1, with probability $1-\frac{1}{2}\delta$, we have

$$ $\displaystyle\left|[\xi_{\text{ff}}]_{N+1}-\eta_{f}\cdot\alpha\right|$ $\displaystyle\leq\eta_{f}\cdot O\left(\sqrt{\frac{\ln\frac{2(N+1)}{\delta}}{m} }\right),$ $\displaystyle\forall~{}k\leq N,~{}~{}\left|[\xi_{\text{ff}}]_{k}-\eta_{f}\cdot \bigg{(}\frac{1-\alpha}{N}-\frac{1}{N+1}\bigg{)}\right|$ $\displaystyle\leq\eta_{f}\cdot O\left(\sqrt{\frac{\ln\frac{2(N+1)}{\delta}}{Nm }}+\frac{\ln\frac{2(N+1)}{\delta}}{m}\right),$ $$

and then triangle inequality finishes the proof for $\xi_{\text{ff}}$.

For $\mathbf{W}_{V}$, since the gradient on $\mathbf{W}_{KQ}$ at initialization is zero, $\mathbf{W}_{KQ}$ being zero after the first step induces a uniform attention over the input sequence.
Consider the input sequence $\{z_{i}\}_{i=1}^{T}$, then the logits will be $[\xi_{\text{attn}}]_{j}=\mathbf{W}_{U}(j)^{\top}\mathbf{W}_{V}\frac{1}{T}\sum_
{t=1}^{T}\mathbf{W}_{E}(z_{t}),~{}\forall~{}j\in[N+1]$.

Then considering the concentration bound of $\mathbf{W}_{V}$ after one-step update in Lemma D.2, denoting $\Gamma(j,k)=\mathbf{W}_{U}(j)^{\top}\mathbf{W}_{V}\mathbf{W}_{E}(k)$, we have

$$ $\displaystyle[\xi_{\text{attn}}]_{j}=\frac{1}{T}\sum_{t\leq T}\Gamma(j,z_{t})= \frac{1}{T}\sum_{k\leq N+1}n_{k}\cdot\Gamma(j,k),$ $$

with concentration bound for each $\Gamma(\cdot,\cdot)$ in Lemma D.2. From Table 2, note that for all $j=N+1,k\leq N$, the expectation and variances are the same, while $k=N+1$ has slightly different expectation and variance (but still in the same order of the others). Hence, denoting $\hat{q}=\frac{1}{T}\sum_{t\leq T}\mathbbm{1}\{z_{t}=N+1\}$ dependent of the test sample $z_{1:T}$, we have

$$ $\displaystyle\left|[\xi_{\text{attn}}(x_{1:T})]_{N+1}-\frac{\eta_{v}}{N}\cdot( \alpha^{2}\hat{q}+\alpha(1-\hat{q}))\right|$ $\displaystyle\leq\eta_{v}\cdot O\left(\sqrt{\frac{(\frac{1}{TN}+\frac{1}{N^{2} })\ln\frac{2(N+1)}{\delta}}{m}}+\frac{\ln\frac{2(N+1)}{\delta}}{m}\right).$ $$

Meanwhile, as the terms in Table 2 for $j\neq N+1$ always have much smaller mean and variance by a factor $1/N$, using the Bernstein’s inequalites for these terms in Lemma D.2 finishes the proof for $\mathbf{W}_{V}$.

∎

In this section, we will present the expectations and variances of $\nabla_{\mathbf{W}_{V}}\hat{L}$ and $\nabla_{\mathbf{W}_{F}}\hat{L}$ with $\mathbf{W}_{V}=\mathbf{W}_{F}=0$ at initialization. The targets are to show:

- 1.
a gap between $\lim_{m\rightarrow\infty}\nabla_{\mathbf{W}_{V}}\hat{L}$ and $\lim_{m\rightarrow\infty}\nabla_{\mathbf{W}_{F}}\hat{L}$ so that a step of GD with large learning rates is enough to learn the noise in $\mathbf{W}_{F}$, and
- 2.
sample complexity of $\nabla_{\mathbf{W}_{V}}\hat{L}$ and $\nabla_{\mathbf{W}_{F}}\hat{L}$ based on expectations and variances.

###### Assumption D.1 (Orthonormal embeddings) .

###### Theorem 4 (Restatement of Theorem 1 ) .

###### Proof.

### D.1 Gradient for the Feed-forward Matrix 𝐖 F subscript 𝐖 𝐹 \mathbf{W}_{F} bold_W start_POSTSUBSCRIPT italic_F end_POSTSUBSCRIPT

###### Lemma D.1 .

Consider zero initialization, $\mathbf{W}_{V}=\mathbf{W}_{F}=\mathbf{W}_{KQ}=0$ and $N\gg 1$.
Then with probability $1-\delta$, for any $j,k\in[N+1]$, it holds

$$ $\displaystyle\left|\mathbf{W}_{U}(k)^{\top}(\nabla_{\mathbf{W}_{F}}\hat{L}) \mathbf{W}_{E}(q)-\mu(k)\right|$ (10) $\displaystyle\leq\sqrt{\frac{4\sigma^{2}(k)\left(\ln(N+1)+\ln(\frac{2}{\delta} )\right)}{m}}+\frac{4R(k)\left(\ln(N+1)+\ln(\frac{2}{\delta})\right)}{m},$ $$

where $\mu(k),\sigma^{2}(k),R(k)$ are expectation, variance and range for different choices of $k\in[N]$ as follows:

|  | $\mu(N+1)=-\alpha,$ | $\sigma^{2}(N+1)=\alpha(1-\alpha),$ | $R(N+1)=\max\{\alpha,1-\alpha\},$ |
| --- | --- | --- | --- |
| $\forall~{}k\leq N:$ | $\mu(k)=\frac{1}{N+1}-\frac{1-\alpha}{N},$ | $\sigma^{2}(k)=\frac{1-\alpha}{N},$ | $R(k)=1.$ |

###### Proof.

Due to zero initialization, i.e., $\mathbf{W}_{V}=\mathbf{W}_{F}=0$, the current predicted probability is $\hat{p}_{\mathbf{W}}(k|x_{i})\equiv\frac{1}{N+1}$ for all $i\in[m]$ and $k\in[N+1]$. Therefore, from Lemma F.1, we have

$$ $\displaystyle\nabla_{\mathbf{W}_{F}}\hat{L}=\frac{1}{m}\sum_{i=1}^{m}\left[ \sum_{k=1}^{N+1}\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{y_{i}=k\}\bigg{)}\mathbf{W} _{U}(k)x_{i,T}^{\top}\right],$ $$

where $x_{i,T}$ $\in\mathbb{R}^{d}=\mathbf{W}_{E}(z_{i,T})+p_{T}$ is the input embedding with input token $z_{i,T}$ at position $T$ in sequence $i$, together with positional encoding $p_{T}$ for position $T$. Since $z_{i,T}$ is set to be the trigger $q$ in the data generation process and $p_{T}$ is assumed to orthogonal to any other vector in $\mathbf{W}_{E}$ in Assumption D.1, we have the following projections for $\nabla_{\mathbf{W}_{F}}\hat{L}$: $\forall~{}k\in[N+1]$,

$$ $\displaystyle\mathbf{W}_{U}(k)^{\top}(\nabla_{\mathbf{W}_{F}}\hat{L})\mathbf{W }_{E}(q)=\frac{1}{m}\sum_{i=1}^{m}\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{y_{i}=k\} \bigg{)}.$ $$

From the data generation process, it is obvious to get

$$ $\displaystyle\mathbb{E}_{(x,y)}\bigg{[}\frac{1}{N+1}-\mathbbm{1}\{y=k\}\bigg{] }=\frac{1}{N+1}-\alpha\cdot\mathbbm{1}\{k=N+1\}-\frac{1-\alpha}{N}\cdot \mathbbm{1}\{k\leq N\}.$ (11) $$

Since $\alpha=\Theta(1)$ is much larger than $\frac{1}{N+1}$ when $N\gg 1$, due to law of large numbers, we have the population gradient $\nabla_{\mathbf{W}_{F}}L$ satisfying

$$ $\displaystyle\mathbf{W}_{U}(N+1)^{\top}(-\nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{ E}(q)$ $\displaystyle\approx\alpha=\Theta(1),$ $\displaystyle\forall~{}k\leq N:~{}~{}~{}~{}~{}~{}\mathbf{W}_{U}(k)^{\top}(- \nabla_{\mathbf{W}_{F}}L)\mathbf{W}_{E}(q)$ $\displaystyle<0,\text{ with absolute value in }O(\nicefrac{{1}}{{N}}).$ $$

The variance of the gradient projection onto $\mathbf{W}_{U}(N+1)\mathbf{W}_{E}(q)^{\top}$ of a single data point follows that of Bernoulli distribution with parameter $\alpha$, which means

$$ $\displaystyle\text{Var}\bigg{[}\frac{1}{N+1}-\mathbbm{1}\{y=N+1\}\bigg{]}= \alpha(1-\alpha).$ (12) $$

Similarly, for any $k\leq N$, the variance of the gradient projection onto $\mathbf{W}_{U}(N+1)\mathbf{W}_{E}(q)^{\top}$ of a single data point follows that of Bernoulli distribution with parameter $\frac{1-\alpha}{N}$, which means

$$ $\displaystyle\text{Var}\bigg{[}\frac{1}{N+1}-\mathbbm{1}\{y=k\}\bigg{]}=\frac{ 1-\alpha}{N}\left(1-\frac{1-\alpha}{N}\right)=\Theta(\nicefrac{{1}}{{N}}).$ (13) $$

The ranges of the gradient projections’ deviation from the expectation are

$$ $\displaystyle\left|\frac{1}{N+1}-\mathbbm{1}\{y=N+1\}-\bigg{(}\frac{1}{N+1}- \alpha\bigg{)}\right|$ $\displaystyle\leq\max\{\alpha,1-\alpha\},$ (14) $\displaystyle\forall~{}k\leq N:~{}~{}~{}~{}~{}~{}\left|\frac{1}{N+1}-\mathbbm{ 1}\{y=k\}-\bigg{(}\frac{1}{N+1}-\frac{1-\alpha}{N}\bigg{)}\right|$ $\displaystyle\lessapprox 1.$ $$

For each choice of $k\in[N+1]$ individually, after having the expectation $\mu(k)$, variance $\sigma^{2}(k)$ and range $R(k)$, by applying Bernstein’s inequality, then: for each $k\in[N+1]$, with probability $1-\delta$, it holds

$$ $\displaystyle\left|\mathbf{W}_{U}(k)^{\top}(\nabla_{\mathbf{W}_{F}}\hat{L}) \mathbf{W}_{E}(q)-\mu(k)\right|\leq\sqrt{\frac{4\sigma^{2}(k)\ln(\frac{2}{ \delta})}{m}}+\frac{4R(k)\ln(\frac{2}{\delta})}{m}.$ $$

Then by the union bound in probability, we need $(N+1)$ events above to hold at the same time, so we can substitute $\delta$ with $\frac{\delta}{N+1}$ to have: with probability $1-\delta$, for any $k\in[N+1]$, it holds

$$ $\displaystyle\left|\mathbf{W}_{U}(k)^{\top}(\nabla_{\mathbf{W}_{F}}\hat{L}) \mathbf{W}_{E}(q)-\mu(k)\right|\leq\sqrt{\frac{4\sigma^{2}(k)\left(\ln(N+1)+ \ln(\frac{2}{\delta})\right)}{m}}+\frac{4R(k)\left(\ln(N+1)+\ln(\frac{2}{ \delta})\right)}{m}.$ (15) $$

∎

###### Lemma D.1 .

###### Proof.

### D.2 Gradient for the Value Matrix 𝐖 V subscript 𝐖 𝑉 \mathbf{W}_{V} bold_W start_POSTSUBSCRIPT italic_V end_POSTSUBSCRIPT

###### Lemma D.2 .

Consider zero initialization, $\mathbf{W}_{V}=\mathbf{W}_{F}=\mathbf{W}_{KQ}=0$. Then with probability $1-\delta$, for any $j,k\in[N+1]$, it holds

$$ $\displaystyle\left|\mathbf{W}_{U}(j)^{\top}(\nabla_{\mathbf{W}_{V}}\hat{L}) \mathbf{W}_{E}(k)-\mu(j,k)\right|$ (16) $\displaystyle\leq\sqrt{\frac{4\sigma^{2}(j,k)\left(2\ln(N+1)+\ln(\frac{2}{ \delta})\right)}{m}}+\frac{4R(j,k)\left(2\ln(N+1)+\ln(\frac{2}{\delta})\right) }{m},$ $$

where $\mu(j,k),\sigma^{2}(j,k),R(j,k)$ are expectation, variance and range for different choices of $(j,k)$ at listed in Table 2.

**Table 2: $\mu(j,k),\sigma^{2}(j,k),R(j,k)$ for different choices of $(j,k)$ in Lemma D.2.**
| $j$ | $k$ | $\mu$ | $\sigma^{2}$ | $R$ |
| --- | --- | --- | --- | --- |
| $N+1$ | $N+1$ | $-\frac{\alpha^{2}}{N}$ | $\frac{\alpha^{2}}{TN}+\frac{\alpha^{3}-\alpha^{4}}{N^{2}}$ | $\frac{1}{2}$ |
| $N+1$ | $q$ | $-\frac{\alpha}{N}$ | $\frac{\alpha}{TN}+\frac{\alpha-\alpha^{2}}{N^{2}}$ | $1$ |
| $N+1$ | $[N]\setminus\{q\}$ | $-\frac{\alpha}{N}$ | $\frac{\alpha}{TN}+\frac{\alpha-\alpha^{2}}{N^{2}}$ | $1$ |
| $q$ | $N+1$ | $\frac{2\alpha-1}{N^{2}}$ | $\frac{1}{TN^{2}}+\frac{\alpha^{2}-\alpha+1}{N^{3}}$ | $\frac{1}{2}$ |
| $q$ | $q$ | $\frac{2\alpha-1}{\alpha N^{2}}$ | $\frac{\alpha^{3}-\alpha^{2}-\alpha+2}{\alpha^{3}TN^{2}}+\frac{\alpha^{2}-<br>\alpha+1}{\alpha^{2}N^{3}}$ | $1$ |
| $q$ | $[N]\setminus\{q\}$ | $\frac{\alpha}{N^{2}}$ | $(2-\alpha)\cdot\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right)$ | $1$ |
| $[N]\setminus\{q\}$ | $N+1$ | $\frac{\alpha^{2}}{N^{2}}$ | $(2-\alpha)\left(\frac{\alpha}{TN^{2}}+\frac{\alpha^{2}}{N^{3}}\right)$ | $\frac{1}{3}$ |
| $[N]\setminus\{q\}$ | $q$ | $\frac{\alpha}{N^{2}}$ | $(2-\alpha)\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right)$ | $\frac{1}{2}$ |
| $[N]\setminus\{q\}$ | $j$ | $\frac{-\alpha^{2}+3\alpha-1}{N^{2}}$ | $\frac{1+(1-\alpha)(2-\alpha)}{TN^{2}}+\frac{1+(1-\alpha)(2-\alpha)^{2}}{N^{3}}$ | $1$ |
| $[N]\setminus\{q\}$ | $[N]\setminus\{q,j\}$ | $\frac{\alpha}{N^{2}}$ | $(2-\alpha)\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right)$ | $1$ |

###### Proof.

Due to zero initialization, i.e., $\mathbf{W}_{V}=\mathbf{W}_{F}=0$, the current predicted probability is $\hat{p}_{\mathbf{W}}(k|x_{i})\equiv\frac{1}{N+1}$ for all $i\in[m]$ and $k\in[N+1]$. Meanwhile, the attention score is uniform as $\frac{1}{T}$ for all context positions due to $\mathbf{W}_{K}=0$. Therefore, from Lemma F.1, we have

$$ $\displaystyle\nabla_{\mathbf{W}_{F}}\hat{L}=\frac{1}{m}\sum_{i=1}^{m}\left[ \sum_{k=1}^{N+1}\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{y_{i}=k\}\bigg{)}\mathbf{W} _{U}(k)\bigg{(}\frac{1}{T}\sum_{t=1}^{T}x_{i,t}\bigg{)}^{\top}\right],$ $$

where $x_{i,t}$ $\in\mathbb{R}^{d}=\mathbf{W}_{E}(z_{i,t})+p_{t}$ is the input embedding with input token $z_{i,t}$ at position $t$ in sequence $i$, together with positional encoding $p_{t}$ for position $t$. With the assumption of orthonormality in Assumption D.1, we have the projection of $\nabla_{\mathbf{W}_{F}}\hat{L}$: $\forall~{}j,k\in[N+1]$,

$$ $\displaystyle\mathbf{W}_{U}(j)^{\top}(\nabla_{\mathbf{W}_{V}}\hat{L})\mathbf{W }_{E}(k)=\frac{1}{m}\sum_{i=1}^{m}\left[\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{y_{ i}=j\}\bigg{)}\bigg{(}\frac{1}{T}\sum_{t=1}^{T}\mathbbm{1}\{z_{i,t}=k\}\bigg{) }\right].$ $$

Since each sample is drawn i.i.d., it suffices to discuss the expectation and variance of

$$ $\displaystyle\Gamma_{i}(j,k)$ $\displaystyle\triangleq\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{z_{i,T+1}=j\}\bigg{) }\bigg{(}\frac{1}{T}\sum_{t=1}^{T}\mathbbm{1}\{z_{i,t}=k\}\bigg{)},$ $\displaystyle\hat{\Gamma}(j,k)$ $\displaystyle\triangleq\frac{1}{m}\sum_{i=1}^{m}\Gamma_{i}(j,k),$ $$

where we use the fact $y_{i}=z_{i,T+1}$.

Recall that, for each sample in the data generation process, the trigger $q$ is fixed while the correct next token $\bar{y}\sim\text{Uniform}([N])$. Hence, conditioning on $z_{i,T}=q$, it has probability $\alpha$ for $z_{i,T+1}=N+1$ and probability $1-\alpha$ for $z_{i,T+1}=\bar{y}$. This leads to the necessity of discussing whether or not $\bar{y}=k$. Meanwhile, a corner case of $\bar{y}=q$ is also necessary to consider, as this implies an event that increases the counting $\frac{1}{T}\sum_{t=1}^{T}\mathbbm{1}\{z_{i,t}=q\}$ than the case of $\bar{y}\neq q$.

Therefore, generally there are 10 cases due to different choices of $(j,k)$ as follows:

- 1.
$j=N+1,k=N+1$,
- 2.
$j=N+1,k=q$,
- 3.
$j=N+1,k\in[N]\setminus\{q\}$,
- 4.
$j=q,k=N+1$,
- 5.
$j=q,k=q$,
- 6.
$j=q,k\in[N]\setminus\{q\}$,
- 7.
$j\in[N]\setminus\{q\},k=N+1$,
- 8.
$j\in[N]\setminus\{q\},k=q$,
- 9.
$j\in[N]\setminus\{q\},k=j$,
- 10.
$j\in[N]\setminus\{q\},k\in[N]\setminus\{q,j\}$.

For each $\Gamma_{i}(j,k)$ individually, if we have its expectation $\mu(j,k)$, variance $\sigma^{2}(j,k)$ and range $R(j,k)$, by applying Bernstein’s inequality, then: for each $j,k\in[N+1]$, with probability $1-\delta$, it holds

$$ $\displaystyle\left|\hat{\Gamma}(j,k)-\mu(j,k)\right|\leq\sqrt{\frac{4\sigma^{2 }(j,k)\ln(\frac{2}{\delta})}{m}}+\frac{4R(j,k)\ln(\frac{2}{\delta})}{m}.$ $$

Then by the union bound in probability, we need $(N+1)^{2}$ events above to hold at the same time, so we can substitute $\delta$ with $\frac{\delta}{(N+1)^{2}}$ to have: with probability $1-\delta$, for any $j,k\in[N+1]$, it holds

$$ $\displaystyle\left|\hat{\Gamma}(j,k)-\mu(j,k)\right|\leq\sqrt{\frac{4\sigma^{2 }(j,k)\left(2\ln(N+1)+\ln(\frac{2}{\delta})\right)}{m}}+\frac{4R(j,k)\left(2 \ln(N+1)+\ln(\frac{2}{\delta})\right)}{m}.$ (17) $$

As a final step of the proof, now we elaborate the expectation, variance and range of $\Gamma_{i}(j,k)$ for these 10 cases.

- Case 1:
$j=N+1,k=N+1$.
There is probability $\frac{1}{N}$ for $\bar{y}=q$ and probability $\frac{N-1}{N}$ for $\bar{y}\neq q$. Hence, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{N-1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q],$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{N-1}{
N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q].$
From Lemma 19 and the independence between $\mathbbm{1}\{z_{i,T+1}=N+1\}$ and $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=k\}$, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-\alpha\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right),$
where the second is from
$\displaystyle\mathbb{E}\bigg{[}\bigg{(}\frac{1}{N+1}-\mathbbm{1}\{z_{i,T+1}=N+
1\}\bigg{)}^{2}\bigg{]}$
$\displaystyle=(1-\alpha)\cdot\bigg{(}\frac{1}{N+1}\bigg{)}^{2}+\alpha\cdot
\bigg{(}\frac{1}{N+1}-1\bigg{)}^{2}\approx\alpha.$
Similarly, from Lemma 22, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]$
$\displaystyle\approx-\alpha\cdot\frac{\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\alpha\cdot\left(\frac{\alpha}{TN}+\frac{\alpha^{2}}{N^{2}
}\right).$
Therefore, it holds
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\frac{-\alpha}{N}+\frac{N-1}{N}\frac{-\alpha^{2}}{N}
\approx-\frac{\alpha^{2}}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{N-1}{
N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\frac{1}{N}\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}
\right)+\frac{N-1}{N}\alpha\cdot\left(\frac{\alpha}{TN}+\frac{\alpha^{2}}{N^{2
}}\right)\approx\frac{\alpha^{2}}{TN}+\frac{\alpha^{3}}{N^{2}},$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{\alpha^{2}}{TN}+\frac{\alpha^{3}-\alpha^{4}}{N^{2}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\leq\frac{1}{2},$
and the extreme case is when half of the sequence is $N+1$ with the rest all being $q$.
- Case 2:
$j=N+1,k=q$.
Similar to Case 1, we have $\mathbbm{1}\{z_{i,T+1}=N+1\}$ is independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=k\}$.
From Lemma 18, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-\alpha\cdot\frac{1}{\alpha N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{\alpha TN}\left(-1+\frac{2}{
\alpha^{2}}\right)+\frac{1}{\alpha^{2}N^{2}}\right).$
From Lemma 21, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]$
$\displaystyle\approx-\alpha\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{N-1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]\approx-\frac{\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{N-1}{
N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]\approx\frac{\alpha}{TN}+\frac{
\alpha}{N^{2}},$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{\alpha}{TN}+\frac{\alpha-\alpha^{2}}{N^{2}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when $\bar{y}=q$ and the sequence is all $q$’s.
- Case 3:
$j=N+1,k\in[N]\setminus\{q\}$.
Similar to Case 1, we have $\mathbbm{1}\{z_{i,T+1}=N+1\}$ is independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=k\}$.
From Lemma 20, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-\alpha\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 24, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]$
$\displaystyle\approx-\alpha\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle\approx-\alpha\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle\approx\alpha\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right),$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{\alpha}{TN}+\frac{\alpha-\alpha^{2}}{N^{2}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when all of the sequence except the last one is $k$.
- Case 4:
$j=q,k=N+1$.
If $\bar{y}\neq q$, we always have $z_{i,T+1}\neq q$ because $z_{i,T+1}\in\{\bar{y},N+1\}$. If conditioning on $\bar{y}=q$, it has probability $1-\alpha$ for $z_{i,T+1}=q$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=N+1\}$.
From Lemma 22, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{\alpha}{TN}+\frac{\alpha^{2}
}{N^{2}}\right).$
From Lemma 19, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{N-1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]\approx\frac{2\alpha-1}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{N-1}{
N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]\approx\frac{1}{TN^{2}}+\frac{
\alpha^{2}-\alpha+1}{N^{3}},$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{1}{TN^{2}}+\frac{\alpha^{2}-\alpha+1}{N^{3}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox\frac{1}{
2},$
and the extreme case is when $\bar{y}=q$ and half of the sequence is $N+1$ with the rest all being $q$.
- Case 5:
$j=q,k=q$.
Similar to Case 4, if $\bar{y}\neq q$, we always have $z_{i,T+1}\neq q$. If conditioning on $\bar{y}=q$, it has probability $1-\alpha$ for $z_{i,T+1}=q$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=q\}$.
From Lemma 21, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 18, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{1}{\alpha N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{1}{\alpha TN}\left(-1+\frac{2}{
\alpha^{2}}\right)+\frac{1}{\alpha^{2}N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{N-1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q]\approx\frac{2\alpha-1}{\alpha N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{N-1}{
N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q]$
$\displaystyle\approx\frac{\alpha^{3}-\alpha^{2}-\alpha+2}{\alpha^{3}TN^{2}}+
\frac{\alpha^{2}-\alpha+1}{\alpha^{2}N^{3}},$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{\alpha^{3}-\alpha^{2}-\alpha+2}{\alpha^{3}TN^{2}}+\frac{\alpha^{2
}-\alpha+1}{\alpha^{2}N^{3}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when $\bar{y}=q$ and all of the sequence are $q$.
- Case 6:
$j=q,k\in[N]\setminus\{q\}$.
Similar to Case 4, if $\bar{y}\neq q$, we always have $z_{i,T+1}\neq q$. If conditioning on $\bar{y}=q$, it has probability $1-\alpha$ for $z_{i,T+1}=q$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=k\}$.
Moreover, we need to consider whether $\bar{y}=k$ or not.
From Lemma 23, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,k=\bar{y}]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{2-\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,k=\bar{y}]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{2-\alpha}{TN}+\frac{(2-
\alpha)^{2}}{N^{2}}\right).$
From Lemma 24, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,k\in[N]\setminus\{q,\bar
{y}\}]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,k\in[N]\setminus\{q,
\bar{y}\}]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 20, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,k=\bar{y}]$
$\displaystyle~{}~{}~{}~{}+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q
,k\in[N]\setminus\{q,\bar{y}\}]$
$\displaystyle\approx\frac{\alpha}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,k=\bar{y}]$
$\displaystyle~{}~{}~{}~{}+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}
\neq q,k\in[N]\setminus\{q,\bar{y}\}]$
$\displaystyle\approx(2-\alpha)\cdot\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}
\right),$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx(2-\alpha)\cdot\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right).$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when all of the sequence except the last one are $k$.
- Case 7:
$j\in[N]\setminus\{q\},k=N+1$.
If $\bar{y}\neq j$, we always have $z_{i,T+1}\neq j$ because $z_{i,T+1}\in\{\bar{y},N+1\}$. If conditioning on $\bar{y}=j$, it has probability $1-\alpha$ for $z_{i,T+1}=j$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=N+1\}$.
Moreover, in the case of $\bar{y}\neq j$, we need to discuss whether or not $\bar{y}=q$.
From Lemma 19, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 22, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{\alpha}{TN}+\frac{\alpha^{2}
}{N^{2}}\right).$
From Lemma 22, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{\alpha}{TN}+\frac{\alpha^{2}}{N
^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j,k)|
y\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{\alpha^{2}}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j
,k)^{2}|y\neq q,\bar{y}\neq j]$
$\displaystyle\approx(2-\alpha)\left(\frac{\alpha}{TN^{2}}+\frac{\alpha^{2}}{N^
{3}}\right),$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx(2-\alpha)\left(\frac{\alpha}{TN^{2}}+\frac{\alpha^{2}}{N^{3}}\right).$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox\frac{1}{
3},$
and the extreme case is when $\bar{y}=j$ and one-third of the sequence are $k$, where the sequence has a repeated pattern like $[q,j,N+1,q,j,N+1,\dots]$.
- Case 8:
$j\in[N]\setminus\{q\},k=q$.
Similar to Case 7, if $\bar{y}\neq j$, we always have $z_{i,T+1}\neq j$. If conditioning on $\bar{y}=j$, it has probability $1-\alpha$ for $z_{i,T+1}=j$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=N+1\}$.
Moreover, in the case of $\bar{y}\neq j$, we need to discuss whether or not $\bar{y}=q$.
From Lemma 18, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{\alpha N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{T}{\alpha N}\left(-1+\frac{2
}{\alpha^{2}}\right)+\frac{T^{2}}{\alpha^{2}N^{2}}\right).$
From Lemma 21, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 21, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j,k)|
y\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{\alpha}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j
,k)^{2}|y\neq q,\bar{y}\neq j]$
$\displaystyle\approx(2-\alpha)\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right),$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx(2-\alpha)\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right).$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox\frac{1}{
2},$
and the extreme case is when $\bar{y}=j$ and half of the sequence are $q$.
- Case 9:
$j\in[N]\setminus\{q\},k=j$.
Similar to Case 7, if $\bar{y}\neq j$, we always have $z_{i,T+1}\neq j$. If conditioning on $\bar{y}=j$, it has probability $1-\alpha$ for $z_{i,T+1}=j$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=N+1\}$.
Moreover, in the case of $\bar{y}\neq j$, we need to discuss whether or not $\bar{y}=q$.
From Lemma 20, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 24, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 23, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{2-\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{2-\alpha}{TN}+\frac{(2-\alpha)^
{2}}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j,k)|
y\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{-\alpha^{2}+3\alpha-1}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]+\frac{N-2}{N}\mathbb{E}[\Gamma_{i}(j
,k)^{2}|y\neq q,\bar{y}\neq j]$
$\displaystyle\approx\frac{1+(1-\alpha)(2-\alpha)}{TN^{2}}+\frac{1+(1-\alpha)(2
-\alpha)^{2}}{N^{3}},$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx\frac{1+(1-\alpha)(2-\alpha)}{TN^{2}}+\frac{1+(1-\alpha)(2-\alpha)^{2}}
{N^{3}}.$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when $\bar{y}=j$ and all of the sequence are $j=k$.
- Case 10:
$j\in[N]\setminus\{q\},k\in[N]\setminus\{q,j\}$.
Similar to Case 7, if $\bar{y}\neq j$, we always have $z_{i,T+1}\neq j$. If conditioning on $\bar{y}=j$, it has probability $1-\alpha$ for $z_{i,T+1}=j$, independent of $\sum_{t\leq T}\mathbbm{1}\{z_{i,t}=N+1\}$.
Moreover, in the case of $\bar{y}\neq j$, we need to discuss whether or not $\bar{y}=q$.
From Lemma 20, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 24, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]$
$\displaystyle\approx-(1-\alpha)\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]$
$\displaystyle\approx(1-\alpha)\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
From Lemma 23, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=k]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{2-\alpha}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=k]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{2-\alpha}{TN}+\frac{(2-
\alpha)^{2}}{N^{2}}\right).$
From Lemma 24, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}\neq q,\bar{y}\neq j,\bar{y}
\neq k]$
$\displaystyle\approx\frac{1}{N+1}\cdot\frac{1}{N},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}\neq q,\bar{y}\neq j,\bar{y
}\neq k]$
$\displaystyle\approx\frac{1}{N+1}\cdot\left(\frac{1}{TN}+\frac{1}{N^{2}}\right).$
Therefore, we have
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)|\bar{y}=j]+\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)|
\bar{y}=k]$
$\displaystyle~{}~{}~{}~{}+\frac{N-3}{N}\mathbb{E}[\Gamma_{i}(j,k)|y\neq q,\bar
{y}\neq j]$
$\displaystyle\approx\frac{\alpha}{N^{2}},$
$\displaystyle\mathbb{E}[\Gamma_{i}(j,k)^{2}]$
$\displaystyle=\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=q]+\frac{1}{N}
\mathbb{E}[\Gamma_{i}(j,k)^{2}|\bar{y}=j]+\frac{1}{N}\mathbb{E}[\Gamma_{i}(j,k
)^{2}|\bar{y}=k]$
$\displaystyle~{}~{}~{}~{}+\frac{N-3}{N}\mathbb{E}[\Gamma_{i}(j,k)^{2}|y\neq q,
\bar{y}\neq j]$
$\displaystyle\approx(2-\alpha)\left(\frac{1}{TN}+\frac{1}{N^{2}}\right),$
$\displaystyle\text{Var}[\Gamma_{i}(j,k)]$
$\displaystyle=\mathbb{E}[\Gamma_{i}(j,k)^{2}]-\mathbb{E}[\Gamma_{i}(j,k)]^{2}
\approx(2-\alpha)\left(\frac{1}{TN^{2}}+\frac{1}{N^{3}}\right).$
The range of $\Gamma_{i}(j,k)$ is
$\displaystyle|\Gamma_{i}(j,k)-\mathbb{E}[\Gamma_{i}(j,k)]|\lessapprox 1,$
and the extreme case is when $\bar{y}=j$ and all of the sequence except the last are $k$.

∎

###### Lemma D.2 .

###### Proof.

## Appendix E Proof for First and Second moments

In this section, we will show the proof of the first and second moments of $\left[\sum_{1\leq t\leq T}\mathbbm{1}\{z_{t}=k\}|\cdot\right]$ for all cases.
Note that we do not consider $z_{T}=q$, but including it will not change the results, as $T\gg 1$ and $z_{T}$ is explicitly fixed as $q$ during data generation in Section 3.
Generally, there are three factors to classify the cases as follows:

- 1.
The i.i.d. uniformly sampled correct token $\bar{y}\in[N]$:
(a)
$\bar{y}=q$,
(b)
$\bar{y}\neq q$.
- 2.
The target token $k\in[N+1]$:
(a)
$k=q$,
(b)
$k=N+1$.
(c)
$k\leq N,k\neq q,k\neq\bar{y}$,
(d)
(if $\bar{y}\neq q$) $k\leq N,k\neq q,k=\bar{y}$,
- 3.
A condition about the token $z_{0}$ before the sequence $\{z_{t}\}_{t\geq 1}$:
(a)
$z_{0}=q$,
(b)
$z_{0}\in[N+1]\setminus\{q\}$.

Note that when $z_{0}$ will be implicitly or explicitly considered. When there is no condition on the first token, which means $z_{1}\sim\text{Uniform}([N])$, this belongs to Case (3b), i.e., $z_{0}\in[N+1]\setminus\{q\}$, following the data generation process.

Table 3 summarizes all lemmas about the seven cases classified by the first two factors. The third factor about $z_{0}$ is explicitly presented in the proof of each corresponding lemma.

**Table 3: All lemmas about the seven cases classified by $\bar{y}$ and $k$.**
|  | (2a) | (2b) | (2c) | (2d) |
| --- | --- | --- | --- | --- |
| (1a) | 18 | 19 | 20 | N/A |
| (1b) | 21 | 22 | 24 | 23 |

### E.1 When y ¯ = q ¯ 𝑦 𝑞 \bar{y}=q over¯ start_ARG italic_y end_ARG = italic_q

###### Lemma E.1 ( $\bar{y}=q,k=q$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}=q$ and $k=q$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k=q\right]$ $\displaystyle\approx\frac{T}{\alpha N},$ (18) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k=q\right]$ $\displaystyle\approx\frac{T}{\alpha N}\left(-1+\frac{2}{\alpha^{2}}\right)+ \frac{T^{2}}{\alpha^{2}N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}=q,k=q$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot(1+Y(T-1))+p(z_{1}=N+1|z_{0}=q)\cdot\hat{ Y}(T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot(1+Y(T-1))+p(z_{1}\in[N]\setminus\{q \}|z_{0}\neq q)\cdot\hat{Y}(T-1).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=(1-\alpha)\cdot Y(T-1)+\alpha\cdot\hat{Y}(T-1)+1-\alpha,$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1)+\frac{1}{ N}.$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Y(T-1)-\hat{Y}(T-1))+1-\alpha-\frac{1}{N},$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\alpha\hat{Y}(T-1)+\frac{1}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{1-\alpha-\frac{1}{N}}{\alpha+\frac{1}{N}}\left(1-\left(1- \alpha-\frac{1}{N}\right)^{T}\right),$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{1}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{1}{\alpha N+1}(T-\alpha N)+\frac{\alpha}{(\alpha+ \frac{1}{N})^{2}}=\frac{1}{\alpha N+1}\left(T-\alpha N+\frac{N^{2}}{\alpha N+1 }\right)$ $\displaystyle\approx\frac{T}{\alpha N}-1+\frac{1}{\alpha^{2}},$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{1}{\alpha N+1}T-\frac{N}{(\alpha N+1)^{2}}+\frac{1}{ \alpha N+1}$ $\displaystyle\approx\frac{T}{\alpha N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k=q\right]=\hat{Y}(T)\approx\frac{T}{\alpha N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot(1+2Y(T-1)+Z(T-1))+p(z_{1}=N+1|z_{0}=q) \cdot Z(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot(1+2Y(T-1)+Z(T-1))+p(z_{1}\neq q|z_{0 }\neq q)\cdot\hat{Z}(T-1),$ $$

where $2Y(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=(1-\alpha)\cdot(1+2Y(T-1)+Z(T-1))+\alpha\cdot\hat{Z}(T-1)$ $\displaystyle=(1-\alpha)Z(T-1)+\alpha\hat{Z}(T-1)+(1-\alpha)(1+2Y(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}\cdot(1+2Y(T-1)+Z(T-1))+\frac{N-1}{N}\cdot\hat{Z}(T-1)$ $\displaystyle=\frac{1}{N}Z(T-1)+\frac{N-1}{N}\hat{Z}(T-1)+\frac{1}{N}(1+2Y(T-1 )).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Z(T-1)-\hat{Z}(T-1))+(1-\alpha-\frac{1}{N })(1+2Y(T-1)),$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\alpha\hat{Z}(T-1)+\frac{1}{N}(1+2Y(T-1)).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=\sum_{t\leq T-1}(1-\alpha-\frac{1}{N})^{T-t}(1+2Y(t))$ $\displaystyle\approx\sum_{t\leq T-1}(1-\alpha-\frac{1}{N})^{T-t}\left(1+\frac{ 2t}{\alpha N}-2+\frac{2}{\alpha^{2}}\right)$ $\displaystyle\approx\left(-1+\frac{2}{\alpha^{2}}\right)\frac{1-\alpha}{\alpha }+\frac{2(1-\alpha)}{\alpha^{2}}\cdot\frac{T}{N}.$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{T}{N}+\frac{2}{N}\sum_{1\leq t\leq T-1}Y(t)$ $\displaystyle\approx\frac{T}{N}+\frac{2}{N}\sum_{1\leq t\leq T-1}\left(\frac{t }{\alpha N}-1+\frac{1}{\alpha^{2}}\right)$ $\displaystyle\approx\frac{T}{N}\left(-1+\frac{2}{\alpha^{2}}\right)+\frac{T^{2 }}{\alpha N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}\left(-\frac{3}{\alpha}+\frac{2}{\alpha^{2}}+ \frac{2}{\alpha^{3}}\right)+\frac{T^{2}}{\alpha^{2}N^{2}}+\frac{1-\alpha}{ \alpha}(\frac{2}{\alpha^{2}}-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{T}{\alpha N}\left(-1+\frac{2}{\alpha^{2}}\right)+ \frac{T^{2}}{\alpha^{2}N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k=q\right]=\hat{Z}(T)\approx\frac{T}{\alpha N}\left(- 1+\frac{2}{\alpha^{2}}\right)+\frac{T^{2}}{\alpha^{2}N^{2}}.$ $$

∎

###### Lemma E.2 ( $\bar{y}=q,k=N+1$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}=q$ and $k=N+1$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k=N+1\right]$ $\displaystyle\approx\frac{T}{N},$ (19) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k=N+1\right]$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}=q,k=N+1$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot Y(T-1)+p(z_{1}=N+1|z_{0}=q)\cdot(1+\hat{ Y}(T-1)),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Y(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Y}(T-1).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=(1-\alpha)\cdot Y(T-1)+\alpha\cdot\hat{Y}(T-1)+\alpha,$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1).$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Y(T-1)-\hat{Y}(T-1))+\alpha,$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\alpha\hat{Y}(T-1)+\frac{\alpha}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{\alpha}{\alpha+\frac{1}{N}}\left(1-\left(1-\alpha-\frac{1} {N}\right)^{T}\right),$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{\alpha}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{T}{N}+1,$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k=N+1\right]=\hat{Y}(T)\approx\frac{T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot Z(T-1)+p(z_{1}=N+1|z_{0}=q)\cdot(1+2\hat {Y}(T-1)+\hat{Z}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Z(T-1)+p(z_{1}\neq q|z_{0}\neq q) \cdot\hat{Z}(T-1),$ $$

where $2\hat{Y}(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=(1-\alpha)\cdot Z(T-1)+\alpha\cdot(1+2\hat{Y}(T-1)+\hat{Z}(T-1))$ $\displaystyle=(1-\alpha)Z(T-1)+\alpha\hat{Z}(T-1)+\alpha(1+2\hat{Y}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}\cdot Z(T-1)+\frac{N-1}{N}\cdot\hat{Z}(T-1).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Z(T-1)-\hat{Z}(T-1))+\alpha(1+2\hat{Y}(T- 1)),$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\alpha\hat{Z}(T-1)+\frac{\alpha}{N}(1+2\hat{Y} (T-1)).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=\sum_{t\leq T-1}\alpha(1-\alpha-\frac{1}{N})^{T-1-t}(1+2\hat{Y}( t))$ $\displaystyle\approx\sum_{t\leq T-1}\alpha(1-\alpha-\frac{1}{N})^{T-1-t}\left( 1+\frac{2t}{N}\right)$ $\displaystyle\approx\frac{2T}{N}+1,$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1}\hat{Y }(t)$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1} \frac{t}{N}$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{\alpha T^{2}}{N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx 3\alpha\frac{T}{N}+\alpha\frac{T^{2}}{N^{2}}+\alpha,$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k=N+1\right]=\hat{Z}(T)\approx\frac{T}{N}+\frac{T^{2} }{N^{2}}.$ $$

∎

###### Lemma E.3 ( $\bar{y}=q,k\leq N,k\neq q$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}=q$ and $k\in[N]\setminus\{q\}$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k\in[N]\setminus\{q\}\right]$ $\displaystyle\approx\frac{T}{N},$ (20) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}=q,k\in[N]\setminus\{q\}$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot Y(T-1)+p(z_{1}=N+1|z_{0}=q)\cdot\hat{Y}( T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Y(T-1)$ $\displaystyle~{}~{}~{}~{}+p(z_{1}\in[N]\setminus\{q\}|z_{0}\neq q)\cdot(p(z_{1 }=k|z_{1}\sim\text{Uniform}([N]\setminus\{q\})+\hat{Y}(T-1)).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=(1-\alpha)\cdot Y(T-1)+\alpha\cdot\hat{Y}(T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot(\hat{Y}(T-1)+\frac{1} {N-1}).$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Y(T-1)-\hat{Y}(T-1))-\frac{1}{N},$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\alpha\hat{Y}(T-1)+\frac{\alpha}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{-\frac{1}{N}}{\alpha+\frac{1}{N}}\left(1-\left(1-\alpha- \frac{1}{N}\right)^{T}\right),$ $\displaystyle\frac{1}{N}Y(T)+\alpha\hat{Y}(T)$ $\displaystyle=\frac{\alpha}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{T}{N},$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}=q,k=N+1\right]=\hat{Y}(T)\approx\frac{T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=p(z_{1}=q|z_{0}=q)\cdot Z(T-1)+p(z_{1}=N+1|z_{0}=q)\cdot\hat{Z}( T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Z(T-1)+p(z_{1}\neq q|z_{0}\neq q) \cdot\hat{Z}(T-1)$ $\displaystyle~{}~{}~{}~{}+p(z_{1}=k|z_{0}\neq q)\cdot(1+2\hat{Y}(T-1)),$ $$

where $2\hat{Y}(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=(1-\alpha)\cdot Z(T-1)+\alpha\cdot\hat{Z}(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}\cdot Z(T-1)+\frac{N-1}{N}\cdot\hat{Z}(T-1)+\frac{1}{ N}(1+2\hat{Y}(T-1)).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})(Z(T-1)-\hat{Z}(T-1))-\frac{1}{N}(1+2\hat{ Y}(T-1)),$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\alpha\hat{Z}(T-1)+\frac{\alpha}{N}(1+2\hat{Y} (T-1)).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=-\frac{1}{N}\sum_{t\leq T-1}(1-\alpha-\frac{1}{N})^{T-1-t}(1+2 \hat{Y}(t))$ $\displaystyle\approx-\frac{1}{N}\sum_{t\leq T-1}(1-\alpha-\frac{1}{N})^{T-1-t} \left(1+\frac{2t}{N}\right)$ $\displaystyle\approx-\frac{1}{\alpha N}\left(\frac{2T}{N}+1\right),$ $\displaystyle\frac{1}{N}Z(T)+\alpha\hat{Z}(T)$ $\displaystyle=\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1}\hat{Y }(t)$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1} \frac{t}{N}$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{\alpha T^{2}}{N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}},$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]=\hat{Z}(T)\approx\frac{T }{N}+\frac{T^{2}}{N^{2}}.$ $$

∎

###### Lemma E.1 ( y ¯ = q , k = q formulae-sequence ¯ 𝑦 𝑞 𝑘 𝑞 \bar{y}=q,k=q over¯ start_ARG italic_y end_ARG = italic_q , italic_k = italic_q ) .

###### Proof.

###### Lemma E.2 ( y ¯ = q , k = N + 1 formulae-sequence ¯ 𝑦 𝑞 𝑘 𝑁 1 \bar{y}=q,k=N+1 over¯ start_ARG italic_y end_ARG = italic_q , italic_k = italic_N + 1 ) .

###### Proof.

###### Lemma E.3 ( y ¯ = q , k ≤ N , k ≠ q formulae-sequence ¯ 𝑦 𝑞 formulae-sequence 𝑘 𝑁 𝑘 𝑞 \bar{y}=q,k\leq N,k\neq q over¯ start_ARG italic_y end_ARG = italic_q , italic_k ≤ italic_N , italic_k ≠ italic_q ) .

###### Proof.

### E.2 When y ¯ ≠ q ¯ 𝑦 𝑞 \bar{y}\neq q over¯ start_ARG italic_y end_ARG ≠ italic_q

###### Lemma E.4 ( $\bar{y}\neq q,k=q$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}\neq q$ and $k=q$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=q\right]$ $\displaystyle\approx\frac{T}{N},$ (21) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}\neq q,k=q\right]$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}\neq q,k=q$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot(1+Y(T-1))+p(z_{1}\in[N]\setminus\{q \}|z_{0}\neq q)\cdot\hat{Y}(T-1).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1)+\frac{1}{ N}.$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=-\frac{1}{N}(Y(T-1)-\hat{Y}(T-1))-\frac{1}{N},$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\hat{Y}(T-1)+\frac{1}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{-\frac{1}{N}}{1+\frac{1}{N}}\left(1-\left(-\frac{1}{N} \right)^{T}\right),$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{1}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{T}{N},$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=q\right]=\hat{Y}(T)\approx\frac{T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot(1+2Y(T-1)+Z(T-1))+p(z_{1}\in[N] \setminus\{q\}|z_{0}\neq q)\cdot\hat{Z}(T-1),$ $$

where $2Y(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\frac{N-1}{N}\hat{Z}(T-1)+\frac{1}{N}(1+2Y(T-1 )).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=-\frac{1}{N}(Z(T-1)-\hat{Z}(T-1))-\frac{1}{N}(1+2Y(T-1)),$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\hat{Z}(T-1)+\frac{1}{N}(1+2Y(T-1)).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=-\frac{1}{N}\sum_{t\leq T-1}(-\frac{1}{N})^{T-1-t}(1+2Y(t))$ $\displaystyle\approx-\frac{1}{N}\sum_{t\leq T-1}(-\frac{1}{N})^{T-1-t}\left(1+ \frac{2t}{N}\right)$ $\displaystyle\approx-\frac{1}{N}-\frac{2T}{N^{2}},$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{T}{N}+\frac{2}{N}\sum_{1\leq t\leq T-1}Y(t)$ $\displaystyle\approx\frac{T}{N}+\frac{2}{N}\sum_{1\leq t\leq T-1}\frac{t}{N}$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}},$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]=\hat{Z}(T)\approx\frac{T }{N}+\frac{T^{2}}{N^{2}}.$ $$

∎

###### Lemma E.5 ( $\bar{y}\neq q,k=N+1$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}\neq q$ and $k=N+1$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=N+1\right]$ $\displaystyle\approx\frac{\alpha T}{N},$ (22) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}\neq q,k=N+1\right]$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{\alpha^{2}T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}\neq q,k=N+1$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1)+p(z_{1}=N+1|z_{0}=q),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Y(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Y}(T-1).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1)+\alpha,$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1).$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=-\frac{1}{N}(Y(T-1)-\hat{Y}(T-1))+\alpha,$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\hat{Y}(T-1)+\frac{\alpha}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{\alpha}{1+\frac{1}{N}}\left(1-\left(-\frac{1}{N}\right)^{T }\right),$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{\alpha}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{\alpha T}{N}+\alpha,$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{\alpha T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=q\right]=\hat{Y}(T)\approx\frac{\alpha T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1)+p(z_{1}=N+1|z_{0}=q)\cdot(1+2\hat{Y}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Z(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Z}(T-1),$ $$

where $2\hat{Y}(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1)+\alpha(1+2\hat{Y}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\frac{N-1}{N}\hat{Z}(T-1).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=-\frac{1}{N}(Z(T-1)-\hat{Z}(T-1))+\alpha(1+2\hat{Y}(T-1)),$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\hat{Z}(T-1)+\frac{\alpha}{N}(1+2\hat{Y}(T-1)).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=\alpha\sum_{t\leq T-1}(-\frac{1}{N})^{T-1-t}(1+2\hat{Y}(t))$ $\displaystyle\approx\alpha\sum_{t\leq T-1}(-\frac{1}{N})^{T-1-t}\left(1+\frac{ 2\alpha t}{N}\right)$ $\displaystyle\approx\frac{2\alpha^{2}T}{N}+\alpha,$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1}\hat{Y }(t)$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{2\alpha}{N}\sum_{1\leq t\leq T-1} \frac{\alpha t}{N}$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{\alpha^{2}T^{2}}{N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}(2\alpha^{2}+\alpha)+\frac{\alpha^{2}T^{2}}{N^{ 2}}+\alpha,$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{\alpha T}{N}+\frac{\alpha^{2}T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]=\hat{Z}(T)\approx\frac{ \alpha T}{N}+\frac{\alpha^{2}T^{2}}{N^{2}}.$ $$

∎

###### Lemma E.6 ( $\bar{y}\neq q,k=\bar{y}$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}\neq q$ and $k=\bar{y}$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=\bar{y}\right]$ $\displaystyle\approx(2-\alpha)\frac{T}{N},$ (23) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}\neq q,k=\bar{y}\right]$ $\displaystyle\approx\frac{(2-\alpha)T}{N}+\frac{(2-\alpha)^{2}T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}\neq q,k=\bar{y}$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1)+p(z_{1}=\bar{y}|z_{0}=q),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Y(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Y}(T-1)+p(z_{1}=\bar{y}|z_{0}\neq q).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1)+(1-\alpha),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1)+\frac{1}{ N}.$ $$

This gives

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=-\frac{1}{N}(Y(T-1)-\hat{Y}(T-1))+(1-\alpha-\frac{1}{N}),$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{1}{N}Y(T-1)+\hat{Y}(T-1)+\frac{2-\alpha}{N}.$ $$

Consider the initialization $Y(0)=\hat{Y}(0)=0$. This implies

$$ $\displaystyle Y(T)-\hat{Y}(T)$ $\displaystyle=\frac{1-\alpha-\frac{1}{N}}{1+\frac{1}{N}}\left(1-\left(-\frac{1 }{N}\right)^{T}\right),$ $\displaystyle\frac{1}{N}Y(T)+\hat{Y}(T)$ $\displaystyle=\frac{2-\alpha}{N}T.$ $$

Then we obtain

$$ $\displaystyle Y(T)$ $\displaystyle\approx(1-\alpha)+(2-\alpha)\frac{T}{N},$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx(2-\alpha)\frac{T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=q\right]=\hat{Y}(T)\approx(2-\alpha)\frac{T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1)+p(z_{1}=\bar{y}|z_{0}=q)\cdot(1+2\hat{Y}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Z(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Z}(T-1)$ $\displaystyle~{}~{}~{}~{}+p(z_{1}=\bar{y}|z_{0}\neq q)\cdot(1+2\hat{Y}(T-1)),$ $$

where $2\hat{Y}(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1)+(1-\alpha)(1+2\hat{Y}(T-1)),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\frac{N-1}{N}\hat{Z}(T-1)+\frac{1}{N}(1+2\hat{ Y}(T-1)).$ $$

This gives

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=-\frac{1}{N}(Z(T-1)-\hat{Z}(T-1))+(1-\alpha-\frac{1}{N})(1+2\hat {Y}(T-1)),$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\hat{Z}(T-1)+\frac{2-\alpha}{N}(1+2\hat{Y}(T-1 )).$ $$

Considering the initialization $Z(0)=\hat{Z}(0)=0$, we have

$$ $\displaystyle Z(T)-\hat{Z}(T)$ $\displaystyle=(1-\alpha-\frac{1}{N})\sum_{t\leq T-1}(-\frac{1}{N})^{T-1-t}(1+2 \hat{Y}(t))$ $\displaystyle\approx(1-\alpha-\frac{1}{N})\sum_{t\leq T-1}(-\frac{1}{N})^{T-1- t}\left(1+\frac{2(2-\alpha)t}{N}\right)$ $\displaystyle\approx(1-\alpha)\left(1+\frac{2(2-\alpha)T}{N}\right),$ $\displaystyle\frac{1}{N}Z(T)+\hat{Z}(T)$ $\displaystyle=\frac{(2-\alpha)T}{N}+\frac{2(2-\alpha)}{N}\sum_{1\leq t\leq T-1 }\hat{Y}(t)$ $\displaystyle\approx\frac{(2-\alpha)T}{N}+\frac{2(2-\alpha)}{N}\sum_{1\leq t \leq T-1}\frac{(2-\alpha)t}{N}$ $\displaystyle\approx\frac{(2-\alpha)T}{N}+\frac{(2-\alpha)^{2}T^{2}}{N^{2}}.$ $$

Then we obtain

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}(2-\alpha)(3-2\alpha)+\frac{(2-\alpha)^{2}T^{2} }{N^{2}}+(1-\alpha),$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{(2-\alpha)T}{N}+\frac{(2-\alpha)^{2}T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]=\hat{Z}(T)\approx\frac{( 2-\alpha)T}{N}+\frac{(2-\alpha)^{2}T^{2}}{N^{2}}.$ $$

∎

###### Lemma E.7 ( $\bar{y}\neq q,k\leq N,k\neq q,k\neq\bar{y}$ ) .

Following the data generation process, assuming $N,T\gg 1$ and $\alpha=\Theta(1)$, if $\bar{y}\neq q$ and $k\in[N]\setminus\{\bar{y},q\}$, it holds

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k\in[N]\setminus\{\bar{y},q\}\right]$ $\displaystyle\approx\frac{T}{N},$ (24) $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}\neq q,k\in[N]\setminus\{\bar{y},q\}\right]$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

###### Proof.

For simplicity, we omit the condition of $\bar{y}\neq q,k\in[N]\setminus\{\bar{y},q\}$ in this proof. Denote

$$ $\displaystyle Y(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Y}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\} \bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Y(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Y}(T-1)+p(z_{1}=k|z_{0}\neq q).$ $$

The iteration becomes

$$ $\displaystyle Y(T)$ $\displaystyle=\hat{Y}(T-1)+(1-\alpha),$ $\displaystyle\hat{Y}(T)$ $\displaystyle=\frac{1}{N}\cdot Y(T-1)+\frac{N-1}{N}\cdot\hat{Y}(T-1)+\frac{1}{ N}.$ $$

Note that these two equations are exactly the same as those in Lemma 21 with same initialization as $Y(0)=\hat{Y}(0)=0$. Therefore, we have

$$ $\displaystyle Y(T)$ $\displaystyle\approx\frac{T}{N},$ $\displaystyle\hat{Y}(T)$ $\displaystyle\approx\frac{T}{N}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg{|}\bar{ y}\neq q,k=q\right]=\hat{Y}(T)\approx\frac{T}{N}.$ $$

To obtain the expectation of the quadratic term, we similarly denote the following terms with different $z_{0}$:

$$ $\displaystyle Z(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}=q\right],$ $\displaystyle\hat{Z}(T)$ $\displaystyle\triangleq\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t }=k\}\bigg{)}^{2}\bigg{|}z_{0}\in[N+1],z_{0}\neq q\right].$ $$

Then the data generation process implies, $\forall~{}T\geq 1$,

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=p(z_{1}=q|z_{0}\neq q)\cdot Z(T-1)+p(z_{1}\in[N]\setminus\{q\}|z _{0}\neq q)\cdot\hat{Z}(T-1)$ $\displaystyle~{}~{}~{}~{}+p(z_{1}=\bar{k}|z_{0}\neq q)\cdot(1+2\hat{Y}(T-1)),$ $$

where $2\hat{Y}(T-1)$ is due to $\mathbb{E}[(1+\sum_{2\leq t\leq T}\cdot)^{2}]=1+2\mathbb{E}[\sum_{2\leq t\leq T
}\cdot]+\mathbb{E}[(\sum_{2\leq t\leq T}\cdot)^{2}]$.

Then the iteration becomes

$$ $\displaystyle Z(T)$ $\displaystyle=\hat{Z}(T-1),$ $\displaystyle\hat{Z}(T)$ $\displaystyle=\frac{1}{N}Z(T-1)+\frac{N-1}{N}\hat{Z}(T-1)+\frac{1}{N}(1+2\hat{ Y}(T-1)).$ $$

Again note that, since $Y(T)\approx\hat{Y}(T)$, these two equations are the same as those in Lemma 21. Therefore, we have

$$ $\displaystyle Z(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}},$ $\displaystyle\hat{Z}(T)$ $\displaystyle\approx\frac{T}{N}+\frac{T^{2}}{N^{2}}.$ $$

Since the data generation process implicitly assumes $z_{0}\neq q$, we have the desired expectation as

$$ $\displaystyle\mathbb{E}\left[\bigg{(}\sum_{t\leq T}\mathbbm{1}\{z_{t}=k\}\bigg {)}^{2}\bigg{|}\bar{y}=q,k\in[N]\setminus\{q\}\right]=\hat{Z}(T)\approx\frac{T }{N}+\frac{T^{2}}{N^{2}}.$ $$

∎

###### Lemma E.4 ( y ¯ ≠ q , k = q formulae-sequence ¯ 𝑦 𝑞 𝑘 𝑞 \bar{y}\neq q,k=q over¯ start_ARG italic_y end_ARG ≠ italic_q , italic_k = italic_q ) .

###### Proof.

###### Lemma E.5 ( y ¯ ≠ q , k = N + 1 formulae-sequence ¯ 𝑦 𝑞 𝑘 𝑁 1 \bar{y}\neq q,k=N+1 over¯ start_ARG italic_y end_ARG ≠ italic_q , italic_k = italic_N + 1 ) .

###### Proof.

###### Lemma E.6 ( y ¯ ≠ q , k = y ¯ formulae-sequence ¯ 𝑦 𝑞 𝑘 ¯ 𝑦 \bar{y}\neq q,k=\bar{y} over¯ start_ARG italic_y end_ARG ≠ italic_q , italic_k = over¯ start_ARG italic_y end_ARG ) .

###### Proof.

###### Lemma E.7 ( y ¯ ≠ q , k ≤ N , k ≠ q , k ≠ y ¯ formulae-sequence ¯ 𝑦 𝑞 formulae-sequence 𝑘 𝑁 formulae-sequence 𝑘 𝑞 𝑘 ¯ 𝑦 \bar{y}\neq q,k\leq N,k\neq q,k\neq\bar{y} over¯ start_ARG italic_y end_ARG ≠ italic_q , italic_k ≤ italic_N , italic_k ≠ italic_q , italic_k ≠ over¯ start_ARG italic_y end_ARG ) .

###### Proof.

## Appendix F Useful Lemmas

###### Lemma F.1 .

Let $p$ be a data distribution on $(x,y)\in\mathbb{R}^{d}\times[N]$. Consider training data as $m$ i.i.d. samples $\mathcal{D}\triangleq\{(x_{i},y_{i})\}_{i=1}^{m}\subset\mathbb{R}^{d}\times[N+1]$ from $p$. Consider the following classification problem, with fixed output embeddings $\mathbf{W}_{U}$:

$$ $\displaystyle\hat{L}(\mathbf{W})=\frac{1}{m}\sum_{i=1}^{m}[l(y_{i},\mathbf{W}_ {U}\mathbf{W}x_{i})].$ $$

The gradients take the following form: denoting $\hat{p}_{\mathbf{W}}(k|x_{i})$ as the current predicted probability of class $k$ in $[N+1]$ classes for input $x_{i}$,

$$ $\displaystyle\nabla_{\mathbf{W}}\hat{L}(\mathbf{W})=\frac{1}{m}\sum_{i=1}^{m} \left[\sum_{k=1}^{N+1}(\hat{p}_{\mathbf{W}}(k|x_{i})-\mathbbm{1}\{y_{i}=k\}) \mathbf{W}_{U}(k)x_{i}^{\top}\right].$ $$

###### Proof.

Recall the form of the cross-entropy loss for classification with $K$ classes:

$$ $\displaystyle l(y,\epsilon)=-\sum_{k=1}^{K}\mathbbm{1}\{y=k\}\log\frac{e^{\xi_ {k}}}{\sum_{j}e^{\xi_{j}}}.$ $$

Its derivatives take the form

$$ $\displaystyle\frac{\partial l}{\partial\xi_{k}}(y,\xi)=s(\xi)_{k}-\mathbbm{1} \{y=k\},$ $$

where $s(\xi)_{k}=\frac{e^{\xi_{k}}}{\sum_{j}e^{\xi_{j}}}$.

The gradient of $L$ is then given by

$$ $\displaystyle\nabla_{\mathbf{W}}\hat{L}(\mathbf{W})$ $\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\left[\sum_{k=1}^{N+1}\frac{\partial l} {\partial\xi_{k}}(y_{i},\mathbf{W}_{U}\mathbf{W}x_{i})\nabla_{\mathbf{W}}( \mathbf{W}_{U}(k)^{\top}\mathbf{W}x_{i})\right]$ $\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\left[\sum_{k=1}^{N+1}(\hat{p}_{\mathbf {W}}(k|x_{i})-\mathbbm{1}\{y_{i}=k\})\mathbf{W}_{U}(k)x_{i}^{\top}\right].$ $$

∎

###### Lemma F.2 .

Consider a sequence $\{S_{t}\}_{t\geq 1}$ with $S_{t}=a^{t}\cdot t$ where $a\neq 1$. Then $\sum_{1\leq t\leq T}S_{t}=\frac{a(1-a^{T})}{(a-1)^{2}}+\frac{a^{T+1}\cdot T}{a
-1}$.

###### Proof.

Denote $X_{t}\triangleq\sum_{1\leq t\leq T}S_{t}$. Then we have
$a\cdot X_{t}=\sum_{2\leq t\leq T+1}a^{t}\cdot(t-1).$
Hence, it holds
$(a-1)X_{t}=-\sum_{2\leq t\leq T}a^{t}-a+a^{T+1}\cdot T=-\frac{a(1-a^{T})}{1-a}
+a^{T+1}\cdot T.$
Therefore, we have

$$ $\displaystyle X_{t}=\frac{a(1-a^{T})}{(a-1)^{2}}+\frac{a^{T+1}\cdot T}{a-1}.$ $$

∎

###### Lemma F.3 .

Consider the following ODE with with $a(0)=b(0)=0$ and $\alpha\in(0.2,0.4)$,

$$ $\displaystyle\dot{a}$ $\displaystyle=\frac{2-2\exp(a)}{\exp(a)+\exp(b)+1}-2+2\alpha,$ $\displaystyle\dot{b}$ $\displaystyle=\frac{2-8\exp(b)}{\exp(a)+\exp(b)+1}-2+10\alpha.$ $$

Then, when $t\rightarrow\infty$, we have

$$ $\displaystyle a\rightarrow-\log(t)-\log(1-\alpha)(4-2\alpha),~{}~{}~{}~{}b \rightarrow\log\frac{\alpha}{1-\alpha}.$ $$

###### Proof.

The ODE can be re-written as

$$ $\displaystyle\dot{a}$ $\displaystyle=2\cdot\frac{(\alpha-2)\exp(a)+(\alpha-1)\exp(b)+\alpha}{\exp(a)+ \exp(b)+1}\triangleq\frac{2D}{\exp(a)+\exp(b)+1},$ $\displaystyle\dot{b}$ $\displaystyle=10\cdot\frac{(\alpha-\frac{1}{5})\exp(a)+(\alpha-1)\exp(b)+ \alpha}{\exp(a)+\exp(b)+1}\triangleq\frac{10E}{\exp(a)+\exp(b)+1}.$ $$

At $t=0$, it holds $\dot{a}(0)<0,\dot{b}(0)<0$ since $D=3\alpha-3<0,E=3\alpha-\frac{6}{5}<0$. Hence, $a$ and $b$ start to decrease from $t=0$. The ending of the decreasing happens when one of $D$ and $E$ gets positive. Let’s show $D$ and $E$ will never be positive when $\alpha\in(0.2,0.4)$ by contradiction.

Assume time $T_{1}$ is when one of $E$ and $E$ equals to 0 for the first time. This means $E=0$, because, for any time $t$, it always holds $D<E$ since $\exp(a)>0$ for any $a\in\mathbb{R}$.
Then at $T_{1}$, we have $\dot{a}<0,\dot{b}=0$, which means $\exp(a)$ will decrease for any small time window $\Delta t>0$ and $\exp(b)$ stays unchanged. Together with $\alpha>0.2$, this means it has $E<0$ again at time $T_{1}+\Delta t$. Therefore, it is possible for $E$ to be 0, but $E$ will never be positive. Meanwhile, this also guarantees $D$ will always be negative because $D<E$.

Then, we make an observation that when $D$ is always negative and $E$ is always non-positive, the decreasing nature of $a$ will have $D\approx E$ when $t\rightarrow\infty$ by $\exp(a)\approx 0$. This implies $b=\log\frac{\alpha}{1-\alpha}$. Then, by taking $\exp(a)=\beta\cdot t^{-\gamma}$, the ODE gives

$$ $\displaystyle-\gamma\frac{1}{t}=\frac{(2\alpha-4)\beta\cdot t^{-\gamma}}{\beta \cdot t^{-\gamma}+\frac{1}{1-\alpha}},$ $$

which gives $\gamma=1,\beta=\frac{1}{(1-\alpha)(4-2\alpha)}$.

Therefore, when $t\rightarrow\infty$, we have

$$ $\displaystyle a\rightarrow\log\bigg{(}\frac{1}{(1-\alpha)(4-2\alpha)}t^{-1} \bigg{)},~{}~{}~{}~{}b\rightarrow\log\frac{\alpha}{1-\alpha}.$ $$

∎

###### Lemma F.1 .

###### Proof.

###### Lemma F.2 .

###### Proof.

###### Lemma F.3 .

###### Proof.

## Appendix G Input Examples for LLMs

### G.1 Examples for Prepositions

For experiments in Appendix B.1, we use two synthetic datasets: inputs are 30 prepositions, and inputs are 40 incomplete sentences ending with a preposition.

The 30 prepositions are:

"about", "above", "across", "after", "against", "along", "around", "at",
"before", "behind", "below", "beneath", "beside", "between", "by",
"during", "for", "from", "in", "inside", "into", "near", "of", "on",
"over", "through", "to", "under", "with", "without".

Generated by Claude 3 [Anthropic, ], the 40 incomplete sentences are:

[
"Inspired painter gazed at pristine canvas, envisioning next creation about",
"Children’s delighted squeals filled yard as they frolicked, stumbling across",
"Singer inhaled deeply, calming nerves before gracing stage before",
"Ominous storm clouds amassed, promising downpour that would soon roll in",
"Awestruck trekker admired breathtaking summit vista, looking over",
"Rich aroma of freshly roasted beans permeated cozy cafe, enticing during",
"With deft sleight of hand, illusionist made coin vanish, leaving spectators in awe without",
"Majestic oak stood tall, branches reaching skyward above",
"Gentle waves caressed shoreline, soothing rhythm lulling along",
"Meticulous investigator scoured crime scene, searching for any evidence left behind",
"Radiant sunbeams filtered through sheer curtains, warming hardwood floor beneath",
"Concert pianist’s nimble fingers glided across ivory keys, room resonating with melody around",
"Crickets’ evening chorus filled silent field from nearby meadow during",
"Jubilant laughter resounded down corridor as jovial group headed towards celebration without",
"Struggling poet tapped pen restlessly, seeking words to capture elusive emotion beneath",
"Soothing patter of raindrops danced on windowpane, inviting serene relaxation with",
"Mouthwatering scent of fresh bread beckoned passersby into cozy bakery without",
"Mighty waves thundered against jagged cliffs, echoing roar along rugged shoreline around",
"Seasoned trekker carefully navigated winding trail, cautiously avoiding exposed roots and rocks beneath",
"Graceful ballerina flowed across stage, movements blending seamlessly with melody during",
"Crackling campfire cast dancing shadows across gathered faces around",
"Vibrant brush strokes danced across canvas, bold hues bursting into life before",
"Photographer framed breathtaking sunset, capturing fleeting beauty over glistening ocean without",
"Stern librarian hushed raucous group, reminding them to stay quiet inside",
"Ink flowed from author’s pen, words brimming with raw passion as page filled during",
"Earthy aroma of freshly steeped tea perfumed air, inviting moment of serenity along",
"Masterful guitarist’s fingers danced nimbly across strings, room alive with haunting melody around",
"Meticulous chef artfully garnished plate, adding delicate finishing touches over",
"Indomitable marathoner pushed through punishing final stretch, fortitude driving every stride before",
"Engrossed scientist examined specimen’s intricate structures through microscope beneath",
"Nervous thespian steadied breathing, striding into dazzling spotlight, delivering flawless performance with",
"Skilled artist’s pencil glided gracefully, deftly capturing subject’s essence without",
"Weary hiker paused to catch breath, marveling at sweeping panorama from lofty peak above",
"Deep in thought, writer drummed fingers, seeking perfect phrasing to convey profound emotion without",
"Lost in reverie, violinist swayed gently, fingers dancing across delicate strings during",
"Painter’s brushstrokes burst into radiant life, canvas ablaze with vivid sunset hues over",
"Adept photographer framed picturesque scene, preserving landscape’s beauty without",
"World-renowned chef meticulously garnished plate, each component strategically placed around",
"Dedicated researcher scrutinized specimen under microscope, documenting minute details beneath",
"Seasoned actor inhaled deeply, embodying character as bright lights engulfed stage with",
].

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: * I would suggest moving the proofs of Lemmas C.1 and C.2 up, before they are invoked in the proof of Theorem 1 / 4. This would make it easier to follow the progression of claims made. 
* In the abstract, the term “gradient noise” is used in line 21, but not referred to in the rest of the paper — please clarify (or remove, if not relevant) this term. 
* The use of the term 'sequel' on line 125 is unclear; this should be clarified or rephrased accordingly. 
* As noted in the weaknesses section, are there any architectural changes that may be proposed, conditioned on the results of the exploration? Is the key takeaway that one should perform low-rank truncation as in LASER on the feed-forward layers since they store noise? Including discussion on this would improve the actionability of the findings of this work and its impact.
soundness: 3
strengths: * The paper sets out to address a clear objective, and designs a 2-layer model that enables a theoretical analysis of the phenomenon of interest, then reduces to a 1-layer model for theoretical analysis, revealing important, mathematically-grounded supporting evidence for said phenomenon. 
* I appreciate that the tasks of interest, and the distinguishment between the two abilities of interest (distributional associations and in-context reasoning) is made clear upfront. 
* Figure 2 is especially helpful to follow the 2-layer setup for the in-context recall task with noise tokens — the finding that the induction head filters out the noisy token is an interesting insight, and reinforced by Theorem 2 and its proof in Appendix D. 
* While the paper is dense, it is also organized quite well and in a manner that emphasizes the key takeaways.
summary: This work theoretically examines a phenomenon regarding the purpose of the feedforward and attention layers of a Transformer model, relative to the reasoning and knowledge learning tasks. It attempts to provide an explanation for the difference through a generic noise token, demonstrating that the feedforward layers store noise while the induction heads in the attention layers are insensitive to the noise tokens.
weaknesses: * The novelty presented the results section (Section 4) is unclear — are there any architectural insights informed by the theoretical analyses that could be applied that have been considered, beyond an application of LASER? Aside from this, this section appears to be a more in-depth, investigative analysis on the behavior of the chosen pre-trained models for the selected tasks; it is also unclear how / whether these insights scale as the models considered are small. 
* I have no major concerns about the theoretical results in the work, which reveal the means through which noise is stored both in the presence and absence of feedforward layers, under a simple 1-layer model. I must acknowledge that while I have largely studied and followed the proofs of Theorems 1 and 2 and the theory in Appendix A, I have not thoroughly verified the correctness of the theory in the rest of the appendix.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. In addition to feed-forward layers and attention layers, there are also wording embedding layers, residual connections, and LayNorm layers in the Transformer architecture. What special role do these layers play in helping the model understand the semantics of the natural language? Did the authors explore?

2. Because the attention module is the interaction between all tokens in the input sequence, it intuitively plays a role in understanding context. It seems obvious that the attention layer mainly focuses on in-context reasoning, right?

3. Are the given conclusions applicable to all models based on the Transformer architecture? Is this applicable to all LLM models?

4. In terms of format:
    Line 1935, Line 1938, Line 1952, Line 2001, Line 2069 exceed the margins and have incorrect formatting.
soundness: 3
strengths: 1. This paper uses a large number of experiments and theoretical proof to draw the following conclusions:
For the Transformer architecture of large language model LLM, feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning.

2. The experimental results and conclusions are relatively solid, and the theoretical proof is quite convincing.

3. This work may be helpful for future researchers studying LLM.
summary: Large language models excel in in-context reasoning tasks like coherent language generation and knowledge storage. Their Transformer architecture relies on feed-forward and attention layers linked to knowledge and reasoning. In controlled experiments, the authors found feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning. Gradient noise explains this difference. Pre-trained models, like Pythia, show similar disparities in simple reasoning tasks.
weaknesses: 1. This work is highly theoretical and lacks certain practical significance and applicability.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: 1) Why did you use such a small model for the experiments? Would it be possible to perform the same kind of analysis with more layers?
2) What is your interpretation of the poorer performance of truncated MLPs after 8-shot on GSM8K (Table 2)?
3) Your results seem to indicate that reducing FFNs greatly improves in-context reasoning, which is often desirable. However, FFNs also play an important role. It would therefore be interesting to study what is lost when they are reduced/removed?
soundness: 4
strengths: - It is very well written and clear.
- Their experimental environment is very clearly established, from the creation of their dataset to the architecture of the model.
- I think that experimenting with LLMs in a controlled setting is an excellent methodology.
- The paper proposes a theoretical justification for their discovery that is more than welcome and very rare in the field.
- The findings are very interesting and help to better understand the behavior of LLMs, in particular FFNs, which are often less studied than self-attention.
summary: In this paper, the authors investigate the role of transformers' self-attention and feed-forward networks (FFNs) both experimentally, in a controlled environment with a simplified model, and theoretically, using an even simpler model that allows them to derive calculations. The authors thus highlight an important distinction: FFNs are biased towards learning distributional associations (i.e. co-occurrence cues), while self-attention focuses on in-context reasoning. 

Concretely, as it is common to have the word “the” after the word “in”, a model using distributional cues only would be tempted to complete the sentence “Madrid is located in” with “the”, whereas a model using in-context reasoning would complete it with “Spain”.

So the authors set up an experiment in which the model is trained on a text containing with equal probability the “a + b” and “a + noise” patterns (where “a” varies from sentence to sentence, while the “noise” token remains the same) and then evaluated on a text containing only the “a + b” pattern. The idea is then to measure in-context recall in such a scenario: indeed, a purely distributional model will perform at 50% accuracy where a perfectly in-context reasoning model will achieve 100%. Using LASER to reduce the rank of the second-layer MLP matrix, until it is completely suppressed (rho = 0), they note that distributional behavior disappears with MLP rank, suggesting that FFNs are biased in this direction. They also note that self-attention manages to focus only on “a + b” patterns, ignoring “a + noise” patterns. 

The authors then propose a theoretical explanation for such behavior, analyzing in particular the first gradient descent step. Finally, the authors propose to test their findings in a more realistic setting, i.e. on real models and real-world reasoning benchmarks.
weaknesses: - The main experiments and theoretical results are based on highly simplified models. However, the observed behavior may not scale with larger models.
- The theorems and their consequences are relatively complicated to understand for a non-initiated reader of this literature, perhaps try to explain the different terms a little more.
- Experimental figures have no error bars.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Is there potential modifications or improvements for the authors to propose to LASER that could strengthen this paper’s contributions?
2. In Section 4.2, why LASER struggles with more shots? 
3. Is it possible to make other modifications beyond those in [2], such as designing an entirely new dataset to better validate your claims?
soundness: 3
strengths: 1. The authors conducts an in-depth study on how attention and feed-forward layers differently process distributional associations and contextual inputs, yielding clear and significant conclusions.
2. The flow of the paper is easy to follow. 
3. The authors employ clear theoretical proofs and a logical structure, supporting the authors' arguments.
summary: The authors studied how attention mechanisms and feed-forward layers handle distributional associations and in-context inputs differently. They considered simple two-layer transformers on an in-context recall task and observed the distinct focuses of feed-forward layers and attention mechanisms through analyzing gradient.
They further employed LLMs like GPT-2 and Pythia, concentrating on IOI dataset . Additionally, the authors validate the effectiveness of the LASER technique.
Overall, the authors underscore the differing roles of attention layers and feed-forward layers in language models.
weaknesses: 1. The experiments in this paper are insufficient to demonstrate authors’ claim, they only evaluate with limited number of datasets and model structures, which might not be representative of the general behavior.

2. The technique authors recommended, LASER[1], is not introduced by this paper, making this paper seems an extension or complement to LASER. Differentiation between this paper’s contributions and [1]’s should be claimed.

3. The methods and datasets this work employed is interesting but mostly appear a simple extension of that in [2]. This makes the research seem a supplementary work.

**Reference**

[1] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction

[2] Birth of a Transformer: A Memory Viewpoint

Meta-review: The papers provide an analysis about feed-forward layers and attention layers. The main observation is that FFN is generally responsible for distributional association, whereas attention layers are responsible for in-context reasoning. The authors back up their claim by both empirical and theoretical analysis, although the setups are a little bit toy.

Reviewers generally recognize contributions of this paper. Concerns are raised regarding the simplified setups, but it is understandable that theoretical analysis on full-sized models could be extremely difficult. On the other hand, the implication of the finding could be further discussed.

Decision: Accept (Poster)