Ptn Masked Mean Pooling

This is the little pooling node that makes transformer classifiers work. After a transformer encoder you've got a 3D tensor - (batch, seq, feature) - and for classification you need to squeeze it to (batch, feature). PtnMaskedMeanPooling does that average, but it uses an attention mask to exclude padding tokens so the mean isn't dragged toward zero by all the [PAD] positions. It's the "mean pooling, but ignoring padding" step the pack's build-a-transformer-from-scratch guide calls for, and it has no parameters at all - you just drop it in.

How it works

The node has no inputs in its schema - you add it to the graph and it emits a PTMODEL that expects to receive a (inputs, mask) tuple from the chaining node in front of it. In the forward pass it takes the inputs (rank 3, batch × seq × features), multiplies by the expanded mask, sums over the sequence axis, and divides by the number of valid tokens per sample (clamped to avoid div-by-zero). That's textbook masked mean pooling, and it's the exact same operation the HF classification nodes run internally when use_mean_pooling is on - this just exposes it as a standalone building block.

The mask convention matters: 1 means a valid token, 0 means padding. That's the inverse of PyTorch's key_padding_mask convention (where True masks out), so keep it straight when you're feeding it from the dataset nodes, which produce masks in this 1=valid form.

Wiring it up

In the mha.json transformer workflow, the chained model that includes the attention layer hands its output plus the mask to this pooling node, which then feeds the final linear classification head. Because it needs the mask, it has to sit in a chain that carries the mask along - the Ptn Chained Model With Attention Mask node is the vehicle that makes sure it gets one. If you see a "mask is not rank 2" error, the mask isn't reaching it, and that's a chaining problem, not a pooling problem.

Installing

Same as every node in the pack. ComfyUI Manager → search "Pt-Wrapper", or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Restart ComfyUI; the pack's requirements install on first launch.

Where people get burned

Two things. First, forgetting that this node requires the mask - connect it to something that only passes the hidden states and you'll hit a rank error. Second, the 1=valid/0=padding convention: if your mask comes from a source that treats 1 as padding (like a reversed key_padding_mask), your "average" silently becomes garbage instead of erroring. Check the mask values once and you'll never be surprised again.