Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Transformer Block Assembly › Causal Masking for Decoder Blocks

Michael BrenndoerferUpdated June 13, 20254 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 110Transformer Block Assembly › Causal Masking for Decoder Blocks

Reading section 7 of 16 · Causal Masking for Decoder Blocks · 4 minSection 7 of 16 · 4 min

Section 7 of 16 · 4 min read

Transformer Block Assembly › Causal Masking for Decoder BlocksOne section · Sections in this chapter7 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Language models like GPT are autoregressive: they generate tokens one at a time, with each token depending only on previous tokens. This requires causal masking in attention, preventing tokens from attending to future positions.

The causal mask modifies the attention score computation. Recall that attention scores between queries and keys are computed as:

\text{score}_{ij} = \frac{\mathbf{q}_i \cdot \mathbf{k}_j}{\sqrt{d_k}}

where:

$\mathbf{q}_i$ : the query vector for token at position $i$
$\mathbf{k}_j$ : the key vector for token at position $j$
$d_k$ : the dimension of the key vectors (used for scaling)

To prevent token $i$ from attending to any future token $j > i$ , we add a mask $M$ to the scores before applying softmax:

\text{masked\_score}_{ij} = \text{score}_{ij} + M_{ij}

where:

$M_{ij} = 0$ if $j \leq i$ (past or current positions, allowed)
$M_{ij} = -\infty$ if $j > i$ (future positions, blocked)

After softmax, positions with $-\infty$ scores become exactly zero, effectively removing them from the weighted average. The result is a triangular attention pattern where each token can only "see" itself and previous tokens:

In[21]:

Code

def create_causal_mask(seq_len):
    """
    Create causal attention mask.

    Returns a matrix where position (i, j) is -inf if j > i (future position),
    and 0 otherwise. Adding this to attention scores before softmax
    ensures tokens cannot attend to future positions.
    """
    mask = np.triu(np.ones((seq_len, seq_len)) * float("-inf"), k=1)
    return mask


# Example mask for 4 tokens
mask_example = create_causal_mask(4)

Out[22]:

Console

Causal Mask (4 tokens):
(0 = can attend, -inf = cannot attend)

  0  -∞  -∞  -∞
  0  0  -∞  -∞
  0  0  0  -∞
  0  0  0  0

Out[23]:

Visualization

Triangular heatmap showing causal mask pattern with lower triangle green (allowed) and upper triangle red (blocked). — Causal attention mask for a 6-token sequence. Green cells indicate allowed attention (token can see that position), red cells indicate blocked attention (future positions). Each token can only attend to itself and previous tokens.

Decoder Block ImplementationLink Copied

A decoder block is identical to the standard block, but with causal masking applied during attention:

In[24]:

Code

class DecoderBlock(TransformerBlock):
    """
    Transformer decoder block with causal masking.

    Identical to TransformerBlock but automatically applies
    causal masking to prevent attending to future positions.
    """

    def __call__(self, x):
        """
        Process input through decoder block.

        Args:
            x: Input tensor of shape (seq_len, d_model)

        Returns:
            Output tensor of shape (seq_len, d_model)
        """
        seq_len = x.shape[0]

        # Create causal mask
        mask = create_causal_mask(seq_len)

        # Use parent class forward with mask
        return super().__call__(x, mask=mask)


# Test decoder block
decoder_block = DecoderBlock(d_model=256, n_heads=4, d_ff=688, use_swiglu=True)

x_test = np.random.randn(8, 256) * 0.02
y_test = decoder_block(x_test)

Out[25]:

Console

Decoder Block Test
========================================
Input shape: (8, 256)
Output shape: (8, 256)
Causal masking: enabled
Parameters: 791,040