Encoder-Decoder Attention Mechanism

The Encoder-Decoder Attention mechanism is a significant component within the Transformer's decoder stack. The decoder employs masked self-attention to consider tokens it has already generated. However, it additionally requires a mechanism to consult the input sequence provided to the encoder. Encoder-decoder attention fulfills this requirement, acting as a bridge that allows the decoder to selectively focus on relevant parts of the encoded input representation when generating each output token.

Think about a task like machine translation. As the decoder generates the translated sentence word by word, it needs to constantly refer back to the source sentence (processed by the encoder) to ensure accuracy and context. The encoder-decoder attention mechanism facilitates this process.

How Encoder-Decoder Attention Works

Structurally, this mechanism is very similar to the multi-head self-attention layers we've already discussed. It typically uses scaled dot-product attention, often implemented with multiple heads running in parallel. The primary difference lies in the origin of the Query (Q), Key (K), and Value (V) vectors.

Unlike self-attention where Q, K, and V are all derived from the same sequence (the output of the previous layer), in encoder-decoder attention:

1. Query (Q) Vectors: These are derived from the output of the preceding sub-layer in the decoder stack. Specifically, they come from the output of the masked multi-head self-attention layer (after the Add & Norm step). This query represents the decoder's current state and what information it needs to generate the next token.
2. Key (K) and Value (V) Vectors: These vectors originate from the output of the encoder stack. The encoder processes the entire input sequence and produces a final set of representations (one for each input position). These representations are used to generate the K and V vectors for the encoder-decoder attention layer in every decoder layer. These keys and values encapsulate the information from the entire input sequence.

The attention calculation then proceeds similarly to what we saw before, often using scaled dot-product attention for each head:

$$\text{Attention}(Q_{\text{decoder}}, K_{\text{encoder}}, V_{\text{encoder}}) = \text{softmax}\left(\frac{Q_{\text{decoder}} K_{\text{encoder}^T}}{\sqrt{d_k}}\right) V_{\text{encoder}}$$

Here, $Q_{\text{decoder}}$ represents the queries from the decoder, while $K_{\text{encoder}}$ and $V_{\text{encoder}}$ represent the keys and values derived from the encoder's output. The scaling factor $d_k$ is the dimension of the key vectors.

The Flow of Information

At each step during decoding, the decoder generates a query based on the output sequence generated so far. This query is then compared (via dot products) with the keys derived from all encoder outputs. The softmax function converts these dot products into attention scores (weights), indicating how relevant each part of the encoded input sequence is to the decoder's current query.

These attention scores are then used to compute a weighted sum of the value vectors (also derived from the encoder outputs). The result is a context vector that summarizes the information from the input sequence most pertinent to generating the next output token.

Flow of information in the Encoder-Decoder Attention mechanism. Queries originate from the decoder's previous sub-layer, while Keys and Values come from the final output of the encoder stack. This mechanism allows the decoder to weigh the importance of different parts of the input sequence.

Integration within the Decoder Layer

This encoder-decoder attention sub-layer sits within each decoder layer, typically positioned after the masked self-attention sub-layer and its corresponding Add & Norm step. Following the encoder-decoder attention calculation, another Add & Norm step is applied, combining the output of this attention mechanism with its input (the output from the masked self-attention sub-layer) and applying layer normalization.

This structured approach, combining multi-head attention (first masked self-attention, then encoder-decoder attention) with residual connections, normalization, and feed-forward networks, allows the decoder to effectively integrate information from both its previously generated tokens and the complete encoded input sequence. This approach is fundamental to the Transformer's ability to handle complex sequence-to-sequence tasks.