Encoder-Decoder

Encoder and Decoder architecture as proposed in “Attention Is All You Need” ¹

Encoder

In an Encoder-Decoder Transformer, the Encoder is responsible for converting the complete source sequence into contextual representations. Given a sequence of $N$ source tokens, Encoder processes all of them simultaneously and generates one contextual representation for each source token.

The source tokens are first converted from token IDS into continuous vectors of dimension $d_{model}$ using a learned embedding matrix $E\in {\mathbb{R}}^{V\times d_{model}}$ where $V$ is the vocabulary size. We get $X_{emb}\in {\mathbb{R}}^{\mathbb{N}\times d_{model}}$.

Since self-attention is position invariant, we have to imbue some kind of positional information. In the original Transformer this was done using Sinusoidal Positional Embeddings (Sinusoidal PE), while modern architecture may follow RoPE (Rotatry PE (RoPE)). After this step the representations are then passed through a stack of $L$ encoder layers.

Each encoder layer contains two principal sub-layers: a multi head self attention sub-layer (Multi Head Attention (MHA)) followed by a position-wise feed forward network(Feed Forward Neural Network (FFNN)). Residual connections are applied around these sub-layers. So, at a high level, an Encoder layer transforms its input representation into a new representation by first allowing each token to “attend” to other tokens and hence share information through self-attention and then applying a non-linear transformation independently to each token representation.

Decoder

The Decoder is responsible for generating the target sequence autoregressively while using the representation produced by the encoder. Unlike Encoder, (at inference) the Decoder doesn’t receive the entire sequence simultaneously and must generate target tokens one by one. Therefore it has two main objectives: use the previously generated tokens to determine what token should come next, and access the relevant information from the source sequence (Encoder).

The first operation is Causal Self-Attention on the Decoder’s own target representation. The Query, Key and Values originate from the decoder input, however, a causal mask is applied so that the representation at position $t$ can’t attend to future tokens at position $>t$. This prevents the model from accessing the information from future tokens while predicting the current one.

After this masked self-attention, the decoder performs Cross-Attention with the encoder output. Here the queries come from the Decoder while the Keys and Values come from the Encoder:

$$Q=X_{dec}{W}_{dec}^Q,K=H_{enc}{W}_{dec}^K,V=H_{enc}{W}_{dec}^V$$

This allows the Decoder to determine which source-token representations are relevant for producing the next target token.

The output of Cross-Attention layer is then passed through the position-wise Feed Forward NN. FFNN applies non-linear transformation individually to each token representation. Hence each Decoder layer has three sub-layers: causal self-attention, encoder-decoder cross attention, and a position-wise feed forward network. Residual connections are applied around these sub-layers.

During training, the target sequence is shifted by one position to form two sequences: a decoder input beginning with $\text{<BOS>}$, and a label sequence ending with $\text{<EOS>}$. Example:

$$\begin{array}{rl}\text{Decoder Input:}& [\text{<BOS>, In, your, dream, lies, an, eternal, truth}]\\ \text{label:}& [\text{In, your, dream, lies, an, eternal, truth, <EOS>}]\end{array}$$

$\text{<BOS> \& <EOS>}$ are special tokens that denote “Beginning Of Sequence” and “End Of Sequence”, respectively.

Using such target-label pair (called standard teacher forcing) we train the model to produce the next token based on the given input sequence. For example, for the above input sequence after seeing $\text{<BOS>}$ model should predict next token as “In”. After going through the whole Decoder stack we get a final hidden representation $h_t\in {\mathbb{R}}^{d_{model}}$ at position $t$ which is then project into the target vocabulary space using an unembedding matrix $W_u\in {\mathbb{R}}^{d_{model}\times V}$ that gives us $z_t\in {\mathbb{R}}^V$. These logits are then converted into a probability distribution using the softmax function:

$$p_{t,i}=\frac{exp(z_{t,i})}{{\sum }_{j=1}^{V}exp(z_{t,j})}$$

$p_{t,i}$ denotes the probability assigned to vocabulary token $i$ at decoder position $t$.

If the ground-truth token at position $t$ is $y_t$ then the Cross-Entropy loss for that position is the negative log-probability assigned to that token:

$${\mathcal{L}}_t=-\log p_{t,y_t}$$

The total training loss is usually the mean of token-level losses over the entire target sequence:

$$\mathcal{L}=-\frac{1}{T}\sum _{t=1}^T\log p_{t,y_t}$$

Minimizing this loss is equivalent to maximizing the probability assigned to the correct target tokens.

Footnotes

1. Attention Is All You Need ↩