SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proposition 3 (Leaky soft attention). Assume the latent variables have bounded norm, i.e.  $\| \pmb{v}\| _2\leq 1$  for any latent vector  $\pmb {v}\in \mathbb{R}^d$ , and let  $\sigma_{\mathrm{max}}$  denote the max singular value of  $W_K^\top W_Q$ . Then for  $T = \Omega (\exp (2\sigma_{\mathrm{max}}))$ , any sequences of latent vectors  $\{\pmb{v}_{\tau}\}_{\tau \in [T]}$ ,  $\| \mathrm{softmax}(\{\pmb{v}_{\tau}\}_{\tau \in [T]})\|_{\infty} = 1 - \Omega (1)$
Proof. The proof follows directly from a simple rewriting.
For any  $\pmb{u},\pmb{v}$  with  $\| \pmb {u}\| _2,\| \pmb {v}\| _2\leq 1$  , the pre-softmax attention score is bounded by  $\pmb{u}^{\top}\pmb{W}_{K}^{\top}\pmb{W}_{Q}\pmb {v}\in$ $[- \sigma_{\mathrm{max}},\sigma_{\mathrm{max}}]$
$$
\frac {\exp (\boldsymbol {v} _ {t} ^ {\top} \boldsymbol {W} _ {K} ^ {\top} \boldsymbol {W} _ {Q} \boldsymbol {v} _ {T})}{\sum_ {\tau \in [ T ]} \exp (\boldsymbol {v} _ {\tau} ^ {\top} \boldsymbol {W} _ {K} ^ {\top} \boldsymbol {W} _ {Q} \boldsymbol {v} _ {T})} \leq \frac {\exp (\sigma_ {\max})}{\exp (\sigma_ {\max}) + (T - 1) \exp (- \sigma_ {\max})} = 1 - \frac {T - 1}{T - 1 + \exp (2 \sigma_ {\max})},
$$
where the last term is  $\Omega(1)$  when  $T = \Omega(\exp(2\sigma))$ .
Attention dilution and failure on dense sequences Strictly speaking, attention dilution caused by an increased sequence length does not necessarily affect the output of the layer. For example, if ignore gets mapped to a subspace orthogonal to that of write, then  $W_V$  can project out the ignore subspace, making the weighted averaged depending only on the number of writes. Hence with the presence of layer norm, attention dilution won't be a problem for the final prediction if the number of write is upper bounded regardless of the sequence length.
For the experiments in Section 5.1, denser sequences (i.e. larger  $p(\text{write})$ ) do increase the number of write compared to the training distribution, hence attention dilution can be a potential cause for the decrease in performance.
C.3 Failure of hard attention: bad margin for positional embeddings
In this section, we look at a failure mode that a 1-layer 1-head Transformer has on the flip-flop automaton simulation task. Why do we care about this setup? Simulating the automaton is in fact a sub-task of FFLM. For example, the second layer of the construction in Proposition 2 reduces to the simulation task.
Consider a 1-layer 1-head Transformer with parameters  $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{k \times d}$ . Write the attention query matrix  $\mathbf{W}_Q$  as  $\mathbf{W}_Q = [\mathbf{W}_{Qe}, \mathbf{W}_{Qp}]$ , where  $\mathbf{W}_{Qe} \in \mathbb{R}^{k \times (d - 1)}$  corresponds to the embedding dimensions, and  $\mathbf{W}_{Qp} \mathbb{R}^k$  corresponds to the dimension for the linear positional encoding. Write  $\mathbf{W}_K = [\mathbf{W}_{Ke}, \mathbf{W}_{Kp}]$  similarly.
Then, we claim that the following must be true, regardless of the choice of the token embedding:
Proposition 4. Consider linear positional encoding, i.e.  $p_i = i / C$  for some (large) constant  $C$ . Then, perfect length generalization to arbitrary length requires  $\mathbf{W}_{Qp}^{\top}\mathbf{W}_{Kp} = 0$ .
Proof. Let  $e^{(i)} \in \mathbb{R}^{d-1}$  denote the embedding vector (without the position encoding) for token  $i \in \{0,1,2\}$ . Let  $\pmb{v}_t = [e_t,p_t]^\top \in \mathbb{R}^d$  denote the embedding for the  $t_{th}$  token, where  $e_t \in \{e^{(0)},e^{(1)},e^{(2)}\} \mathbb{R}^d$  is the embedding of the token itself, and  $p_t := i / C$  is the linear positional encoding.
Let  $s_{i\rightarrow j}$  denote the pre-softmax attention score that the  $i_{th}$  token puts on the  $j_{th}$  token, which is given by
$$
\begin{array}{l} s _ {i \rightarrow j} = \left\langle \boldsymbol {W} _ {Q} \boldsymbol {v} _ {i}, \boldsymbol {W} _ {K} \boldsymbol {v} _ {j} \right\rangle (C.2) \\ = e _ {i} ^ {\top} W _ {Q e} W _ {K e} e _ {j} + e _ {i} ^ {\top} W _ {Q e} ^ {\top} W _ {K p} \cdot p _ {j} + \left(e _ {j}\right) ^ {\top} W _ {K e} W _ {Q p} \cdot p _ {i} + W _ {Q p} ^ {\top} W _ {K p} \cdot p _ {i} p _ {j} (C.3) \\ = e _ {i} ^ {\top} \boldsymbol {W} _ {Q e} \boldsymbol {W} _ {K e} e _ {j} + \frac {\boldsymbol {e} _ {i} ^ {\top} \boldsymbol {W} _ {Q e} ^ {\top} \boldsymbol {W} _ {K p}}{C} \cdot j + \frac {\left(\boldsymbol {e} _ {j}\right) ^ {\top} \boldsymbol {W} _ {K e} \boldsymbol {W} _ {Q p}}{C} \cdot i + \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}} \cdot i j. (C.4) \\ \end{array}
$$
We will prove the proposition in two cases, which respectively require  $\mathbf{W}_{Qp}^{\top}\mathbf{W}_{Kp} \leq 0$  and  $\mathbf{W}_{Qp}^{\top}\mathbf{W}_{Kp} \geq 0$ .
Case 1:  $W_{Qp}^{\top}W_{Kp} \leq 0$  required Consider the case of long-term dependency, where the input sequence consists of an initial write and a series of reads, i.e.  $\sigma_1 = 1$  and  $\sigma_t = 0$  for  $t > 1$ . Then for the  $T_{th}$  position, the score for the first write token is
$$
\begin{array}{l} s _ {T \rightarrow 1} = \left\langle \boldsymbol {W} _ {Q} \boldsymbol {v} _ {T}, \boldsymbol {W} _ {K} \boldsymbol {v} _ {1} \right\rangle (C.5) \\ = e ^ {(0) ^ {\top}} \boldsymbol {W} _ {Q e} \boldsymbol {W} _ {K e} e ^ {(1)} + \frac {e ^ {(0) ^ {\top}} \boldsymbol {W} _ {Q e} ^ {\top} \boldsymbol {W} _ {K p}}{C} + \frac {\left(e ^ {(1)}\right) ^ {\top} \boldsymbol {W} _ {K e} \boldsymbol {W} _ {Q p}}{C} \cdot T + \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}} \cdot T (C.6) \\ = \left(\frac {\left(\boldsymbol {e} ^ {(1)}\right) ^ {\top} \boldsymbol {W} _ {K e} \boldsymbol {W} _ {Q p}}{C} + \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}}\right) \cdot T + O (1) = O (T), (C.7) \\ \end{array}
$$
and the score for the last write token is
$$
\begin{array}{l} s _ {T \rightarrow T} = \left\langle \boldsymbol {W} _ {Q} \boldsymbol {v} _ {T}, \boldsymbol {W} _ {K} \boldsymbol {v} _ {T} \right\rangle (C.8) \\ = e ^ {(0) ^ {\top}} \boldsymbol {W} _ {Q e} \boldsymbol {W} _ {K e} e ^ {(0)} + \frac {e ^ {(0) ^ {\top}} \boldsymbol {W} _ {Q e} ^ {\top} \boldsymbol {W} _ {K p}}{C} T + \frac {e ^ {(0) ^ {\top}} \boldsymbol {W} _ {K e} \boldsymbol {W} _ {Q p}}{C} \cdot T + \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}} \cdot T ^ {2} (C.9) \\ = \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}} \cdot T ^ {2} + O (T). (C.10) \\ \end{array}
$$
Think of  $C$  as going to infinity. If  $W_{Qp}^{\top}W_{Kp} > 0$ , then there exists a sufficiently large  $T$  such that  $s_{T\to T} > s_{T\to 1}$ . Hence we need  $W_{Qp}^{\top}W_{Kp}\leq 0$ .
Case 2:  $W_{Qp}^{\top}W_{Kp} \geq 0$  required Consider the input sequence where  $\sigma_1 = 1$ ,  $\sigma_{T - 1} = 2$ , and  $\sigma_t = 0$  for  $t \in [T] \setminus \{1, T - 1\}$ . Similar to the above, calculate the pre-softmax attention scores for  $\sigma_1, \sigma_{T - 1}$  as
$$
s _ {T \rightarrow 1} = O (T) \tag {C.11}
$$
$$
s _ {T \rightarrow T - 1} = \frac {\boldsymbol {W} _ {Q p} ^ {\top} \boldsymbol {W} _ {K p}}{C ^ {2}} \cdot T ^ {2} + O (T). \tag {C.12}
$$
Since we need  $s_{T\to T - 1} > s_{T\to 1}$ , it must be that  $\mathbf{W}_{Qp}^{\top}\mathbf{W}_{Kp}\geq 0$