wenmin-wu/nlp-mean-pooling-with-mask

Mean Pooling with Attention Mask

Overview

Instead of using the CLS token or naive mean, mask-aware mean pooling averages only the real token embeddings by weighting with the attention mask. This prevents padding tokens from diluting the representation, especially important for variable-length inputs.

Quick Start

import torch
import torch.nn as nn

class MeanPooling(nn.Module):
    def forward(self, last_hidden_state, attention_mask):
        mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
        sum_embeddings = torch.sum(last_hidden_state * mask, dim=1)
        sum_mask = torch.clamp(mask.sum(dim=1), min=1e-9)
        return sum_embeddings / sum_mask

Workflow

1. Get last_hidden_state from transformer (shape: B x S x H)
2. Expand attention mask to match hidden state dimensions
3. Multiply hidden states by mask, sum along sequence dimension
4. Divide by sum of mask values (count of real tokens)

Key Decisions

• Clamp min=1e-9: Prevents division by zero for empty sequences
• vs CLS token: Mean pooling captures more context; CLS can be undertrained
• vs max pooling: Mean is smoother; max captures salient features. Concatenate both for best results
• Normalization: Optionally L2-normalize output for cosine similarity tasks

References

• Feedback Prize - English Language Learning (Kaggle)
• Source: fb3-deberta-v3-base-baseline-train