Sentence Pooling

The problem, from first principles

A transformer outputs one vector per token, but classification and retrieval need one vector per sentence. Three standard ways to pool: take the vector of the special first token ([CLS]), which the model was trained to summarize the input; take the mean over the tokens; or take the elementwise max. In a batch, sequences are padded, and padding positions hold meaningless vectors, so mean and max must use only the real tokens as marked by the attention mask. A mean that includes padding is silently biased toward zero and changes with how much padding the batch happens to have.

From theory to code

Implement pool_sentence.

Constraints

• hidden has shape (B, T, d) and mask has shape (B, T) with 1 for real tokens and 0 for padding. Every row has at least one real token, and real tokens come first (right padding).
• pool_sentence(hidden, mask, mode) returns (B, d). 'cls': hidden[:, 0]. 'mean': the average over positions where mask == 1. 'max': the elementwise maximum over positions where mask == 1. Any other mode raises ValueError.

Hints