Encoder-Decoder Cross-Attention

While the masked self-attention mechanism allows the decoder to consider previous tokens within the sequence it's generating, it lacks direct access to the information encoded from the input sequence. To bridge this gap and enable the decoder to selectively focus on relevant parts of the source information, the Transformer architecture incorporates a second attention mechanism within each decoder layer: encoder-decoder cross-attention.

Unlike self-attention where Queries (Q), Keys (K), and Values (V) all originate from the same sequence (either the input sequence in the encoder or the partially generated output sequence in the decoder's masked self-attention), cross-attention draws its components from different sources:

1. Queries (Q): These vectors come from the output of the preceding sub-layer in the decoder. Specifically, they are derived from the output of the masked self-attention layer after the residual connection and layer normalization step. These queries represent the decoder's current state and what information it needs to generate the next token.
2. Keys (K) and Values (V): These vectors come directly from the output of the final layer of the encoder stack. The encoder has processed the entire input sequence, and its output representation contains contextualized information about the source. The keys and values derived from this representation allow the decoder to "look back" at the input sequence.

This structure allows the decoder, at each step, to pose a query based on its current context (the generated sequence so far) and match it against the keys representing the entire input sequence. The resulting attention scores determine how much weight to give to the corresponding values from the encoder's output when constructing the decoder's output for that step.

Mathematical Formulation

The computation itself uses the same scaled dot-product attention function discussed previously. However, the inputs clearly distinguish the sources:

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

Here:

• $Q_{dec}$ represents the query matrices derived from the decoder's previous sub-layer output.
• $K_{enc}$ and $V_{enc}$ represent the matrices derived from the encoder stack's output.
• $d_k$ is the dimension of the vectors, used for scaling.

The softmax function ensures the weights assigned to the encoder's value vectors sum to one, creating a weighted average based on relevance as determined by the query-key interactions.

Multi-Head Cross-Attention

Just like self-attention, cross-attention benefits from multiple attention heads operating in parallel. Each head applies separate linear projections to the incoming $Q_{dec}$, $K_{enc}$, and $V_{enc}$ to map them into different representation subspaces.

$$\text{head}_i = \text{Attention}(Q_{dec}W_i^Q, K_{enc}W_i^K, V_{enc}W_i^V)$$

Where $W_i^Q$, $W_i^K$, and $W_i^V$ are the learned projection matrices for head $i$.

The outputs of these independent heads are then concatenated and passed through a final linear projection layer, identical to the process in multi-head self-attention:

$$\text{MultiHead}(Q_{dec}, K_{enc}, V_{enc}) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O$$

Using multiple heads allows the decoder to simultaneously attend to information from the encoder output based on different criteria or from different representational perspectives. For instance, in translation, one head might focus on syntactic alignment while another focuses on semantic correspondence.

Purpose and Integration

The primary purpose of encoder-decoder cross-attention is to condition the decoder's generation process on the relevant parts of the input sequence. Without it, the decoder would only have access to the input information through the initial decoder state, lacking the ability to dynamically focus on specific source tokens as generation progresses.

Consider translating "The black cat" (input) to "Le chat noir" (output).

1. When generating "Le", the decoder's masked self-attention looks at the start-of-sequence token. Its cross-attention might weakly attend to the whole input or focus on "The".
2. When generating "chat", the decoder's masked self-attention looks at "Le". Its cross-attention mechanism, using a query influenced by "Le", should strongly attend to the encoder's representation of "cat".
3. When generating "noir", the masked self-attention looks at "Le chat". The cross-attention should now strongly attend to the encoder's representation of "black".

This dynamic focusing is fundamental to the Transformer's success in sequence-to-sequence tasks.

Data Flow within the Decoder Layer

The cross-attention mechanism sits between the masked self-attention sub-layer and the position-wise feed-forward sub-layer within each decoder layer. Residual connections and layer normalization are applied around it, just as with the other sub-layers.

Simplified data flow for the Encoder-Decoder Cross-Attention sub-layer within a Transformer Decoder layer. Queries originate from the decoder's state, while Keys and Values come from the encoder's final output.

Understanding this cross-attention mechanism is important for grasping how the Transformer decoder effectively utilizes the encoded representation of the input sequence to guide the generation of the output sequence. It acts as the primary bridge connecting the encoder's processing of the source sequence to the decoder's step-by-step generation of the target sequence.