SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 10: Average self-attention map entropy over the validation sets in 1-layer transformer after training, when the learning rate  $\eta_{Y}$  and  $\eta_{Z}$  changes. Note that higher learning rate  $\eta$  leads to higher  $B_{n}(t)$  and thus low entropy (i.e., more sparsity), which is consistent with our theoretical finding (Sec. 6). All the experiments are repeated in 5 random seeds. Error bar with 1-standard is shown in the figure.
equal to 1 to the total number of distinct tokens. In the figures, the horizontal coordinates denote the threshold values around  $1 / M$  for different last/next tokens setting, and the vertical coordinates denote the recall rate. The dataset is Syn-Medium mentioned in Section 8, and every data point in the figures is the mean value over 10 seeds. It's obvious that a sparser attention map will result in a slower descent rate of the recall-threshold line since sparser attention corresponds to fewer distinct tokens with higher attention weights, and the results of Fig. 11 and Fig. 12 match that of Fig. 6.
H Technical Lemma
Lemma 18. Let  $\pmb{h} = [h_1, h_2, \dots, h_M]^\top \in \mathbb{R}^M$  is some  $M$ -dimensional vector,  $\pmb{h}_X := [h_{x_1}, \dots, h_{x_{T-1}}]^\top \in \mathbb{R}^{T-1}$  is a vector selected by input sequence  $X$ , then given event  $x_T = m, x_{T+1} = n$ , there exists some  $\pmb{q}_{m,n} = [q_{1|m,n}, q_{2|m,n}, \dots, q_{M|m,n}]^\top \in \mathbb{R}^M$  so that  $\pmb{q} \geq 0$  and
$$
\frac {1}{T - 1} X ^ {\top} \boldsymbol {h} _ {X} = \sum_ {l = 1} ^ {M} q _ {l | m, n} h _ {l} \boldsymbol {e} _ {l} = \boldsymbol {q} _ {m, n} \circ \boldsymbol {h} \tag {134}
$$
$$
\frac {1}{T - 1} X ^ {\top} \operatorname {d i a g} \left(\boldsymbol {h} _ {X}\right) X = \sum_ {l = 1} ^ {M} q _ {l | m, n} h _ {l} \boldsymbol {e} _ {l} \boldsymbol {e} _ {l} ^ {\top} = \operatorname {d i a g} \left(\boldsymbol {q} _ {m, n} \circ \boldsymbol {h}\right) \tag {135}
$$
where  $q_{l|m,n}$  satisfies  $\sum_{l=1}^{M} q_{l|m,n} = 1$ . And with probability at least  $1 - \delta$  we have
$$
\max  \left(0, \mathbb {P} (l | m, n) - \sqrt {\frac {\ln (2 / \delta)}{2 (T - 1)}}\right) \leq q _ {l | m, n} \leq \mathbb {P} (l | m, n) + \sqrt {\frac {\ln (2 / \delta)}{2 (T - 1)}} \tag {136}
$$
And thus  $q_{l|m,n} \to \mathbb{P}(l|m,n)$  when  $T \to +\infty$ .
Proof. Given that  $x_{T} = m$  and  $x_{T + 1} = n$ , then we have
$$
\frac {1}{T - 1} X ^ {\top} \boldsymbol {h} _ {X} = \frac {1}{T - 1} \sum_ {t = 1} ^ {T - 1} h _ {x _ {t}} \boldsymbol {x} _ {t} = \sum_ {l = 1} ^ {M} \left(\frac {1}{T - 1} \sum_ {t = 1} ^ {T - 1} \mathbb {I} [ x _ {t} = l ]\right) h _ {l} \boldsymbol {e} _ {l} =: \sum_ {l = 1} ^ {M} q _ {l | m, n} h _ {l} \boldsymbol {e} _ {l} \tag {137}
$$
And similar equations hold for  $\frac{1}{T - 1} X^{\top}\mathrm{diag}(\pmb{h}_X)X$ . Then we consider the case that the previous tokens are generated by conditional probability  $\mathbb{P}(l|m,n)$  as the data generation part, so  $\mathbb{I}[x_t = l], \forall t \in [T - 1]$  are i.i.d. Bernoulli random variables with probability  $\mathbb{P}(l|m,n)$  and  $Tq_{l|m,n}$  satisfies binomial distribution. By Hoeffding inequality, we get
$$
\mathbb {P} \left(\left| q _ {l \mid m, n} - \mathbb {P} (l \mid m, n) \right| \geq t\right) \leq 2 \exp (- 2 (T - 1) t ^ {2}) \tag {138}
$$
Then we get the results by direct calculation.
Figure 11: Recall value of attention on all distinct tokens versus threshold with changing learning rate ratio  $\eta_{Y} / \eta_{Z}$ . Smaller  $\eta_{Y} / \eta_{Z}$  corresponds to a smaller descent rate and thus sparser attention.
Figure 12: Recall value of attention on all distinct tokens versus threshold with changing learning rate  $\eta_Z$ . Larger  $\eta_Z$  corresponds to a smaller descent rate and thus sparser attention.