SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Derivation for single head eq 18:
Consider two matrices  $\mathbf{X}_1, \mathbf{X}_2 \in \mathcal{X} = \{\mathbf{X} \in \mathbb{R}^{d \times m}; \| \mathbf{X}\|_2 \leq D\}$ . Denote with  $\mathbf{A}_1, \mathbf{A}_2$  the corresponding attention matrices respectively, which can be defined as:
$$
\mathbf {A} _ {1} = \sigma \left(\left(\mathbf {W} _ {k} \mathbf {X} _ {1}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {1}\right)
$$
$$
\mathbf {A} _ {2} = \sigma \left(\left(\mathbf {W} _ {k} \mathbf {X} _ {2}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2}\right) \tag {20}
$$
The output of the attention head, denoted with  $Att(\cdot)$  admits the following:
$$
\left\| \operatorname {A t t} \left(\mathbf {X} _ {1}\right) - \operatorname {A t t} \left(\mathbf {X} _ {2}\right) \right\| _ {2} = \left\| \mathbf {W} _ {v} \mathbf {X} _ {1} \mathbf {A} _ {1} - \mathbf {W} _ {v} \mathbf {X} _ {2} \mathbf {A} _ {2} \right\| _ {2} \tag {21}
$$
$$
\leq \left\| \mathbf {X} _ {1} \mathbf {A} _ {1} - \mathbf {X} _ {2} \mathbf {A} _ {2} \right\| _ {2} \left\| \mathbf {W} _ {v} \right\| _ {2} \tag {22}
$$
$$
= \left\| \mathbf {X} _ {1} \mathbf {A} _ {1} - \mathbf {X} _ {2} \mathbf {A} _ {1} + \mathbf {X} _ {2} \mathbf {A} _ {1} - \mathbf {X} _ {2} \mathbf {A} _ {2} \right\| _ {2} \left\| \mathbf {W} _ {v} \right\| _ {2} \tag {23}
$$
$$
\leq \left(\| \mathbf {A} _ {1} \| _ {2} \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} + \| \mathbf {X} _ {2} \| _ {2} \| \mathbf {A} _ {1} - \mathbf {A} _ {2} \| _ {2}\right) \| \mathbf {W} _ {v} \| _ {2} \tag {24}
$$
$$
\leq \left(\| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} + \| \mathbf {A} _ {1} - \mathbf {A} _ {2} \| _ {2} D\right) \| \mathbf {W} _ {v} \| _ {2} \tag {25}
$$
where  $(a)$  holds from the spectral norm properties and in  $(b)$  we use the bounded input spectral norm assumptions.
We now focus on the second term  $\| \mathbf{A}_1 - \mathbf{A}_2\| _2$  in eq 25. From the bound in lemma 9, we have:
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq 2 \sqrt {m} \| \mathbf {G} \| _ {2} \tag {26}
$$
where  $\mathbf{G}$  is the diagonal matrix with entires described in lemma 8
We can now invoke lemma 10 to obtain the following:
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq 2 \sqrt {m} \times 2 \times 2 \| \mathbf {W} _ {k} ^ {T} \mathbf {W} _ {q} \| _ {2} D \times \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} \tag {27}
$$
Combining the previous inequality with eq 25, we have the following bound:
$$
\left\| A t t \left(\mathbf {X} _ {1}\right) - A t t \left(\mathbf {X} _ {2}\right) \right\| _ {2} \leq \left(1 + 8 \sqrt {m} \| \mathbf {W} _ {k} ^ {T} \mathbf {W} _ {q} \| _ {2} D ^ {2}\right) \| \mathbf {W} _ {v} \| _ {2} \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} \tag {28}
$$
Derivation for the entire block eq 11:
The proof follows simply by leveraging the following property:
Property: for a matrix  $\mathbf{C} = [\mathbf{A},\mathbf{B}]$ , the spectral norm of  $\mathbf{C}$  admits the bound:
$$
\| \mathbf {C} \| _ {2} \leq \sqrt {\| \mathbf {A} \| _ {2} ^ {2} + \| \mathbf {B} \| _ {2} ^ {2}}
$$
We then simply combine the definition of the attention block and the lipschitz constant bound in eq 18 with the above property in order to obtain the desired bound.
Lemma 8 (Dong et al. [2021] Lemma A.1). For the column stochastic matrix  $\mathbf{A}_1$  obtained by performing column-wise softmax of some matrix  $\mathbf{Z}_1$  (where in our setting  $\mathbf{Z}_1 = (\mathbf{W}_k\mathbf{X}_1)^\top \mathbf{W}_q\mathbf{X}_1$ , and another row stochastic matrix  $\mathbf{A}_2$  obtained by performing column-wise softmax of some matrix  $\mathbf{Z}_2$ , where  $\mathbf{Z}_2 = \mathbf{Z}_1 - \mathbf{E}$  (for some  $\mathbf{E}$ , which need not belong to  $\mathcal{X}$ ), we have the following bound:
$$
\mathbf {A} _ {2} (\mathbf {I} - \mathbf {G}) \leq \mathbf {A} _ {1} \leq \mathbf {A} _ {2} (\mathbf {I} + 2 \mathbf {G}) \tag {29}
$$
where the inequality is elementwise and  $\mathbf{G}$  is a diagonal matrix with entries as  $\mathbf{G}_{ii} = \max_{j,j'} |\delta_i^T \mathbf{E}(\delta_j^T - \delta_{j'}^T)|$ . Here  $\delta_k$  is a one-hot vector with the entry 1 in the  $k^{th}$  dimension.
Lemma 9. Following the notations of lemma 8, we have the following spectral norm bound:
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq 2 \sqrt {m} \| \mathbf {G} \| _ {2} \tag {30}
$$
Proof. We begin by noting the following entry-wise inequality from eq 29:
$$
\mathbf {A} _ {2} \mathbf {G} \leq \mathbf {A} _ {1} - \mathbf {A} _ {2} \leq 2 \mathbf {A} _ {2} \mathbf {G} \tag {31}
$$
which ensures that  $\| \mathbf{A}_1 - \mathbf{A}_2\| _F\leq 2\| \mathbf{A}_2\mathbf{G}\| _F$
We also have the following using matrix norm equivalence:
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq \left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {F} \tag {32}
$$
Invoking the matrix norm equivalence again, we have that
$$
2 \| \mathbf {A} _ {2} \mathbf {G} \| _ {F} \leq 2 \sqrt {\operatorname {r a n k} \left(\mathbf {A} _ {2} \mathbf {G}\right)} \| \mathbf {A} _ {2} \mathbf {G} \| _ {2} \tag {33}
$$
where  $\text{rank}(\cdot)$  is the matrix rank.
Combining the inequalities, we attain the bound :
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq 2 \sqrt {m} \left\| \mathbf {A} _ {2} \right\| _ {2} \| \mathbf {G} \| _ {2} \tag {34}
$$
since  $\mathbf{A}_2$  is column-stochastic,  $\| \mathbf{A}_2\| _2 = 1$
Lemma 10. The term  $\mathbf{G}$  in lemma 9 admits the following spectral norm bound:
$$
\left\| \mathbf {G} \right\| _ {2} \leq 2 D \left\| \mathbf {W} _ {q} \mathbf {W} _ {k} ^ {T} \right\| _ {2} \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} \tag {35}
$$
here  $D$  is the previously stated spectral norm bound of the inputs  $\mathbf{X}^l\in \mathcal{X}^l$
Proof. We begin by noting that since  $\mathbf{G}$  is a square diagonal matrix with non-negative real values, the singular values of  $\mathbf{G}$  are the corresponding diagonal elements.
We thus have that  $\| \mathbf{G}\|_{max} = \| \mathbf{G}\| _2$  , where  $\| \cdot \|_{max}$  is the max norm.
Since  $\mathbf{G}$  admits the form described in lemma 8, it is trivial to note that:
$$
\left\| \mathbf {G} \right\| _ {\max } = \max  _ {i, j, i ^ {\prime}, j ^ {\prime}} \left| \mathbf {E} _ {i, j} - \mathbf {E} _ {i ^ {\prime}, j ^ {\prime}} \right| \tag {36}
$$
$$
\leq 2 \| \mathbf {E} \| _ {\max } \leq 2 \| \mathbf {E} \| _ {2} \tag {37}
$$
where the second inequality follows from the matrix norm equivalence.
Now, we can bound the last term  $\| \mathbf{E}\| _2$  by noting that the inputs  $\mathbf{X}$  belong to a bounded set. This allows us to provide the following bounds:
$$
\begin{array}{l} \left\| \mathbf {E} \right\| _ {2} = \left\| \left(\mathbf {W} _ {k} \mathbf {X} _ {1}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {1} - \left(\mathbf {W} _ {k} \mathbf {X} _ {2}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2} \right\| _ {2} (38) \\ = \left\| \left(\mathbf {W} _ {k} \mathbf {X} _ {1}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {1} - \left(\mathbf {W} _ {k} \mathbf {X} _ {1}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2} + \left(\mathbf {W} _ {k} \mathbf {X} _ {1}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2} - \left(\mathbf {W} _ {k} \mathbf {X} _ {2}\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2} \right\| _ {2} (39) \\ \leq \left(\| \mathbf {X} _ {1} \| _ {2} \| \mathbf {W} _ {k} ^ {T} \mathbf {W} _ {q} \| _ {2} + \| \mathbf {X} _ {2} \| _ {2} \| \mathbf {W} _ {k} ^ {T} \mathbf {W} _ {q} \| _ {2}\right) \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} (40) \\ \leq 2 D \left\| \mathbf {W} _ {k} ^ {T} \mathbf {W} _ {q} \right\| _ {2} \left\| \mathbf {X} _ {1} - \mathbf {X} _ {2} \right\| _ {2} (41) \\ \end{array}
$$
C.11 Extension of Lemma 6
Lemma 6 and theorem 4 operate over functions from  $\mathcal{P}^1\times \mathcal{X}^1\to \mathcal{P}^{L + 1}\times \mathcal{X}^{L + 1}$ . We can relax the requirement of the prompt and provide the Lipschitz constant upper bound in consideration to functions of the form  $\mathcal{X}^1\rightarrow \mathcal{X}^{L + 1}$  by using the following assumption:
Assumption 3. Assume for simplicity that  $\| \mathbf{P}_1^l -\mathbf{P}_2^l\| _2\leq \alpha^l\| \mathbf{X}_1^l -\mathbf{X}_2^l\| _2;\forall l\geq 1.$ $\alpha^l = 0$  when  $l = 1$
Note: A recursive expression for  $\alpha^l$  in the above assumption can be provided, but the expression does not admit a simplified form and we thus omit it here.
We will use  $\mathcal{D}_X^l$ , akin to eq 9, to denote the compactness corresponding to the input matrix across the layers.
Based on this assumption, we have the following Lipschitz constant upper bound:
Lemma 11. The Lipschitz constant of the single head  $A$  at  $t^{i,l}$  admits the following bound w.r.t the input part,  $\mathbf{X}^1$  of length  $m_X$ , of the input sequence:
$$
L i p \left(\boldsymbol {A} t \boldsymbol {t} ^ {i, l} (\cdot , \cdot)\right) \leq \left(\sqrt {1 + (\alpha^ {l}) ^ {2}} + 8 \sqrt {m _ {X}} \left(D _ {X} ^ {l}\right) ^ {2} (1 + (\alpha^ {l}) ^ {2}) \| \left(\mathbf {W} _ {k} ^ {i, l}\right) ^ {T} \mathbf {W} _ {q} ^ {i, l} \| _ {2}\right) \| \mathbf {W} _ {v} ^ {i, l} \| _ {2} \tag {42}
$$
For  $l = 1$ ,  $\alpha^l = 0$  in the above bound.
The Lipschitz constant of the entire attention block in layer  $l$  follows similarly.
Proof. For some first layer input  $\mathbf{X}_1^1$  and prompt  $\mathbf{P}$ , let us denote the direct output of the attention head in the  $l$ -th layer with  $\vec{\mathbf{X}}_1^l$ . We have the following update rule for  $\vec{\mathbf{X}}_1^l$ :
$$
\overrightarrow {\mathbf {X}} _ {1} ^ {l} = \mathbf {W} _ {v} ^ {i, l} \left[ \mathbf {P} _ {1} ^ {l}, \mathbf {X} _ {1} ^ {l} \right] \cdot \sigma \left(\left(\mathbf {W} _ {k} ^ {i, l} \left[ \mathbf {P} _ {1} ^ {l}, \mathbf {X} _ {1} ^ {l} \right]\right) ^ {\top} \mathbf {W} _ {q} ^ {i} \mathbf {X} _ {1} ^ {l}\right) = \mathbf {W} _ {v} ^ {i, l} \left[ \mathbf {P} _ {1} ^ {l}, \mathbf {X} _ {1} ^ {l} \right] \cdot \mathbf {A} _ {1} ^ {l} \tag {43}
$$
Here,  $\mathbf{P}_1^l$  is the updated prompt matrix w.r.t the input. For two different inputs  $\mathbf{X}_1^1$  and  $\mathbf{X}_2^1$  at the first layer,  $\mathbf{P}_1^1 = \mathbf{P}_2^1 = \mathbf{P}$ , since the prompt is same across all inputs.  $\mathbf{A}_1^l$  is then simply the corresponding column-stochastic matrix.
With the context clear, we now drop the superscripts  $i,l$ , as done previously. For  $\| \vec{\mathbf{X}}_1 - \vec{\mathbf{X}}_2\| _2$ , we have:
$$
\begin{array}{l} \left\| \overrightarrow {\mathbf {X}} _ {1} - \overrightarrow {\mathbf {X}} _ {2} \right\| _ {2} \leq \left\| \left[ \mathbf {P} _ {1}, \mathbf {X} _ {1} \right] \mathbf {A} _ {1} - \left[ \mathbf {P} _ {2}, \mathbf {X} _ {2} \right] \mathbf {A} _ {2} \right\| _ {2} \left\| \mathbf {W} _ {v} \right\| _ {2} \\ \leq \left(\| \mathbf {A} _ {1} \| _ {2} \sqrt {\| \mathbf {P} _ {1} - \mathbf {P} _ {2} \| _ {2} ^ {2} + \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} ^ {2}} + \sqrt {\| \mathbf {P} _ {2} \| _ {2} ^ {2} + \| \mathbf {X} _ {2} \| _ {2} ^ {2}} \| \mathbf {A} _ {1} - \mathbf {A} _ {2} \| _ {2}\right) \| \mathbf {W} _ {v} \| _ {2} \tag {44} \\ \end{array}
$$
where the second inequality is attained using the property of spectral norm of concatenated matrices.
We now consider the term  $\| \mathbf{A}_1 - \mathbf{A}_2\| _2$ . By invoking lemmas 9 and 10, we have that:
$$
\left\| \mathbf {A} _ {1} - \mathbf {A} _ {2} \right\| _ {2} \leq 2 \sqrt {m _ {X}} \times 2 \| \mathbf {E} \| _ {2}
$$
$$
\text {w h e r e} \mathbf {E} = \left(\mathbf {W} _ {k} \left[ \mathbf {P} _ {1}, \mathbf {X} _ {1} \right]\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {1} - \left(\mathbf {W} _ {k} \left[ \mathbf {P} _ {2}, \mathbf {X} _ {2} \right]\right) ^ {\top} \mathbf {W} _ {q} \mathbf {X} _ {2} \tag {45}
$$
Invoking assumption 3,  $\| \mathbf{E} \|_2$  can further be bounded as:
$$
\left\| \mathbf {E} \right\| _ {2} \leq 2 D _ {X} \sqrt {1 + \alpha^ {2}} \left\| \left(\mathbf {W} _ {k}\right) ^ {T} \mathbf {W} _ {q} \right\| _ {2} \tag {46}
$$
Finally, by combining the bound for  $\| \mathbf{E}\| _2$  and assumption 3 with eq 44, we obtain:
$$
\begin{array}{l} \left\| \overrightarrow {\mathbf {X}} _ {1} - \overrightarrow {\mathbf {X}} _ {2} \right\| _ {2} (47) \\ \leq \left(\sqrt {1 + \alpha^ {2}} + D _ {X} \sqrt {1 + \alpha^ {2}} \times 8 \sqrt {m _ {X}} D _ {X} \sqrt {1 + \alpha^ {2}} \| \left(\mathbf {W} _ {k}\right) ^ {T} \mathbf {W} _ {q} \| _ {2}\right) \| \mathbf {W} _ {v} \| _ {2} \| \mathbf {X} _ {1} - \mathbf {X} _ {2} \| _ {2} (48) \\ \end{array}
$$
which provides us the desired bound.
By setting  $\alpha^l = 0$ , the case when there is no prompt, we obtain a similar bound as lemma 6