wenmin-wu/nlp-learned-attention-pooling

Learned Attention Pooling

Overview

Mean pooling treats every token equally, which dilutes signal from key tokens. A lightweight trainable attention layer learns which tokens matter most for the task. The network computes a scalar weight per token via Linear → Tanh → Linear → Softmax, then produces a weighted sum of hidden states. Adds minimal parameters (~500K) but can significantly improve regression and classification tasks.

Quick Start

import torch.nn as nn

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

    def forward(self, hidden_states, attention_mask=None):
        # hidden_states: (batch, seq_len, hidden_size)
        weights = self.attention(hidden_states)  # (batch, seq_len, 1)
        if attention_mask is not None:
            weights = weights * attention_mask.unsqueeze(-1)
            weights = weights / (weights.sum(dim=1, keepdim=True) + 1e-8)
        return (weights * hidden_states).sum(dim=1)  # (batch, hidden_size)

class Model(nn.Module):
    def __init__(self, model_name):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        h = self.encoder.config.hidden_size
        self.pool = AttentionPooling(h)
        self.head = nn.Linear(h, 1)

    def forward(self, **inputs):
        hidden = self.encoder(**inputs).last_hidden_state
        pooled = self.pool(hidden, inputs.get('attention_mask'))
        return self.head(pooled)

Workflow

1. Replace mean/CLS pooling with AttentionPooling module
2. Initialize attention weights with Xavier/He init
3. Apply attention mask to zero out padding tokens before softmax
4. Use the weighted sum as input to the task head

Key Decisions

• Attention hidden size: 512 is standard; smaller (128) for small models, larger for XL models
• Mask handling: always mask padding tokens — without masking, attention leaks to pad positions
• Init: use nn.init.xavier_uniform_ on linear layers for stable training
• vs multi-head: this is single-head; use nlp-attention-head-pooling for multi-head variant

References

• PPPM / Deberta-v3-large baseline w/ W&B [train]
• PPPM / Deberta-v3-large baseline [inference]