Transformer encoder, decoder, and fixed-width expressivity blocks

Fix positive integers $d,d_{\mathrm{f}\mathrm{f}},n,m$, where $d$ is the model width, $d_{\mathrm{f}\mathrm{f}}$ is the feedforward hidden width, and $n,m$ are the encoder and decoder sequence lengths. A row-token matrix $X\in {\mathbb{R}}^{n\times d}$ has one token in each row. Define $\mathrm{ReLU}(t)=\max (0,t)$ and apply it coordinatewise to vectors and matrices. A positionwise feedforward network with parameters $W_1\in {\mathbb{R}}^{d\times d_{\mathrm{f}\mathrm{f}}}$, $b_1\in {\mathbb{R}}^{d_{\mathrm{f}\mathrm{f}}}$, $W_2\in {\mathbb{R}}^{d_{\mathrm{f}\mathrm{f}}\times d}$, and $b_2\in {\mathbb{R}}^d$ maps each row $x$ to $\mathrm{ReLU}(x{W}_1+b_1)W_2+b_2$, using the same parameters at every position. For $z\in {\mathbb{R}}^d$, let $\bar{z}=d^{-1}{\sum }_{r=1}^d{z}_r$ and $s^2(z)=d^{-1}{\sum }_{r=1}^d(z_r-\bar{z}{)}^2$. Layer normalization with parameters $\gamma ,\beta \in {\mathbb{R}}^d$ and $ϵ>0$ is

$${\mathrm{LN}}_{\gamma ,\beta ,ϵ}(z)=\gamma \odot \frac{z-\bar{z}\mathbf{1}}{\sqrt{s^2(z)+ϵ}}+\beta ,$$

applied rowwise; $\odot$ is coordinatewise multiplication. The positive $ϵ$ makes the denominator nonzero.

Every multi-head attention sublayer in the following post-normalized encoder and decoder layers has query, key, value, and output feature widths $d_q=d_k=d_v=d_o=d$ in def-multihead-self-attention, with its own projection parameters. The head count $h$ and head widths $d_h,d_{vh}$ remain arbitrary positive integers, and the output projection has shape $W^O\in {\mathbb{R}}^{h{d}_{vh}\times d}$. Thus each attention output has the same row count and feature width as its residual input. Any attention mask must leave at least one allowed key in every query row; positive key length supplies this for unmasked attention, and the decoder causal mask supplies it through its diagonal.

A post-normalized encoder layer is the composition

$S={\mathrm{LN}}_{{\gamma }_1,{\beta }_1,{ϵ}_1}\text{\hspace{0.17em}⁣}(X+\mathrm{MHA}(X,X,X)),Y={\mathrm{LN}}_{{\gamma }_2,{\beta }_2,{ϵ}_2}\text{\hspace{0.17em}⁣}(S+\mathrm{FFN}(S))\mathrm{.}$

A post-normalized decoder layer takes decoder states $U\in {\mathbb{R}}^{m\times d}$ and encoder states $E\in {\mathbb{R}}^{n\times d}$. It first applies masked self-attention to $U$ using the causal mask, then cross-attention with queries from the resulting decoder states and keys and values from $E$, and then a positionwise feedforward network; each sublayer output is added to its own input and followed by layer normalization. The decoder's causal mask admits keys at positions $j\le i$ for query position $i$. Encoder and decoder stacks are finite compositions of their respective layers, with positional encodings added to token embeddings before the stacks.

The diagram records the dataflow and residual branches of one encoder layer and one decoder layer. Each arrow is a tensor passed to the indicated sublayer; each “Add & LN” node adds its sublayer input to its output and applies the stated rowwise layer normalization. The encoder output $E$ supplies the key and value rows for decoder cross-attention. The diagram illustrates the definition and is not used in any proof.

The fixed-width expressivity result uses a different, explicitly defined block without layer normalization. In the column-token convention $X\in {\mathbb{R}}^{d\times n}$, one Yun attention sublayer with two heads of size one is

$${\mathrm{Attn}}_Y(X)=X+\sum _{i=1}^2{W}_O^i{W}_V^{i}X{\mathrm{softmax}}_{\mathrm{c}\mathrm{o}\mathrm{l}}\text{\hspace{0.17em}⁣}((W_K^{i}X{)}^{\mathsf{T}}(W_Q^{i}X+b_Q^i{\mathbf{1}}_n^{\mathsf{T}})),$$

where $W_Q^i,W_K^i,W_V^i\in {\mathbb{R}}^{1\times d}$, $W_O^i\in {\mathbb{R}}^{d\times 1}$, and $b_Q^i\in \mathbb{R}$ is a shared scalar query bias; ${\mathrm{softmax}}_{\mathrm{c}\mathrm{o}\mathrm{l}}$ normalizes each column over its key positions. Its tokenwise feedforward sublayer is

$${\mathrm{FF}}_Y(Z)=Z+W_2\mathrm{ReLU}(W_{1}Z+b_1{\mathbf{1}}_n^{\mathsf{T}})+b_2{\mathbf{1}}_n^{\mathsf{T}},$$

with $W_1\in {\mathbb{R}}^{4\times d}$ and $W_2\in {\mathbb{R}}^{d\times 4}$. A Yun block is ${\mathrm{FF}}_Y\circ {\mathrm{Attn}}_Y$, and ${\mathcal{T}}^{2,1,4}$ is the set of all finite compositions of such blocks. Its positional class consists of maps $X\mapsto g(X+E)$ with $g\in {\mathcal{T}}^{2,1,4}$ and a trainable $E\in {\mathbb{R}}^{d\times n}$. This expressivity block has no layer normalization; its widths and head sizes are exactly as displayed.

Multi-head attention and self-attention

Vaswani et al., §3.1: encoder and decoder residual sublayers and post-normalization

Vaswani et al., §§3.2.3 and 3.3: causal masks and shared positionwise ReLU feedforward maps

Yun et al., §2: the width-4, two-head, head-size-one block without layer normalization

Dive into Deep Learning, §§11.7.1–11.7.5: encoder, decoder, residual, and positionwise sublayers

Multi-head attention and self-attention: projected query, key, and value heads and featurewise concatenation.

Positional encodings and causal masks: rowwise positional addition and the lower-triangular decoder mask.

Finite feedforward computation graphs and layered MLPs: finite feedforward computation graphs and layered MLP evaluation.