Research map / Architectures and efficiency / Attention and transformers
Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives closed-form expressivity bounds for Transformers with Average, Sum, Last-token and Max pooling, and empirically compares pooling strategies across vision, NLP, and time-series tasks with frozen backbones.
Problem
Pooling, the step that aggregates token embeddings into a fixed-size vector, is widely used in Transformers but has little theoretical analysis; its effect on expressivity and downstream performance is unclear.
Approach
Model a single-layer Transformer plus pooling as a Lipschitz map and bound its (epsilon, sigma, gamma)-expressivity per pooling type, where the pooling operator's norm (1/sqrt(n) for average, sqrt(n) for sum, 1 for last-token, sqrt(min(n,d)) for max) sets contractive, non-expansive or expansive behavior; then validate empirically with frozen backbones.
Key results
- Average pooling outperforms other fixed pooling methods on inpainting, segmentation and MiniPlaces, while last-token (CLS) pooling yields the best classification accuracy for ViT-small on most classification datasets (e.g., 90.35% on CIFAR-10).
- On TinyStories next-token prediction, last-token pooling reaches about 90.6% top-10 accuracy on Llama3-8B versus roughly 56% for average pooling, consistent with the expansive-vs-contractive theory.
- Weighted Average pooling matches or exceeds the best fixed pooling on larger NLP models such as Mistral-7B and Llama3-8B, and attention pooling is strongest on several time-series imputation and classification tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Average rather than Last-token pooling | ViT-small · MiniPlaces | accuracy: 54.51 ± 0.38 -> 56.94 ± 0.66 | Helped |
| use Average rather than Last-token pooling | ViT-small · CUB-200-2011 | accuracy: 77.45 ± 0.13 -> 71.48 ± 0.18 | Hurt |
| use Average rather than Last-token pooling | BERT · STSB | Spearman: 0.587 ± 0.009 -> 0.713 ± 0.008 | Helped |
| use Last-token rather than Average pooling | Mistral-7B · TinyStories | Top-10 Acc: 61.047 ± 0.228 -> 89.948 ± 0.092 | Helped |
| use Attention rather than Average pooling | MOMENT-small · Traffic imputation | MSE: 1.583 ± 0.016 -> 0.265 ± 0.015 | Helped |
| use Sum rather than Average pooling | MOMENT-base · classification | accuracy: 65.19 ± 0.00 -> 83.30 ± 1.51 | Mixed |
| learn scalar weights over token positions instead of using fixed pooling | Transformer models · downstream tasks across vision and NLP | task-specific performance | Mixed |
| use Attention pooling rather than fixed pooling | MOMENT models · time-series classification | accuracy | Helped |
| use Last-token rather than global pooling | Transformer models · next-token prediction | Top-10 Acc | Helped |
| use Average rather than Sum pooling | Transformer models · global-context vision tasks | task performance | Helped |
Surprising findings
- Attention pooling often underperforms in vision and NLP despite its extra flexibility, attributed to added parameterization complexity.
- Sum pooling sometimes beats average pooling on time-series classification, attributed to larger norms yielding stronger gradients.
Limitations
- Evaluation is limited to frozen backbones; joint end-to-end training of pooling and backbone is not studied.
- The expressivity bounds are necessary but not sufficient for downstream performance.
- Theory analyzes a single-layer Transformer; extension to multi-layer models is argued by composition rather than proven.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet-100 CUB-200-2011 MiniPlaces CelebA Oxford Flowers-102 Oxford-IIIT Pet PascalVOC STSB HellaSwag TinyStories
Benchmarks: CIFAR-10 CIFAR-100 ImageNet-100 CUB-200-2011 MiniPlaces CelebA Oxford Flowers-102 Oxford-IIIT Pet PascalVOC STSB HellaSwag TinyStories
Metrics: accuracy MSE Spearman correlation F1 top-10 accuracy
Baselines: Last-token pooling Average pooling Sum pooling Max pooling Attention pooling Weighted Average pooling
Builds on: L2-MHA LipsFormer Swin Transformer MOMENT
Closest papers by meaning
- Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs (ICML 2025 · similarity 0.87)
- Talking Heads: Understanding Inter-Layer Communication in Transformer Language Models (NeurIPS 2024 · similarity 0.87)
- Effective Context in Transformers: An Analysis of Fragmentation and Tokenization (NeurIPS 2026 · similarity 0.87)
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.87)
- TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters (ICLR 2025 · similarity 0.87)
- Core Context Aware Transformers for Long Context Language Modeling (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.