Google

In autoregressive decoders, a specific type of self-attention called causal or masked self-attention is employed. This mechanism restricts the model's attention, ensuring that the prediction for a token at position `i` can only depend on the tokens at previous positions (from 0 to `i-1`). It is forbidden from attending to tokens at or after position `i`. This is fundamental to maintaining the autoregressive property of generating a sequence one token at a time, as the model cannot 'see' into the future.

Causal Self-Attention in Autoregressive Decoders

In an autoregressive decoder, the probability for the next token is calculated at each step `i` by conditioning on the input `x` and all previously generated tokens `y_{<i}`. The process begins by concatenating `x` and `y_{<i}` and passing them through an embedding layer. This sequence of embeddings is then processed by a stack of decoder layers (which typically include self-attention and feed-forward networks) to produce a sequence of hidden states, `H`. A final linear transformation (using an output weight matrix `W^o`) is applied to these hidden states to get logits, followed by a Softmax function. The probability distribution for the next token, `y_i`, is the probability vector taken from the final position of the output sequence. The formula is: `Pr(·|x, y_{<i}) = (\text{Softmax}(H W^o))_{\text{last}}` where `H = \text{Decoder}([x, y_{<i}])`.

Next-Token Probability Calculation in Autoregressive Decoders

In a causal self-attention mechanism, the query vector for a given token `i`, denoted as $\mathbf{q}_i$, only computes dot products with the key vectors of tokens up to and including its own position. This ensures that the prediction for token `i` does not depend on future tokens. The sequence of dot products computed is as follows:
- For token 0: $\mathbf{q}_0 \mathbf{k}_0^T$
- For token 1: $\mathbf{q}_1 \mathbf{k}_0^T, \mathbf{q}_1 \mathbf{k}_1^T$
- For token 2: $\mathbf{q}_2 \mathbf{k}_0^T, \mathbf{q}_2 \mathbf{k}_1^T, \mathbf{q}_2 \mathbf{k}_2^T$
This pattern continues for subsequent tokens, where for any token `i`, the dot products $\mathbf{q}_i \mathbf{k}_j^T$ are calculated for all $j \le i$.

Enumeration of Dot Products in Causal Self-Attention

An engineer is debugging a language model designed to generate text one word at a time. They observe that when prompted with 'The cat sat on the', the model's prediction for the next word is heavily influenced by a future, unseen word 'rug' that was accidentally included in the decoder's input sequence during a specific training step. Which fundamental principle of sequential data processing is being violated, and how is this principle typically enforced in the model's internal architecture to prevent such 'cheating'?

Debugging a Generative Language Model

In a causal self-attention mechanism, a query at a given position `i` is restricted to attending only to keys at positions `j` that are less than or equal to `i` ($$j \\le i$$). This ensures that the prediction for a token depends only on the preceding tokens. The following list illustrates the specific query-key dot products that are computed for a sequence, demonstrating the lower-triangular pattern of attention scores before masking:

- For query $$\\mathbf{q}_0$$: $$\\mathbf{q}_0 \\mathbf{k}_0^T$$
- For query $$\\mathbf{q}_1$$: $$\\mathbf{q}_1 \\mathbf{k}_0^T$$, $$\\mathbf{q}_1 \\mathbf{k}_1^T$$
- For query $$\\mathbf{q}_2$$: $$\\mathbf{q}_2 \\mathbf{k}_0^T$$, $$\\mathbf{q}_2 \\mathbf{k}_1^T$$, $$\\mathbf{q}_2 \\mathbf{k}_2^T$$
- For query $$\\mathbf{q}_3$$: $$\\mathbf{q}_3 \\mathbf{k}_0^T$$, $$\\mathbf{q}_3 \\mathbf{k}_1^T$$, $$\\mathbf{q}_3 \\mathbf{k}_2^T$$, $$\\mathbf{q}_3 \\mathbf{k}_3^T$$

This pattern continues for the entire sequence length.

Example of Causal Attention Dot Products

An AI development team is building two different language models. Model A is a chatbot designed to generate human-like conversational responses one word at a time. Model B is designed to analyze legal documents and classify them into categories based on the full content of the document. Evaluate whether the same self-attention mechanism should be used for both Model A and Model B. Justify your reasoning by explaining how the information flow requirements differ between the two tasks.

Choosing the Right Attention Mechanism

A language model is designed to generate text one token at a time, predicting the next token based only on the ones that came before it. The image below shows four possible heatmaps (A, B, C, D) representing the attention scores between tokens in a 4-token sequence. The token making the query is on the vertical axis, and the token providing the key is on the horizontal axis. A darker square indicates that a query token is paying more attention to a key token. Which heatmap correctly illustrates 

Learn Before

Related