Practical NLP for Risk Modeling Supplement - Mean Pooling With an Attention Mask
supplement to practical NLP for risk modeling part I - turning text into embeddings with pretrained transformers
Python
Machine Learning
Transformers
Published
September 7, 2026
In Part I of the Practical NLP for Risk Modeling series, we used mean pooling to turn the token-level output of a pre-trained transformer into a single vector representing an a single text sample. In this supplement, we’ll walkthrough how any padding tokens added to a sample are excluded from the mean pooling operation.
where \(h_t\) is the transformer embedding for token \(t\), and \(m_t\) is the corresponding attention-mask value: 1 for a real token and 0 for padding.
It is easiest to see why masking matters when a batch contains texts of different lengths. We’ll use three short examples and ask the tokenizer to pad every sequence to the length of the longest one.
import torchimport pandas as pdfrom transformers import AutoTokenizer, AutoModelmodel_name ="distilbert-base-uncased"tokenizer = AutoTokenizer.from_pretrained(model_name)model = AutoModel.from_pretrained(model_name)texts = ["The price I have paid, to live and to learn","It's too late","The windows are filled with frost",]enc_input = tokenizer( texts, return_tensors="pt", padding=True, truncation=True,)with torch.no_grad(): out = model(**enc_input)input_ids = enc_input["input_ids"]attention_mask = enc_input["attention_mask"]last_hidden_state = out.last_hidden_stateprint(f'input_ids.shape : {input_ids.shape}')print(f'attention_mask.shape : {attention_mask.shape}')print(f'last_hidden_state : {last_hidden_state.shape}')
c:\Users\jtriv\miniforge3\envs\gnlp\Lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
Loading weights: 100%|██████████| 100/100 [00:00<00:00, 457.12it/s, Materializing param=transformer.layer.5.sa_layer_norm.weight]
DistilBertModel LOAD REPORT from: distilbert-base-uncased
Key | Status | |
------------------------+------------+--+-
vocab_projector.bias | UNEXPECTED | |
vocab_layer_norm.bias | UNEXPECTED | |
vocab_transform.weight | UNEXPECTED | |
vocab_transform.bias | UNEXPECTED | |
vocab_layer_norm.weight | UNEXPECTED | |
Notes:
- UNEXPECTED :can be ignored when loading from different task/architecture; not ok if you expect identical arch.
For this batch, input_ids and attention_mask both have shape (B, T), and the transformer output last_hidden_state has shape (B, T, H), where:
B is the number of texts in the batch.
T is the padded sequence length.
H is the transformer embedding dimension.
For our example, B = 3, T = 13 and H = 768.
At this point we do not yet have one embedding per text. We have one embedding per token. The objective of mean pooling is to take the one embedding per token and collapse it to one embedding per text.
The tokenizer pads the shorter examples so that all three sequences have the same length. The attention mask records which positions contain actual tokens. A value of 1 means keep this position. A value of 0 means this position is padding.
Each entry in the DataFrame above is shown as token (attention_mask).
Conceptually, the second example looks like this:
tokens: [CLS] it ' s too late [SEP] [PAD] [PAD] ...
mask: 1 1 1 1 1 1 1 0 0 ...
|--------------- include ----------------| |----- ignore -- ... --|
The model receives the attention mask as part of its forward pass, which prevents padded positions from being treated as meaningful during self-attention. However, the model can still return non-zero hidden-state vectors at padded positions. That is why we use the same mask again when we pool the token embeddings.
H_out = last_hidden_stateM = attention_maskB, T, H = H_out.shapeprint(f"H_out.shape : {tuple(H_out.shape)}")print(f"M.shape : {tuple(M.shape)}")print(f"Valid tokens per text : {M.sum(dim=1).tolist()}")print(f"\nM:\n{M}")
The three texts contain different numbers of valid tokens even though they all occupy the same padded length T. Notice that the counts include special tokens such as [CLS] and [SEP]. The pooling described here includes DistilBERT’s special tokens ([CLS] and [SEP]) because their attention mask values are 1. If a particular embedding recipe calls for averaging only the ordinary text tokens, those special token positions can also be excluded by modifying the pooling mask.
H_out has shape (B, T, H), while the attention mask M has shape (B, T). We add a trailing dimension to the mask so PyTorch can broadcast it across all H embedding dimensions.
H_out * M_expanded represents element-by-element multiplication, not matrix multiplication.
For a real token, the mask value is 1:
\[
1 \times h_t = h_t
\]
For a padded token, the mask value is 0:
\[
0 \times h_t = \mathbf{0}
\]
So every padded 768-dimensional token vector becomes a vector of zeros.
padding_vectors = H_masked[M ==0]print(f"Number of padded token positions: {padding_vectors.shape[0]}")print(f"All padded embeddings are zero : {torch.all(padding_vectors ==0).item()}")
Number of padded token positions: 11
All padded embeddings are zero : True
11 is the total number of padded token positions across the entire batch of 3, which is what we expect.
Now we sum across the token dimension, T. Since the padded vectors have already been zeroed out, they contribute nothing to the sum.
S = H_masked.sum(dim=1) # (B, H)print(f"S shape: {tuple(S.shape)}")
S shape: (3, 768)
We’ve summed over the sequence dimension. We now have one 768-dimensional vector per text. To get the mean, each row needs to be divided by the number of valid tokens in that text.