Hard and soft attention sequence models

Let $\mathrm{\Sigma }$ be a finite token alphabet, let $\mathtt{e}\mathtt{o}\mathtt{s}\mathrm{\notin }\mathrm{\Sigma }$ be an end-of-sequence symbol, and let $d,L,H\ge 1$ be fixed integers. For each length $n\ge 1$, an input is $(x_1,\dots ,x_{n-1},\mathtt{e}\mathtt{o}\mathtt{s})$. A symbol embedding assigns a vector $v_a\in {\mathbb{R}}^d$ to each $a\in \mathrm{\Sigma }\cup \{\mathtt{e}\mathtt{o}\mathtt{s}\}$; a positional embedding sequence assigns $p_i\in {\mathbb{R}}^d$ independent of the input symbols. A fixed shared input map gives $y_i^{(0)}=f_{\mathrm{i}\mathrm{n}}(v_{x_i},p_i)$. For each layer $k=1,\dots ,L$ and head $h=1,\dots ,H$, a fixed score map $s_{k,h}:{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$ gives scores $a_{i,j}^{(k,h)}=s_{k,h}(y_i^{(k-1)},y_j^{(k-1)})$ for every query position $i$ and key position $j$. In the soft model, ${\widehat{a}}_{i,\cdot }^{(k,h)}=\mathrm{softmax}(a_{i,\cdot }^{(k,h)})$. In the hard model, ${\widehat{a}}_{i,j}^{(k,h)}=1$ at the smallest index $j$ attaining ${\max }_{r\le n}{a}_{i,r}^{(k,h)}$ and is $0$ at other indices. With fixed value maps $V_{k,h}:{\mathbb{R}}^d\to {\mathbb{R}}^d$ and fixed positionwise maps $F_k:({\mathbb{R}}^d{)}^{H+1}\to {\mathbb{R}}^d$, the update is

$$b_{i,k,h}=\sum _{j=1}^n{\widehat{a}}_{i,j}^{(k,h)}{V}_{k,h}(y_j^{(k-1)}),y_i^{(k)}=F_k(y_i^{(k-1)},b_{i,k,1},\dots ,b_{i,k,H})\mathrm{.}$$

Each score, value, and positionwise map is shared across positions and is fixed independently of $n$. In the hard-attention limit theorem, the score, value, and positionwise maps may be arbitrary fixed maps; no activation-function or parameter-norm bound is imposed. For the soft-influence theorem below, each score is a scaled dot product of fixed linear query and key projections, each value map is a fixed linear map, the token and positional embeddings are uniformly bounded over all positions and lengths, the combined initial states satisfy ${\sup }_{i\ge 1,a\in \mathrm{\Sigma }\cup \{\mathtt{e}\mathtt{o}\mathtt{s}\}}\mathrm{\parallel }{f}_{\mathrm{i}\mathrm{n}}(v_a,p_i){\mathrm{\parallel }}_2<\mathrm{\infty }$, and every positionwise map is globally Lipschitz. The initial-state bound holds for addition or concatenation of bounded embeddings, but must be imposed separately for an arbitrary input map. A hard-attention recognizer classifies the input from $y_n^{(L)}$ after reading $\mathtt{e}\mathtt{o}\mathtt{s}$. A soft next-token predictor applies a fixed linear logit map and softmax to $y_n^{(L)}$ to obtain a probability vector on a fixed finite alphabet.

Hahn, §3: positionwise activations, softmax attention, hard attention, and end-of-sequence prediction

Vaswani et al., §3.2: query, key, value, and scaled dot-product attention conventions

Finite feedforward computation graphs and layered MLPs: finite feedforward graphs and layered MLPs specify positionwise maps when those implementations are chosen; the abstract hard-attention model above also permits arbitrary fixed maps.

Scaled dot-product attention: row softmax and scaled query-key products.