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 transformer paper includes elements such as positional encoding, masked attention, and 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.