SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Long-tailed data. For direct training on long-tailed sequences, we train on an equal mixture of FFL(0.8), FFL(0.98), and FFL(0.1), ensuring that the sparse and dense tails are both adequately represented in training.
Scaling up data and compute. For both the default and long-tailed training distributions, we consider increasing the number of training iterations by a factor of 3 or 10, either with freshly generated samples (" $N \times$  data") or looping over the same training dataset for multiple epochs (" $N \times$  steps").
Scaling up the model. We also perform these exhaustive tail error evaluations on all 64 of the architecture scale settings shown in Figure 6, as well as a limited number of larger architectures (shown in Figure 9).
(R4) Training on rare sequences works best, by a wide margin. See the teal violins in Figure 9 (left); training for longer (either with fresh data, or for more epochs on the same data) reduces the fraction of unlucky error-prone outliers. Recall that extrapolation is possible without such favorable coverage, via using a recurrent model.
(R5) Resource scaling (in-distribution data, training steps, network size) helps. Training on more data from the same distribution, as well as for more steps on the same examples, both improve sparse-sequence performance, at the expense of dense-sequence performance (blue violins in Figure 9 (left). As best seen through Figure 9 (right), there is no discernible monotonic relationship between any of the Transformer's standard architecture size parameters (i.e. number of layers, embedding dimension, and number of parallel self-attention heads per layer) and extrapolative performance (navy violins).
B.4 Indirect algorithmic controls for extrapolation (Section 5.2)
As shown in Figure 4 in the main paper, various architectural, algorithmic and regularization choices can help improve error rates compared to the baseline Transformer. The various settings of weight decay, {attention, MLP, embedding} dropout, position embedding, activation function, attention-sharpening penalty are all found in Figures 10 and 11.
Details for architecture variants. There is no clear consensus on the advantages and drawbacks of various positional encodings, but it has been known Dai et al. (2019) that the choice of positional symmetry-breaking scheme modulates long-sequence performance on natural tasks. We evaluate
Regularizer effects on o.o.d. performance
Figure 10: Detailed comparisons of standard regularizers (weight decay, and 3 forms of dropout). While some regularizer choices reduce rare-sequence error rates (in particular, large embedding dropout reduces sparse-sequence errors by 2 orders of magnitude), nothing eliminates the glitches entirely.
Architecture & attention-sharpening effects on o.o.d. performance
Figure 11: Detailed comparisons of a selection of architectural changes, attention-sharpening losses, and combinations of indirect algorithmic controls. Our best runs come from jointly tuning these interventions, including an annealing schedule for the attention-sharpening loss; however, even the best models have nonzero glitch rates. Figure 13 provides an exhaustive view of these results.
Figure 12: Supplementary plots. Left: Additional comparisons: modifying the softmax temperature  $\beta$  inside the attention heads (multiplying the attention alignment scores by  $1 / \beta$ ), and generalizations of the FFLM task to larger vocabularies (i.e. the "data" tokens are uniformly drawn from  $\{0,1,\dots,M-1\}$ ). In the rightmost violin, " $2*2"$  refers to a token space consisting of length-2 bitstrings. In all of these settings, attention glitches remain (and mostly worsen). Right: Distribution of dependency lengths (number ignore operations since last read or write operation) of wrong answers on the union of the o.o.d. test sets, for a canonical baseline model. This serves as a quick check that the errors are diverse.
various choices which appear in high-profile LLMs: sinusoidal, learned, ALiBi (Press et al., 2021), and RoPE (Su et al., 2021). We also try the zero positional encoding (which breaks symmetry via the causal attention mask; see Haviv et al. (2022). We find that non-trainable position encodings help on dense sequences (FFL(0.1)), but have no clear benefit on sparse ones (FFL(0.98)) which require more handling of long-term dependency. We also try the gated activation units considered by (Shazeer, 2020).
Details for attention sharpening. There are many possible choices of continuous regularization terms which can promote sparsity in an attention head's weights—we consider entropy, negative  $L_{2}$  loss, and negative  $L_{\infty}$  loss. These terms are averaged across every attention head in the Transformer, and added as a surrogate objective during training. We perform a large grid sweep over coefficients  $\{0.01, 0.03, \dots, 0.1, 0.3, 1, 10, 30\}$ , annealing schedules (linear and oscillating, starting from 0, 2000, and 5000 steps), and display in Figure 11 the 3 choices which appear on the Pareto front.
Optimizer hyperparameters. We also varied the optimizer parameters for AdamW ( $\beta_{1} \in \{0.85, 0.9, 0.95\}$ ,  $\beta_{2} \in \{0.95, 0.99, 0.999\}$ , learning rate  $\eta \in \{0.0001, 0.0003, 0.001, 0.003\}$ ) and found no significant improvements to extrapolation performance.
We restate the main findings:
(R6) Many algorithmic choices influence extrapolative behavior. We sweep over various forms of implicit and explicit regularizers; see Figures 10 and 11. Details are provided below.
(R7) Despite many partial mitigations, nothing eliminates attention glitches entirely. Refer to the scatter plot in Figure 4 (left) for a visualization of every training run.
Figure 12 provides a few supplementary experiments, verifying that attention glitches persist: scaling the attention softmax temperature, and larger vocabularies for the memory register. We also provide a quick verification that errors are not concentrated at the same dependency lengths (i.e. distances between write and read).
Finally, to provide a bird's-eye view of all of our experiments, Figure 13 provides a scatter plot of o.o.d. error rates for all models trained in this paper, color coded by intervention category.
B.5 Preliminary mechanistic study and challenges (Section 5.3)
We continue the discussions in Section 5.3 and provide preliminary mechanistic interpretability results on simulating the flip-flop automaton (Definition 1). Recall the main takeaway:
Figure 13: Error scatter plots for all models evaluated in this work (a larger version of Figure 4 (left) from the main paper), color-coded by intervention category. Left: All runs on the same plot. Right: Segmented by intervention type. The best runs result from an exhaustive joint hyperparameter search over dropout, position embeddings, and annealed attention-sharpening regularizers.
(R8) Attention-sharpening regularizers successfully promote hard attention, but errors persist.
Sparsity regularization helps sharpen the attention. Figures 16a and 16b compare the attention patterns of 1-layer 1-head models with or without attention-sharpening regularization. While both types of models give correct results, the attention-sharpened model puts all attention weights to the most recently write position, which is the solution given according to the definition of the task, whereas the attention patterns of the non-regularized model (Figure 16a) are much less clean.
Are there solutions other than the "ideal" solution? There is a solution naturally associated with the definition of the flip-flop automaton (i.e. the sparse pattern shown in Figure 16b), but it is not necessarily the only solution. For example, an equally valid (dense) solution is for the model to attend to every write token of the correct type.
Are attention patterns reliable for interpretability? Prior work has pointed out the limitations of interpretations based solely on attention patterns (Jain and Wallace, 2019; Bolukbasi et al., 2021; Wen et al., 2023). The intuition is that attention patterns can interact with other components of the network in various ways; for example,  $W_V$  can project out certain dimensions even though they may have contributed to a large attention score. Hence, for multi-layer multi-head non-sparse models, the magnitude of attention weights may not have an intuitive interpretation of "importance" (Meister et al., 2021). This is observed in our experiments as well; for instance, Figure 17 shows examples where the attention on an incorrect token may be higher than that of the correct token. However, in a 1-layer 1-head model, 1-sparse attention (Figure 16b) indeed offers interpretability, since if zero
However, if we consider the "importance / influence" as measured by the norms of the attention-weighted value vectors, then the max norm still corresponds to the correct token, which helps explain why the final output is correct.
Figure 14: Attention drifts as the length increases. The model is trained on length-500 sequences with  $p(\sigma \neq \bot) = 0.5$ . The testing sequences are (a)  $[2, \underbrace{0 \cdots, 0}_{800}]$ , and (b)  $[1, \underbrace{0 \cdots, 0}_{32}, \underbrace{0 \cdots, 0}_{800}]$ . We
sample every 32 positions for visualization.
attention weight  $^{18}$  necessarily means the absence of dependency, which greatly reduces the set of possible solutions implemented.
Sporadic errors persist. Section Section 5.1 (R5) showed that none of the mitigations was successful at making Transformers reach  $100\%$  accuracy. One common failure mode is long-range dependency, where the input sequences contain very few writes. The failure could be attributed to multiple factors; we will explore one aspect related to attention patterns, demonstrated with a 1-layer 1-head Transformers with linear position encoding, on a length-834 sequence with 2 writes. As shown in Figure 14, the attentions for positions early in the sequence correctly attend to the most recent write. However, attention starts to "drift" as we move to later positions, and the positions at the end of the sequence attend entirely  $^{19}$  to the recent read tokens, which contains no information for solving the task. This may be because the attention weights are undesirably affected by the position encodings, as discussed in Proposition 4.
Optimization hurdles. While sparse solutions may be preferred for various reasons, sparsity itself is not sufficient to guarantee good performance: As shown in Figure 16d, sparsity regularization can lead to bad local minima, where the model tends to (incorrectly) rely on earlier positions. This is observed across different types of sparsity regularization. While we do not yet have a full explanation of the phenomenon, a possible explanation for this bias is that earlier positions show up more often during training, due to the use of the causal attention: a valid flip-flop solution is for the model to attend to every write token of the correct type; positions earlier in the sequence get observed in more subsequences because of the causal mask, and are hence more likely to be attended to. We also observe that the phenomenon seems to be closely related to the training distribution. For example, the model is much more likely to get stuck at a bad local minima when  $p(\perp) = 0.5$  (denser sequences) compared to  $p(\perp) = 0.9$  (sparse sequences). We leave a more thorough study on the training dynamics (Jelassi et al., 2022; Li et al., 2023; Ahn et al., 2023) for future work.
Effect of sparsity regularization on training dynamics An interesting future direction is to understand the learning dynamics of flip-flop tasks with attention-sharpening regularization, as suggested by the (quantitively and qualitatively) different results and optimization challenges. As some initial empirical evidence that the regularization indeed have a large impact on the dynamics, we found that sharpened attention seems to have a regularization effect on the weight norms (Figure 15), and also lead to different behaviors of the attention heads (Figure 18).
<sup>18</sup>By “zero” we mean an attention score on the magnitude of 1e-8 empirically.
19The attention weights that are not on the most recent write sum up to around 1e-7.
non-regularized attention-sharpened attention-sharpened from epoch 30
Figure 15: Frobenius norms of weight matrices in 1-layer 1-head models, trained without regularization (blue), with attention-sharpening regularization (yellow), or first without regularization and then adding regularization from epoch 30 (red; epoch 30 marked by the dashed lines). The solid curve and the shadow shows the median and the standard deviation calculated on 8 models.
(c)
Figure 16: Attention-sharpening regularization on 1-layer 1-head models. Compared to a nonregularized model (16a), the sparsity-regularized model (16b) shows clear attention at the last write position. However, sparse attention does not have to align with the "ideal" pattern (16c), and can even be wrong (16d). Positions with yellow borders are where the max attention in each row occurs; errors are marked in red.
More examples of attention patterns Figure 19 shows the full set of attention patterns of two 6-layer 8-head models trained with and without attention-sharpening regularization, corresponding to Figure 5 (a,b). Attention-sharpening regularization can be applied in different ways; for example, Figure 20 shows results of a model for which only the first layer is regularized. The attention patterns of subsequent layers remain sharp, even though there is no explicit regularization.
Figure 17: Non-sparse attention pattern can be misleading: a non-sparse model may put more attention on an incorrect token (i.e. a token that is not the write with the right type), while making the correct predictions. Yellow boxes mark the position of the max attention of each row.
(a) Model trained without sparsity regularization.
(b) Model trained with entropy sparsity regularization with  $\lambda = 0.01$
Figure 18: Examples of the  $\ell_2$  difference in attention patterns from two 6-layer 8-head 512-dimension models. Differences are calculated between all pairs of heads in the same layer.
(a) Without regularization.
Figure 19: Attention patterns for 6-layer 8-head 512-dimension models on the input sequence  $[\sigma_1, \bot, \sigma_0, \bot, \bot, \sigma_0, \sigma_1, \bot]$ : attention-sharpening regularization lead to cleaner attention patterns. 1 attention head in the first layer of the regularized model (marked by the purple box) matches the "ideal" attention pattern Figure 5c.
(b) With attention-sharpening regularization.
(a)  $\ell_2$  differences between pairs of attention heads in the same layer, throughout training ( $x$ -axis).
Figure 20: Attention heads and attention patterns for a 6-layer 8-head 512-dimension model, trained with attention-sharpening regularization (entropy regularization with strength 0.01) on the first layer only. 1 attention head in the first layer (marked by the purple box) matches the "ideal" attention pattern Figure 5c.
(b) Attention patterns on the input sequence  $[\sigma_1, \bot, \sigma_0, \bot, \bot, \sigma_0, \sigma_1, \bot]$ .
B.6 Software, compute infrastructure, and resource costs
GPU-accelerated training and evaluation pipelines were implemented in PyTorch Paszke et al. (2017). For the FFLM experiments, we used the x-transformers $^{20}$  implementations of the Transformer architecture and variants. For the fine-grained mechanistic interpretability experiments on the pure flip-flops, we used the "vanilla, GPT-2"-like Transformer implementation published by HuggingFace (Wolf et al., 2019). Our benchmarks have been released at https://huggingface.co/datasets/synthseq/flipflop.
Each training run was performed on one GPU in an internal cluster, with NVIDIA P40, P100, V100, and RTX A6000 GPUs, with at least 16GB of VRAM. Each (6-layer, 512-dimensional, 8-head) baseline model took  $\sim 10$  minutes to train (and evaluate online) for  $10^{4}$  steps. A nontrivial fraction of the compute time ( $\sim 20\%$ ) was spent on fine-grained evaluation through the course of training. The vast majority of training runs are close to these specifications; consequently, one set of replicates under identical conditions (i.e. each violin plot in each figure) is the product of  $\sim 4$  GPU-hours of training time.
We hope that this computational investment will aid in understanding how to build robust Transformer models and training pipelines at much larger scales.
C Proofs for Section 4.1
Transformer recap. A Transformer (Vaswani et al., 2017) consists of multiple self-attention layers. Given  $d$ -dimensional embeddings of a length- $T$  sequence, denoted as  $\mathbf{X} \in \mathbb{R}^{T \times d}$ , a self-attention layer  $f$  computes
$$
f (\boldsymbol {X}) = \phi \left(\boldsymbol {W} _ {V} \operatorname {s o f t m a x} \left(\boldsymbol {X} \boldsymbol {W} _ {Q} \boldsymbol {W} _ {K} ^ {\top} \boldsymbol {X} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {V} \boldsymbol {W} _ {C}\right). \tag {C.1}
$$
where  $W_{Q}, W_{K} \in \mathbb{R}^{d \times k}$  for  $k \leq d$  are the query and key matrix;  $W_{V}, W_{C}^{\top} \in \mathbb{R}^{d \times k}$  project the representations from and back to  $\mathbb{R}^d$ . softmax calculates row-wise softmax.  $\phi: \mathbb{R}^d \to \mathbb{R}^d$  is a 2-layer fully-connected network. Residual links and layer norm can be optionally included at different places of a self-attention layer.
C.1 Realizability of FFL by small Transformers
Proposition 2. A 2-layer 1-head Transformer with residual connections can represent "deterministic" FFL.
$^{20}$ https://github.com/lucidrains/x-transformers
Proof. Let us consider predicting in the deterministic mode (Section 3.1). Then we need to predict  $x_{t+1}$  given  $x_{1:t}$  with  $x_t = \mathbf{r}$ . In order to do this, we need to find the largest  $\tau < t$  such that  $x_{\tau} = \mathbf{w}$  and output  $x_{\tau+1}$ . There are multiple ways to implement this, we will consider the following: (1) layer 1 converts FFL to the flip-flop automaton (Definition 1), (2) layer 2 implements the flip-flop construction. For layer 2, we can use the construction described in Liu et al. (2023). Here we present the full construction for completeness.
We will consider a two-layer Transformer with one head in each layer followed by a 2-layer MLP and a residual connection. In particular, for  $x \in \{\mathsf{w}, \mathsf{r}, \mathsf{i}, 0, 1\}^T$ :
$$
f (x) = \phi_ {2} (\boldsymbol {W} _ {V} ^ {(2)} \operatorname {s o f t m a x} (f _ {1} (x) \boldsymbol {W} _ {Q} ^ {(2)} \boldsymbol {W} _ {K} ^ {(2) ^ {\top}} f _ {1} (x) ^ {\top}) f _ {1} (x) \boldsymbol {W} _ {V} ^ {(2)} \boldsymbol {W} _ {C} ^ {(2)})
$$
$$
\text {w h e r e} f _ {1} (x) = E (x) + \phi_ {1} (\boldsymbol {W} _ {V} ^ {(1)} \operatorname {s o f t m a x} (E (x) \boldsymbol {W} _ {Q} ^ {(1)} \boldsymbol {W} _ {K} ^ {(1) ^ {\top}} E (x) ^ {\top}) E (x) \boldsymbol {W} _ {V} ^ {(1)} \boldsymbol {W} _ {C} ^ {(1)})
$$
where  $E(x) \in \mathbb{R}^{T \times d}$  is the encoding for the input sequence  $x$  given some encoding function  $E$ .
Our construction is as follows:
- Select  $d = 7, k = 2, H = 1$  (recall from Equation C.1 that  $d, k$  are the dimensions of  $W_{Q}, W_{K}$ ). Among the  $d = 7$  embedding dimension, two dimensions are for the operations (w versus r, i), two for the two write values, one for the positional embedding, one for padding, and the final dimension is for storing whether the previous position is the most recent write, as calculated by the first layer.
- Select input symbol encodings such that for the token at position  $t$ , denoted as  $x_{t}$
$E(x_{t})\coloneqq \mathbb{1}[x_{t} = \mathbf{w}]e_{1} + \mathbb{1}[x_{t} = \mathbf{r}\lor x_{t} = \mathbf{i}]e_{2} + \mathbb{1}[x_{t} = 0]e_{3} + \mathbb{1}[x_{t} = 1]e_{4} + e_{5} + P_{t}\in \mathbb{R}^{7},$  where  $P_{t}$  is the positional encoding. We use the linear positional encoding  $P_{t}\coloneqq (t / C)\cdot e_{6}$ , for some (large) constant  $C$ . For a fixed sequence length  $T$ , we can set  $C = T$ .
-  $\mathbf{W}_Q^{(1)} := \left[ \begin{array}{cc} e_5 & e_5 \end{array} \right] \in \mathbb{R}^{7 \times 2}, \mathbf{W}_K^{(1)} := \left[ \begin{array}{cc} 3c\frac{e_1}{2T} & ce_6 \end{array} \right] \in \mathbb{R}^{7 \times 2}$  for  $c = O(T\log (T))$ ,  $\mathbf{W}_V^{(1)} := \left[ \begin{array}{cc} e_1 & 0 \end{array} \right] \in \mathbb{R}^{7 \times 2}$ , and  $\mathbf{W}_C^{(1)}^\top := \left[ \begin{array}{cc} e_7 & 0 \end{array} \right] \in \mathbb{R}^{7 \times 2}$ .
-  $\mathbf{W}_Q^{(2)} := \left[ \begin{array}{cc} e_5 & e_5 \end{array} \right] \in \mathbb{R}^{7 \times 2}$ ,  $\mathbf{W}_K^{(2)} := \left[ \begin{array}{cc} ce_7 & ce_6 \end{array} \right] \in \mathbb{R}^{7 \times 2}$  for  $c = O(T\log (T))$ ,  $\mathbf{W}_V^{(2)} := \left[ \begin{array}{cc} e_4 & 0 \end{array} \right] \in \mathbb{R}^{7 \times 2}$ , and  $\mathbf{W}_C^{(2)} := \left[ \begin{array}{cc} e_1 & 0 \end{array} \right] \in \mathbb{R}^{7 \times 2}$ .
In layer 1, the unnormalized attention score for query position  $i$  to key position  $j$  is
$$
\left\langle \boldsymbol {W} _ {Q} ^ {(1) ^ {\top}} x _ {i}, \boldsymbol {W} _ {K} ^ {(1) ^ {\top}} x _ {j} \right\rangle = \left\langle \frac {c}{T} \cdot \left[ \frac {3}{2} \cdot \mathbb {1} [ x _ {j} = \mathbf {w} ], j \right], [ 1, 1 ] \right\rangle = \frac {c}{T} \cdot \left(\frac {3}{2} \mathbb {1} [ x _ {j} = \mathbf {w} ] + j\right).
$$
Note that the max attention value for position  $i$  is achieved at  $i$  if  $x_{i-1} \neq \mathbf{w}$ , else the max is achieved at position  $i - 1$ .
In the setting of hard attention, the output for the  $i_{th}$  token after the attention module is  $\mathbb{1}[x_{i-1} = \mathbf{w} \vee x_i = \mathbf{w}] e_7$ . Now similar to the constructions in Liu et al. (2023) (Lemma 6), with a appropriate choice of  $c = O(T \log T)$ , we can approximate hard attention by soft attention, and subsequently use the MLP to round the coordinate corresponding to  $e_7$ . The MLP otherwise serves as the identity function. Together with the residual link, the first layer output (i.e. the second layer input) at position  $i$  takes the form
$$
f _ {1} (x _ {i}) = E (x _ {i}) + \mathbb {1} [ x _ {i - 1} = \mathbb {w} \vee x _ {i} = \mathbb {w} ] e _ {7}.
$$
In layer 2, the unnormalized attention score computed for position  $i$  attending to  $j$  is
$$
\begin{array}{l} \left\langle \boldsymbol {W} _ {Q} ^ {(2) ^ {\top}} f _ {1} (x _ {i}), \boldsymbol {W} _ {K} ^ {(2) ^ {\top}} f _ {1} (x _ {j}) \right\rangle = \frac {c}{T} \left\langle [ 1, 1 ], \left[ \mathbb {1} [ x _ {j - 1} = \mathbf {w} \vee x _ {j} = \mathbf {w} ], \frac {j}{T} \right] \right\rangle \\ = c \cdot \left(\mathbb {1} \left[ x _ {j - 1} = \mathrm {w} \vee x _ {j} = \mathrm {w} \right] + \frac {j}{T}\right). \\ \end{array}
$$
Note that the max attention value is achieved at the position right after the closest w to  $x_{i}$ . Let us denote this position by  $\tau \leq i$ , then with hard attention, the output at the  $i_{th}$  position is  $x_{\tau}e_{1}$ , as desired. Now similar to before, we can approximate this with soft attention and use the MLP to do the appropriate rounding to get our final construction.
Remark: The construction in Proposition 2 is a construction, but it is not the only construction. For example, for the second layer implementation for the flip-flop automaton, there could be an equally valid dense solution, where the model uniformly attends to all write tokens of the correct type.
C.2 Failure of soft attention: attention dilution with bounded Lipschitzness
Consider any attention layer with weight matrices  $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{k \times d}$ . If  $\| \mathbf{W}_K^\top \mathbf{W}_Q \|_2$  is bounded, then the attention cannot be sparse as the sequence length increases: