Transformer
Definition1 2
For a matrix $\mathbf{X} \in \mathbb{R}^{d \times n}$ that represents a sequence of length $n$ with embedding dimension $d$, let the function $\operatorname{Attention} : \mathbb{R}^{d \times n} \to \mathbb{R}^{d \times n}$ be defined as follows, and call it attention.
$$ \operatorname{Attention}(\mathbf{X}) = \mathbf{X} + \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \left( \mathbf{W}_{h}^{V}\mathbf{X} \right) \operatorname{Softmax} \left[ (\mathbf{W}_{h}^{K}\mathbf{X})^{\mathsf{T}} (\mathbf{W}_{h}^{Q}\mathbf{X}) \right] $$
Here $H$ is the number of heads, $m$ is the dimension of a single head, and for each $h = 1, \dots, H$ we have $\mathbf{W}_{h}^{Q}, \mathbf{W}_{h}^{K}, \mathbf{W}_{h}^{V} \in \mathbb{R}^{m \times d}$, $\mathbf{W}_{h}^{O} \in \mathbb{R}^{d \times m}$. And $\operatorname{Softmax}$ denotes the function that, for a given matrix $\mathbf{X} = \begin{bmatrix} \mathbf{x}_{1} & \cdots & \mathbf{x}_{n}\end{bmatrix}$, applies the softmax $\operatorname{softmax}$ to each column vector so that the entries of each column sum to $1$.
$$ \operatorname{Softmax}(\mathbf{X}) := \begin{bmatrix} \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{1})}} & \cdots & \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{n})}} \end{bmatrix} $$
Let the feedforward neural network $\operatorname{FF} : \mathbb{R}^{d \times n} \to \mathbb{R}^{d \times n}$ be defined as follows.
$$ \begin{align*} \operatorname{FF}(\mathbf{X}) &= \mathbf{X} + \mathbf{W}_{2} \operatorname{ReLU} \left( \mathbf{W}_{1} \mathbf{X} + \mathbf{b}_{1} \mathbf{1}^{\mathsf{T}} \right) + \mathbf{b}_{2} \mathbf{1}^{\mathsf{T}} \end{align*} $$
$$ \mathbf{W}_{1} \in \mathbb{R}^{r \times d}, \quad \mathbf{W}_{2} \in \mathbb{R}^{d \times r} $$
$$ \mathbf{b}_{1} \in \mathbb{R}^{r}, \quad \mathbf{b}_{2} \in \mathbb{R}^{d} $$
Here $r$ is the hidden layer dimension of $\operatorname{FF}$, and $\operatorname{ReLU}$ is the ReLU applied entrywise. Also, $\mathbf{1} \in \mathbb{R}^{n}$ is the vector all of whose entries are $1$, so $\mathbf{b}_{k} \mathbf{1}^{\mathsf{T}}$ means the broadcast that adds the bias $\mathbf{b}_{k}$ equally to every column of $\mathbf{X}$.
Let a transformer block be defined as below.
$$ \operatorname{Block} = \operatorname{FF} \circ \operatorname{Attention} $$
Transformer
The composition of $N$ transformer blocks is called a transformer.
$$ \operatorname{Transformer} := \overbrace{\operatorname{Block}_{N} \circ \operatorname{Block}_{N-1} \circ \cdots \circ \operatorname{Block}_{1}}^{N} $$
Explanation
In the definition, $\operatorname{Attention}$ is multi-head self-attention including a residual layer. Mathematically, dot-product attention is nothing but the special case $H=1$ of multi-head attention, so in mathematical papers attention is usually defined as above. For details, refer to the multi-head attention post.
In practice, the architecture introduced in the 🔒(26/08/27)transformer paper includes elements such as positional encoding, 🔒(26/08/31)masked attention, and 🔒(26/08/29)layer normalization, but in mathematical papers it is usually defined as simply as above.
Properties
(a) For any permutation matrix $P$, the following holds.
$$ \operatorname{Transformer}(\mathbf{X}P) = \operatorname{Transformer}(\mathbf{X})P $$
(b) The number of parameters of a single block is $4Hmd + 2rd + r + d$, which does not depend on the length $n$ of the input sequence.