Model Architectures › Transformer · Note 1

Attention Masks in the Transformer

Xinheng Ding

Oct 18, 2026

While reading The Annotated Transformer, I kept tripping over its two masks. The source mask has shape (batch, 1, 1, m). The decoder mask has shape (batch, 1, n, n). Padding appears on both sides of attention, yet only some of it is masked, and it shows up once more in the loss. This note follows a single batch through all three places and explains every shape with one principle.

Where the mask acts

Inside multi-head attention, the scores form a tensor of shape (batch, h, n, m). Entry scores[b, h, i, j] is how strongly query position i attends to key position j, for sentence b and head h. The mask is applied right before the softmax:

def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    ...
    return torch.matmul(p_attn, value), p_attn

Wherever the mask is 0, the score becomes −10⁹ and its softmax weight becomes effectively 0. The softmax runs over the key axis (dim=-1), so the keys that remain visible in each row still sum to 1. The order matters. Zeroing weights after the softmax would leave rows that no longer sum to 1. The value −10⁹ is used instead of −∞ because a fully masked row would turn −∞ into 0/0 = NaN.

The principle

PyTorch aligns shapes from the right and stretches any axis of size 1. A mask therefore stores only the axes along which its value actually changes. For any mask, ask which of b, h, i, j the answer depends on. The shape follows directly.

Masks hide columns, not rows

A column j is a key/value pair, a source of information. A row i is a query, a reader. Padding masks only ever hide columns. Which columns to hide is decided by the padding of the sequence that supplies K and V.

• Padding as a key must be hidden. Otherwise real tokens would mix <blank> vectors into their outputs. A sentence would then encode differently depending on how much padding its batch happened to need.
• Padding as a query can be left alone. Its output row is garbage, but nothing downstream uses it. In the encoder, these rows become memory rows, which are masked again as keys. In the decoder, padded positions are dropped from the loss, as shown below. Masking a whole row would not help anyway: a row of all −10⁹ becomes a uniform distribution, not "nothing".

The source mask: (batch, 1, 1, m)

Whether key j may be seen depends only on whether source token j is padding:

mask[b, h, i, j] = (src[b, j] != pad)        # depends on b and j only

The batch axis is full, because each sentence pads differently. The head axis and the query axis are 1, because every head and every query obeys the same rule. The key axis is full. The mask is built once per batch in Batch, and multi-head attention adds the head axis:

class Batch:
    def __init__(self, src, tgt=None, pad=2):  # 2 = <blank>
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)  # (batch, 1, m)
        ...

class MultiHeadedAttention(nn.Module):
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)  # (batch, 1, 1, m)
        ...

The same src_mask is used twice: in the encoder's self-attention, and in the decoder's cross-attention, discussed below.

The decoder mask: (batch, 1, n, n)

The decoder's self-attention hides two things: padding, and the future.

mask[b, h, i, j] = (tgt[b, j] != pad)  AND  (j <= i)
                    depends on b, j          depends on i, j

The causal rule j ≤ i involves the query position. Position 0 sees one token and position 3 sees four, so the rows now differ and the query axis has to be written out. Both axes have length n, because queries and keys both come from the target sentence. The code builds each factor with only its own axes and lets & broadcast them together:

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)
    return subsequent_mask == 0  # (1, n, n), lower triangle

class Batch:
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)  # (batch, 1, n)
        tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
        return tgt_mask  # (batch, n, n)

The extra n in the shape comes entirely from the causal rule, not from padding. A padding-only target mask, as in a bidirectional encoder like BERT, would be (batch, 1, 1, n): the same form as the source mask.

Cross-attention reuses the source mask

In the decoder's second sub-layer, queries come from the target and keys/values come from the encoder memory. The keys are source tokens, so the mask is src_mask:

class DecoderLayer(nn.Module):
    def forward(self, x, memory, src_mask, tgt_mask):
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))  # target ↔ target
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))  # target → source
        return self.sublayer[2](x, self.feed_forward)

The score tensor is now (batch, h, n_tgt, m_src), which is not square. The mask still broadcasts, because its query axis is 1. There are two kinds of padding here:

• Source padding consists of keys, so those columns are hidden.
• Target padding consists only of queries (rows), so no column is hidden. The padded rows attend normally to the real source tokens, and their outputs are discarded in the loss.

Padding meets the loss

The decoder is trained with teacher forcing. Its input tgt is the target without its last token, and the labels tgt_y are the target shifted left by one. Positions whose label is padding must contribute nothing. Batch counts the real labels:

class Batch:
    def __init__(self, src, tgt=None, pad=2):
        ...
        self.tgt = tgt[:, :-1]  # decoder input
        self.tgt_y = tgt[:, 1:]  # labels, shifted by one
        self.tgt_mask = self.make_std_mask(self.tgt, pad)
        self.ntokens = (self.tgt_y != pad).data.sum()  # real labels only

The criterion builds a smoothed target distribution, then removes padding twice. The <blank> class never receives probability mass. Every position whose label is <blank> gets an all-zero target row. With KLDivLoss(reduction="sum"), an all-zero row contributes exactly 0 loss and 0 gradient:

class LabelSmoothing(nn.Module):
    def forward(self, x, target):
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (self.size - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0  # never predict <blank>
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)  # drop <blank> positions
        return self.criterion(x, true_dist.clone().detach())  # KLDivLoss(reduction="sum")

class SimpleLossCompute:
    def __call__(self, x, y, norm):
        x = self.generator(x)
        sloss = self.criterion(x.contiguous().view(-1, x.size(-1)),
                               y.contiguous().view(-1)) / norm
        return sloss.data * norm, sloss

# in run_epoch
out = model.forward(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)
loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)

This completes the picture. Padding is masked where it would act as a source: key columns, in both self-attention and cross-attention. It is dropped where it would act as a reader: in the loss.

When is a causal mask needed?

A causal mask prevents exactly one thing: peeking at tokens that have not been generated yet. That can only happen when queries and keys lie on the same sequence that is being produced step by step. In the standard encoder–decoder, that is only the decoder's self-attention.

• Encoder self-attention. The source is fully given, so attention is bidirectional.
• Cross-attention. The whole source is available before decoding starts. Moreover, j ≤ i would compare positions in two different sentences that share no timeline.
• Variants. Streaming and simultaneous translation (e.g. wait-k) limit cross-attention to the source prefix received so far. The idea is the same, but the rule is not j ≤ i. Decoder-only LMs have only causal self-attention. Prefix-LMs are bidirectional over the prefix and causal over the continuation.

Summary

A mask stores only the axes its answer depends on. The causal rule is the only rule that depends on the query position, which is why the decoder mask is the only one with a full query axis.

Appendix: try it yourself

A self-contained script that reproduces every shape and number in this note. The helpers mirror the_annotated_transformer.py. I ran it with PyTorch 1.11, the version pinned by the repository.

"""Self-contained check of the mask shapes and loss masking used in the note.

The helper definitions mirror the_annotated_transformer.py.
"""
import math

import torch
import torch.nn as nn

torch.manual_seed(0)
torch.set_printoptions(precision=2, sci_mode=False, linewidth=120)
PAD, BOS, EOS = 2, 0, 1


def subsequent_mask(size):
    attn_shape = (1, size, size)
    return torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8) == 0


class Batch:
    def __init__(self, src, tgt, pad=PAD):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        self.tgt = tgt[:, :-1]
        self.tgt_y = tgt[:, 1:]
        tgt_mask = (self.tgt != pad).unsqueeze(-2)
        self.tgt_mask = tgt_mask & subsequent_mask(self.tgt.size(-1)).type_as(tgt_mask)
        self.ntokens = (self.tgt_y != pad).data.sum()


def attention(query, key, value, mask=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    return torch.matmul(p_attn, value), p_attn


class LabelSmoothing(nn.Module):
    def __init__(self, size, padding_idx, smoothing=0.0):
        super().__init__()
        self.criterion = nn.KLDivLoss(reduction="sum")
        self.padding_idx, self.size = padding_idx, size
        self.confidence, self.smoothing = 1.0 - smoothing, smoothing

    def forward(self, x, target):
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (self.size - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, true_dist.clone().detach())


# Sentence 0: 3 real source tokens + 2 pads; sentence 1: 4 real + 1 pad.
src = torch.tensor([[5, 7, 9, PAD, PAD],
                    [6, 8, 3, 4, PAD]])
# Targets: <s> w w w </s> <pad>   and   <s> w w </s> <pad> <pad>
tgt = torch.tensor([[BOS, 11, 12, 13, EOS, PAD],
                    [BOS, 14, 15, EOS, PAD, PAD]])
b = Batch(src, tgt)

print("== shapes")
print("src_mask           ", tuple(b.src_mask.shape))
print("tgt_mask           ", tuple(b.tgt_mask.shape))
print("src_mask.unsqueeze ", tuple(b.src_mask.unsqueeze(1).shape))
print("tgt_mask.unsqueeze ", tuple(b.tgt_mask.unsqueeze(1).shape))
print("ntokens            ", int(b.ntokens))

print("\n== decoder mask, sentence 1 (tgt = <s> 14 15 </s> <pad>)")
print(b.tgt_mask[1].int())

h, d_k = 8, 64
n, m = b.tgt.size(1), b.src.size(1)
q = torch.randn(2, h, n, d_k)  # from the decoder (target side)
kv = torch.randn(2, h, m, d_k)  # from the encoder memory (source side)

print("\n== cross-attention weights, sentence 0, head 0  (rows = target queries, cols = source keys)")
_, p_cross = attention(q, kv, kv, mask=b.src_mask.unsqueeze(1))
print(p_cross[0, 0])
print("row sums:", p_cross[0, 0].sum(-1))

print("\n== cross-attention weights, sentence 1, head 0  (last row is a <pad> query)")
print(p_cross[1, 0])

print("\n== loss: LabelSmoothing ignores <pad> targets")
V = 16
crit = LabelSmoothing(size=V, padding_idx=PAD, smoothing=0.1)
logits = torch.randn(2, n, V).log_softmax(-1)
y = b.tgt_y
print("tgt_y:\n", y)
loss = crit(logits.view(-1, V), y.reshape(-1))
rows = crit.true_dist.sum(-1).view(2, n)
print("true_dist row sums (0 = ignored):\n", rows)
print("loss / ntokens = %.4f" % (loss / b.ntokens).item())

== shapes
src_mask            (2, 1, 5)
tgt_mask            (2, 5, 5)
src_mask.unsqueeze  (2, 1, 1, 5)
tgt_mask.unsqueeze  (2, 1, 5, 5)
ntokens             7

== decoder mask, sentence 1 (tgt = <s> 14 15 </s> <pad>)
tensor([[1, 0, 0, 0, 0],
        [1, 1, 0, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 1, 0],
        [1, 1, 1, 1, 0]], dtype=torch.int32)

== cross-attention weights, sentence 0, head 0  (rows = target queries, cols = source keys)
tensor([[0.67, 0.11, 0.22, 0.00, 0.00],
        [0.08, 0.41, 0.51, 0.00, 0.00],
        [0.72, 0.06, 0.23, 0.00, 0.00],
        [0.20, 0.74, 0.06, 0.00, 0.00],
        [0.50, 0.26, 0.23, 0.00, 0.00]])
row sums: tensor([1.00, 1.00, 1.00, 1.00, 1.00])

== cross-attention weights, sentence 1, head 0  (last row is a <pad> query)
tensor([[0.10, 0.22, 0.11, 0.57, 0.00],
        [0.54, 0.05, 0.10, 0.32, 0.00],
        [0.06, 0.12, 0.61, 0.20, 0.00],
        [0.21, 0.09, 0.28, 0.41, 0.00],
        [0.26, 0.33, 0.32, 0.09, 0.00]])

== loss: LabelSmoothing ignores <pad> targets
tgt_y:
 tensor([[11, 12, 13,  1,  2],
        [14, 15,  1,  2,  2]])
true_dist row sums (0 = ignored):
 tensor([[1.00, 1.00, 1.00, 1.00, 0.00],
        [1.00, 1.00, 1.00, 0.00, 0.00]])
loss / ntokens = 2.5445

Code excerpts are from harvardnlp/annotated-transformer, lightly trimmed. Shape comments are mine.