Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs

TL;DR

This paper introduces an attention-based adaptive pooling method for transformer outputs that maintains robust performance across varying signal-to-noise ratios, outperforming standard pooling techniques in noisy environments.

Contribution

The work proposes a novel adaptive pooling approach that approximates signal-optimal vector quantization, improving transformer robustness to noise in diverse applications.

Findings

Adaptive pooling outperforms AvgPool, MaxPool, and ClsToken in noisy conditions.

Theoretical bounds show the method approximates optimal vector quantization.

Experimental results confirm robustness across vision, reinforcement learning, and reasoning tasks.

Abstract

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.