Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Authors: Greyson Brothers
Venue: ICML 2025 spotlightposter (2025)
Keywords: Pooling
Attention
Vector Quantization
AdaPool
Transformer
Robustness
Primary area: deep_learning->attention_mechanisms
TL;DR: Existing methods for aggregating transformer embeddings (MaxPool, AvgPool, ClsToken) are vulnerable to noisy inputs, we show that an attention-based pooling mechanism is provably robust across the spectrum of input signal-to-noise ratios.

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.

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Greyson Brothers
Abstract
Pooling
AdaPool
We investigate the design of pooling methods used to summarize the outputs of transformer embed- ding models, primarily motivated by reinforce- ment learning and vision applications. This work considers problems where a subset of the input vectors contains requisite information for down- stream task (signal) while the rest distractors (noise). By framing pooling vector quantiza- tion with the goal of minimizing signal loss, we demonstrate that the standard methods used to ag- gregate transformer outputs, AvgPool, MaxPool, and ClsToken, are vulnerable to performance col- lapse as the signal-to-noke to-soise ratio (SNR) of inputs fluctuates. We then show that attention-based adaptive pooling method signal-optimal vector quantizer within derived er- ror 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.
Transformer
Transformer
Figure 1: Given mix of signal (blue) and noise (red) in- puts, standard pooling methods result in unwanted interfer- ence in the aggregate representation of transformer outputs. AdaPool learns to adaptively attenuate different ratios of noise and obtain cleaner signal.
approximate the
Complex reinforcement learning (RL) environments pose similar challenges, especially those seeking to bridge the gap with the real world. Modern approaches have adopted relational architectures, like transformers (Vaswani et al., 2017) and graph neural networks (GNNs) (Battaglia et al., 2018), to process their observations (Zambaldi et al., 2019; Hu et al., 2021; Nayakanti et al., 2023; Huang et al., 2023; Nayak et al., 2023). By breaking each observation into vector-based tokens representing entities, memories, multi-modal sensory streams, these networks can achieve the following benefits over standard feedforward policy ar- chitectures: (1) variable-sized observation spaces enabling scalability (Iqbal & Sha, 2019; Hsu al., 2021), (2) sample efficiency via permutation invariance and/or equivariance over the tokens (depending on architectural choices) (Liu et al., 2020; McClellan al., 2024; Hao al., 2023), and (3) parameter efficiency by learning pairwise relations between tokens explicitly rather than implicitly via redundant feed- forward subnetworks (Santoro et al., 2017). In light of the benefits relational architectures over their feedforward counterparts, we examine a critical implementation detail for transformers that is often overlooked.
1. Introduction
Autonomous systems, both artificial and biological, require diverse and redundant sensors to capture a reliable picture of their environment. Limited processing resources and com- petitive environments pressure biological systems to process their abundant sensory information efficiently. To satisfy this constraint, only the subset of information relevant to the current task is retained while the rest is considered noise and attenuated the process of selective attention (Broadbent, 1958; Driver, 2001).
'Johns Hopkins University Applied Physics Laboratory, Maryland, USA. Correspondence to: Greyson Brothers wrwpsoobbroottrsssuuual....
Each inference of a standard transformer encoder produces as many output embeddings as input vectors. Under the sequence-to-sequene training paradigm that they were de- signed for, the next token in the sequence is an obvious target for each output embedding. However, in domains such as computer vision RI,, deciding how to map a set
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
of output embeddings to a class action distribution is less straightforward. The goal is discard irrelevant embed- dings while condensing the most salient information from those remaining into a representation that can be used for downstream tasks. Notably, this mirrors the broader aim of pooling in general (Boureau et al., 2010), Indeed, for non-sequential transformer applications, this usually ac- complished via average pooling (Karamcheti et al., 2023: Zhai et al., 2022; Vinyals al., 2019), max pooling (Zam- baldi et al., 2019; Chen et al., 2021), or the use of a learned class token (Devlin et al., 2019; Dosovitskiy et al., 2021; He et al., 2022). However, many prior works treat this step as an arbitrary design choice without theoretical justifications.
Networks (Hopfield, 1982), are models that store sets of memory vectors. They utilize an update rule analogous to pooling method that summarizes and returns information from the set of memories according to their relationship with a retrieval cue. Such networks have a storage capacity under
which the pooling method is guaranteed to return one of the memory vectors and beyond which may return a spu- rious mixture of interfering memories (Pham et al.. 2024). This capacity is generally exceeded when the number of memories far exceeds the dimensionality of each vector. A growing line of work on Dense Asociativve Memories (Kro- tov & Hopfield, 2016; Demircigil et al., 2017; Hoover et al., 2024) has shown the design of this update/pooling function has significant ramifications on the number of memories that can be pooled without destructive interference. By introduc- ing increasingly sharp non-linear activation functions into the update rule, Dense AMs are able to achieve a memory capacity that is exponential with respect to the dimension of the memories. Ramsauer et al. (2021) showed that an expo- nential form of this update rule is equivalent to the attention mechanism, going as far as to propose a standalone Hopfield Pooling layer utilizing attention. Taking inspiration from attention- based recall methods, our goal is to extract the most task-relevant information from set of abundant and diverse sensory embeddings while minimizing destructive interference, as conveyed in Figure 1.
In this paper, we derive a framework to evaluate the perfor- mance of differentiable vector pooling methods on inputs composed of both "signal" and "noise" vectors. Our analysis reveals that popular approaches like average and max pool- ing can suffer catastrophic performance drops when their inductive biases are misaligned with the noise level. Moti- vated by this finding, we explore an attention- pooling (AdaPool) mechanism that dynamically weights rel- evant vectors during aggregation, mitigating interference from distractors. We then show that AvgPool, MaxPool, and ClsToken are special cases of AdaPool. Crucially, we prove that AdaPool can approximate the signal-optimal vec- tor quantizer for any signal-to-noise ratio under explicit error bounds derived from the data distribution. We validate these theoretical insights through synthetic supervised exper- iments and then demonstrate their practical impact on multi- agent reinforcement learning (MPE), relational reasoning (Box World), and vision (CIFAR) benchmarks. Across all tasks, the adaptive pooling method consistently outperforms standard baselines and avoids their failure modes in the presence of significant noise.
adaptive
2.2. Attention-based Pooling
A number of works have recently utilized cross-attention with a single query vector as a pooling method. Research on memory augmented neural networks yielded the Neural Tur- ing Machine (Graves et al., 2014) and Diffcrentiable Neural Computer (Graves et al., 2016), both using attention to ex- tract information from a set of memory vectors, but without the recall guarantees afforded by Hopfield networks. These methods used the output of recurrent controller as the query vector. Ilse et al. (2018) 8) use attention based pooling with a learned query for multi-instance learning. Stergiou & Poppe (2023) apply attention- pooling to computer vision problems, terming the technique adaptive pooling (AdaPool), which we adopt. They use the centroid of the input set as the query for pixel-level pooling. Several works utilize learned query for vision-based transformers, which is similar to the ClsToken approach but only introduces the learned query in a final aggregation layer (Touvron et al., 2021; Przewięzlikowski et al., 2024; Torres et al., 2024). Fi- nally, Perceiver (Jaegle et al., 2021), Set Transformer (Lee et al., 2019), and Universal Physics Transformers (Alkin et al., 2024) use multiple learned queries as inducing points to reduce the cardinality of an input set for more efficient self-attention blocks. We further discuss the implementation details of adaptive pooling in Section 3.5, and motivate a novel choice of query based on our analysis.
The contributions of this work can be summarized as fol- lows:
1. We provide a theoretical framework for analyzing the robustness of vector pooling methods to noisy inputs.
2. We show that attention can robustly pool observations across the full spectrum of noisy inputs, and derive error bounds on its ability to optimally retain signal.
3. We perform extensive experiments on supervised and RL benchmarks to corroborate our theoretical results.
2. Related Work
2. Interference and Associative Memories
The problem of noise robustness posed by this paper is most closely related to the problem of noisy recall in the context of associative memories (AMs). AMs, such as Hopfield
2

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
3. Methods
form:
N C(X) Σ W. W.
In this section, we formally define vector pooling methods and establish a theoretical framework under which they can be evaluated analytically for signal loss under various signal-to-noise regimes. Critically, we demonstrate that attention-based pooling can approximate signal-optimal vector quantizer for inputs with any noise ratio.
(I)
where @ indicates the Hadamard product with a weight vector W, = [Wt,1. wr,a whose elements Way E R. all elements of w, are identical, then the result is equivalent to using a scalar weight w/- We label this function C' it represents compressor.
3.1. Data and Noise
In the VQ literature, the information loss incurred by quan- tizer is referred to as quantization error or distortion. The standard metric used is mean squared error (MSE) between each point and its cluster's representation. Gray (1984) de- fines the optimal quantizer as that which minimizes MSE over the input set, yielding the centroid in the case of a single cluster. For our purposes, we care only about signal distortion; rolsterelated quantization error is of no conse- quence. Thus, we define related measure of information loss specifically for our use case:
We define our input domain as set of vectors containing sensory information. It is often helpful to imagine these vec- tors as points, particles, or entities. We represent these sets as matrices X E RNx, where N indicates the cardinality of the set and d reflects the dimensionality of each vector. In a multi-agent RL setting, each vector might contain the x-y positions and velocities of different entities, as well as relevant features including entity type, team, health, etc. In multi multi-modal perception example, the set of input vectors may be split between encoded image patches, encoded audio tokens, and embedded text tokens. Abstractly, each vector in the set represents a snapshot from the different sensory streams an agent uses to perceive its environment.
Definition 3.2. The Signal Loss IL of a vector pool C' on a noisy set X is the MSE between the compressed representa- tion C(X) = Xx and the subset of signal vectors X, S X.
For any given inference, some number l s N of those input vectors will contain information that is relevant to the learning task. We label the k task-relevant vectors as the signal subset X, and the rest as noise subset X.,. such that the input set X = X, U X,, and X, n nX,, = Ø. Under this framework, we define a signal-to-non to-se ratio SNR= = N-
L(X,xo) = Σ (x. - (x)) X,EX,
Corollary 3.3. The point x² that minimizes signal loss is the centroid of the signal subset (see proof A. We say that the global vector pool C that computes this point is
A vector x; belongs the signal subset when the learning target y a function of that vector. Formally, x4 E X, = By A and E X,, dx.
From Definition 3.2 and Corollary 3.3, it follows that the signal-optimal global vector pool assigns weights
By O. The following sections utilize this notation to lay the groundwork for our theoretical analysis of vector pooling methods.
if E X. if E X,
K
3.2. Vector Quantization
3.3. Global Average Pooling
We extend analytical tools developed for vector quantization (VQ) (Gray, 1984) and lossy compression to evaluate the theoretical noise-robustness of pooling methods. It is worth noting that vector quantization should not be confused with the recent weight quantization techniques used to reduce the memory footprint of language models (Jacob al., 2018). Instead, VQ is concerned with evenly dividing large set
Definition 3.4. Global Average Pooling is a vector quantizer of the following form:
N
Avol X)
Σ
X;
(2)
of points into discrete clusters for data compression. The compressed set uses a single vector representation for each cluster in place of the original points. well-known exam- ple is the k-means clustering algorithm (MacQueen, 1967). We frame global vector pooling methods as a degenerate case of vector quantization with single cluster.
where weights w; are are scalars given by
Wi
N
(3)
Corollary 3.5. AvgPool is signal-optimal if the input set contains no noise.
X,, x, = - AugPool = CC
Definition 3.1. A Global Vector Pool is any differentiable vector quantizer C(X) RNxd - R of the following
See proof A.2
3

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Corollary 3.6. When the input contains noise, X, + O, AvgPool is signal-optimal if and only if the centroid of X is equivalent to the centroid of Xny
functions and inner product spaces study the set of rela- tional functions that attention can approximate (T'ai et al., 2019; Altabaa & Lafferty, 2024). Following this line of work, we let x.) IR x R - IR be parametrizable, asymmetric kernel used to compute the relation between a query and a sensory vector. Given parametrizable weight matrices We, WK Rdxdd and a query vector X₂ E R we define the relation kernel as the scaled dot product
AugPool= -C- AugPool(X.) = AugPool(X)
See proof A.3
We note that these two cases are extremely limiting. Average pool is only signal-optimal in the unlikely cases that the entire input set is strictly composed of task-relevant vectors or the signal and noise vectors are identically distributed. When that is not the case, AvgPool tends to yield higher signal loss with each additional noise vector.
r r(xg= x,) (ibo (x,) (o0))))(()))) = WoWW Vd
For our purposes, the query X, is fixed for all x,, so we use the shorthand (x. = r; for notational convenience.
Definition 3.9. Given parametrizable weight matrices Wq, WK, Wy R and query vector X₂ E R Global Adaptive Pooling is vector quantizer of the following form:
3.4. Global Max Pooling
Definition 3.7. Global Max Pooling is a vector quantizer of the following form:
N
AdaPool (X) = Σ
wi x, W.W
(6)
N MaxPPoo(X) = Σ w; xx
(4)
where each weight wy is given by the softmax of relation ry
where the elements of the weight vector W; [wi,1, 1, .. Wi,a] are given by
exp (ri)
W/ Σ3 exp (15)
wi,d { if af Tz.d Vj
(5)
The choice of the query is thus an inductive bias that controls the nature of the relation kernel, which in turn affects the attenuation of the input set.
otherwise
Any elements whose values are the maximum along a fea- ture column of X are given a weight of one and all others a weight of zero. For notational simplicity, we make the assumption that a unique maximum exists. If a given col- umn has m elements that share the maximum value, then the MaxPool can still be obtained with this summation by giving each weight of
Corollary 3.10. AvgPool is special case of AdaPool. See proof A.6
Corollary 3.11. MaxPvol is special case of AdaPool. See proof A.7
Corollary 3.8. MaxPool is signal-optimal if and only if the input set X. contains single signal vector (k 1) taking the maximal value along each feature dimension.
M Mark Pool =C -
X,1 +
1, MaxPool(X) Xa
Dat Product
See proof A.4
& Softmax
Weighted Average
We observe that MaxPool acts complement to AvgPool, where signal loss tends to increase with each additional sig- nal vector added to the input set. They both have inductive biases that result in best performance at opposite ends of the signal-to-note noise spectrum and experience increasing signal loss when the SNR changes.
W K
W Q
W
Query Xq
X
Input Set
3.5. Global Adaptive Pooling
Adaptive pooling utilizes cross-attention with a single query to pool set of vectors. Prior works have utilized kernel
Figure 2: Attention as a vector pooling method (AdaPool).
4

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Next, construct set of error bounds on AdaPool's ap- proximation of the signal-optimal quantizer. The general intuition is that the output of the softmax should equally distribute weights amongst the signal vectors and give zero weight to the noise vectors, following from Corollary 3.3. Since softmax normalizes inputs by the sum of their expo- nentiated values, there needs to a meaningful margin between signal and noise relation scores to push the noise weights zero, and the signal scores need to be similar in magnitude such that one does not dominate the others. To formalize this intuition, we introduce the following notation to describe the neighborhood widths e for signal and noise:
and upper bounds respectively:
-1
Ln 1-== h k.e +1+ +1+(N- -N-
Un k. +1+ (NN - 6)
See proof A.5
Remark 3.13. Crucially, as the margin M increases and the signal and noise neighborhoods c, and ey shrink, the bounds squeeze the approximation error to zero. This is influenced by the data distribution, the expressiveness of the relation kernel, and any preprocessing of the input set, such as embedding via transformer.
Eg max (rx - min (r,} } 2 0
(8)
E = max(r,) min(r,) 2 0
We additionally define the minimum margin M and max distance D between signal and noise neighborhoods as
To make this likely in practice, we propose selecting the query from the signal subset X, X., Since the dot- product relation kernel computes notion of similarity, when the query is signal vector, the dot product with other signal vectors should tend to be higher than the dot product with noise vectors. We argue that using the cen- troid of the whole set as query (Stergiou & Poppe, 2023) is not robust, as it is influenced by the ratio of signal and noise vectors in given input. We also argue that selecting learned embedding (Lee et al., 2019) is not ideal, as is fixed regardless of changes to the distributions of signal and noise per sample. For example, a simple rotation of all signal and noise points about the origin could severely alter their dot products with a fixed learned embedding, while dot products with a signal vector would be invariant to such transformations.
M = min } - max(r,)
(9)
D max(r.} min(r,) = M + c, +
These values are illustrated in Figure 3 using an example set of relation scores for better intuition. These are used to bound the worst-case weights for any given set of inputs.
En
{rn)
M
(rs)
Figure 3: An example displaying the neighborhoods of signal (blue) and noise (red) relation values, along with the margin M between the two neighborhoods (ra} and {r,, } The axis represents the codomain of the relation kernel.
While choosing a known signal vector for every input might seem like nebulous task, provide numerous examples here to ground and inspire such choices. For entity-based RL problems, the network is controlling one of the N entities in its observation space. We use the transformer embedding of that entity's 's state token as the query since it always contains task-relevant information. If the inputs are memory vectors, the current state of the environment query vector. For vision tasks with structured images, like video games, one might choose a particular image patch to be the query, such as a patch covering a mini-map map or status indicators. For real-world vision tasks, a patch from the center of the image may be desirable, as will generally contain focal content coineiding with the gaze of the agent. We explore the consequences of such choices in Section 4.5,
Theorem 3.12.
For any signal-to-nois toise ratio, AdaPool can approximate the signal-optimal vector pool within an error bound determined by the distribution of signal and noise re- lations. Explicitly, for any signal vector x; E X,, AdaPool approximates the optimal weight of w;: = within error bounds L,, S " Wi S Us. which are the following lower and upper bounds respectively:
signal-rich
1
L,
1 1 (k- -k- e 1e) + (N (N- -k) -k)e--l
k
U,
1 (k - e
x + (N- - k) e- e
By using a query from the input set, we also have a natural way of preserving the residual stream by adding the output of AdaPool back to the query. This helps with gradient propagation, and we observe better empirical performance than without extending a skip connection to the pooling layer. Finally, for an analysis of time complexity, see B.1.
k
for any noise vector X; E X.y, AdaPool approximates the optimal weight of w. = 0 within the ervor bounds Tr S w.: - wi S Un. which are the following lower
5

---

We designed variety of experiments evaluate the- oretical findings. Using large synthetic dataset, we first construct supervised learning task to isolate the robustness of transformers with each pooling method across the full spectrum of noise levels. The same architectures are then applied to RI. tasks in the Multi-Particle Environment and BoxWorld, first mirroring the supervised task and then look- ing at performance in standard benchmark scenarios with increasing noise. Finally, we test the generalization of our analysis to real-world data with the CIFAR image classifi- cation dataset. Critically, we show that AdaPool displays superior robustness across noise levels on all baselines, vali- dating our theoretical results. Code is publicly available at Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
3.6. ClsToken
AdaPool
Avg MaxPool
Centroid Baseline
0.12
Clstoken
Target Baseline
Class tokens are common alternative the above pooling methods, introduced by BERT (2019) and frequently used in vision applications (Dosovitskiy et al., 2021; He et al., 2022). For this method, a learned parameter vector is ap- pended to the input set before being fed into the transformer encoder. The corresponding embedding is taken from the output and used for downstream tasks. If one ignores the final feedforward sublayer and discards the N other out- put embeddings produced by the transformer, the ClsToken and AdaPool methods differ only by the choice of query vector x ClsToken takes x, to be the output embedding of the learned parameter vector, whereas we construct AdaPool to select the output embedding of a signal vector.
0.10
Signal Loss
0.08
0.06
0.02
0.00
0.0
0.2
0.4
0.6
0.8
1.0
Signal-to-Noise Ratio
Figure 4: Signal Loss (MSE) on the KNN-Centroid Task (N=128, d=16). Lower is better, data is provided in Table 3.
4. Experiments
can explicitly control the signal-to-noise ratio of each obser- vation, allowing us to gauge robustness by measuring signal loss under each noise regime. As baselines, we display the signal loss incurred by naively predicting the centroid of the entire input set, as well as naively predicting the target vector used source the nearest neighbor subset.
For these experiments, we use 5-fold cross- voss-valldation to train a 1 12-layer transformer encoder capped by a pooling method. All models were on the order of 600k parame- ters with the same initial weight configuration for the base transformer. Additional hyperparameters are listed in the appendix C.2. We report the signal loss test set com- prised of 100k holdout samples in Figure 4, averaged across the model checkpoints with the lowest validation loss on each fold. Additional ablations are reported in B.5, varying the dimensionality and cardinality of the inputs as well as the dimensionality and depth of the networks.
4.1. Synthetic Dataset
For our first experiments, we generated a synthetic dataset with 1 million samples. Each sample is a set containing N 128 vectors with d 16 features per vector, repre- sented by 2D array. Each feature column is drawn from unique distribution, evenly split between randomly param- eterized exponential, gaussian, and uniform distributions. The ordering of distributions is shuffled for each sample to ensure diversity and prevent overfitting.
AdaPool exhibits the lowest signal loss and most consistent performance across noise regimes. In partieular, it exceeds all other approaches by an order of magnitude in the low SNR regimes (0.03-0.25). Predictably, MaxPool has its best relative performance in the lowest SNR regime and rapidly declines as signal increases. Conversely, AvgPool performs best for high SNR (0.5-1.0) and deteriorates predictably as signal becomes sparse. AvgPool also closely follows the naive centroid baseline, while MaxPool closely follows the naive target baseline until SNR of 1.0, at which point it improves slightly. The CosToken method performs slightly worse than AvgPool at each noise level. All methods exhibit their worst absolute performance on the highest noise KNN- task, with the exception of MaxPool. These findings directly support our theoretical analysis.
4.2. Noise Robustness Experiments
To assess our theoretical findings, we design a supervised experiment to mirror our analytical framework. Given a set of N' input vectors representing points in space, the network must predict the centroid of the k-Nearest Neighbors (the signal subset) to an arbitrarily chosen target point. The tar- get point is indicated to the network by adding learned embedding to it at inference time, similar to the use of learned positional embeddings used in language modeling tasks (Radford et al., 2018). As we vary < & < N. we
4.3. Multi-Agent Experiments
To generalize our findings to reinforcement learning, we use the standard ParticParticle Environment (MPE) benchmark.
6

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
0.0
ASsPez! Nachool
AePeel Maxtool
Adaho Cabsinn
-05-
Miofool
-1.0
Namtel
-20
-25
AepPeol
MooPtel
0.2
0.4
0.5
1.0
0.2
0.5
1.0
0.0
0.4
2.0
300
Traniin Steae
Los
Thanaa Steps
(a) 1-31-0
(b) 1-15-16
(c) 1-3-28
(d) 1-1-30
Figure 5: Performance on the simple centroid environment, with the number of particles labeled Soetts-Signl--Nooie The (b) 32 (c) 2 (d). Lines represent the mean value across seeds, with standard error
SNR decreases from 32 (a) 32
16
shaded and the Y-axis fixed across plots to show performance decay.
back to each agent. We use 10 seeds per method for each simple centroid experiment and 20 per method for simple tag. The policy architecture mirrors the architecture used in the previous experiments (4.2) with the addition of pair of 3-layer MLPs that map the pooled output to action logits and value predictions respectively. For each trial, all four meth- ods are run with the same seeds, base transformer weights, and environmental initial conditions. Training batches con- sist of 8192 samples drawn from 128-timestep episodes, with all experiments training for 4 million timesteps total. Additional training hyperparameters can be found in C.4.
Figure 6: Left, the simple centroid environment with an agent (green), signal entities (blue), and noise entities (red). Right, the simple tag environment with predator (red), prey (green), and obstacles (black).
4.3.1. 1. SIMPLE CENTROID
We first design custom MPE scenario to mirror the su- pervised KNN-Centroid task in multi-agent RL setting, We sample a set of signal agents that move according to a random policy, along with a set of noise agents whose positions are fixed, as shown on the left of Figure 6. We train single additional agent with a reward to minimize its distance to the centroid of the signal agents. We use a continuous action space, where the actions accelerations in the cardinal directions. Results are displayed in Figure 5,
The environment provides a simple 2D world with baseline scenarios and an API to implement custom scenarios. We utilize the baseline simple tag environment and implement custom simple centroid environment, both described below. While the default environment returns flat observation vector, we process the observation into entity state tokens, such that each entity is represented by an 8-dimensional state vector containing (x.y) position, (x.y) velocity, and a 4-dimensional learned embedding corresponding to the entity type: [self, predator, prey, obstacle]. These tokens are each passed through the same projection layer map them to the hidden dimension of the network. Expanding the size of the hidden dimension relative to the inputs alleviates the burden on the pooling method to effectively compress information, as found our ablation studies in B.5. For these experiments, we thus limit the hidden dimension to at most the size of the input tokens.
AdaPool consistently reaches the same level of performance across noise levels, while all other methods deteriorate as noise increases. MaxPool was the worst performer by a wide margin across all levels, even when signal was sparse. As in the synthetic experiment, we observe that AdaPool performs the best relative to other methods in the mid-low SNR regime 4 ) With all methods, we observe an intuitive 3z
correlation between sample efficiency and noise level; more noise requires more samples to reach the same performance.
We use the Proximal Policy Optimization (PPO) (Schulman et al., 2017) reinforcement learning algorithm in online, centraliiee--rainnngg decentralizzed-execuii (CTDE) setup: each agent is controlled by copy of the same policy net- work but observes and acts independently. The training batch compiles experience from all learning agents to up- date the weights of the single policy, which is then copied
4.3.2. SIMPLE TAG + NOISE
For this experiment, we examine the impact of increas- ing the volume and ratio of noise in the default simple tag benchmark provided by MPE. Simple tag is a predator-pre scenario, where predators receive +10 reward for colliding with prey and prey receive -10 reward for colliding with
7

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
AcePocl Okthan Anghos
Adahoo Cetheen sugeel
Adutcal
Ochoon Aajtico
Cartean
Naghool
52-
0.0
1.5
2.
2.5
4.5
1.0
1.5
2.0
2.5
3.0
1.5
2.0
2.5
4.0
3.0
3.5
2.0
2.5
4.
Thaeeng Stom
109
Thanin May
166
Thamne Meem
105
Thankng Meeas
l0s
(a) 1v3v0
(b) 1v3v4
(c) 1v3v12
(d) 1v3v28
Figure 7: Scaling the number of distractor obstacles (noise) while keeping reward dynamics fixed. Subplots are labeled Predators-Prey Prey--bstaclless and the Y-axis is shared across plots.
predators, and collidable obstacles are scattered about, as shown on the right of Figure 6, For this experiment, we turn off obstacle collisions such that they have no influence on the reward functtion or dynamics of the game. We use single predator and prey agents, and train with increas- ing numbers of obstacles acting as noise in the observation space. To enable objective evaluation and prevent compe- tition from introducing non-stationarity into the problem, we only train the predator and control the prey agents with a heuristic that randomly samples a direction to move in. Unlike the simple centroid experiment, we use a discrete action space here.
Box World is a vision-based relational reasoning task intro- duced by Zambaldi et al. (2019). It involves picking up keys to unlock a sequence of boxes in the correct order to reach goal gem. There are also distractor paths that lead to dead ends, requiring the agent to reason over which key to pick up or which box unlock next. Keys, Locks, and Boxes are represented by pixels on a 2D grid, as depicted in Figure 10. We train on an instance with a goal sequence of length 2, 2 distractors, and 7x7 grid under two observation regimes. One presents tokens containing the normalized RGB color values and relative pixel coordinates only for the entities (max 8 tokens), while the other presents a token for each pixel plus the key currently held by the agent (50 tokens). Like 4.3.2, the underlying learning dynamics are the same, but the pixel observations contain a significantly higher vol- ume and ratio of noise. We use 5 seeds per method and train for 40 million steps. Additional training parameters can be found the appendix C.6.
As shown in Figure 7, the final mean reward of all methods declines significantly as both the volume and ratio of noise increase, indicating major decrease in sample efficiency within the fixed training budget of 4 million timesteps. Since the reward dynamics are identical across all scenarios and the training time fixed, this performance decrease can directly linked to the difficulty of learning from observa- tions with sparse signal. Predictably, AvgPool suffers the worst from increase in noise, with the final mean reward dropping 77.4 17.4% Cis'Token drops a similar 70.4%, MaxPool falls 60.7%, and AdaPool drops 50.9%. AdaPool attains the highest mean reward and lowest performance decline across all noise levels.
In Figure 8, we observe that MaxPool is able achieve su- perior sample efficiency and performance the entity-level observations, but then collapses to the worst performance under pixel-level observations. As in the previous experi- ment, we observe the smallest decline from AdaPool, which achieves superior performance in the high noise regime. AvgPool and ClsToken perform comparatively worse in both regimes. Regarding the initial high performance of MaxPool, the goal gem is always represented by white pixel, giving it a normalized color value of all ones the upper bound of the observation space. Corollary 3.8 im- plies that MaxPool is uniquely suited for problems of this form, hence the superior performance in 10a. However, the non-entity pixels (empty space) are a light grey color, relatively close in numerical representation white. This likely exacerbated the performance drop, despite MaxPool's helpful inductive bias. Overall, this demonstrates that our theoretical framework can be used an analytical tool for neural network designers to map their problem domains to the best pooling methods for the job, rather than relying on intuition or trial and error.
4.4. BorWd
10
Cbiftie
Aoyivll
Taikiig Drom
tor
(a) Entity-Based Obs
(b) Pixel-Based Ohs
Figure& 8: Box BorWo performance on entities only (8) vs all
pixels (50). The Y-axis is shared across plots.
8

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
4.5. CIFAR
METHOD
CIFAR-10
CIFAR-
Prior experiments carefully controlled for signal and noise, but left the following unanswered: (1) Does our analysis hold up on real-world data where signal and noise are less clearly defined, and where individual vectors may contain a mix of both? (2) How do you employ AdaPool when the choice of query is not obvious? To answer these questions, we conducted additional studies on image classification us- ing the CIFAR 10 and 100 benchmark datasets. We adopted the Vision Transformer (ViT) approach (Dosovitskiy al., 2021), partitioning the 32x32 pixel RGB images into 64 separate 4x4 pixel patches as shown in Figure 9. These are then flattened, projected, and fed into the transformer for embedding. Standard ViT implementations use ClsToken to pool output embeddings, with a minority favoring AvgPool.
ClsToken
84.52 1:02//
55.56 1:.13
AvgPool
87.15 10.35
59.63 1.023
MaxPool
87.65 1:0.7
60.55 1:028
Ada-Focal
87.98 1:042
61.22 1:0.3
Ada-Mean
87.84 1:030
61.23 1:0.20
Ada-Corner
87.00 1:030
57.08 1:0.31
Table 1: VIT Top-1 test accuracy the CIFAR image classification dataset using different pooling methods.
Additionally, the standard ClsToken approach underper- forms all other methods by a significant margin. AvgPool and MaxPool perform relatively well, with MaxPool having the edge in both experiments. In Figure 4, we previously ob- served that this occurred in the very low SNR regime (SNR <0.1). If that trend generalizes beyond our synthetic dataset, then it is quantitative indication that, after embedding. CIFAR samples tend to be signal-sparse.
One of our major theoretical results was that the selection of query critical for approximating the optimal quantizer. This could be a practical limitation, depending on AdaPool's sensitivity to non-ideal queries. A key takeaway from this experiment is that, while AdaPool is indeed sensitive to poor choices like the Corner query, the Mean query is strong default option that can outperform the standard pooling methods, even if the data is more noise than signal.
Figure 9: Splitting the CIFAR images into patches. We ex- periment with different choices of query for AdaPool, with the Corner query using the red patch embedding, the Focal query averaging the four central yellow patch embeddings, and the Mean query averaging all patch embeddings.
5. Conclusion
Most samples in the CIFAR dataset tend contain the object being classified in the center of the image. While pe- riphery patches may include some useful information, like grass or water, we hypothesize that the central patches are more likely to contain discriminating signal on average than any other patches. We thus examine 3 choices of query: Cor- ner, the embedding of the upper-left-nne lestmoostt periphery patch; Focal, averaging the embeddings of the four central patches; and Mean, using the average of all patch embeddings as the query. These query choices are also highlighted in Figure 9,
In this work, we investigated the selection and design of global pooling methods for aggregating embeddings pro- duced by transformers. We drew connections between pool- ing, vector quantization, and associative memory to reframe pooling as lossy compression problem rather than a trivial operation for aligning dimensions. This reframing helped formalize the limitations of common approaches like Avg- Pool and MaxPool, and we showed that ittention-based adaptive pooling, a niche approach, approximates the opti- mal compressor within derived bounds. These theoretical findings were first evaluated with carefully designed super- vised experiments, and then in more general reinforcement learning and vision tasks using the Multi-Particle Environ- ment, Box World, and CIFAR. The results confirmed that AvgPool and MaxPool fail in predictable ways when the hignal-to-noise ratio of inputs changes, validating the theo- retical analysis. Additionally, ClsToken had similar noise sensitivity to AvgPool, with less predictable relative perfor- mance between tasks. Crucially, we found that AdaPool was predictably robust to different ratios and quantities of noise, resulting superior performance across tasks.
We report the Top-1 accuracy scores on the holdout test set in Table 1, averaged from the best models across 5 folds after 300 training epochs each; see additional details in C.7. We observe that AdaPool Focal and Mean queries out- perform all other methods in both cases, reaching similar accuracies. The Corner query underperforms in both experi- ments, aligning with the prediction that a noisy query should lead to worse performance. By evaluating the efficacy of different queries, one could discover how signal tends to be distributed. This evidence supports the hypothesis that signal tends to be more concentrated in the central patches.

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Acknowledgements
Huge Storage Capacity. Journal of Statistical Physics, 168()))288-22))) 299, July 2017.
We thank Mark Fleischer, Joshua McClellan, Zachary God- dard, John Winder, and Jovanna Aragon for their helpful feedback.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of Deep Bidirectional Transformers for Lan- guage Understanding. North American Chapter of the Association for Computational Linguistics, 2019.
Impact Statement
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR, 2021.
The underlying motivation of this paper is to enhance the design of neural network architectures used for controlling autonomous systems. We focus on formalizing and im- proving one particular aspect, vector pooling, which commonly used and broadly applicable technique in many machine learning domains. In terms of specific impacts, there is rich body of discussion on the very real potential benefits and harms of intelligent autonomous systems, none of which we feel merits particular discussion in this work.
Driver, A selective review of selective attention research from the past century. British Journal of Psychology, 92 (1):53-78, 2001.
Graves, A., Wayne, G., and Danihelka, I. Neural Turing Machines, December 2014. arXiv:1410.5401 [es].
References
Alkin, B., Fürst, A., Schmid, S. I.., Gruber, L.., Holzleitner, M., and Brandstetter, Universal Physics Transformers: Framework For Efficiently Scaling Neural Operators. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, 2024.
Graves, A., Wayne, G., Reynolds, M. Harley, T., Dani- A., Colmenarejo, S. G., Grefenstette, E., Ramalho, T., Agapiou, J., Badia, A. P., Hermann, K. M., Zwols, Y., Ostrovski, G., Cain, A., King, H., Summerfield, C., Blunsom, P., Kavukcuoglu, K., and Hassabis, D. Hybrid computing using neural network with dynamic external memory. Nature, 538(7626):471- 476, October 2016. Publisher: Nature Publishing Group.
helka, I.,
Altabaa, A. and Lafferty, J. Approximation of rela- tion functions and attention mechanisms, June 2024. 088855200200886 [es].
Gray, R. Vector quantization. JEEE ASSP Magazine, I (2):4-29, April 1984. Conference Name: IEEE ASSP Magazine.
Battaglia, P. W., Hamrick, B., Bapst, V., Sanchez- Gonzalez, A., Zambaldi, V. F., Malinowski, M., Tac- chetti, A., Raposo, D., Santoro, A., Faulkner, R., Çaglar Gülçehre, Song, H. F., Ballard, A. J., Gilmer, J., Dahl, G. E., Vaswani, A., Allen, K. R., Nash, C., Langston, V., Dyer, C., Heess, N. M. O., Wierstra, D., Kohli, P., Botvinick, M. M., Vinyals, O., Li, Y., and Pascanu, R. Relational inductive biases, deep learning, and graph net- works. ArXiv, abs/1 Abs//880.011 266,,0221, 2018.
Hao, J., Hao, X., Mao, H., Wang, W., Yang, Y., Li, D., Zheng, Y., and Wang, Z. Boosting Multiagent Reinforce- ment Learning via Permutation Invariant and Permutation Equivariant Networks. In The Eleventh International Conference on Learning Representations, 2023.
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, Masked Autoencoders Are Scalable Vision Learners. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15979-15988, 2022.
Boureau, Y.-.., Ponce, J., and LeCun, Y. A theoretical analysis of feature pooling in visual recogmition. In Pro- ceedings of the 27th International Conference on Inter- national Conference on Machine Learning, pp. 111-118, Madison, WI, USA, 2010. Omnipress.
Hoover, B., Chau, D. H., Strobelt, H., Ram, P., and Krotov, D. Dense Associative Memory Through the Lens of Ran- dom Features. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Broadbent, D. E. Perception and communication. Pergamon Press, Elmsford, NY, US, 1958.
Chen, J., Hu, H., Wu, H., Jiang, Y., and Wang, C. Learning the Best Pooling Strategy for Visual Semantic Embed- ding. In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15784-15793. June 2021.
Hopfield, J. Neural networks and physical systems with
emergent collective computational abilities. Proceed- ings of the National Academy of Sciences, 79:2554-2558, April 1982. Publisher: Proceedings of the National Academy of Sciences.
Demircigil, M., Heusel, J., Löwe, M., Upgang, S., and Vermet, F. On a Model of Associative Memory with
Hsu, C. D., Jeong, H., Pappas, G. J., and Chaudhari, P. Scal- able Reinforcement Learning Policies for Multi-Agent
10

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Control. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 4795-4791, September 2021.
Liu, I.J.,, Yeh, R. A., and Schwing, A. G. PIC: Permuta- tion Invariant Critic for Multi-Agent Deep Reinforcement Learning. In Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learn- ing Research, pp. 590-602. PMLR, 30 Oct-01 Nov 2020.
Hu, S., Zhu, F., Chang, X., and Liang, X. UPDeT: Universal Multi-agent RL via Policy Decoupling with Transformers. In International Conference on Learning Representations, 2021.
MacQueen, J. Some methods for classification and analy- sis of multivariate observations. In Proceedings of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, Volume 1: Statistics, volume 5.1, pp. 281- 298. University of California Press, January 1967.
Huang, Z., Liu, H., and Lv, C. GameFormer: Game theoretic Modeling and Learning of Transformer-basee Interactive Prediction and Planning for Autonomous Driv- ing. In Proceedings of the IEEE/CVE International Con- ference on Computer Vision (ICCV). pp. 3903-3913, Oc- tober 2023.
McClellan, J., Haghani, N., Winder, J., Huang. F,, and Tokekar, P. Boosting Sample Efficiency and Generaliza- tion in Multi-agent Reinforcement Learning via Equivari- ance. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Ilse, M., Tomezak, J., and Welling, M. Attention-based deep multiple instance learning. In Dy, and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Ma- chine Learning Research, pp. 2127-2136. PMLR, 10-15 Jul 2018.
Nayak, S., Choi, K., Ding, W., Dolan, S., Gopalakrishnan, K., and Balakrishnan, H. Scalable Multi-Agent Reinforce- ment Learning through Intelligent Information Aggrega- tion. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 25817-25833. PMLR, 23-29 Jul 2023.
Iqbal, S. and Sha, F. Actor-Attention--Ce Critic for Multi-Agent Reinforcement Learning. In Proceedings of the 36th International Conference on Machine Learning, pp. 2961- 2970. PMLR, May 2019.
Nayakanti, N., AI-Rfou, R., Zhou, A., Goel, K., Refaat, K. S., and Sapp, B. Wayformer: Motion Forecasting via Simple & Efficient Attention Networks. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 2980-2987, May 2023.
Jacob, B., Kligys, S., Chen, B., Zhu, M., Tang, M., Howard, A., Adam, H., and Kalenichenko, D. Quantization and Training of Neural Networks for Efficient Integer- Arithmetic-Only Inference. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2018
Pham, B., Raya, G., Negri, M., Zaki, M. J., Ambrogioni, L., and Krotov, D. Memorization to Generalization: The Emergence of Diffusion Models from Associative Mem- ory. In NeurlPS 2024 Workshop on Scientific Methods for Understanding Deep Learning, 2024.
Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, J. Perceiver: General Perception with Iterative Attention. In Proceedings of the 38th Interna- tional Conference on Machine Learning, pp. 4651-4664.
Przewięžlikowkki, M., Balestriero, R., Jasinssii, W., Smieja, M., and Zieliński, B. Beyond [cls]: Exploring the true po- tential of masked image modeling representations, 2024.
PMLR, July 2021. ISSN: 2640-3498.
Karamcheti, S., Nair, S., Chen, A., Kollar, T., Finn, C., Sadigh, D., and Liang, P. Language-Driver Driven Represen- tation Learning for Robotics. Robotics: Science and Systems XIX. Robotics: Science and Systems Foundation, July 2023.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre- training. OpenAl, 2018
Krotov, D. and Hopfield, J. Dense Associative Memory for Pattern Recognition. In Advances in Neural Informa- fion Processing Systems, volume 29. Curran Associates, Inc., 2016.
Ramsauer, H., Schäff, B., Lehner, J.. Seidl, P., Widrich, M., Gruber, L., Holzleitner, M., Adler, T., Kreil, D., Kopp, M. K., Klambauer, G., Brandstetter, J., and Hochreiter, S. Hopfield Networks is All You Need. In International Conference on Learning Representations, 2021.
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y. W. Set Transformer: A Framework for Attention-basee Iommutation-Invariinn Neural Networks. In Proceedings of the 36th International Conference on Machine Learn- ing, volume 97 of Proceedings of Machine Learning Re- search, pp. 3744-377777 3353. PMLR, 09-15 Jun 2019.
Santoro, A., Raposo, D., Barrett, D. G., Malinowski, M., Pascanu, R., Battaglia, P., and Lillicrap, T. A simple neu- ral network module for relational reasoning. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal Policy Optimization Algorithms, August 2017. HcXiv:1707.06344 [es].
Computer Vision and Pattern Recognition (CVPR), pp. 1204-12133 June 2022.
Stergiou, A. and Poppe, R. AdaPool: Exponential Adaptive Pooling for Information-Retainne Retaining Downsampling. IEEE Transactions Image Processing, 32:251-266, 2023.
Torres, F., Zhang, H., Sicre, R., Ayache, S., and Avrithis, Y. CA-Stream: Attention- Pooling for Interpretable Image Recognition. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR) Workshops, pp. 8206-8214 June 2024.
Touvron, H., Cord, M., Sablayrolles, A., Synnaeve, G., and Jégou, H. Going Deeper With Image Transformers. Proceedings of the IEEE/CVF International Conference Computer Viston (ICCV), pp. 32 42, October 2021.
Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P., and Salakhutdinov, R. Transformer Dissection: Unified Understanding for Transformer's Attention via the Lens of Kernel. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Lan- guage Processing (EMNLP-IJCNLP), pp. 4344-4353, Hong Kong, China, November 2019. Association for Computational Linguisties.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is All you Need. In Advances in Neural Information Processing Systems, 2017.
Vinyals, O., Babuschkin, I., Czarnecki, W. M., Mathieu, M.,
Dudzik, A., Chung. J., Choi, D., Powell, R., Ewalds, T., Georgiev, P., Oh, J., Horgan, D., Kroiss, M., Danihelka, I., Huang, A., Sifre, L., Cai, T., Agapiou, J. P., Jaderberg, M., Vezhnevets, S., Leblond, R., Pohlen, T., Dalibard, V., Budden, D., Sulsky, Y., Molloy, J., Paine, T. L., Gulcehre, C., Wang, Z., Pfaff, T., Wu, Y., Ring, R., Yogatama, D., Wünsch, D., McKinney, K., Smith, O., Schaul, T., Lillicrap, T. P., Kavukcuoglu, K., Hassabis, D., Apps, C., and Silver, D. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575:350- 354, 2019.
Zambaldi, V., Raposo, D., Santoro, A., Bapst, V., Li, Y., Babuschkin, I., Tuyls, K., Reichert, D., Lillicrap, T., Lockhart, E., Shanahan, M., Langston, V., Pascanu, R., Botvinick, M., Vinyals, O., and Battaglia, Deep rein- forcement learning with relational inductive biases. In International Conference on Learning Representations, 2019.
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling Vision Transformers. In 2022 IEEE/CVF Conference on
12

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
A. Proofs
1
Σ X,
Σ X,
Xi
Xi
Xi
N
N
h
A.I. Proof of Corollary 3.3
x,
The point x that minimizes signal loss is the centroid of the signal subset X, S X.
1
Σ
Xj
X
x;EX,
x;EX,
Proof. Given a vector set X as presented in Section 3, a global vector pool C yielding C(X) = x E RR, and the following definition for signal loss given in Section 3.2
1 Σ X
N- 1
Σ Xi
N X;EX;
N k
X,EX,
L L(X,, xx)
1 k
Σ X;
(xl - xx))²
N
1
K
Σ X:EX,
X1
k.
1
Σ
X:EX,
AugPool(X,) = Aug Aoo(((o(X.))
the signal-loss minimizing point x is given by
OL
2
Σ
(xi
- xx) =
Case 2. (5) Suppose X, +0 and AugPool(X, = AcgPool(Xn). Then
0xc
K
X.EX,
AvgP (X.))
= Aug
k x
Σ Xi = 0
x,EX,
1 x
Σ X.EX.
1
Σ
1 Σ xi
N
N
Xi
x X;EX,
X,EX,
The remainder of the proof is symmetric to Case
yielding the centroid of the signal subset.
A.4. Proof of Corollary 3.8
A.2. Proof of Corollary 3.5
MaxPool is signal-optimal if and only if the input set X contains a single signal vector (k = 1) taking the maximal value along each feature dimension.
AvgPool is an optimal vector compressor if the input set contains no noise.
x, = Ø = AugPool C*
M PaxP = C - X = 1, Pool = X,
Proof. Suppose X,, = Ø. Then X,, = X, = N, and
Proof. Case 1. ( (=) ) Suppose MaxPool C When x, E X, is a signal vector, the weight vector yielded by MaxPool is w; = [w w = w/, where wi,' =
AugPool(X) =
1 N X, X,eX,
Σ Xi
1 Σ xi -
g for all features d.
Since the weights assigned by MaxPool can only take values of I or O, this implies w = for all features d which, in turn, implies that = X,I| = Furthermore, we know that MaxPool only assigns weight of 1 when Fi,d takes the max value over the set along dimension d. Since w, = [1,... ..... 11. it implies that x; = MarPool By our assumption, = X, = MarPool(X)
A.3. Proof of Corollary 3.6
When the input contains noise, X,, # Ø, AvgPool is an optimal vector compressor if and only if the centroid of X, equivalent to the centroid of X,,
AugPool= = C 0
Case 2. ( (=) ) Suppose X,| = 1, Mar (X) = x, When x; E X, is the only signal vector, by definition of MaxPool, the weight vector must be w; [w....... W;,a]
AugPool(X.) = AugPool(Xn)
where w = 1 for all features. By our assumption, |X,| = h. so the optimal weight vector and the MaxPool weight vector coincide, w; = w.
Proof. Case 1. (=) Suppose X,, # Ø and AogPool- = C², Since AvgPool assigns weight w, = + to all vectors while the optimal compressor C" assigns weights w, = \ all x, and wi all X,y we have
Now instead, when x; E X., is a noise vector. Since MaxPool(X) = Xg, by definition, MaxPool assigns a weight of widd = 0 for all features of all noise vectors,
Pool C =
13

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
which coincides with the weight vector yielded by the signal- optimal global vector pool w [0. ....... 0]. Since Max- Pool assigns the optimal weights to all vectors in the X, MaxPool = C
Case 1. Weight for a signal vector: Given x; E X,, it follows that ry belongs the set of signal relations {ra)- We can obtain an upper bound on the corresponding weight Wi via
-1
A.5. Proof of Theorem 3.12
Σ -e +Eer-n
W
1
For any signal-to-noise ratio, AdaPool can approximate the signal-optimal vector pool within an error bound determined by the distribution of signal and noise relations.
##1
1 + (k- - 1) e + Σe e'inrr
Explicitly, for any signal vector 24 Xg. AdaPool approxi- mates the optimal weight of wh within error bounds
1+ (k- 1).e e - + (N - k)-e""
L S w ws
U
v w
where L and U are the following lower and upper bounds respectively:
Intuitively, we use the lower bound distance -6, between any two signal relation values along with the maximal mar- gin D = M Eg + Ey between signal and noise to bound the value for wy. We can similarly obtain a lower bound on wy with the following:
L 1+ (ks 1).e e- + (N - -k) k)k.e -))
k
1 k U - 1 (k- e (N - k). e A5
and for any noise vector vi E X,, AdaPool approximates the optimal weight of : = 0 with error bound
-1
Σ
W/
1
L S w w; S
8#1
7
where L and U are the following lower and upper bounds respectively:
1 1 (k- - 1) e + Eenr-n
L L= - k. k.e M +1+ (N 1)
1 (k- 1) E + (N - k)) ee-Me
w
U U- - ke +1 - -N 1)
Proof. First, we let Wv = la. We then start by expanding the equation for the attention weight corresponding x; as follows:
In this case, the largest possible distance +e, between any two signal relation values and the minimum margin M are used. Note that these bounds are independent of any particular sample from the input set, instead relying purely on characteristics of the input distributions defined in the assumptions.
P
JEX
1
With these bounds on wa, it follows that the upper and lower bounds on the error from the optimal weight w = k are
Σ
-
1
JeX
-1
Σ 8EX,
2 + Σ r
k w 5 2 - w.
5 K
-w =
nEX,
As a general technique, we substitute the maximal and mini- mal neighborhood sizes well as maximal and minimal mar- gins between signal and noise relations as defined in Section 3.5 bound the range of the softmax operation.
Case 2. Weight for a noise vector: Given X; E X,, it follows that r, belongs to the set of noise relations {r,}- We can obtain an upper bound on the corresponding weight
14

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
W via
A.7. Proof of Corollary 3.11
MaxPool special case of AdaPool.
Σ 5en-n++4 +1+
Eeron Nift
Note:
The following proof assumes multi-head implemen- tation of AdaPool (attention), which was not discussed in the main text for brevity. We outline the details of multi-head attention in B.2.
-1
k cdd +1+ Eeti-ti TT
7
k- k MM + 1 (N - 1) -1
Proof. We will show this by construction. First, we sume a multihead adaptive pooling mechanism with many heads as input dimensions, or h = d. This is done to com- pute unique attention weight for each feature of each vector. Let WK. K Wy = Id, and let Wi = a La be the iden- tity matrix scaled by some large constant a. As a - oc, the softmax of I-dimensional head approaches the maximum over some column d:
Way
Similar to Case 1, we use the lower bound distance -C, between any two noise relationship scores along with the minimum margin M between signal and noise to the value for w, from above. We can similarly obtain a lower bound on Wi with the following:
exp
Fq.a
Iiid
vi
lim widd a0500
0/4
Σ
ΣS exp
Va
aggid Bjjdd
MF
exp B (-Ii,a)
k. D + +++ Σ TT
Ej exp (33 -j,a)
71
1 if ai, > Ljdd Vjj E
10 otherwise
IV
k e + 1 (N - 1)
where we substitute B = Va Lq.d- By then concatenating all of the heads, we obtain weight vectors w; E {0, 1] corresponding to each input vector x, and the formula for AdaPool then becomes
Again, the largest possible distance +e, between any two signal relationship scores and the maximum margin D are used to bound wi from below.
N AdaPool(X)) Σ w; x; MarPool(X)
With these bounds on w;, it follows that the upper and lower bounds on the error from the optimal weight w = are
L 0 w < W wi <
: U
B. Additional Results
A.6. Proof of Corollary 3.10
B.1. Time Complexity of AdaPool
AvgPool a special case of AdaPool.
AvgPool and MaxPool can be implemented with O(n d) algorithms, both require each of the d features of each of the vectors to be visited during aggregation.
Proof. We will show this by construction. Let Wq = 0₂ zero matrix, then the formula for the weights reduces to
Self-attention famously experiences quadratic computa- tional costs with respect to the context size n. In particular, the attention weight computation is O(n2 as each vector in the input set computes attention weights for each other vector in the input set, i.e. n² dot products of dimensional vectors. However, AdaPool uses cross-attention with sin- gle query, requiring the computation of only single set of attention weights rather than n sets. This reduces the weight computation and pooling time down to O d), the same as AvgPool and MaxPool.
exp (0) 1 Wt N exp N Σ
By letting Wy = Ta be the identity, the formula for AdaPool then becomes and
AdaPool Σ w. x, AugPool(X)
15

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
However, AdaPool still retains the overhead of the QKV projections of the input set, which involves standard matrix multiplications for each input vector, running (n d²). Thus, the overall time complexity of AdaPool is O(n d n : d²). The key takeaway is that it scales linearly with the number of inputs but quadratically with the number of features (like a standard linear layer).
burden on the query, ideally requiring it contain signal in all subspaces used in downstream processing.
B.3. Aggregation Function Approximation
Before we evaluate the noise robustness problem, we first ask the following question: Does the choice of output pooling even matter when the transformer is composed of a bunch of learned weighted averaging (attention) lay- ery? As a simple test, we examine how well a trans- former with each pooling method can approximate each other. Using the synthetic dataset, we generate 3 sets of targets: y MaxPool(X)) y (XnPoolX)), and y Aug (X). If the pooling method was of little consequence, we would expect all models to achieve the same performance regardless of the method used.
In isolation, AdaPool is slower the increased expressivity comes at computational cost. However, in practice, these pooling layers are placed at the end of a multi-layer trans- former or similarly sized encoder network, and the compute time added by the pooling layer is marginal. For our partic- ular applications, when using a transformer with 3+ layers, the time difference between AdaPool and Avg/MaxPool networks was negligible, regardless of the size of d.
Also of note, the ClsToken implementation is significantly slower than all other methods. This is due to the fact that the learned class embedding must be copied and concatenated along the batch dimension of the input for every inference, which is quite expensive in terms of both time and memory during training. This is a known hindrance investigated by Zhai et al. (2022), which led them explore the use of AvgPool for efficiency reasons when scaling vision trans- formers.
METHOD
MAX
MEAN
MIN
MaxPool
0.000
0.026
0.297
10.000
10.046
+:0239
AvgPool
0.003
0.000
0.003
1.0,000
10.000
1.0000
ClsToken
0.011
0.008
0.011
#0,000
10,000
+00000
AdaPool
0.002 1.00000
0.000 1.0000
0.002
1.0.000
B.2. Multi-Head AdaPool
In the theoretical portions of this work, we primarily consid- ered the use of a single-headed attention for simplicity and conciseness of notation and proofs. However, we note that implementing multi-head attention, as proposed by Vaswani et al. (2017), is quite straightforward. This can be obtained by partitioning the input vector space RRd into h partitions (i.e. heads) belonging Rd/h, computing a separate set of adaptive weights over the inputs per head, and then concate- nating the resulting weighted averages from each head. We assume this form of AdaPool in Proof A.7.
Table 2: Aggregation Approximation MSE (N=128, d=16)
Unsurprisingly, pooling methods can approximate them- selves perfectly, with MaxPool achieving near-zero MSE on the max targets and AvgPool achieving near-zero MSE on the average targets. Interestingly, MaxPool performance deteriorates on the average and min targets. AdaPool also able to perfectly fit the average targets and has the best per- formance on min targets, with a similarly good performance max targets. AvgPool performs just marginally worse than AdaPool on min and max targets, while Cis con- sistently performs worse than AdaPool and AvgPool on all targets. This shows that the inductive bias of the pooling method chosen has a clear impact on a transformer encoder's ability to approximate functions. It is also worth noting that, although we have proven MaxPool to be special case of AdaPool, AdaPool is not able to perfectly fit the max targets. This is likely because the approximation requires that one of the weight matrices obtain extremely high values, which is unlikely to occur in practice via gradient descent.
Multi-head adaptive pooling presents one key potential ben- efit that we did not mention in the main text- - it allows different subspaces of the input to be pooled in parallel. This means that, with further work, our analysis could po- tentially be generalized beyond vectors that purely signal or purely noise, more akin to what we expect from real-world data. Consider a vector that contains meaningful signal in one subspace, but serves as noise in another. In an au- tonomous system, one could imagine that faulty sensor providing accurate position but bad velocity measurements might create such an input. By using multihead attention with the right partitioning, AdaPool could assign it high weight in the position subspace while suppressing it in the velocity subspace, where it is noisy. While this multi-hear head partitioning allows us to relax some of the strict assump- tions made about the input data, it does place more relative
B.4. Noise Robustness Table
We provide the data for Figure 4 in Table 3. This cor- responds to the signal loss of each method for different signal-to-noise ratios of the input data.
16

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
METHOD
KNN-1 SNR 0.01
KNN-2 SNR 0.02
KNN-4 SNR 0.03
KNN-8 SNR = 0.06
KNN-16 SNR 0.13
KNN-32 SNR 0.25
KNN-64 SNR = 0.50
KNN- 128
SNR 1.00
MaxPool
0.066
0.061
0.050
0.045
0.047
0.055
0.068
0.046
10.702
=0.110
+0.093
+0.089
+0,005
+0.096
=0,096
#0.000
AvgPool
0.093
0.071
0.055
0.043
0.031
0.016
0.001
0.000
+0.000
+00000
+00000
#:0000
+0,008
=00000
#:0000
ClsToken
0.102
0.079
0.063
0.050
0.038
0.027
0.016
0.008
+0,000
=0,000
+0,000
+0,000
+0,000
=0,000
AdaPool
0.037
0.019
0.006
0.004
0.003
0.002
0.001
0.000
+00006
=0.055
+0.001
#0,000
#00000
+0,000
=0,000
100000
Baseline
0.093
0.071
0.055
0.043
0.031
0.020
0.008
0.000
(Centroid)
Baseline 0.058
0.044
0.040
0.041
0.048
0.060
0.080
0.126
(Target)
Table 3: Signal Loss (MSE) the KNN-Centroid Task (N=128, d=16) presented in Figure 4.
B.5. KNN Centroid Ablations
the network grows relative to the input data dimension, the burden on pooling mechanism compress information is relaxed.
We performed several ablation studies to examine how key data and network hyperparameters changed the relative per- formance of each pooling method. The ablations were per- formed on much smaller cardinality datasets than the main experiment to enable faster iteration. We perform initial study to use baseline with which compare each abla- tion, using synthetic dataset with million samples, each of size N=32 and d=16. We use the exact same network architecture as the main experiment, with 12 layers and a hidden dimension of We perform evaluations using k- the ONN-Centroid task as described in section 4.2, and we report results for the baseline in Table 4. The results reflect the main experiment, with AdaPool exhibiting superior performance across different SNRs on dataset with samples that have a quarter of the cardinality as those in the main experiment (32 vs 128), with all other hyperparameters being the same.
Finally, look increasing the depth of the network from 12 layers to 24 layers. The results in Table do not show clearly dominant method in any particular SNR regime. AdaPool appears to be most consistent, but is edged out by AvgPool the high signal regimes. AvgPool does predictably poorly in the high noise regime, while MaxPool does predictably poorly in the high signal regime. Finally, ClsToken performs the worst by far in all regimes.
fold cooss-validation
C. Experiment Parameters
C.1. Synthetic Dataset Construction Parameters
To construct a dataset full of sets of points with diverse distributions, we use the following sampling scheme. For each of the d features, all N vectors sample values from one of the following distributions:
Next, we examine the effect of increasing the dimensionality of the input vectors from d=16 to d=64. As shown in Table 5, MaxPool slightly outperforms AdaPool in the low signal regimes, while AdaPool outperforms the other methods in the medium to high signal regimes. AvgPool and ClsToken perform very poorly in the low signal regime, with AvgPool getting better as signal increases much more rapidly than ClsToken.
Gaussian with randomly sampled mean and standard devia- tion:
I Uni (orm(m- (-3,
a ~ form (1,m113)) (1,3) a. ~ Normal(y,
We then look expanding the width of the network from hidden dimension of 16 to 32 to 64. As displayed in tables 8 and 7, AdaPool outperforms all methods at all SNRs for all widths. However, we observe that the absolute signal loss for all methods reduces as the network width increases. In particular, AvgPool and MaxPool become much more competitive in the higher signal regimes. One possible interpretation of this phenomenon is that the width of
Exponential with randomly sampled rate parameter, sign, and shift:
sign {-1,1}
shi,ft ~ (0, * sign
A Uni 2)
Z. Exponentiall)) # sign shift
17

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
hid
L
METHOD
SNR
SNR
SNR
SNR
SNR
32
16
MaxPool
0.022
0.009
0.006
0.009
AvgPool
0.086
0.049
0,010
0.005
0.003
ClsToken
0.094
0.066
0.046
0.030
0.017
AdaPool
0.018
0.011
0.007
0.004
0.002
Table 4: Ablation baseline study.
hid
T,
METHOD
SNR J
2 SNR
4 SNR
SNR
&
16 SNR
32
32
32
32
32
32
64
16
12
MaxPool
0.0056
0.0025
0.0025
0.0018
0.0018
AvgPool
0.0204
0.0104
0.0053
0.0023
0.0006
ClkToken
0.0229
0.0126
0.0069
0.0036
0.0017
AdaPool
0.0064
0.0031
0.0010
0.0005
0.0002
Table 5: Ablation increasing the data dimension x4 from 16 64.
N
hid
L
METHOD
SNR
SNR
SNR
SNR
SNR
16
32
MaxPool
0.022
0.011
0.0023
0.0006
0.0003
AvgPool
0.048
0.022
0.0012
0.0006
0.0002
ClsToken
0.086
0.057
0.0371
0.0214
0.0082
AdaPool
0.003
0.001
0.0006
0.0003
0.0001
Table Ablation increasing the network width x2 from 16 to 32.
hid
L
METHOD
SNR 32
SNR 32
SNR 32
SNR 32
SNR
16
32
32
16
64
12
MaxPool
0.022
0.0012
0.0005
0.00024
0.00010
AvgPool
0.004
0.0016
0,0005
0.00021
0.00009
ClsToken
0.086
0.0574
0.0373
0.02147
0.00825
AdaPool
0.001
0.0010
0.0003
0.00019
0.00007
Table 7: Ablation increasing the network width x4 from 16 to 64.
N
hid
L
METHOD
SNR 32
SNR
SNR 4
SNR
SNR 16
37
32
32
16
16
28
MaxPool
0.020
0.011
0.007
0.005
0.0056
AvgPool
0.061
0.013
0.006
0.003
0.0019
ClsToken
0.094
0.066
0.046
0.030
0.0165
AdaPool
0.017
0.014
0.014
0.008
0.0021
Table Ablation increasing the network depth x2 from to 24.
Uniform with randomly sampled lows and highs:
We scale each resulting vector by v reduce the impact of the curse of dimensionality the scale of the loss when different dimensionalities set are drawn evenly from these 3 distributions, with each individual feature having a uniquely parametrized distribu- tion according to the above scheme. The ordering of these
low Uniform( 3).3)
used. The features for given
high (0.mforrm(0.... + low
F.dd mniifrr((low,hig]
18

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
feature distributions is shuffled randomly for each sample. For our main supervised experiment, we generate million samples, 128 vectors per sample, and 16 features per vector. Our dataset was generated using NumPy version 2.0.2 with a seed of 42.
C.3. Supervised Aggregation Approximation
We use the same architecture and training setup outlined in the previous section. Given the main dataset gener- ated above according to C.I, we produced three sets of targets y using standard pooling functions to see if the underlying transformer was good enough to approximate all aggregation types regardless of the method used to ag- gregate its outputs. Targets include y = PinPoolX y MarPool(XX) and y = Aug (X).
C.2. Supervised KNN Centroid
The base network consisted of transformer encoder with 12 layers, a hidden dimension of 16 (same as the input data dimension), a feedforward dimension of 64, and attention heads. We follow the standard practice of moving the layer norm to the input of the attention and feedforward sublayers (pre-norm). We use PyTorc for all model implementations. Weight matrices and embedding vectors are initialized ac- cording to N (frmal 0.0, = 0.02) with bias set to 0.0. Layer norms are initialized with weight set to 1.0 and bias set to 0.0.
C.4. MPE Simple Centroid
For the multi-agent experiments on the simple centroid task, network architecture hyperparameters are displayed in Table 10 and training hyperparameters in Table 11.
Hyperparameters
Value
Obs Shape
N X
Num Layers
C.2.1. MODEL ARCHITECTURE
Num Heads
8
Dim Inpul
Dim Hidden
Hyperparameters
Value
Dim FF
32
Drupout Alin
0.0
Input Shape
x 16
Dropout FF
0.1
Num Layers
Bias Alin
Fabe
Heads
8
Bias
True
Dim Hidden
16
Action MLP Layers
Dim FF
64
Action MLP Dim
32
Dropout Atta
0.0
Value MLP Layers
Dropout FF
0.1
Value MLP Dim
32
Bias Artn
False
Continuous Actions
True
Bias FF
True
Table 10: Network Hyperparameters for MPE Simple Centroid
Table 9: Network Hyporparameters for KNN-Centroid
Hyperparameters
Value
Algorichm
PPO
We train on the main dataset described above (N=128, d=16). For each desired level of noise, we generated X.y pairs by picking an arbitrary target vector from the input set X. find- ing the k-nearest neighbors to that target point, and comput- ing y to be the centroid of those k neighbors (excluding the target). These k neighbors are considered the signal vectors, while the rest are considered noise. We add learned em- bedding to the target vector at inference time to indicate it to the transformer. We compute X.y pairs for the following values of 1, 8, 16, 32, 64, 128. For each of these pairs, we use 5-fold cross-validation with holdout test set of 100k samples. Each transformer pooling method (AdaPool, ClsToken, AvgPool, MaxPool) was trained for 100 epochs with learning rate of 5.0e 4 and a batch size of 750. The base transformer model initialized with the same weights for each method, and different seeds are used on each fold.
Number of Seeds
Number Episodes
of
500
Training Batch
8192 Agent Steps
Minihatch
8192 Agent Steps
Its/BBtth
Episode Length
128 Fiw Steps
Leaming Rate
2.5e
3
7
0.99
X
0.95
KL Coeff
0.0
Coeff'
0.05
Entropy Coeff
0.0
PPO Clip
0.2
Clip
10.0
Table 11: MPE Simple Centroid Training Hyperparameters
19

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
C.5. MPE Simple Tag Noise
For the multi-agent experiments on the simple tag task, network architecture hyperparameters displayed in Table 12 and training hyperparameters in Table 13.
Hyperparameters
Value
Obs Shape
N x 8
Num Layers
Num Heads
3
Dim Hiddon
3
(a) Loose Key
(b) Two Distractor Boxes
Dim FF
Drogout Alin
0.0
Figure 10: (Left) The Box World environment showing the agent, the initial loose key, and the box composed of blue lock and the gem. (Right) A scenario containing two distractor paths, with the agent holding the orange key as visible the top left.
Dropout FF
0.1
Bins FF
True
Action MLP Layers
3
Value MLP Layers
3
Value MLP Dim
to unlock the box with the white pixel referred to as the gem. In all training scenarios, we provide two distractor paths, shown by the additional boxes in Figure 10b. When it picks up the loose key, carns +1 reward; when it unlocks the box with the gem, it earns +10 reward; opens distractor box, it gets penalty of
Table 12: MPE Simple Tag Network Hyperparameters
Hyperparameters
Value
Algorithm
PPO
Num Scods
20
Num Episodes
500
Hyperparameters
Value
Training Batch
8192 Agent Steps
Shape
Minibatch
8192 Agont Steps
Num Layers
SGD Iter/Batch
16
Num Heads
Episode Length
128 Env Steps
Learning
Dim FF
24
0.99
Dropout Attn
0.0
0.95
Dropout FF
0.1
Bias Attn
False
KL Target
0.01
Bias
True
VF Coeff
0.0005
Action MLP Layers
Entropy Coeff
Action MLP Dim
32
PPO Clip
0.2
Value MI P Layers
Grad Clip
10.0
Value MLP
32
Continuous Actions
False
Table 13: MPE Simple Tag Training Hyperparameters
Table 14: Box World Network Hyperparameters
C.6. BoxWorld
Hyperparameters
Value
For the relational reasoning experiments in the Box World environment, network architecture hyperparameters are dis- played in Table 14 and training hyperparameters in Table 15. We also display images taken from the rendered envi- ronment in Figure 10. The learning agent controls the grey pixel. The agent's actions move it up, down, left, or right by one pixel on each step. Each round starts with a loose key, depicted by the lone blue pixel on the left side of Figure 10a. "Boxes" are indicated by pairs of colored pixels, with lock on the right and the item in the box the left. The agent must first pick up the loose key, which it can then use to open locks of the same color by moving onto them. The currently held key displayed in fixed position at the top left pixel, as shown in Figure 10b. The goal of the agent is
Algorithm
PPO
Num Sceds
Num Episodes
3000
Training Batch
14410 Agent Steps
Minibatch
14400 Agent Steps
SGD Iter/Bachh
32
Episede Length
30 Eav Steps
Leaming
0.99
X
0.95
Coeff
VF Coeff
0.01
Entropy Coeff
0.0
PPO Clip
Grad Clip
10.0
Table 15: World Training Hyperparameters
20

---

Robust Noise Attenuation via Adaptive Pooling of Transformer Outputs
C.7. CIFAR Training Hyperparameters
All experiments trained vision transformers (ViT) from scratch on the CIFAR dataset using 5-fold cross- voss-vaaliation. As in the original paper, note that vision transformers trained from scratch are less performant than similarly sized convolutional neural networks, as they lack spatial inductive biases, such as translation invariance (Dosovitskiy et al., 2021). Thus, performance reported in Table is expected to be lower than similarly sized ResNet trained from scratch on the same task, for example.
Training hyperparameters and network architecture details are outlined in Table 16. Each fold was trained for 300 epochs and, as in previous experiments, the network for each method was initialized with the same underlying transformer weights and seeds. Following the ViT paper, we apply gradient clipping at global norm of 1.0, as well AdamW with weight decay value of 0.1. Finally, rate scheduler with a linear warmup from 0.0 to 0.001 after 30 epochs (10% of training) and cosine annealing back down to 0.0 by epoch 300.
learning
Hyperraaraeeerrs
Value
Image Shape
32
32 3
Shape
Input Shape
64 48
(Flattened Forches)
Num Layers
6
Num Heads
8
Hidden
Dim FF
Dropout Artn
0.1
Dropout
Bias Artn
False
Blas FF
True
Training Epaccs
Warmup Epochs
Learning Rate
1.0e
Batch Size
512
Weight Decay
0.1
Grad Clip
1.0
Table 16: CIFAR network and training hyperparameters. Both experiments trained models from scratch and differed only by the number of output classes, 100 respectively.
21


Reviews:
[{"review_id": "idHwKCWUa1", "reviewer": "Reviewer_PXAR", "summary": "The paper studies pooling methods for aggregating transformer embeddings—particularly in settings where only a subset of input vectors (signal) is task-relevant while the remainder (noise) may deteriorate performance. The authors reframe pooling as a vector quantization (or lossy compression) problem and show that common methods like Average Pooling (AvgPool) and Max Pooling (MaxPool) can fail under fluctuating signal-to-noise ratios. They introduce an attention-based adaptive pooling method (AdaPool) and provide theoretical error bounds that demonstrate its ability to approximate the signal-optimal compressor across noise regimes. These theoretical results are validated with experiments on synthetic supervised tasks as well as on reinforcement learning benchmarks, where AdaPool consistently exhibits greater robustness to noise.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work and provide thorough feedback.\n\nThe work from Przewiezlikowski et al. is indeed interesting and highly relevant to our discussion, and we will update our related works to acknowledge it. Their method appears equivalent to AdaPool with a learned embedding as the query. We believe that our analytical framework for analyzing pooling methods and our more general form of AdaPool distinguishes us as novel from this concurrent work, while the results of our analysis led to a similar solution of weighted pooling.\n\nDue to character limits, we address question 3 in our response to GbuG's review below. As to 1 & 2, we realize that by seeking to explicitly control the noise level, our experiments failed to address the following: What happens with real-world data where vectors are not purely signal/noise and the choice of query for AdaPool is not obvious?\n\nTo address this concern, we conducted additional studies on image classification using the CIFAR 10 & 100 benchmark datasets. We use this experiment to explore various choices of query on real-world, less cleanly defined input vectors. We adopted the Vision Transformer (ViT) approach, partitioning the 32x32 pixel RGB images from CIFAR into 64 separate 4x4 pixel patches which are then flattened and projected to the dimension of the transformer (see Fig 1 in https://arxiv.org/abs/2010.11929). We show the arrangement of those 4x4 pixel patches by their indices below:\n\n```\n[ 0,  1,  2,  3,  4,  5,  6,  7]\n[ 8,  9, 10, 11, 12, 13, 14, 15]\n[16, 17, 18, 19, 20, 21, 22, 23]\n[24, 25, 26, 27, 28, 29, 30, 31]\n[32, 33, 34, 35, 36, 37, 38, 39]\n[40, 41, 42, 43, 44, 45, 46, 47]\n[48, 49, 50, 51, 52, 53, 54, 55]\n[56, 57, 58, 59, 60, 61, 62, 63]\n```\n\nAs in our previous experiments, we feed these patches through a transformer and pool the outputs. Most CIFAR images contain the object being classified (i.e. the signal) in the center patches. We thus study 3 choices of query: (CORNER) - the embedding of patch 0 which appears *less likely* to include signal, (FOCAL) - averaging the embeddings of the center four patches 27, 28, 35, and 36, which appear *highly likely* to contain signal, and (MEAN) - averaging all patch embeddings as the query, as proposed by Stergiou & Poppe (2023). \n\nBoth experiments use a transformer with layers=6, dim=512, heads=8, & dropout=0.1. We use 5-fold cross-validation, 150 epochs/fold, batch=512, and LR=1e-4. We report the Top 1 accuracy scores on the holdout test set, averaged across the best models from each fold. Our ClsToken scores are consistent with common open-source implementations of ViT trained from scratch on the CIFAR benchmarks. \n\n```\nMETHODS         CIFAR 10          CIFAR 100\nClsToken   | 0.7863 ±0.0041 | 0.4892 ±0.0018 |\nMaxPool    | 0.8301 ±0.0041 | 0.5547 ±0.0044 |\nAvgPool    | 0.8265 ±0.0023 | 0.5633 ±0.0070 |\nAda-CORNER | 0.8240 ±0.0016 | 0.5385 ±0.0044 |\nAda-FOCAL  |*0.8325 ±0.0035*|*0.5730 ±0.0051*|\nAda-MEAN   | 0.8317 ±0.0022 |*0.5731 ±0.0073*|\n```\n\nWe observe that Ada-FOCAL and -MEAN queries outperform in both cases. Interestingly, Ada-FOCAL and MaxPool have relatively better performance on C-10, while Ada-MEAN and AvgPool perform relatively better on C-100. With C-10, there are only 10 highly distinct classes, so the object at the center of the image may be relatively more important for discriminating between classes. For C-100, the data is the same, but there are 100 highly similar classes, and border patches may contain subtle features that are important for discriminating between related classes where they were otherwise distracting before. \n\nInterestingly, the Ada-CORNER query underperforms in both experiments, which aligns with the prediction that choosing a query with little to no signal will lead to worse performance. Since a noise query would yield higher dot product relations with other noise vectors, this likely results in higher attention weights for irrelevant patches and lower weights for signal-containing patches, a sort of worst-case for AdaPool. This violates our assumption of linearly separable signal and noise relation scores, implying the error from an optimal weighting would not be guaranteed to fall within the bounds we presented *for this choice of query*. It is worth noting that the standard ClsToken approach still underperforms this and all other methods by a significant margin.\n\nThis experiment demonstrates how one might use intuition, data analysis, and experimentation to find a competitive query. If one cannot be found, Avg/Max may be strong alternatives *if* the noise level of the dataset is constant. While the explicit SNR level in CIFAR is not known, the key takeaways are that the choice of AdaPool’s query (1) can be used to investigate the way signal and noise vectors are distributed, (2) can lead to superior performance on real-world data, and (3) may underperform with a poor choice of query. \n\nWe really appreciate your comments and insightful questions!", "claims_and_evidence": "The paper makes two core claims:\n\n- AdaPool can approximate the optimal vector quantizer for any signal-to-noise ratio with quantifiable error bounds.\n- Standard pooling methods (AvgPool, MaxPool, ClsToken) fail or degrade predictably as noise increases, while AdaPool remains robust.\n\nThe claims seem to be supported by a theoretical framework (with the primary result being Theorem 3.11) and are accompanied by proofs provided in the appendix. On the experimental side, the evidence is drawn from controlled synthetic datasets as well as from RL environments where noise is systematically varied.", "methods_and_evaluation_criteria": "The authors:\n\n- Formally define pooling as a differentiable vector quantization task with a focus on minimizing “signal loss”.\n- Analyze AvgPool and MaxPool as special cases, and then derive AdaPool based on an attention mechanism.\n- Use both theoretical error bounds and empirical evaluations to compare the pooling methods.\n\nThe evaluation criteria—especially the use of synthetic datasets where the signal-to-noise ratio can be explicitly controlled—provide a clear measure of robustness. In reinforcement learning tasks, using both entity-based and pixel-based observations further tests the method in settings with varying noise levels.", "theoretical_claims": "The paper’s theoretical contributions include:\n\n- A derivation of the optimal pooling strategy as one that minimizes the signal loss.\n- A series of corollaries showing under what conditions AvgPool and MaxPool become optimal.\n- Theorem 3.11, which provides error bounds for the proposed adaptive pooling method.\n\nWhile the theory & proofs appear correct I am not familiar enough with the techniques and line of work to check this in detail.", "experimental_designs_or_analyses": "The experiments are synthetic and controlled in nature and include:\n\n- A supervised KNN-centroid task where noise levels are precisely controlled allows for direct measurement of signal loss across varying signal-to-noise ratios.\n- The authors evaluate performance in both a custom “simple centroid” scenario and a “simple tag” task in the Multi-Particle Environment, demonstrating how increased noise leads to degradation in performance for standard pooling methods.\n- A vision-based relational reasoning task that further challenges the pooling methods by introducing high-dimensional, pixel-based noise.\n\nThe analysis of performance degradation is convincing. However, due to the synthetic nature of these experiments, it's difficult to extrapolate whether or not the proposed pooling method improves performance in more realistic tasks and environments.", "relation_to_broader_scientific_literature": "Most works in this space don't approach the problem of aggregating representations in sequence models from a signal-to-noise lens. Given my unfamiliarity with how the authors approached this problem and how it's somewhat disconnected from the literature in the vision and text space I'm unable to definitively comment on how the paper relates to the broader literature.", "essential_references_not_discussed": "There is one prior approach [1] that also replaced the common CLS token representation with a similar learned pooling mechanism. This seems close enough that it's worth discussing.\n\n[1] Marcin Przewiezlikowski, Randall Balestriero, Wojciech Jasinski, Marek Smieja, Bartosz Zielinski. Beyond [cls]: Exploring the true potential of Masked Image Modeling representations.", "other_strengths_and_weaknesses": "Strengths\n\n- The paper provides a novel framework that reinterprets pooling as vector quantization.\n- It offers comprehensive empirical validation across multiple domains.\n- The clear exposition of failure modes for AvgPool and MaxPool under varying noise regimes is a valuable contribution.\n\nWeaknesses\n\n- While the experiments are extensive, additional tests in real-world or more diverse noisy environments could further establish generality.\n- It's unclear how novel the approach is given related work in this space, e.g., Przewiezlikowski et al.", "other_comments_or_suggestions": "No further comments.", "questions_for_authors": "1. How sensitive is AdaPool’s performance to the choice of query vector? Have you explored different strategies for selecting the query in settings where the signal vector is not obvious?\n2. Is the linear separability assumption (Assumption 3.10) going to hold in practical, real-world problems?\n3. What is the computational overhead of AdaPool compared to AvgPool or MaxPool in terms of training time and inference speed? Although it should be negligible compared to the cost of a forward pass through a large model it's still more expensive than a simple reduction.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6U0YCCKtjm", "reviewer": "Reviewer_cS6k", "summary": "This work analyzes the various pooling methods used in deep neural architectures. They establish a connection between pooling and vector quantization and demonstrate adaptive pooling is more robust to signal-to-noise ratio. They provide experimental results with a carefully created synthetic dataset and multi-agent RL environments.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read and review our work! Regarding your questions - \n\nUpon further empirical analysis and review of our proof of Theorem 3.11, we realized that our assumption that the margin M must be greater than zero (i.e. signal and noise must be linearly separable) was not necessary for the bounds to hold. That is, the bounds are still effective when the neighborhoods of signal and noise relations overlap. We also address the case where there may be an overlap between signal and noise vectors (i.e. a single vector contains both relevant and irrelevant information) by using multiple attention heads. Please see our response to 3DVE below for more details on both of these topics. \n\nAs to your question about applications to CNNs, we believe this line of work would be highly relevant to CNN-based image retrieval tasks, where global pooling of output feature maps is heavily utilized. We are less certain about the effectiveness of AdaPool in replacing Avg/Max pooling for pixel-level filter downsampling in standard convolutional architectures, but our analysis should be generally relevant to any pooling applications. \n\nIf you have any further questions, please let us know!", "claims_and_evidence": "The experimental results supports the superiority of the adaptive pooling method. The experiments to assess the robustness are thorough and includes varying signal-to-noise ratio, network width, network depth, data dimension.", "methods_and_evaluation_criteria": "While this work doesn't propose any new method, it systematically analyzes different pooling methods and their outputs. The designed evaluation metric \"signal loss\" seems reasonable to me.", "theoretical_claims": "The formal connection established between pooling and vector quantization is pretty interesting. However, I would like to see the case if the margin M vanishes.", "experimental_designs_or_analyses": "The experiments are well designed and considered different factors of variations. Also, the experiments with multi-agent environments are nicely designed to capture the problem at hand.", "relation_to_broader_scientific_literature": "The authors discussed the relevance of the work with works related to associative memories. Also, it puts a different perspective based on vector quantization. However, adaptive pooling is already known as a robust approach compared to avg. pool or max pool.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "This paper assumes that there is no overlap between the task relevant vectors and noise vectors. However, in reality, this may not be the case. Some channels may have important information mixed with noise.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "While the paper's discussion is for transformer outputs, I am curious whether the same analysis may be valid for CNN output?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FqkfTIraVa", "reviewer": "Reviewer_3DVE", "summary": "This paper investigates pooling methods for transformer embeddings in tasks where only a subset of inputs carries signal and the remainder are noise. It shows that standard methods like average and max pooling can collapse in performance as the signal-to-noise ratio fluctuates. The authors propose an attention-based adaptive pooling method that approximates a signal-optimal vector quantizer, providing theoretical error bounds that guarantee robustness across various noise levels. Their approach is validated through supervised experiments on synthetic datasets and reinforcement learning benchmarks, demonstrating superior performance in noisy settings. Overall, the work establishes a framework for enhancing transformer pooling in non-sequential, noise-prone applications.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your comprehensive review. We saw the weak reject rating as a great opportunity to improve the quality of the paper, and hope to address the questions you raised regarding performance outside of our theoretical assumptions.\n\n**1 & 3. Idealistic Assumptions:** As questions 1 & 3 are related, we address them jointly here. We empirically evaluated the theoretical error bounds from Theorem 3.11 on the KNN-centroid dataset (N=32, d=16) for SNRs 0.5-0.03. We know which vectors are signal and noise by design, so we can save off the relations rₛ and rₙ after computing dot products in the attention mechanism. With these, we derived ϵₛ , ϵₙ , M, and D, computed the error bounds, and checked if the weights for each vector fell between them. We note that this type of analysis is only possible when the signal and noise vectors are known and will typically not be feasible on a real-world dataset where the line between signal and noise is blurry.\n\nSurprisingly, the error bounds held for all 1 million samples at each noise level, even when the assumption of linearly separable signal and noise relations was violated (M < 0). This was observed both with trained networks and randomly initialized networks. Upon revisiting the proof of Theorem 3.11, **we realized that Assumption 3.10 was overly restrictive and NOT necessary for the error bounds to hold**. We can instead define M more loosely as M = min( rₛ ) – max( rₙ ), allowing this quantity to be positive or negative rather than strictly positive as assumed before. Theorem 3.11 thus applies more generally than was claimed in the original submission, which is a pleasant surprise. Also worth noting, we observed empirically that M is nearly always less than zero on the KNN-centroid dataset, meaning AdaPool still outperformed other methods with some degree of overlap between signal and noise neighborhoods in those experiments.\n\nIt is also worth taking a step back and discussing why we derived and presented the bounds in the first place. The sole purpose was to show that with AdaPool, it is possible to approximate the optimal vector quantizer exactly for any SNR. The bounds are also helpful in illustrating which factors influence the approximation error most and how a signal-rich query can shrink those bounds. Since it is impractical to compute the bounds exactly on most real-world datasets, we did not present them to guarantee modeling performance on a particular problem.\n\nAnother point we did not discuss is that using multiple attention heads allows the network to pool chunks of the input vectors separately. If the input vectors each contain a mix of signal and noise, the heads can weight partitions of the vector space separately. In fact, when using the same number of heads as features, AdaPool can be viewed as pooling each feature dimension separately. This would enable it to optimally pool such inputs, given a good choice of query, and may be a better modeling choice for real-world data. We experiment on real-world data in response to PXAR above.\n\n**2. Hyperparameters:** We performed a number of ablation studies on model hyperparameter choices in the Appendix; please see Tables 3-7.\n\n**4. Computational Complexity:** Due to character limits, we ask that you see our response to reviewer GbuG regarding this below.\n\n**5. Signal/Noise Overlap:** Suppose the signal and noise relation distributions overlap completely. If variance is low, applying softmax results in relatively uniform weights, yielding something closer to AvgPool. If variance is high, one relation score will dominate the others in the softmax, yielding weights closer to MaxPool.\n\nAs a worst case, suppose the query is more similar to noise than signal. The dot product relations will be higher for all noise vectors, and the resulting aggregation will suppress signal and accentuate noise. However, if such a query vector exists, then it implies that the noise and signal are easily distinguished via dot product. A learned linear transformation (the Q projection) could rotate the query away from the noisy direction and towards the signal, given that dot products largely reflect directional alignment. So, practically speaking, the worst case for AdaPool would instead be when signal and noise are indistinguishable, in which case it reverts towards AvgPool or MaxPool as discussed above.\n\n**Related Works:** Can you provide full titles or links to the papers you referenced by Lee et al., 2024 and Xu et al., 2022? We also reviewed the line of work on noise-robust graph neural networks. However, those approaches focus on robust learning in the presence of mislabeled graph classification targets. While the titles and themes are similar, we felt that the actual research questions were not related enough to justify inclusion.\n\nAs a result of your questions, we were able to identify an oversight and make a significant improvement to our theoretical claims. We really appreciate your review!", "claims_and_evidence": "The paper’s primary claims—that standard pooling methods collapse under variable signal-to-noise ratios and that adaptive, attention-based pooling can approximate the signal-optimal vector quantizer—are generally supported by a combination of rigorous theoretical analysis and empirical validation. The authors derive clear error bounds and show mathematically that AvgPool and MaxPool are specific instances of the more general AdaPool method, which is capable of handling noise across any SNR regime. Their experimental results, conducted on both synthetic datasets and complex reinforcement learning benchmarks, provide convincing evidence that AdaPool consistently achieves superior robustness and performance compared to traditional pooling methods. However, some of the theoretical guarantees rely on idealized assumptions regarding the separability of signal and noise in the relation space, which may not always hold in real-world settings and could benefit from further investigation.", "methods_and_evaluation_criteria": "The paper's methods and evaluation criteria are well-suited to the problem of noise attenuation in transformer pooling. The authors establish a rigorous theoretical framework using vector quantization and signal loss minimization, which logically underpins their proposed adaptive pooling method. Their experimental design includes a synthetic dataset to isolate the noise impact, followed by evaluations on established reinforcement learning benchmarks (MPE and BoxWorld) that effectively mimic real-world noisy conditions. Together, these methods and benchmarks provide a comprehensive and convincing validation of the approach for non-sequential tasks where noise is a critical factor.", "theoretical_claims": "Overall, the proofs appear correct within their stated assumptions, but the reliance on these idealized conditions may limit the direct applicability of the theoretical guarantees in real-world settings.\n\n- The proof that the signal-optimal pooled vector is the centroid of the signal subset (Corollary 3.3), which is mathematically sound given the mean squared error definition.\n- The derivations showing that AvgPool and MaxPool are special cases of the more general adaptive pooling framework (Corollaries 3.5, 3.6, and 3.8). These proofs correctly demonstrate the limitations of traditional pooling methods under varying noise conditions.\n- The error bounds established in Theorem 3.11 for AdaPool's approximation to the signal-optimal quantizer are derived rigorously under the assumption of linearly separable relation neighborhoods. While the derivations are mathematically consistent, they rely on idealized assumptions—such as the clear separation (margin) between signal and noise relations—which might be challenging to guarantee in practical scenarios.", "experimental_designs_or_analyses": "The experimental designs are generally sound and align well with the paper’s goals. The authors first construct a synthetic dataset where the noise-to-signal ratio is carefully controlled, allowing for a focused evaluation of how each pooling method manages signal loss. They then validate their findings on established benchmarks such as the Multi-Particle Environment (MPE) and BoxWorld, which are relevant for multi-agent reinforcement learning and relational reasoning tasks. This multi-tiered approach strengthens the empirical evidence for the proposed adaptive pooling method. One potential issue is that the synthetic dataset, while useful for isolating noise effects, might oversimplify the complexities of real-world noise; however, the inclusion of realistic RL benchmarks helps mitigate this concern.", "relation_to_broader_scientific_literature": "The paper’s contributions build on the foundational transformer architecture (Vaswani et al., 2017) by addressing a critical yet underexplored aspect: the pooling of transformer outputs for non-sequential tasks. Its adaptive pooling method extends prior work in vector quantization and selective attention, drawing parallels with research on associative memories such as Hopfield Networks and Dense Associative Memories, which also focus on robustly retrieving relevant information amidst noise. Moreover, the work connects with recent advances in computer vision and reinforcement learning, where attention-based pooling and class token strategies have been successfully employed to handle complex, noisy data. By offering rigorous theoretical error bounds and empirical validation, the paper unifies traditional pooling methods (like AvgPool and MaxPool) with modern attention mechanisms, thereby enriching the broader literature on robust representation learning and relational reasoning.", "essential_references_not_discussed": "While the paper comprehensively cites many foundational works, some recent studies could provide additional context for its key contributions. For instance, adaptive pooling strategies in transformer architectures—such as the dynamic pooling method proposed in [Lee et al., 2024] at ICML, which adjusts to varying noise levels—offer alternative approaches that are highly relevant to the paper's focus but are not currently discussed. Additionally, although the authors reference classical vector quantization and associative memory literature, they omit recent work on noise-robust pooling in graph neural networks (e.g., Xu et al., 2022), which addresses similar challenges in aggregating noisy, high-dimensional data. These omitted references highlight complementary ideas and alternative mechanisms for robust aggregation that could enrich the theoretical and empirical framework of the current work. Including a discussion of these studies would situate the paper more firmly within the broader landscape of robust representation learning.", "other_strengths_and_weaknesses": "The paper is notable for its originality in creatively combining vector quantization techniques with attention-based pooling mechanisms, which provides a fresh perspective on mitigating noise in transformer outputs. Its strengths include a solid theoretical framework with detailed proofs that derive meaningful error bounds, and an empirical evaluation that spans both synthetic and realistic RL benchmarks, underscoring the practical relevance of the approach. The work significantly advances our understanding of how traditional pooling methods like AvgPool and MaxPool fail under varying noise conditions and demonstrates the potential of adaptive pooling in overcoming these limitations. However, the reliance on idealized assumptions regarding the separability of signal and noise in relation space may restrict the direct applicability of the theoretical guarantees in all real-world scenarios. Additionally, while the paper is generally clear, some of the denser theoretical derivations could benefit from more intuitive explanations to enhance accessibility for a broader audience.", "other_comments_or_suggestions": "The paper is overall very strong, but a few suggestions might help improve its clarity and accessibility. For instance, some of the theoretical derivations are quite dense—adding more intuitive explanations or illustrative examples could make these parts more digestible for a broader audience. There are also a few minor typos and notation inconsistencies in both the main text and supplementary material that should be corrected. Additionally, while the experimental setup is comprehensive, it might be beneficial to include a discussion of scenarios where the assumptions (e.g., the clear separation between signal and noise) might not hold, and how the proposed method would perform in those cases.", "questions_for_authors": "1. Could you provide more insight into how robust the adaptive pooling method is when the assumption of clear separability between signal and noise relations (as per Assumption 3.10) is relaxed or violated? A more detailed discussion or empirical analysis in scenarios with overlapping distributions would help assess the method's applicability in real-world settings.  \n2. In your synthetic experiments, how sensitive is AdaPool to the choice of hyperparameters in the attention mechanism? An ablation study highlighting the impact of key parameters would clarify the robustness and ease of tuning the method.  \n3. Have you empirically validated the theoretical error bounds derived in Theorem 3.11? If so, could you compare the theoretical predictions with actual observed performance on benchmark datasets? This connection would strengthen the confidence in your theoretical claims.  \n4. Since AvgPool and MaxPool are special cases of AdaPool, can you discuss any potential trade-offs in computational complexity or performance when scaling AdaPool to larger datasets or more complex models? Understanding these trade-offs would inform its practical deployment.  \n5. Could you elaborate on AdaPool's performance in scenarios where the noise is highly correlated or where the noise and signal distributions significantly overlap? Clarification on these cases would help determine the method's limitations and potential areas for further improvement.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LHwl675v8b", "reviewer": "Reviewer_GbuG", "summary": "This paper studies the problem of pooling methods in transformers for non-sequential tasks where only a subset of input tokens (signal) is relevant for downstream decision-making, while the rest (noise) may degrade performance.\n\nThe authors formulate a theoretical framework that formalizes pooling as a vector quantization problem, leading to error bounds for different pooling methods.\nUnder this view, the authors analyze common pooling methods—average pooling, max pooling, and class token-based pooling—and show that each fails under certain signal-to-noise conditions. They then investigate adaptive attention-based pooling (AdaPool), a technique adapted from Stergiou & Poppe (2023), and demonstrate that it can approximate an optimal pooling function across varying noise conditions.\nThey provide empirical validation for their theoretical findings on synthetic supervised tasks and reinforcement learning (RL) environments (Multi-Particle and BoxWorld), demonstrating AdaPool’s robustness to varying signal-to-noise ratios.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your time and feedback. Your questions raise some important points that we should have discussed in the original submission. \n\n**1. Computational Complexity**: AvgPool and MaxPool can be implemented with `O(n * d)` algorithms, as both require each of the `d` features of each of the `n` vectors to be visited during aggregation. \n\nSelf-attention famously experiences quadratic computational costs with respect to the context size `n`. In particular, the attention weight computation is ` O(n^2 * d)`, as each vector in the input set computes attention weights for each other vector in the input set, i.e. `n^2` dot products of `d` dimensional vectors. However, AdaPool uses cross-attention with a single query, requiring the computation of only 1 set of attention weights rather than `n` sets. This reduces the weight computation and pooling time down to `O(n * d)` - the same as Max & Avg. \n\nHowever, AdaPool still retains the overhead of the QKV projections of the input set, which involves standard matrix multiplications for each input vector, running `O(n * d^2)`. Thus, the overall time complexity of AdaPool is ` O(n * d + n * d^2)`. The key takeaway is that it scales linearly with the number of inputs (like max/avg) but quadratically with the number of features (like a standard linear layer). \n\nIn isolation, AdaPool is slower - the increased expressivity comes at a computational cost. However, in practice, these pooling layers are placed at the end of a multi-layer transformer or similarly sized encoder network, and the compute time added by the pooling layer is marginal. For our particular applications, when using a transformer with 3+ layers, the time difference between AdaPool and Max/AvgPool networks was negligible, regardless of the number of features `d`. \n\nAlso of note, ClsToken is significantly slower than all other methods. This is due to the fact that the learned class embedding must be copied and concatenated along the batch dimension of the input for every inference, which is quite expensive in terms of both time and memory during training. This is a known hindrance investigated by Zhai et al. (https://arxiv.org/abs/2106.04560), which led them to explore the use of AvgPool for efficiency reasons when scaling vision transformers.\n\n**2. Handling Ambiguous Inputs:** To address this question, we conducted additional experiments on an image classification task (CIFAR 10/100) where each input vector may contain nebulous amounts of signal or noise, in contrast to the other experiments we presented. Due to character count limits, please see our response to reviewer PXAR above for details about these experiments and their results. In summary, the findings were consistent with our theory and other experiments. \n\n**3. Multi-Query Extension:** Using multiple queries or even employing multiple different pooling methods and concatenating the results may result in even better downstream representations. We have not tested this, but it is certainly a viable approach that is worth looking into in the future.  \n\nIncorporating this additional discussion into our revisions will help make the paper much more robust. Thank you for your review!", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "The proofs and theoretical claims seem sound, however, I did not check all the math in detail.", "experimental_designs_or_analyses": "- Synthetic dataset experiments are well-structured to explicitly test signal loss across varying SNRs.", "relation_to_broader_scientific_literature": "Connects well with prior work in attention-based pooling, vector quantization, and associative memories.", "essential_references_not_discussed": "I am not very familiar with the related work, as such I cannot answer this question.", "other_strengths_and_weaknesses": "The paper is very well written and nicely presented. Even without deep knowledge of the literature I had no problem following the authors.", "other_comments_or_suggestions": "\\-", "questions_for_authors": "- How does AdaPool’s computational cost compare to avg/max pooling, particularly for large-scale models?\n- How does AdaPool handle cases where some tokens are partially informative but not strictly signal or noise?\n- Have you considered extending AdaPool to multi-query setups for more robust feature aggregation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper focuses on problems where the network inputs contain a mix of signal and noise components. The authors demonstrate that traditional pooling methods (like MaxPool and AvgPool) are affected by high levels of input noise, and propose AdaPool as an attention-based pooling mechanism to better deal with noisy inputs. The authors provide theoretical analyses bounding the error of their method, and demonstrating that MaxPool and AvgPool are special cases of AdaPool. The authors further demonstrate the efficacy of their method on a number of synthetic experiments.

With the exception of 3DVE (whose review is suspect and I am not considering in my decision), all reviewers are supportive of this work, and I agree with them.
Decision: Accept (spotlight poster)