PyTorch / Training state and architecture interfaces / Lesson 11

Attention masks and batch layout

Read each PyTorch attention API’s shape and mask polarity instead of transferring assumptions by name.

About 5 min readContent reviewed 2026-10-28

On this page

Lesson goals

Distinguish MultiheadAttention module layout from direct scaled-dot-product layout.
Translate Boolean blocked and allowed masks deliberately.
Control dropout and inspect valid attention dependencies.

The module and the functional interface have different contracts

nn.MultiheadAttention projects inputs into query, key, and value spaces, combines heads, and applies an output projection. Its batch_first=True sequence interface uses (N,L,E), while the direct scaled_dot_product_attention function exposes a head dimension such as (N,H,L,E). Matching embedding width and head count requires the documented dimensions and divisibility. A batch-first flag does not remove the head axis from the functional operation. The fixture uses one head and explicitly assigned identity projections so both calculations represent the same mathematical function; arbitrary learned projections would not justify that comparison.

Boolean mask polarity is API-specific

For the module’s Boolean attn_mask, True marks a query-key position that is not allowed to attend. Its key_padding_mask likewise marks excluded key positions under its own shape contract. In direct scaled_dot_product_attention, a Boolean True marks a position that participates. The illustrative program therefore passes blocked to the module and its logical inverse to the functional call. A float mask uses additive score semantics instead of the Boolean rule. Verify mask dtype, shape, and broadcast scope, and follow documented type requirements when combining masks. A variable called mask is not enough to determine what True means.

python Illustrative CPU PyTorch 2.14 API program; torch unavailable, expected output derived

import torch
from torch import nn
from torch.nn import functional as F
attention = nn.MultiheadAttention(2,1,dropout=0.0,bias=False,batch_first=True,dtype=torch.float64,device="cpu")
with torch.no_grad():
    attention.in_proj_weight.copy_(torch.cat([torch.eye(2,dtype=torch.float64,device="cpu")]*3))
    attention.out_proj.weight.copy_(torch.eye(2,dtype=torch.float64,device="cpu"))
x = torch.tensor([[[1.0,0.0],[0.0,1.0],[1.0,1.0]]],dtype=torch.float64,device="cpu")
blocked = torch.ones(3,3,dtype=torch.bool,device="cpu").triu(diagonal=1)
output, weights = attention(x,x,x,attn_mask=blocked,need_weights=True,average_attn_weights=False)
q = x.unsqueeze(1)
direct = F.scaled_dot_product_attention(q,q,q,attn_mask=~blocked,dropout_p=0.0)
print("weight shape", tuple(weights.shape))
print("first output", output[0,0].tolist())
print("future weights", weights[0,0,0,1:].tolist())
print("same after mask conversion", torch.allclose(output,direct.squeeze(1),atol=1e-12,rtol=1e-12))

Expected output

weight shape (1, 1, 3, 3)
first output [1.0, 0.0]
future weights [0.0, 0.0]
same after mask conversion True

Identity projections make the module and direct attention mathematically comparable; their Boolean masks have opposite meanings. No attention kernel, projection, or allclose check was executed in torch.

Inspect the permitted dependency set

The upper-triangular blocked mask excludes future keys for each query. The first query can attend only to the first value, so the first output is [1,0] and both future weights are zero. Returned per-head weights have shape (1,1,3,3) because averaging across heads is disabled. The companion independently enumerates allowed keys and computes one weighted context. A padding key mask does not automatically define what to do with invalid query outputs or exclude them from a loss. An all-excluded query needs a supported explicit policy rather than assuming every implementation returns useful zeros.

python Executed deterministic standard-library companion

from math import exp, sqrt
rows = [[1.0,0.0],[0.0,1.0],[1.0,1.0]]
blocked = [[j>i for j in range(3)] for i in range(3)]
allowed = [[not value for value in row] for row in blocked]
keys = [[j for j,ok in enumerate(row) if ok] for row in allowed]
scores = [sum(a*b for a,b in zip(rows[0],key))/sqrt(2) for key in rows]
maximum = max(score for score,ok in zip(scores,allowed[0]) if ok)
raw = [exp(score-maximum) if ok else 0.0 for score,ok in zip(scores,allowed[0])]
weights = [value/sum(raw) for value in raw]
context = [sum(weight*row[j] for weight,row in zip(weights,rows)) for j in range(2)]
print("allowed key indices", keys)
print("first context", context)

Expected output

allowed key indices [[0], [0, 1], [0, 1, 2]]
first context [1.0, 0.0]

The executed companion verifies the causal allowed set and one weighted context, not PyTorch mask broadcasting or implementation equivalence.

Functional dropout does not read module mode automatically

The documented direct attention function applies dropout according to dropout_p, so an enclosing module’s eval state does not by itself change that argument. The fixture explicitly passes zero and uses no stochastic mask. A wrapper that wants mode-dependent dropout must pass an appropriate value from its own contract. Kernel choice, precision, and floating-point operation order can affect numerical results, so compare under justified tolerances rather than promising bitwise equality across implementations. The 2.14 documentation identifies the functional API’s change-sensitive status; check the pinned contract when migrating code to another version.

Validate attention as a component before a larger model

Test a tiny sequence with identifiable values, a nontrivial causal boundary, padding, and unequal query and source lengths where supported. Check whether batch and sequence axes were confused and whether a context includes unavailable future information. Attention weights show an internal combination, not automatically a causal explanation of the final prediction. The framework program is illustrative because torch is absent; the companion executes list and scalar arithmetic only. No MultiheadAttention module, direct attention kernel, optimized path, gradient, model, or device-performance measurement was tested during review.

Common errors and fixes

A converted attention mask exposes future positions

Why it happens: Boolean polarity from one API was reused unchanged in another.
How to fix it: Map the allowed dependency set to each documented mask convention.
Prevent it next time: Assert excluded future weights on a tiny identifiable sequence.

Evaluation remains stochastic inside a custom attention wrapper

Why it happens: The direct dropout_p argument remained positive despite module.eval.
How to fix it: Set the functional dropout argument according to the intended mode.
Prevent it next time: Test the wrapper’s actual argument and repeated behavior in a supported runtime.

Attention mixes different observations

Why it happens: The batch or sequence axis was interpreted under the wrong layout.
How to fix it: Restore the exact module or functional input shape contract.
Prevent it next time: Use distinct batch identities and unequal axis lengths in a fixture.

Check your understanding

1. Why is ~blocked passed to direct attention instead of blocked?

Show worked solution

The module’s Boolean True means excluded, while direct scaled-dot-product attention uses True for participation. Inverting this explicitly constructed Boolean blocked mask preserves the same allowed query-key pairs. This conversion does not imply that arbitrary float masks or differently shaped padding masks can be inverted without further interpretation.

2. Why is the first output [1,0] regardless of the later values in this fixture?