Scaled dot-product attention

Let $n_q,n_k,d_k,d_v$ be positive integers, let $Q\in {\mathbb{R}}^{n_q\times d_k}$ and $K\in {\mathbb{R}}^{n_k\times d_k}$ be query and key matrices, and let $V\in {\mathbb{R}}^{n_k\times d_v}$ be a value matrix. For a real matrix $S$ with $n_q$ rows and $n_k$ columns, define its row softmax by ${\mathrm{softmax}}_{\mathrm{r}\mathrm{o}\mathrm{w}}(S{)}_{ij}=\exp (S_{ij})\mathrm{/}{\sum }_{\mathrm{\ell }=1}^{n_k}\exp (S_{i\mathrm{\ell }})$. The scaled dot-product attention output is

$$\mathrm{Attn}(Q,K,V)={\mathrm{softmax}}_{\mathrm{r}\mathrm{o}\mathrm{w}}\text{\hspace{0.17em}⁣}\left(\frac{Q{K}^{\mathsf{T}}}{\sqrt{d_k}}\right)V\in {\mathbb{R}}^{n_q\times d_v}\mathrm{.}$$

Each output row is the weighted average of the value rows, with weights given by the softmax of that query's scaled dot products with all keys. A mask is an array $M\in (\mathbb{R}\cup \{-\mathrm{\infty }\}{)}^{n_q\times n_k}$ whose entries are $0$ for allowed query-key pairs and $-\mathrm{\infty }$ for disallowed pairs; masked attention replaces the score matrix by $Q{K}^{\mathsf{T}}\mathrm{/}\sqrt{d_k}+M$ before row softmax and is defined only when every row has at least one allowed key. Here $\exp (-\mathrm{\infty })=0$.

Vaswani et al., §3.2.1: scaled dot-product formula and query/key/value dimensions

Dive into Deep Learning, §11.3.3: matrix formula and batched query/key/value shapes