Configuring Padding Masks in a Non-Causal Transformer Encoder

Welcome back. Your input module now produces a vector for each position in a corrupted sequence. Those vectors still have not exchanged information. Today you’ll put a small Transformer encoder after that module and give it two instructions: ignore [PAD] as a source of information, but allow a masked position to use visible tokens on either side. This should take about 15–20 minutes, including a Colab check.

Padding and missing text are different

Consider the corrupted sequence c[MASK]t[PAD]. To predict the middle character, the denoiser should be able to use both c and t. The [MASK] position is a real position whose answer is hidden; [PAD] merely fills unused space in the batch.

The padded-sequences image shows why this distinction matters when examples have different lengths:

A causal mask would block earlier positions from attending to later ones. That is appropriate for left-to-right next-token prediction, but it would prevent the masked position in c[MASK]t from seeing t. A padding mask does something else: it prevents attention from drawing information from [PAD] positions. We need the padding mask, not a causal mask.

Match the encoder to your existing tensor

Your DenoiserInputs module returns a tensor shaped [batch_size, length, d_model]. PyTorch’s TransformerEncoderLayer accepts that layout when batch_first=True. Skim the TransformerEncoderLayer documentation’s Parameters list, from d_model through batch_first; concentrate on the input width, number of attention heads, and tensor layout.

The layer performs self-attention and additional processing while preserving that overall tensor shape. TransformerEncoder stacks such layers. Its src_key_padding_mask argument handles padding for each batch item; its separate mask argument is for a sequence-level attention restriction that we do not want to supply here.

Read the TransformerEncoder documentation’s constructor example and forward arguments. The distinction between mask and src_key_padding_mask is the important part.

There is one easy-to-miss polarity change from your batching code. Your earlier valid-token mask was True for usable positions. PyTorch’s Boolean src_key_padding_mask instead expects True where a key must be ignored. Thus we pass x_t.eq(pad_id), or equivalently the inverse of the valid-token mask. A [MASK] token is not [PAD], so it remains eligible for attention.

Wire it up in Colab

Run the cell defining DenoiserInputs from the previous lesson first. This example keeps the same [batch, length, width] convention throughout:

import torch
from torch import nn

torch.manual_seed(0)
pad_id = 0
mask_id = 1

inputs = DenoiserInputs(
    vocab_size=8, max_length=6, d_model=16, pad_id=pad_id
)

layer = nn.TransformerEncoderLayer(
    d_model=16,
    nhead=4,
    dim_feedforward=64,
    dropout=0.0,
    batch_first=True,
)
encoder = nn.TransformerEncoder(
    layer,
    num_layers=2,
    enable_nested_tensor=False,
)

# The two rows differ only in the visible token after [MASK].
x_t = torch.tensor(
    [[2, mask_id, 3, pad_id],
     [2, mask_id, 4, pad_id]],
    dtype=torch.long,
)
t = torch.tensor([0.5, 0.5])

valid_mask = x_t.ne(pad_id)
padding_mask = ~valid_mask        # True means "do not attend to this key"

h = inputs(x_t, t)
out = encoder(
    h,
    mask=None,                    # no sequence-level restriction
    src_key_padding_mask=padding_mask,
    is_causal=False,              # no left-to-right restriction
)

print("padding mask:\n", padding_mask)
print("input shape:", h.shape)
print("output shape:", out.shape)

assert padding_mask.tolist() == [
    [False, False, False, True],
    [False, False, False, True],
]
assert out.shape == (2, 4, 16)

d_model=16 matches the width produced by DenoiserInputs, and four heads divide that width evenly. batch_first=True means no transpose is needed. Two layers are enough for a small demonstration, while enable_nested_tensor=False keeps this toy example in the ordinary padded-tensor representation.

The call to encoder expresses the essential configuration. src_key_padding_mask excludes the fourth position as a key that other positions might consult. mask=None supplies no triangular attention mask, and is_causal=False makes the intended non-causal setting explicit.

A subtle limitation: masking padding keys does not necessarily make the output vectors at padding positions zero. Those positions still exist in the output tensor. Later, when computing the training loss, you will exclude padding positions separately.

Check what “bidirectional” permits

The two rows above have the same [MASK] at position 1 but different visible tokens at position 2. With dropout off, compare their encoder outputs at the masked position:

inputs.eval()
encoder.eval()

with torch.no_grad():
    out = encoder(
        inputs(x_t, t),
        mask=None,
        src_key_padding_mask=padding_mask,
        is_causal=False,
    )

print("Difference at [MASK]:",
      (out[0, 1] - out[1, 1]).abs().max().item())

You should see a nonzero difference. The position-1 representation can respond to a token to its right, as a bidirectional denoiser needs. This is only a wiring check: the randomly initialized encoder does not yet know which character belongs behind [MASK].

What to carry forward

You now have the middle of the mini diffusion language model: corrupted-token, position, and noise-level vectors enter a non-causal Transformer encoder; a per-example padding mask prevents [PAD] from serving as context. Keep the two mask meanings distinct: [MASK] marks text to predict, while src_key_padding_mask=True marks padding to ignore.

Next, you’ll turn each encoder output vector into vocabulary logits, giving the denoiser candidate characters at every position.