wenmin-wu/nlp-attention-head-pooling

Attention Head Pooling

Overview

Instead of CLS token or mean pooling, add a learnable attention layer that scores each token position and computes a weighted sum. This lets the model focus on the most informative tokens in the sequence — e.g., complex vocabulary in readability tasks or key entities in classification.

Quick Start

import torch
import torch.nn as nn

class AttentionPooling(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, hidden_states, attention_mask=None):
        # hidden_states: (B, S, H)
        weights = self.attention(hidden_states).squeeze(-1)  # (B, S)
        if attention_mask is not None:
            weights = weights.masked_fill(attention_mask == 0, -1e9)
        weights = torch.softmax(weights, dim=1)  # (B, S)
        return torch.sum(hidden_states * weights.unsqueeze(-1), dim=1)  # (B, H)

Workflow

1. Get full sequence hidden states from transformer (not just CLS)
2. Apply learned attention layer to score each token
3. Mask padding positions before softmax
4. Compute weighted sum as the sequence representation
5. Feed into regression/classification head

Key Decisions

• vs mean pooling: Attention pooling learns which tokens matter; mean treats all equally
• vs CLS: CLS depends on pre-training; attention adapts to your task
• Architecture: Tanh + linear is standard; single linear also works
• Combine: Concatenate attention-pooled + mean-pooled for best results

References

• CommonLit Readability Prize (Kaggle)
• Source: commonlit-two-models