A padded batch can run through an attention layer without an error and still produce misleading results. The tensor shapes look correct, but placeholder positions may receive attention or keep nonzero output unless the masks match PyTorch’s rules.

Use nn.MultiheadAttention with batch_first=True. Pass a Boolean key_padding_mask shaped (batch, key_length), with True at every padded key that queries must not read. Set average_attn_weights=False to inspect each head separately, then clear padded query rows because a key-padding mask blocks padded keys, not padded queries.

This tutorial makes those rules visible with two short equipment-alert sequences. We will use PyTorch’s attention layer rather than rebuild the calculation, inspect executed weights from two heads, and add checks that catch a reversed or incomplete mask. If the score calculation itself is new to you, read scaled dot-product attention in NumPy first.

Before the Layer: Setup and Teaching Data

You need Python 3.11 or newer and basic experience with Python lists and arrays. You do not need to know Transformer architecture or train a model.

Create and activate a virtual environment, then install the three packages used by the executed lesson. On macOS or Linux, run:

python -m venv .venv
source .venv/bin/activate
python -m pip install torch pandas matplotlib

On Windows PowerShell, use .venv\Scripts\Activate.ps1 for activation. PyTorch packages differ by operating system and accelerator, so use the official PyTorch installation selector when your computer needs a platform-specific command.

The published artifacts were executed with Python 3.13.2, PyTorch 2.12.0, pandas 3.0.3, NumPy 2.5.1, and Matplotlib 3.11.0.

Download equipment_alert_vectors.csv and save it beside your Python file. It is an original synthetic dataset released under CC0 1.0. The alert text and feature values were created for this lesson; they do not represent real equipment data or vectors from a trained model.

Load the file and inspect the first sequence. Each row is one token position, and <pad> fills unused positions so both sequences have the same length:

import pandas as pd
import torch
from torch import nn

alerts = pd.read_csv("equipment_alert_vectors.csv")
print(alerts[alerts["sequence_id"] == "alert_a"].to_string(index=False))

sequence_id  position    token  is_padding  asset_feature  signal_feature  alert_feature  action_feature
    alert_a         0     pump       False            2.0             0.1            0.0             0.0
    alert_a         1 pressure       False            1.7             1.4            0.2             0.0
    alert_a         2   rising       False            0.3             2.0            2.0             0.4
    alert_a         3  inspect       False            0.1             0.3            1.5             2.0
    alert_a         4    <pad>        True            0.0             0.0            0.0             0.0
    alert_a         5    <pad>        True            0.0             0.0            0.0             0.0

An embedding is a numeric vector that represents one token. Models usually learn token embeddings and attention projections from data. Here, each token already has four named teaching features so that the two heads are easy to compare. The first pair describes asset and signal information; the second pair describes alert and action information.

How Multi-Head Attention Works Across Four Axes

Self-attention means the same sequence supplies the queries, keys, and values. A query asks what information one position needs, a key is used to score a possible source position, and a value carries the information collected from that position.

Multi-head attention runs several attention calculations in parallel. Learned matrices first project the inputs into queries, keys, and values. PyTorch then splits each projection across the heads, lets each head produce its own attention pattern, joins the head results, and applies one output projection.

For this lesson, the important shapes are:

input and output:       (batch, sequence, embedding)
key padding mask:       (batch, key sequence)
per-head weights:       (batch, heads, query sequence, key sequence)

Our batch contains two sequences, each stored in six positions with four features. The input shape is therefore (2, 6, 4). We choose two heads, so PyTorch splits each four-feature query, key, and value projection into two features per head. embed_dim must be divisible by num_heads; the current nn.MultiheadAttention documentation defines the head width as embed_dim // num_heads.

Rows and columns have different jobs inside the weight matrix. A row is a query asking where to read. A column is a key that may be read. This difference explains why a key-padding mask hides columns but does not remove query rows.

Build the Batch and Its Padding Mask

First, place the four feature columns into a batch-first tensor. Build a second Boolean tensor from is_padding; in nn.MultiheadAttention, True means that the key position must be ignored:

import numpy as np

feature_columns = [
    "asset_feature",
    "signal_feature",
    "alert_feature",
    "action_feature",
]

groups = [
    group.sort_values("position")
    for _, group in alerts.groupby("sequence_id")
]
vectors = torch.from_numpy(np.stack([
    group[feature_columns].to_numpy(np.float32)
    for group in groups
]))
key_padding_mask = torch.from_numpy(np.stack([
    group["is_padding"].to_numpy(bool)
    for group in groups
]))

print("input shape:", tuple(vectors.shape))
print("padding mask shape:", tuple(key_padding_mask.shape))
print(key_padding_mask)

input shape: (2, 6, 4)
padding mask shape: (2, 6)
tensor([[False, False, False, False,  True,  True],
        [False, False, False, False, False,  True]])

The first sequence has four real tokens followed by two padding positions. The second has five real tokens followed by one padding position. The mask has no feature axis because it answers one question for each key position: should attention ignore this key?

Do not infer this mask from feature values. A real token can have a valid all-zero embedding, and a padded vector might become nonzero after another layer. Carry sequence lengths or a padding indicator from the data pipeline instead.

Configure Two Heads You Can Inspect

A new attention layer starts with random projection weights. That is appropriate for a layer that will be trained, but random heads would make this shape lesson harder to inspect. We will set the query, key, value, and output projections to identity matrices. The projected vectors will therefore equal the input vectors: head 0 will use the first two features, and head 1 will use the final two.

Create the layer with batch_first=True so inputs follow the (batch, sequence, feature) order used above:

attention = nn.MultiheadAttention(
    embed_dim=4,
    num_heads=2,
    dropout=0.0,
    bias=False,
    batch_first=True,
)

identity = torch.eye(4)
with torch.no_grad():
    attention.in_proj_weight.zero_()
    attention.in_proj_weight[0:4].copy_(identity)    # query projection
    attention.in_proj_weight[4:8].copy_(identity)    # key projection
    attention.in_proj_weight[8:12].copy_(identity)   # value projection
    attention.out_proj.weight.copy_(identity)

attention.eval()

in_proj_weight stores three projection matrices in one tensor: query first, key second, and value third. This direct overwrite is only a teaching device. In a trainable model, leave the parameters alone and let optimization learn them. You can still use the same masking, shape checks, and weight inspection after training.

Dropout is zero here so every run has the same weights. eval() also puts the module in evaluation mode. If you later train a surrounding model, call train() during training and eval() during validation or inference.

Run Attention and Keep Each Head Separate

Pass vectors three times because this is self-attention. Request weights and turn off head averaging so the returned tensor keeps a separate head axis:

with torch.inference_mode():
    output, head_weights = attention(
        vectors,
        vectors,
        vectors,
        key_padding_mask=key_padding_mask,
        need_weights=True,
        average_attn_weights=False,
    )

print("output shape:", tuple(output.shape))
print("per-head weight shape:", tuple(head_weights.shape))

output shape: (2, 6, 4)
per-head weight shape: (2, 2, 6, 6)

The output keeps the same batch size, sequence length, and four-feature width as the input. The weight tensor has two extra sequence axes: six query rows by six key columns for each of two heads in each batch item.

The average_attn_weights option matters only when need_weights=True. Its default value, True, would return shape (2, 6, 6) and hide differences between heads. For a production path that does not need weights, the PyTorch documentation recommends need_weights=False so that the layer can use its best optimized attention path.

Read One Query Across Both Heads

The word rising is at query position 2 in the first sequence. Print its full key row for each head:

rising_weights = head_weights[0, :, 2, :]

print("head 0:", rising_weights[0].numpy().round(3))
print("head 1:", rising_weights[1].numpy().round(3))
print("rising output:", output[0, 2].numpy().round(3))

head 0: [0.055 0.327 0.568 0.049 0.    0.   ]
head 1: [0.028 0.037 0.527 0.408 0.    0.   ]
rising output: [0.843 1.615 1.674 1.028]

Read the first four numbers in each row; the final two belong to padding. Head 0 uses asset and signal features. It gives rising 56.8% and pressure 32.7% of its weight. Head 1 uses alert and action features. It gives rising 52.7% and inspect 40.8%.

The two heads are not choosing labels. Each one creates a different weighted mixture of its value features. PyTorch joins the two mixtures and passes the result through the output projection, which leaves it unchanged in this lesson because that projection is an identity matrix.

Each chart row is one query and each column is one key. For example, follow the rising row horizontally to compare where that query reads. Compare the two panels only after fixing the row: head 0 and head 1 use different feature pairs, so their weight patterns need not match.

Attention weights show how this layer mixed its value vectors. They do not prove what a whole model understands, and they are not a complete explanation of a prediction.

Verify That Padding Cannot Be Read

A zero vector is not automatically invisible to softmax. A zero score becomes exp(0) = 1 inside softmax, so it can receive positive weight. To see the problem, run the same layer without the key-padding mask and total the two padding columns for the pump query:

with torch.inference_mode():
    _, unmasked_weights = attention(
        vectors,
        vectors,
        vectors,
        need_weights=True,
        average_attn_weights=False,
    )

pump_padding_share = unmasked_weights[0, :, 0, 4:].sum(dim=-1)
print(pump_padding_share.numpy().round(3))

[0.058 0.333]

Without the mask, padding receives 5.8% of head 0’s weight and 33.3% of head 1’s weight for this query. The pump query is all zeros in head 1, so its dot product with every key is zero. All six keys tie, and softmax gives each one-sixth of the weight. Two padded keys therefore collect two-sixths, or one-third.

Now test the masked result across the complete batch. The first check selects every weight whose key column is padded. The second confirms that each query row still totals 1:

masked_key_weights = head_weights.masked_select(
    key_padding_mask[:, None, None, :]
)
row_sum_error = (head_weights.sum(dim=-1) - 1.0).abs().max()

print("largest masked-key weight:", masked_key_weights.abs().max().item())
print("largest row-sum error:", row_sum_error.item())

largest masked-key weight: 0.0
largest row-sum error: 1.1920928955078125e-07

Every padded key receives exactly zero weight. Every query row sums to 1 within about 0.00000012; the tiny difference is normal floating-point rounding. These numeric checks are stronger than looking at one heatmap because they cover both heads, both sequences, and every query.

Clear Padded Query Outputs Separately

The key-padding mask has done its job, but it has not removed padded query rows. PyTorch defines key_padding_mask as positions in key to ignore. Our self-attention call also uses padded positions as queries, so the layer still creates outputs for those rows by reading the valid keys.

Measure the largest absolute value at any padded query position, then clear those rows explicitly:

padded_query_values = output.masked_select(
    key_padding_mask.unsqueeze(-1)
)
clean_output = output.masked_fill(
    key_padding_mask.unsqueeze(-1),
    0.0,
)
clean_padded_values = clean_output.masked_select(
    key_padding_mask.unsqueeze(-1)
)

print("before cleanup:", padded_query_values.abs().max().item())
print("after cleanup:", clean_padded_values.abs().max().item())

before cleanup: 1.024999976158142
after cleanup: 0.0

This is not a defect in the key mask. Keys are sources; queries are destinations. The mask prevents every query from reading padded sources. The final masked_fill enforces a separate rule: downstream layers should not use outputs whose destination position is padding.

Keep both rules when a residual connection or feed-forward block follows attention. You can reapply the valid-position mask after a block, or use a later operation that already ignores padded rows. What matters is that padded query outputs do not silently enter a loss, pooling operation, or metric.

Troubleshooting and Quick Checks

The mask meaning is reversed. For a Boolean nn.MultiheadAttention key-padding mask, True means “ignore this key.” Print one sequence beside one mask row before calling the layer.

The input axes are swapped. With batch_first=True, use (batch, sequence, embedding). Without it, the layer expects sequence first. A shape can be accepted yet represent the wrong axes when batch size and sequence length happen to match.

embed_dim is not divisible by num_heads. Each head receives an equal-width slice. For embed_dim=4 and num_heads=2, each head has width 2. Choose compatible values before building the layer.

Head weights disappear. need_weights=True is not enough. Also pass average_attn_weights=False, or PyTorch averages the head axis before returning weights.

Padding gets nonzero weight. Confirm that the mask shape is (batch, key_length), its dtype is Boolean, and padded locations are True. Apply it in the attention call, not after softmax.

Padded output rows remain nonzero. This is expected from a key-padding mask. Clear padded query rows or ensure every downstream operation ignores them.

All weights become nan. A query cannot form an attention distribution when every key is masked, and some execution paths may return nan values. Ensure every sequence has at least one real key. Reject empty sequences or define a separate representation for them.

The two mask APIs seem inconsistent. PyTorch mask meaning depends on the API. In nn.MultiheadAttention, Boolean key_padding_mask=True ignores a key. Do not copy a Boolean mask into a lower-level attention function without checking its current documentation. For a separate order constraint, see causal self-attention in NumPy; causal masks block future keys, while padding masks block placeholders.

Inspection code slows inference. Returning full weight matrices uses memory and may prevent the fastest path. Inspect weights during development, but use need_weights=False when production code does not consume them.

Recap: Split, Mask, Inspect, Clear

Multi-head attention becomes easier to debug when you keep four actions separate:

1. Split: choose an embedding width divisible by the head count, and know the width each head receives.
2. Mask: pass (batch, key_length) with True at padded keys.
3. Inspect: request unaveraged weights when you need shape and mask evidence.
4. Clear: remove outputs at padded query positions before downstream code can use them.

The executed batch confirmed all four points. PyTorch returned output shape (2, 6, 4) and per-head weight shape (2, 2, 6, 6). Every masked key had zero weight, every row summed to 1 within floating-point precision, and separate cleanup reduced the largest padded-query output from about 1.025 to zero.

For your next step, replace the four teaching features with embeddings from your model and remove the identity-weight override. Keep the same assertions around shapes, masked key weights, row sums, and padded outputs. Those checks test the data path even after the heads become learned and much harder to inspect by eye.