# Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs

arXiv: https://arxiv.org/abs/2506.09215 · HTML: https://mlsanity.com/paper/2506.09215/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 26 of its 45 references that are in the map.
- First posted on arXiv 10 Jun 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [A simple neural network module for relational reasoning](https://mlsanity.com/paper/1706.01427/) (2017)
- [Wayformer: Motion Forecasting via Simple & Efficient Attention Networks](https://mlsanity.com/paper/2207.05844/) (2022)
- [Actor-Attention-Critic for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1810.02912/) (2018)
- [Dense Associative Memory for Pattern Recognition](https://mlsanity.com/paper/1606.01164/) (2016)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Language-Driven Representation Learning for Robotics](https://mlsanity.com/paper/2302.12766/) (2023)
- [GameFormer: Game-theoretic Modeling and Learning of Transformer-based Interactive Prediction and Planning for Autonomous Driving](https://mlsanity.com/paper/2303.05760/) (2023)
- [Learning the Best Pooling Strategy for Visual Semantic Embedding](https://mlsanity.com/paper/2011.04305/) (2020)
- [AdaPool: Exponential Adaptive Pooling for Information-Retaining Downsampling](https://mlsanity.com/paper/2111.00772/) (2021)
- [Scalable Multi-Agent Reinforcement Learning through Intelligent Information Aggregation](https://mlsanity.com/paper/2211.02127/) (2022)
- [Memorization to Generalization: Emergence of Diffusion Models from Associative Memory](https://mlsanity.com/paper/2505.21777/) (2025)
- [Scalable Reinforcement Learning Policies for Multi-Agent Control](https://mlsanity.com/paper/2011.08055/) (2020)
- [Boosting Sample Efficiency and Generalization in Multi-agent Reinforcement Learning via Equivariance](https://mlsanity.com/paper/2410.02581/) (2024)
- [Dense Associative Memory Through the Lens of Random Features](https://mlsanity.com/paper/2410.24153/) (2024)
- [Approximation of relation functions and attention mechanisms](https://mlsanity.com/paper/2402.08856/) (2024)
- [CA-Stream: Attention-based pooling for interpretable image recognition](https://mlsanity.com/paper/2404.14996/) (2024)

## Built on by (most cited first, 3 of 3)

- [AdaJudge: Adaptive Multi-Perspective Judging for Reward Modeling](https://mlsanity.com/paper/2601.08097/) (2026)
- [Towards Improved Sentence Representations using Token Graphs](https://mlsanity.com/paper/2603.03389/) (2026)
- [The Hidden States Cookbook: A Large-Scale Ablation Study for Noise-Robust Conversational Intent Classification in Industry](https://mlsanity.com/paper/2610.05394/) (2026)

## Abstract (arXiv)

We investigate the design of pooling methods used to summarize the outputs of transformer embedding models, primarily motivated by reinforcement learning and vision applications. This work considers problems where a subset of the input vectors contains requisite information for a downstream task (signal) while the rest are distractors (noise). By framing pooling as vector quantization with the goal of minimizing signal loss, we demonstrate that the standard methods used to aggregate transformer outputs, AvgPool, MaxPool, and ClsToken, are vulnerable to performance collapse as the signal-to-noise ratio (SNR) of inputs fluctuates. We then show that an attention-based adaptive pooling method can approximate the signal-optimal vector quantizer within derived error bounds for any SNR. Our theoretical results are first validated by supervised experiments on a synthetic dataset designed to isolate the SNR problem, then generalized to standard relational reasoning, multi-agent reinforcement learning, and vision benchmarks with noisy observations, where transformers with adaptive pooling display superior robustness across tasks.