Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models
Authors: Sofiane ENNADIR
Levente Zólyomi
Oleg Smirnov
Tianze Wang
John Pertoft
Filip Cornell
Lele Cao
Venue: NeurIPS 2025 poster ()
Keywords: Transformer-based Models
Pooling
Expressivity
Primary area: deep_learning
TL;DR: We theoretically and empirically analyze the effect of the Pooling Operation in the downstream performance.

Abstract: Transformer models have become the dominant backbone for sequence modeling, leveraging self-attention to produce contextualized token representations. These are typically aggregated into fixed-size vectors via pooling operations for downstream tasks. While much of the literature has focused on attention mechanisms, the role of pooling remains underexplored despite its critical impact on model behavior. In this paper, we introduce a theoretical framework that rigorously characterizes the expressivity of Transformer-based models equipped with widely used pooling methods by deriving closed-form bounds on their representational capacity and the ability to distinguish similar inputs. Our analysis extends to different variations of attention formulations, demonstrating that these bounds hold across diverse architectural variants. We empirically evaluate pooling strategies across tasks requiring both global and local contextual understanding, spanning three major modalities: computer vision, natural language processing, and time-series analysis. Results reveal consistent trends in how pooling choices affect accuracy, sensitivity, and optimization behavior. Our findings unify theoretical and empirical perspectives, providing practical guidance for selecting or designing pooling mechanisms suited to specific tasks. This work positions pooling as a key architectural component in Transformer models and lays the foundation for more principled model design beyond attention alone.

Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models
Sofiane Ennadir*
Levente Zólyomi*i
King AI Labs, Microsoft Gaming sofiane. ennadir@king com
NXAI GmbH
levente eeventezzolyen zolyomi Cnx- att.ai..oo. com
Oleg Smirnov
Tianze Wang
King Al Labs, Microsoft Gaming oleg. wmwrnov@miccommmmmmmmmmmo crosoft. com
Kreditz AB
tianze. wang@kreditz.com
John Pertoft
Filip Cornell Amazon
King Al Labs, Microsoft Gaming john pertoferooff weattoot@king.co com
com
Lele Cao
King AI Labs, Microsoft Gaming wel.ecao@iicroooof..ooo com
Abstract
Transformer models have become the dominant backbone for sequence modeling. leveraging self-attention to produce contextualized token representations. These are typically aggregated into fixed-size vectors via pooling operations for downstream tasks. While much of the literature has focused on attention mechanisms, the role pooling remains underexplored despite critical impact on model behavior. In this paper, introduce theoretical framework that rigorously characterizes the expressivity of Transformer based models equipped with widely used pooling methods by deriving closed-form bounds their representational capacity and the ability to distinguish similar inputs. Our analysis extends to different variations attention formulations, demonstrating that these bounds hold across diverse archi- tectural variants. We empirically evaluate pooling strategies across tasks requiring both global and local contextual understanding, spanning three major modalities: computer vision, natural language processing, and time-series analysis. Results reveal consistent trends in how pooling choices affect accuracy, sensitivity, and optimization behavior. Our findings unify theoretical and empirical perspectives, providing practical guidance for selecting or designing pooling mechanisms suited to specific tasks. This work positions pooling as key architectural component in Transformer models and lays the foundation for more principled model design beyond attention alone.
Introduction
The profound impact of the Transformer [39] architectures across different modalities and in cross- modal modeling cannot be underestimated. These models have become the building stones of
"quual contribution.
Work conducted while at King Al Labs.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

large-scale systems capable of successfully addressing a multitude of downstream tasks in com-
puter vision [.. natural language processing (NLP) [37 14,, and time series analysis [22] 12] Since these models typically require substantial training data to perform effectively, pre-trained large-scale models trained with self-supervised objectives such as autoregressive, autoencoding. contrastive, or hybrid formulations have become the standard. These models are first trained to capture rich, contextualized representations and are later fine-tuned by attaching a classifica- tion regression head specific to the downstream task, thereby serving as feature extractors.
Transformer-basee based models produce sequence of token- level embeddings, which are typically aggregated into a single vector that captures tesk-relevant information. This operation, commonly referred to as pooling, has been widely studied in domains such as multi-modal learning (10) and
Lest Thkre
Graph Neural Networks [23], where it is recognized as a key architectural component. The choice of pooling function directly affects the content of the final representation and thus the model's performance on downstream tasks. Re- cent studies have increasingly focused on understanding Transformer models from both empirical and theoretical per- spectives. However, much of this research has concentrated on the backbone encoder, which processes the imput into con- textual embeddings, often neglecting the final pooling step that aggregates these into single representation for predic- tion. Although some empirical work (35] has explored the effects of various pooling strategies, systematic theoretical analysis is still largely absent.
Atnestica
Weadned Avetage
Clossidc
Clesifrsiko Tweet Seattment
Tilaes Prodbcton
Bonking
Figure 1: Performance of different pooling strategies using a GPT-2 pre- trained model.
In this work, we close this gap by providing end-to-end analysis of Transformer-based models that explicitly incorporates the pooling stage. We introduce a theoretical framework to quantify Transformer expressivity, measuring the model's 's ability distinguish dissimilar inputs while pre- serving similarity. Applying this framework, we analyze how common pooling functions influence expressivity by encoding different token-level properties. From these insights, derive practical guidelines for choosing pooling methods based on a task's need for local versus global context. We then validate our theory with experiments in vision, natural language, and time-series domains using standard models and datasets. The results confirm that no single pooling strategy dominates all tasks, underscoring the importance of task-tailored pooling design. To our knowledge, this is the first theoretical examination of pooling mechanisms, offering a unified treatment of standard strategies found in the literature.
2 Related Work
A growing body of work has provided theoretical insights into Transformer-based models and their attention mechanisms. Prior research has explored broad range of topics, including training dynamics (36). inductive biases [18], and in-context learning [41] Much of this literature focuses on the core Transformer backbone, proposing architectural improvements and refined training procedures. However, in practical scenarios, particularly those involving large pretrained models, the backbone is typically frozen, and a lightweight classification or regression head appended after the final pooling layer. This common design choice limits the direct applicability of many of the previously proposed modifications.
Prior studies have demonstrated that different pooling strategies can significantly influence down- stream performance. For example, (6) showed that, in Vision Transformers (ViT), the choice between using CLS token and average pooling leads to measurable differences in classification accuracy. Those findings were reaffirmed a follow-up study (28). which examined the influence of the classifi- cation token across different Transformer layers. In the language domain, [9 conducted an empirical analysis of BERT-based [5] embedding models and decoder-only models from the GPT 333 family, along with their standard pooling layers, and proposed new technique to mitigate information dilution and recency bias commonly observed in the respective pooling operations. In addition, (35) systematically evaluated combinations of attention mechanisms and pooling methods in large lan-

---

guage models, and introduced an attention baaeet pooling approach that aggregates representations from multiple hidden layers. However, the performance of these novel pooling mechanisms has been shown to vary considerably across different types of tasks.
Multiple studies have further examined the impact of pooling strategies across variety of settings [38 43.. This line of research concludes that the optimal pooling method generally depends on the specific downstream task, well as other factors such as model size. Although these works provide extensive empirical evidence, to the best of our knowledge, none have proposed a principled theoretical framework explain the behavior of pooling operations in a broader context.
More recently, increased attention has been directed toward studying the theoretical properties of Transformer architectures through the lens of Lipschitz continuity, with the goal of understanding its behavior and dynamics. For instance, [.5] propose an 1.2-based attention mechanism that replaces the standard dot-product, providing both theoretical and empirical evidence of its Lipschitz continuity. In a similar direction, [4] introduce Lipschitz.Norm, a normalization technique designed to enforce Lipschitz continuity within the self-attention mechanism. Building on these efforts, [32] further redesign the full Transformer architecture to ensure that the model remains Lipschitz continuous throughout. However, these analyses typically do not extend to the final pooling operation, despite its
widespread use in practical applications.
Our work extends both lines of inquiry by closing the gap between empirical findings and theoretical understanding, and contributing to a deeper comprehension of how pooling functions influence the performance of Transformer--asee models.
3 Preliminaries
We start by reviewing the Transformer architecture and its key components, forming the basis for the concepts of our theoretical study.
Transformer Architecture. Let X E X S Rxdd denote sequence of n tokens, where each token I; E Rdd The backbone of a transformer the self-attention mechanism, which computes a weighted combination of all token representations. Specifically, given learnable query, key, and value parameter matrices WQ, WK, WV E Rdx(d/H), the output single attention head AH for input X is defined as
X S Rand + Z S Rxxd as introduced in [39], is
(XWX)(XW*) WOWKY
AH(X) softmax
(XWV),
(I)
Va/H
where H denotes the number of parallel attention heads and d/H is the dimension per head. In practice, multiple attention heads AH, are computed in parallel, then concatenated and projected using a learnable weight matrix /0 E &Rdxd, yielding the Multi-Head Attention (MHA) operation:
MHA(X) = concat (AH,(X), AH2 ... ,H (X)
(2)
Attention Block. In addition to MHA, each Transformer attention block (AB) incorporates a residual
connection, layer normalization (20) and position-wii wise norward network (FFN), and can be written the following two steps:
X' LN(X MHA(X)); AB(X) = LN(X' + FFN(X')).
with (.) denoting the layer normalization operation. FFN(-) is a feed-forward network, formulated as FFN(X') a (X'WErN) with being a non-linear activation function. While different placements of normalization layers, commonly called Pre-LN and Post-LN, have been examined in prior work (11) this study focuses on the original Post-LN setup. We note that our main findings are transferable to other configurations.
Given the output of Transformer backbone Z E 2 C Raxd, pooling function g: 2 Y C T2d produces a fixed-size embedding for downstream tasks. Common choices are Average pooling, which computes the mean over tokens; Sum pooling, which sums the token embeddings; Max pooling, which takes the elementwise maximum; and Last-token pooling, which selects a designated token (for example the final or CLS Thhen).These These operations can be formally defined as:
Pooling,
n
12
(Av = n
o]. DSum (2)
= Σ Zi
9Max (Z) = max Z Zii,3]: Mtas (Z) = Z(n, :).
1=1
3

---

Problem Setup. Let f: X' S RTxd + + J Rd be Transformer- based model incorporating final pooling layer. For our theoretical analysis, we model f as single MHA block with H heads, followed by one-layer FFN and the layer-normalization variant of [32. which is provably stable under small input perturbations. We assume all activation functions are I-Lipschitz (e.g. ReLU, LeakyReLJU, H)n [40], and that the input space is bounded, X' C [0, B Bjnxd This bound is realistic: in vision and time-series applications B = after normalization, and NLP the embedding process usually results in bounded input representations due to the initialization of the embedding matrix.
4 On the Expressivity of Transformer-Based Models
In this section, we introduce the notion of expressivity for Transformer-aseee models (TBMs). Building upon this definition, we develop a theoretical analysis of several attention mechanisms and commonly used pooling strategies.
4.1 Expressivity of TBMs
Inspired by work in graph representation learning, we define exprexsivity as the capacity of a model to distinguish between similar and dissimilar inputs 44 27 26.. Specifically, by being able to distinguish between such cases, model is able to produce meaningful representation that could be used by a classification or regression head to produce the final downstream task, For instance, in natural language processing, two semantically similar sentences should yicld representations that are closer in the output embedding space than those produced by two semantically disparate sentences. In this perspective, defining across diverse domains such as NLP and computer vision, is crucial and fundamental to evaluating expressivity. Let X' and ) denote the input and output spaces, respectively. we consider both spaces to be measurable and equipped with a measures |x and . ly. With a well-designed embedding function, semantically similar elements from the input space mapped to proximate points in the output space. Consequently, the distances in y are expected to accurately reflect the semantic relationships present in X.
accurate measure of semantic similarity that is applicable
Let f: S Radd + y S Rd be a TBM as defined in Section 3 For a given input X E X, we define its neighborhood with respect to the input distance metric and threshold by:
B(X, e) = (X e X |
X x Ec..
Since the desired behavior is for inputs in close proximity to yield similar output representations, we consider the following measure:
{f P PX x B(X,e) ((x) oj.
where Dx represents the underlying distribution over the input space X, and dy is distance metric on V. The quantity [f] encapsulates the probability that two inputs, which are similar (within the same neighborhood) in the input space X, are mapped to outputs that differ by more than a threshold Intuitively, small input distance should yield correspondingly small output distance, while larger differences in the input should result in more pronounced variations in the output. A model that appropriately distinguishes these nuances is said to exhibit stronger expressive power, which is as motivated previously a vital attribute for achieving robust downstream performance. Definition 4.1 reflects such expressivity in the context of TBMs.
Definition 4.1. Let f: X C Raxd + y C d be TBM. The model f is said to be (E, expressive if & [f] 7.
Definition depends on several hyperparameters. The threshold e specifies when two inputs are considered semantically similar and is inherently application-specinn perturbation in an image may be negligible, while the same in a financial time series could be meaningful. The parameter a defines the acceptable variation in the output space for representations to be considered similar. By setting e based on domain knowledge, the interaction between and o allows us to capture the model's 's expressive capacity in a way that reflects the semantic structure of the data. This formulation highlights the model's adaptability and expressivity application-specffic contexts.
For instance, a minor

---

4.2 Expressivity of Pooling Strategies
Based on Definition 4.1 and given fixed values of and o, our objective is to quantify the correspond- ing expressivity parameter for different pooling strategies. This analysis allows us to assess how the choice of pooling influences the model's ability to distinguish semantically meaningful variations in the input. Throughout the remainder of this paper, ||- denotes the operator norm.
Theorem 4.2. Let I: X C Raxd y C IRd be a TBM following the frumework introduced in Section In respect to Definition 4. we have:
( ( [
If employs Average pooling, then I is (e, -) with 7
_
C1C2
ovR
2
If f employs Sum pooling, then f is (e, - with 7 vice
C₁C2
2-1
If f employs Last-token pooling, then S is o, expressivve with = ( C1O2
If f employs Max pooling, then f is (e, o, 7)) with y symiulndd)
( CiC,,
2
with: C1 = =+ max [ ( n
(170)T B² W
W 1 +1)] C2 =2=1+ 1+ Werx
Theorem shows that expressivity bounds across pooling strategies depend on shared architectural parameters, including the number of attention heads H, embedding dimension d, and sequence length n. captured through constants C1 and C2. These constants reflect the norms of key model components, such attention weights, projection layers, and feed-forward networks. Each pooling function introduces distinct scaling effects, shaping how these elements combine to influence the model's ability to separate similar from dissimilar inputs.
For Average pooling, the bound scales with 1/ 1/v smoothing the output by evenly distributing token contributions. This favors tasks where global structure matters more than individual token details. In contrast, Sum pooling scales with v vT,, amplifying token-level variation. This is useful for tasks where localized information is essential. Using a single token (e.g.. the last or CLS token) leads to a scaling of 1, preserving variations without change. This suits scenarios where a specific token encodes the most relevant context, such as in sentiment amalysis. Max pooling introduces bound that scales / /minnn,, d). When d is large relative to n. it behaves similarly to Sum pooling, capturing fine-grained differences. When is smaller, it emphasizes broader context. This flexibility enables Max pooling to shift between local and global focus based on model size and sequence length.
Overall, the theoretical results emphasize that pooling a key factor in how Transformer models aggregate local token information into a global representation. Theorem 4.2 formalizes how this choice affects model expressivity across different settings.
On the generalization to multi-layer TBMs. single-layer Transformer-bastt based model; nonetheless, the results naturally extend to the multi-layer case. Specifically, a Transformer model with L layers, denoted as f(L), be expressed
We note that the current theoretical analysis focuses on
composition of L single-layer functions: f()) (z) (I.-1) f(r.-2)
f()) (x). Under this formulation, and following standard results from Lipschitz continuity, the overall expressivity bound bound for each pooling becomes multiplicative composition of the bounds for each individual layer. As a result, our theoretical study remains applicable to deeper architectures, as confirmed by experiments involving exclusively multi-layer models.
4.3 Expressivity of Alternative Attention Mechanisms
Recent studies have proposed alternative formulations of the scaled dot-product self-attention mech- anism to improve model behavior and facilitate theoretical analysis. For example, L2 -ead Attention (L2-MHA) [15] employs an L.2-kernel attention function, while LipsFormer [2) replaces the doo-product with scaled cosine similarity. The theoretical results from the previous section are general and extend to these variants. In the following, we consider the same problem setup, with the only change being the use of an alternative attention mechanism in place of the standard formulation.
5

---

Lemma 4.3. Let f: 2' S Raxd + y S Rd be L2-MHA-based TBM 15
In respect to
Definition 4.1 the following holds:
2
If f employs Average pooling, then f is (c,a, (.7 expressive with ry ovn c
C₁C2
If f employs Sum pooling, then f is (e, exxressiive with vice d
C1C2
1-1
If employs Last-token pooling, then I is (t, eexpresssvv with 2= T-
d CiO²
If employs Max pooling, then f is (c,0, -xxppressvee with symmn((nd) (211) C.C2.
with Ci
n H
4Wo
12 (7)
(V>
Σ | www.ii/iaaiii. || c==1+ WWFNN
15 in which the query and key matrices are tied. This constraint influences the constant C12 in the expressivity bound, reflecting the interaction among shared parameters, sequence length n., number of heads H, and embedding dimension d. Compared to the standard -ro-prrduct formulation, this structure alters how the L.2-kernel shapes the bound, resulting slightly different expressivity dependencies.
Lemma 4.3 analyzes an L2-based attention mechanism
The pooling-related terms in Lemma are consistent with those derived under standard self- attention, and the same trade-offs between local and global context remain applicable. Similar behavior is observed in Swin 24 and LipsFormer (32), which employs scaled cosine similarity and normalizes the key, query, and value matrices maintain Lipschitz-continuity Lemma provides the corresponding bound for single-layer model with H attention heads and window size w.
In both cases, the bounds reveal comparable structure, reinforcing that pooling remains critical component in controlling the balance between local preservation and global aggregation in TBMs.
Lemma 4.4. Let f: X + y to be function based on the LipsFormer [32] framework, with corresponding hyper-parameters V, >> 0 and window size w. In respect to Definition 4.1 we have:
If f is based on Average pooling, then f is (,, ( with 7 oxxn (2)) Gc
If f is based on Sum pooling, then f is expressive with y -Txe x d-1 )
d If f is based Last-token pooling, then f is (e, o, -77-eexexsssiv with 7 = T=- C22
If f is based on Max pooling, then f is (e, () exprresssv with 7 s/min(ndd) x
( 2 C,
where
|| C 1+ [Woll H max 2w(w L)urV [ [|+(( 2, 1 1 |
C2 (1 [Wrexll)
5 Experimental Validation
Our theoretical analysis shows that the pooling strategy influences the classifier's expressivity bound via a leading multiplicative factor, which be contractive (1 / /v), non-expansive (1), or expansive (yn or / /min(n, d)). Contractive methods like Average pooling enhance stability by smoothing small variations, whereas expansive methods such as Last-token and Sum pooling increase expressivity, but can be more sensitive to minor perturbations. Therefore, we posit that pooling should be selected based on task requirements: tasks emphasizing global context (e.g., image inpainting or text classification) benefit from contractive pooling, while those relying on local detail (e.g.. next-token prediction) may perform better with expansive alternatives. In this section, validate

---

these theoretical insights through empirical evaluation, to demonstrate the applicability of our findings and provide practical guidance for choosing pooling strategies different domains.
Experimental Setup. We evaluate how pooling choice affects downstream performance across three domains where TBMs have shown strong results: (a) computer vision, (b) natural language processing, and (c) time series analysis. For each modality, we select diverse set of established benchmarks with tasks requiring global and local contexts. Across all settings, examine commonly used pooling methods: (i) Last-token pooling (or CLS/EOS, depending on the task), (ii) Average (Avg) pooling, (iii) Sum pooling, and (iv) Max pooling. We also include two learnable strategies: (v) Attention (Attn) pooling, which uses learnable latent dictionary attended by the model output [19] 35|, and (vi)
Weighted Average (W-Avg) pooling, which learns scalar weights over token positions. Further details on training and evaluation protocols are provided in Appendix [D] Our code and implementation to reproduce the results is available the following link: http:///hwuuu.oooo////n///nnnnee
5.1 Expressivity Analysis
We begin by empirically analyzing the expressivity of the pooling strategies under study, in accordance with the theoretical bounds introduced earlier. Using the framework in Section 4.1 we define local neighhorhoods by injecting Gaussian noise into input samples, scaled chosen We then compute the average distance between the resulting pooled outputs, yielding an empirical estimate of
[1) CV: VIT om CIFARLOD Dataset
similar
11F-
Pooled Outpur. Discace (log soale)
10
similag-2
Aue
Sun.
Mas
non-simular-1 non-sinular-2
17.0
0.2
0.4
10.0
Pooled Distance (log
(2) NLP: GPT-2 on Sample Tea
10"
10
Max
10
11.0
0.2
0.4
10.0
0.8
(3) TS. MOMENT- s-mall on Electricly Dalese
10'-
Avg
Max
Luci
Sumn
Pooling Method
10-1
original:
saw significant fluctuotions
sinilar-1: The stock market encountered notable swinggs last week.
Stet Mas.
sinilar-2: The equity market experienced marked valatility last week.
0.4
0.8
Lo
non-similar The stock market encountered notable castle last week. non non-simi swi-ssiilla---. 2: The stock market encountered notable oar last week.
Noorhoor Toeshold
Figure 2: Empirical analysis of the expressivity power across modalities and pooling strategies. Left: Mean pooled-output distance y versus perturbation across modalities highlighting the behavior of various methods. Right: pooled-output distances for similar and dissimilar inputs, exemplifying expressivity of different strategies.
Figure 2 (left part) presents results across different values and modalities, confirming the theoretical contrast between contractive and expansive pooling. Sum pooling shows high sensitivity to even small perturbations;
increases, its Y grows rapidly. In contrast, Average pooling remains stable. To further illustrate this, Figure 2 (right part) shows how replacing a word with either a synonym or a semantically different term in NLP settings affects the pooled output across different strategies. This supports our hypothesis that expansive pooling better captures subtle variations, as required in tasks like sentiment analysis.
5.2 Effect on the Downstream Performance
In line with our theoretical analysis, we empirically evaluate tasks with varying dependence on local versus global context to assess how different pooling strategies perform. This allows us to validate the extent to which each method's empirical behavior aligns with its expected theoretical properties.
7

---

Table Mean and standard deviation of test metrics for computer vision tasks. Best performance per dataset and model is indicated in bold. Best performance among non-leanee non-learnable pooling methods is underlined,
Chonliction (Accarac)
Iepoieting (MSE)
Segencotiaion (Avonmse))
Pooling
CIFAR-100
Imrashks-100 100
Miniflac
Caktas
OctooRRowcs-002
Osfosk-UIT Pot
Pacc PaadVOCCC
Laan (13.5)
00:13
13.45
26.12 1:0
ST.NG E000
27.45
1.33
11200 2 111611
0.275 1.275-100.00
03305 2
TILB 22 14:35
21.11 1::33
Ang
90:11
0.11
76.26 16.6-1-0.0 5-0.3
65.5544 0.10
71.48
02209 6 0.001
0.261 0.001
0.900 0
71 66 4 0.00
35. 10:4-0055
85.45 10:33
75-72 1.999
12
16.37
11.711667
01001
0:151 0.012
0.303 1:000
08.02 0.97
35:31 1.01
0.03
8T
00:96:0.11
1286-200006
0.077 000007
11.88 0.00
0:022
74.51 +0001
$8.8
0 0.003
0.273
24:22:00.00
07:20 0013
1618
18.17 5.3
0.255 5 8:00
0:309:1-00003
70.00 8:01
53.35
75.05
0.01
50.76
72:25:022
0.:06
11200:11000
0.237 0.257 0.M6
TUL11 0.36
Sun
96.82 0.005
72.70=00113
6:24
66,96 0 0.13
$7.08++0006
0.368 e 0.00
0.064:20006
0.305 0.0000
T2.02 e: 0:06
24.27 +.229
0:035
9:155
16:66
16.15
0:30-4 0/913
0:317 0:0.00
T1.
0:000
7..00
00:08:0-:11
05 16:21:2
0.90
0.351 4:0.00
16.08 0.16
31.39-50..00
77.37
Computer Vision. Results for image-bases tasks are shown in Table 1 As predicted by the analysis, Average pooling outperforms other fixed pooling methods in inpainting, segmentation, and in the MiniPlaces classification dataset, which benefits from modeling global structure. In contrast, Last-token (CLS) pooling generally yields the best results on classification tasks, particularly those involving large-scale fine-grained datasets where local information is more critical. Max and Sum pooling consistently perform worse across all tasks. These trends also hold for alternative attention mechanisms, such as LipsFormer (32), which employs scaled cosine similarity attention (see Lemma 4.4. In this Swin-based model, which does not include a CLS token, Average pooling again achieves the best performance among fixed strategies, further emphasizing its strength in capturing localized context when dedicated classification token is absent.
Among learnable pooling methods, Weighted Average pooling performs competitively across tasks, likely due to its capacity to adaptively weight tokens based on task-specific context. Attention-based pooling often underperforms, except in high-resource settings such ImageNet-1000 where sufficient supervision allows it to learn effective attention patterns. Its reduced reliability in low-resource or fine-grained tasks may be attributed the additional complexity introduced by its parameterization.
NLP. The impact of pooling across downstream NLP tasks and models is summarized in Table The results support our theoretical analysis: no single pooling method is optimal across all tasks, and the best strategy depends on the task's 's contextual requirements. For tasks requiring global context, such as classification and semantic similarity, global pooling methods like Average Sum significantly outperform Last-token pooling. Conversely, Last-token pooling yields superior performance in next-token prediction tasks.
These trends hold across models, although the performance gap between local and global pooling narrows for larger architectures (e.g.. Mistral-7B [14] and Llama (37)). In such models, Weighted Average pooling matches or exceeds the best-performing non-lcarnable methods, due to its ability to adaptively approximate effective pooling strategies. In smaller models, particularly the GPT-2 (33 family, Attention pooling performs well on global-context tasks, often outperforming fixed global methods like Average or Sum. However, this advantage does not consistently generalize to larger models or all task types. Additional results on larger models are provided in Table (Appendix E We additionally analyze the effect
Time Series. As shown in T for time series classification tasks, Last-token and Max pooling generally yield the worst results, as they focus on local features and fail to capture the broader context required for accurate classification. In contrast, Attention-basee pooling consistently achieves the highest performance, due to its ability to assign task-specific weights to different input segments during joint training. Sum pooling outperforms Average pooling in several cases, that can be attributed to the larger norm of summed representations, which results in stronger gradients and faster learning under fixed training hyperparameters.
Forecasting shows that Last-token pooling yields the best results overall. This is consistent with its ability to retain fine-grained temporal details, which are critical for predicting future values based on recent history. In imputation tasks, results indicate that Attention-based pooling again performs best, while Sum pooling performs the worst. This supports the hypothesis that given sufficient data and training time, Attention pooling can adaptively focus on the most relevant parts of the sequence,
8

---

Table 2: Mean and standard deviation of test metries for NLP tasks. Best performance per dataset and model is indicated in bold. Best performance among non-learnable pooling methods is underlined. (-) indicates non-applicable, as the model uses bidirectional attention mechanism.
Dataset
Pooling
BERT
12-GPT2
CPT2
Q.c
Misual-BB
Llanu3-58
(8peammni
0-587 0.0000
0-378 0.006
1.602-2
D. 0.005
0-286 d: 006
0-534 0.001
0.017 +: 0.085
5758
0.004
0.670
000003
0.619
0.036 0.005
Similanity
0.002
0.653
0.002
0.017
W--aeg
0.727
+: 002
0.582 0.002
0.588
D. 0.003
0.871 = 0 00.00
0,458
0.673 + 1
0.088
Last
0:307
0.001
0.231
022
0.264
D 0.024
0.344 1 0.001
0.770 0.002
675 1: D D02
Hellsoog
Avg
0.315
0.000
0.297
00
0. 305
0 0.000
0-452
0.00
0..69 0.000
0.784 + 0.000
(F1)
0.001 0:816
0:209
00T
0.008
0.451
0.00T
0.769 0.000
0.001
0.003
0.002
0.293
0.002
0.364 0.008
0,709 0.001
0.005
W-Are
0.316 + 0. 001
0.264 0.001
0.278
0.001
0.452 + 0.000
0,801
0.000
0.763 + 0. 0.000
Ami
0.205
0.004
0.012
0.260
0.088
0.410 0.009
0.787 0.018
0.460 0.271
Last
77-175 +: 0.112
17 0144 1: 0.323
45. 528 1: 244
23.253
0.378
74.847 +: 1.584
45 107 1 0-403
Baankigg
(Accursc))
Sum
Avg
85.142 + 0.002
83.865 2: 0.200
86. 852 + 0.076
130133222 4: 0.807
83 131-724-2 4: 1.047
-0.497 +.0 0.103
78.161
$3.486
0.776
0.402
88. 183 + 0.402
86-42 0.708
$7 65.688 + 0.523
1: 22.324 1.002
Chavdhoniee
Max
50.755
0,008
$3.00 0330091 0.315
$3.023 0.226
74.007
1.106
74.890 2.476
W-Arg
84.967 +: 0.153
67. 253 4: 0.275
73.389 4: 0.223
85.613
0.221
89 271 0.426
88 928 0 296
All
858 0.589
27.3402 + 0.302
86.968 0.604
63.792 = 1:813
73:352 + 1,731
51.1 + $4.316
Lasi
67.621 + 0.085
48. 583 + 0.204
63.8800 899 1: 0.145
51.758
0.713
56, 603 + 0.192
80- 466 2: 0.538
Twret
(Ausmme))
Sam
AVg
6O 244
65:351
0. 1228
0:275
67 055 1 0.103
+ 0.195
G9 5735
1.0 200
0.317
59.625
08.961
0 0.330
2:702
G7 231 4 0.214
11.151 1.623
04:231
67 128 022
67.070
0.177
0.254
21.200 0:178
02:971
0.433
2.401
1,899
W-Arg
560 0.121
622 d5is d: 0.123
7118 224
06.031
0.303
67 293 0.228
67 477 0 1.16
Alli
69.455
+: 0.879
69.181 +: 0.529
70.844
0.621
70.627
0.668
46.684 + 6.908
56. 850 + 13.252
Timy Stories
(Top 10 No)
Laxi
Avg
82, 170 + 0.225
SY-215 4 0.223
$4.589 + 0.001
55.3226 7: 0.455
86.634
10:054
0.428
0.327
60 948 + 0.002
05.0017 + 0.225
90, 808 t: 0 0.227
11: a: 0.537
Sum
24 8012 d: o 12442
3912 8933
28.911
0.630
50-481 0.625
7:81 1: 1,2188
Neat
W-AAg
3100 d: 0.018
53
0.348
86 850
212
89: 160 4: 0.118
3est 0.039
50. 60:364 0.287
53.299
0.335
65.970
0.302
14.210 5.239
11.188 8,980
Table 3: Mean and standard deviation test metrics in time series tasks. Best performance per dataset and model in bold. Best performance among non-learnable pooling methods is underlined.
Chosiiccaiion (Accarasy)
Forccasting (MSE)
Imputation (MSS)
Model
Foeling
BCC2200
Devioo
Appliances
Eactrocy
Thalfc
17Th
Blurriity
Traffc
Last
12.21 2
00.45 a 01:8
70.30 2: 0.15
65005 2
00002 2 0.000
0.0003
11001
6,213
0.901
1101 1: 0002
0.35 5: (L000
1 2 0.155
MOMENT
12
0.106 2 0,030
0.700
1.734 0.000
0.000 2,000
0,374 0.013
2 0.015
Sam
20 10:35-1
00.52
[LB0 0.003
0,82
0.01
1,422 0.096
0.005 0008
L/72
2: 0L142
2 2
02.04
0.300 0.003
0.800 0.001
1,300 0.002
0.022 0.002
6 0.012
0005
Wi-oo
60.11 = 0.00
02:54 2: 0.07
67.02 2: 0.64
12 01
2: 1.01
0,100 0.000
0.519 0.008
0.964
0.014
[0000 2: 0.002
0,774 0.013
1.582 00006
02.00
67
2:00
0.106 0.000
0.475 0.000
0:268
0.000
0:379 0.008
0.260 0.005
0.00
03.45:00.40
0.55
0.897 0.001+000000
0.205 0.000
0.002
0.750 0.011
1,008
N0MENTIboe
65 10
00000000 0.00
03.71 + 0:00
09.75
173---037
02 1:00
0.105 0.000
0.755 0:001
1.719 0.001
010074-0000
0.718 + 0.016
1.494 0.70
0.00
02.51
0:13
06:0 1.45
0.301 0.0000
0 0.000
1,009
0.023
0.103
L/000 0.161
1,501
Max
65
78
+ 1:37
07:31 6: 1.14
05:33
4 0.50
64
7T 9:19
0.106 4 0.030
0.796 0:001
1.717 0.003
0.007 4 0.000
0,707 0.00
1.000 0.113
0.00
61. 68
0.09
99.03-4--36
00:00
1:71
0.105 4.0000
0.513
0.006
0.857
0.014
0.007
4- 0.001
0,750
0.005
1. 1.455
Alm
52.57
1.19
01.35 L 1:37
M2.74 0.35
06.76
1:3M
(LED6) 0.000
0.907 0.108
0/007 0.316
0.072
0.003
0:370
0.013
0.273
0.00
MOMENT-RR
Lae
72.97
0.33
61.10 1. 0.33
79.01 0:311
89:45 1:22
0.080 1 L..000
0.319 8.000
0.272
0.001
(8083 1 0.001
0.732 1. 0.01-
1.800 1 0.0T
0001
Sam
2.92
03.48
64922
0.301 0.0000
0,685 0.007
0.752 = 0.003
a0005:00000
0,857 2 0.075
1.00
12:27
W-Aog
65.19 0.00
01.48 0:00
86007 0:38
00:01:00033
(L100 2.0000
0.31 B00G
0/851 0.005
[1,81 00002
0.738 a: 0L008
1.MG 0011
Alaa
78.73
5.00
62.55 : 1.88
00,02 0.30
06.64
(.097 0.034
0.68 2.00
0.458 0.122
0.072 s: 0.003
0.205 :: 0.000
0,381 :: 0.009
whereas Sum pooling may amplify irrelevant noise and obscure important local patterns required for accurate imputation. Additional results for all datasets are provided in Appendix
5.3 Positional Weighting in Weighted Average Pooling
We investigate the role of learnable weights in the Weighted Average pooling layer through an ablation-style comparison, while keeping all other model components and training settings constant.
Figure 3 depicts results for the Mistral [4] backbone on representative tasks. In each case, the trainable variant converges to weight distributions that closely resemble the expected optimal pooling strategy for each task: near-uniform weights on text classification, strong focus on the final token for next-token prediction, and intermediate patterns for tasks that require both local detail and global context. Consistent with our theoretical analysis, it shows that the benefits of Weighted Average
9

---

Sum
Nett Tabe Prodake
Last Token
Avarage
Bonking
Classihaation
Similority
Semantic
HellaSwogg
Next Tolom
Prodiction
Figure 3: Left: Cosine similarity between W-Avg pooling and other pooling methods, showing task-dependent alignment. Right: The distribution of the learned weights in the W-avg pooling, illustrating the adaptability of the pooling mechanism.
pooling arise from its ability to mimic the best-performing fixed pooling strategy which depends on the task's contextual demands. Full results for all models and datasets are provided in Appendix
Conclusion
We presented an end-to-end study of pooling in Transformer models by introducing formal expres- sivity framework and deriving closed-form bounds for standard pooling mechanisms. Our theory extends to alternative attention variants and shows that pooling is the key factor balancing local detail and global aggregation. Extensive experiments in vision, language, and time-series tasks validate these bounds: contractive pooling excels global-context tasks, expansive pooling captures fine-grained distinctions, and learnable methods converge to the best-balanced strategy when given sufficient training data and time. No single pooling method dominates universally, highlighting the need for lask-specific pooling design. In data-scarce regimes, our guidelines enable principled selection of pooling methods based task demands and inductive biases. These contributions bridge theory and practice, enhance understanding of Transformer expressivity, and inform the design of adaptive pooling schemes for diverse downstream applications.
Limitations and future work. While our theoretical and empirical findings provide solid foundation for understanding pooling in Transformer architectures, several limitations remain. Our evaluation. though broad, is limited to frozen backbones, leaving the effect of jointly adapting pooling and the backbone under end-to-end training less explored. The expressivity bounds establish necessary but not sufficient for optimal performance; bridging the gap between theoretical capacity and practical effectiveness remains open challenge. In future work, aim to develop hybrid pooling methods that dynamically balance global smoothing with token-level sensitivity while adapting to the TBM's inherent smoothing behavior. In addition to the analysis provided in Appendix E..3 also plan to investigate further how pooling impacts robustness under perturbations, derive scaling laws that govern pooling performance datasets grow, and refine theoretical framework with tighter bounds that identify when specific strategies are provably optimal.
Acknowledgments and Disclosure of F Funding
This work partially supported by Wallenberg Autonomous Systems Program (WASP). L.Z. gratefully acknowledges NXAI GmbH for supporting his participation in NeurlPS 2025.
References
[1] Muhammad Awais, Muzammal Naseer, Salman Khan, Rao Muhammad Anwer, Hisham Cholakkal, Mubarak Shah, King-Hsuan Yang, and Fahad Shahbaz Khan. Foundation models
10

---

defining new era in vision: survey and outlook. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2025.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
[3] Daniel Cer, Mona Diab, Eneko Agirre, Inigo pazeeoo and Lucia Specia. Semeval- 2017 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 1 1th International Workshop on Semantic Evaluation (SemEvva-20117) page Association for Computational Linguistics, 2017.
[4] George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self-attention layers with application to graph neural networks. In International Conference on Machine Learning, pages 2456-2466. PMLR, 2021.
[5] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pages 4171 14186, 2019.
[6] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[7] Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint 2.050::900000779,, 2023.
[8] Sofiane ENNADIR, Johannes F. Lutzeyer, Michalis Vazirgiannis, and El houcine Bergou. If you want to be robust, be wary of initialization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[9] Mark Everingham, SM Ali Eslami, Lue Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes challenge: A retrospective. International journal of computer vision, 111:98-12: 136, 2015.
[10] Akira Fukui, Dong Huk Park, Daylen Yang, Anna Rohrbach, Trevor Darrell, and Marcus Rohrbach. Multimodal compact bilinear pooling for visual question answering and visual grounding. In Conference on Empirical Methods in Natural Language Processing, pages 457-468. ACL, 2016.
[11] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http:://kkllooo07..ggg wwt.:o....tthuu..... 1o/ OpeniebTeettCorpus 2019.
[12] Mononito Goswami, Konrad Szafer, Arjun Choudhry, Yifu Cai, Shuo Li, and Artur Dubrawski. Moment: A family of open time-series foundation models. arXiv preprint 240.0380. 2024.
[13] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788, 2024.
[14) Albert Q. Jiang, Alexandre Sablayrolles. Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
[15] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The Lipschitz constant of self-attention. In International Conference on Machine Learning, pages 5562-5571. PMLR, 2021.
[16] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv.1412.6980, 2014.
11

---

[17] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[18] Itay Lavie, Guy Gur-Ari, and Zohar Ringel. Towards understanding inductive bias in trans- formers: view from infinity. In Forsy-first International Conference on Machine Learning. 2024.
[19] Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, and Wei Ping. NV-embed: Improved techniques for training LLMs as generalist embedding models. In The Thirteenth International Conference on Learning Representations. 2025.
[20] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ArXiv e-prints, pages arXiv-1607, 2016.
[21] Pengxiang Li, Lu Yin, and Shiwei Liu. Mix-ln: Unleashing the power of deeper layers by combining pre-In and post-In. In International Conference on Learning Representations, 2025.
[22] Yuxuan Liang, Haomin Wen, Yuqi Nie, Yushan Jiang, Ming Jin, Dongjin Song, Shirui Pan, and Qingsong Wen. Foundation models for time series analysis: A tutorial and survey. In Proceedings of the 30th ACM SIGKDD conference on knowledge discovery and data mining, pages 3555-6565, 2024.
[23] Chuang Liu, Yibing Zhan, Jia Wu, Chang Li, Bo Du Wenbin Hu, Tongliang Liu, and Dacheng Tao. Graph pooling for graph neural networks: progress, challenges, and opportunities. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, UJCAI '23, 2023.
[24] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[25] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(2018):11, 2018.
[26] Christopher Morris, Gaurav Rattan, and Petra Mutzel. Weisfeiler and leman go sparse: Towards scalable higher-order graph embeddings. Advances in Neural Information Processing Systems, 33:21824-21840, 2020.
[27] Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference artificial intelligence, volume 33, pages 4602-4609, 2019.
[28] Muhammad Muzammal Naseer, Kanchana Ranasinghe, Salman H Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Intriguing properties of vision transformers. Advances in Neural Information Processing Systems, 34:23296-23308, 2021.
[29] Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian conference on computer vision, graphics & image processing, pages 722-729. IEEE, 2008
[30] Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. In 2012 IEEE conference on computer vision and pattern recognition, pages 3498-3505. IEEE, 2012.
[31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Bue, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
12

---

[32] Xianbiao Qi, Jianan Wang. Yihao Chen, Yukai Shi, and Lei Zhang. Lipsformer: Introducing Lipschitz continuity vision transformers. In The Eleventh International Conference on Learning Representations, 2023.
[33] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[34] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 1 11::211-252, 2015.
[35] Yixuan Tang and Yi Yang. Pooling and attention: What are effective designs for Ilm-based embedding models? arXiv preprint arXiv:2409.02720 2024.
[36] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Sean and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in neural information processing systems, (66711111-79477 2023.
[37] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro. Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977 2023.
[38] Hayato Tsukagoshi, Ryohei Sasano, and Koichi Takeda. Defsent: Sentence embeddings using definition sentences. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguisties and the 1 11th International Joint Conference on Natural Language Processing (Volume 2: Short Papers), pages 411-418, 2021.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[40] Aladin Virmaux and Kevin Scaman. Lipschitz regularity of deep neural networks: analysis and efficient estimation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. and R. Garnett, editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[41] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[42] C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The caltech-ucsd birds-200-2011 dataset. Technical Report CNS-TR-2011-00 California Institute of Technology, 2011.
[43] Jinming Xing, Dongwen Luo, Chang Xue, and Ruilin Xing. Comparative analysis of pooling mechanisms in Ilms: A sentiment analysis perspective. arXiv preprint vxXiv:2411.14655 2024.
[44] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanic Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
[45] An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint wrXiv:2412.15115. 2024.
[46] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791-4800, 2019.
[47] Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE transactions on pattern analysis and machine intelligence, 40(6))1452-1444 2017.
13

---

Supplementary Material:
A Proof of Theorem 4.2
Theorem. Let f: X C goxd S IR be a TBM following the framework introduced in Section
In respect to Definition 4.7 we have:
f employs Average pooling, then f is (c,o,"))exppee expressiiv with - ovn (di))
0.22
If f employs Sum pooling, then f is (e, expressiive with q v d 2
d-1 C1C2
If f employs Last-token pooling, then f is (e, o, -xxpresssivee with
/- C,C²
I employs Max pooling, then f is - with y svmin(n.d)
( CiC2,
where
C1=1+ [ max |
n Va/H B*|l | || ]
B²
C2 C2=1+ WEFN|
Proof. Let the input X E X consist of tokens æ; Rdd We consider Transformer model using scaled dot-product attention as defined in Equation 1 and formulated as:
(Xx9) (XWX))
VD/H ) (XXW)
AH(X) softmax
PXW = h(X)WV
where Wº, WK WV are learnable projection matrices. The attention matrix computed from the softmax of pairwise scores:
((x) PX softmax XxxlmmaxxXAAXXx XAA
with
A WKWQ WQT
dd/H E Rddd
Each row of the output f(X) can be expressed as:
(x))
f(X)= =
E RRdd with 4(X) = Σ
Por
hn (X)
j=1
where P. = Cextmax(XAArgr
To analyze the Jacobian of h, derive its partial derivatives:
Ju = x XT p()
EAXA X
(XT p()
Payla
where:
p()) diag(P:))) P.T P. is the Jacobian of the softmax,
Eji is an matrix with at position (j. i) and zeros elsewhere.

---

From this, we have:
If #: Jú = X x FX +
Pijla
(5)
ji=j: J4 = X pl) BAXA + X + Pala.
(6)
Assuming the input space is bounded, ie., l0ell2 S B for all i, we get:
| = Σ | v nB²
|X| < | S VTB.
Since Pt. is probability distribution and o man (diag(p)) S 1, we have || S 2.
Case 1: i +
| $ || p() 11
- ||||i|00|||1
+ +1 || +1.
Case 2: i==) j
| S |X p() E + | ( X|I| + IPP.aall
B A|l Bn A +1= 4n ||| +1.
Thus, the Jacobian of bounded:
[2nBa|Al| +1, ifi [ BABB||||++) +1, ifi=j.
||ssllop S
Therefore, the function h is bounded with constant:
La S 4n B² B +
Attention Head Bound.
Including the value projection:
Chaad
||WVA 4 Va/H
Ba W
|||l K.j +1
Multi-Head Attention Bound. Since is represented by H attention head, their concatenated
output as explained in Equation satisfies:
Csu < [Wollv H max Liead
{ [INNAA 4 Va/H
WollV max
B² | & K. +
Full Transformer Block.
Incorporating FFN and layer norm (with y = 3 1):
Cs < 12 (1 + (1 + [[FFen]))
d
(1 + LMmI + [[Wanl])
2
C2C2,
d-1
where:
C 1 + |Woll max max |
4
1 +1
B² B'll
Vd/H
C2 1 1+ |FeN]

---

Impact of Pooling Strategies
Given a final representation 2 = g(f(X)) using pooling function 9, we evaluate its effect on the bound.
Average pooling.
We recall that this pooling method can be written as a linear layer with weights W Aug therefore:
Avg 12
1n
WAygl
Vn
1
d
[If(X)
f(8)1
S
ClCse.
vn
d
Sum pooling.
Similarly to Average pooling:
Wsum 1n
| || vn.
d
2
( s(8)
vn
CiC2e.
d d-1
Last-Token pooling,
Considering the last token as the output as the pooling operation:
Waatt = .01 + || = 1.
d ) C.C2c.
lI((X) - f(Xi1l <
d d
Note that the same treatment can be applied to CLS or any other chosen token.
Max pooling.
Using norm bounds:
d
- - 2(X))12= = Σ
(f(X)), - -(((x)) |
j=1
d
$ Σ max Xi,s - Xi. 12
j=1
|| X1 p
For spectral norm:
[If(x) - s(8) S X 1
S V rank(X -X]1IX X1
< min(n, d)]li, X11
From those two results, we have:
[If(X) - f(8)1l s Vmin(n, d)llX - xX1
s Vmin(n, d) (--1 )² C2 xe
Applying the Markov inequality concludes the proof.
3

---

B Proof of Lemmal 4.3
Lemma. Let f: X C Bund y C Rd be L22MMHA-bsed TBM 775.. In respect Definition 4.1 the following holds:
If f employs Average pooling, then f is (e,a, expressivv with ovn (z))² C12 9c0
If f employs Sum pooling, then f is (c,a, expressiiv with y yue d
C2
[-1
If f employs Last-token pooling, then f is (e, 7) exxprsssive with = - C1C²
f employs Max pooling, then f is (c,0,7))expresee)) with y= <ymmlu(n.d) (a-1) C1C2.
where
( Σ
-1=1
Vn Vd/H
(Wo ( +1)
||
C2 = 1+ [Weeel]
Proof. Let the input X E X be composed of n tokens 2; E Rd In this proof, we consider the Transformer-basee model built using the L2 Multi-Head Attention (L2-MHA) mechanism, where the attention weights are computed as:
|xx - -x,W -
Po σ exp
Vd/H
with Wº, WK being learnable projections.
From Theorem 3.2 of [5], the L2-MHA operator is bounded by:
(F) < Vn
(4W,
12 +1)
I |aak |lol
Va/H
where Wo(.) denotes the Lambert W-function.
Following the Transformer architecture defined in Section 3 account for the additional effects of LayerNorm (LN) and the Feed-Forwadd For-ard Network (FFN). As in previous derivations, we obtain:
C S Li + LMHA) + LFEN).
where we now substitute the bound for L2-MHA:
L
d
2
1
Vn
(Wo
n (")
+1)
(oall] Iwaaj//wwvaji
+
[Wrxx]))
dd/H
d
C2C2,
d
with constants defined as:
(1>
C 1+ Vn
WWo
(0) +1)
Wiw.ariiiiii |llil|
Vd/H
C2 = 1+ W|
Following the same steps as in the Theorem 4.2 proof, we get the following:
For Average pooling:
1 x vn
d
Livg S
CrC2.
d-1

---

For Sum pooling:
d
2
Ssu S SVx
CiC²
d-1
For Last-token pooling:
Last S d-1
C2O
For Max pooling:
d 2
LMax S min(n. d)
(d-1) C2C2.
C Proof of Lemma 4.4
Lemma. Let f: X y be a function based the LipsFormer [32] framework, with cor- responding eyper-parameters V,V,T T >> 0 and window size w. In respect to Definition 4.1 we have:
If f is based on Average pooling, then f is (c, (c,0,)) -.)orre)eeeessiveee with y oxva (141 Cic
c
If f is based Sum pooling, then f is - with y vaxe x d C2C
If f is based on Last-toked pooling, then f is (e,or, -))e)))eesssiii with 7 - ( C202
If is based on Max pooling, then f is (e, 7) expressivee with min(n,dd)
1 C1C2,
where
G 1+ [Woll max 2w(w 1) I || 1)urV
|| 11 | I)
C2 =
Proof. Before delving in the specific analysis of the Lipsformodel model, start the proof by providing some preliminary elements about the Swin Transformer 24 which is different from the original Transformer-bases Model defined in Section 3
Swin Transformer block's ' input is similar to the one from TBM, specifically, the input X Redd can be viewed as n tokens (or patches) each of dimension d. Rather than applying global self-attention to all n tokens, the model partitions complexity. Between successive Swin stages, there "patch merging" step. which consists of linear downsampling that reduces the number of tokens while increasing their dimension.
into small "local windows" of size w, thereby reducing
Let W denote the total number of windows, X E Rexd be the slice of input corresponding to window and WºO W W are the query/key/value projection matrices, within each block, window-based self-attention is computed follows:
LocalA(tn(X)
softmax
(XX (XeW)) Va/H
In our analysis, rather focus the LipsFormer 32 model, which is adaptation of the previous equation. Specifically. the model modifies the Swin Transformer architecture by replacing the standard dot dot-product attention with scaled cosine similarity attention mechanism. Given input X e IR' ux (e.g., the tokens in a window), define the following:
X; WQ
X, WK
X, WV
1x
h; Ix, WK
Mj
X +
5

---

where X, E Rd is the i-th token row of X, and V > O is small constant to avoid division by 0. Then accordingly write the usual attention matrices as:
Q 1
K
k
01
which are all in Rxxd The scaled cosine similarity attention (SCSA) can be then formulated as:
SCSA(X) v/ sofomaa =
where > are scalars that scale the argument of the softmax and the final output. As can be seen in the formulation, the key difference from standard attention is that each query/key vector is row-normalized to unit &g length (up V).
Similar TBM, H attention heads are used with each one using separate projection matrices W w wX, forming Q, K, for each head, then concatenates the outputs and multiplies by Wo:
MultiHead SCSA(X) =
Let's now derive the upper-bounds of this model. Similar to the previous proofs, let's consider that our model f is built using the scaled cosine similarity attention, with // attention heads and one layer. From Appendix H.2 the original paper [32], we have the following for single head of attention:
LSOSA & 2n(n - V) 1
W 1|+2(n - T) | |WVI]
with n being the number of tokens within a local window, and v, 0 and V the chosen hyper-parameters of in SCSA.
When considering the multi-head attention framework, we get:
MMH-SCSA S WollVH max IShSH
CSCSA.]
Similar to previous proofs and since we consider the same the aspect, we directly get the following result:
and Layer Normalization
C, S (
(1)
1+ WollVE VV
max Am I..,,,, H
(w 2w
- 1)2T ||
+ 2( - 1)LTV 2
| | []] [|F
which could written as:
d 2 Cj S d d CiC2.
with C1 = 1+ WollV max 20( 1)ur
l. 2
1)rrV 2 1 +200 [w ]|}
C== = (1 + [WFew])).
For the pooling operation, similar analogy that was used in the case of dot- dot-product attention can be used, and find therefore the final results:
For Average pooling:
d
Lavg <
C1C2.
/n
d-1

---

For Sum pooling:
d sum S Vix
CiC2.
d-1
For Last-token pooling:
d Linst S
C2
d-
For Max pooling:
Cistt S v min(n, d)
C12²
d=
By applying the Markov inequality get the desired result.
D Experimental Details
We start by noting that the necessary code to reproduce the results is publicaly available in the following link: https://gitthbb. https://gitub..om//ingg// wow//ingtteassttormrrrreeee pooolngg
In what follows, we provide experimental details and hyper-parameters choices.
D.1 Computer Vision.
All computer vision experiments used frozen Transformer backbone (ViT-base (6), ViT-small, or LipsFormer [32] built on Swin Transformer (24] architecture), with randomly initialized heads fine-tuned on each task. We optimized with Adam [16] learning rate of 1 10-3, All the tasks were trained for 10 epochs; all of which yielded stable convergence. Images were resized and either padded or center-cropped the model's input resolution. No data augmentation was applied during
training.
We evaluated tasks that require both local and global context. For classification (CIFAR 10/100 177. ImageNet-100 100 [34], MiniPlaces (47), Caltech-UCSD Birds (CUB) (4]). the head's output dimension matched the number of classes and was trained with cross-entropy loss. For inpainting (CelebA (25), Oxford Flowers [29]. Oxford-IIIT Pet (30)), the head predicted pixel values in masked regions and trained with mean squared error. For segmentation (Pascal-VOC (), linear per-pixel classifier trained with cross- cross-entropy loss was used and report mean pixel accuracy. All the experiments were run using a single NVIDIA L4 GPU and took 25 GPU hours obtain all results.
D.2 NLP
In all experiments involving LLMs, the Transformer backbone was kept frozen, and only a randomly initialized linear head and parametric pooling parameters were fine-tuned. We used the provided pretrained (non instruction-tunede checkpoints for Llama3 [13] Mistral 7B [14), Qwen 2.5 [45] and BERT (5), and we pre-trained GPT-2 and I.2-GPT-2 (see details below). Optimization was performed using the Adam [6] optimizer with a learning rate x 10-3 Ten epochs fine-tuning
consistently yielded stable convergence across tasks. Each experiment was repeated five times with fixed random seeds to improve robustness and ensure reproducibility.
All pooling methods, including those with learnable components, were trained using the same config- uration. Experiments were conducted on an instance with 2xNVIDIA L4 GPUs using Py ToTch [1 with the Distributed Data Parallel framework and batch size of 32 per GPU. Running all experiments took 1832 GPU hours on L4 GPUs.
Fach dataset's training split was used for fine-tuning. Hyperparameters were selected based on validation performance (where available), and final results were reported on the held-out test set. To maintain consistent input dimensions, all sequences were padded or truncated to predefined
7

---

maximum length. Tokenization was done using each model's 's default tokenizer, and [PAD] tokens were used for padding.
For classification tasks, we used a linear head with output dimensionality matching the number of classes, trained by minimizing the cross-entropy loss. In semantic similarity tasks, the pooled embeddings were linearly projected without changing dimensionality. Cosine similarity between embedding pairs was used as the main metric. For STS Benchmark (STSB) [3], similarity scores (rescaled from [1, 5] to [0, ID were predicted by minimizing mean squared error. In the HellaSwag 44]] task, the goal was to match a given context to its correct ending. The context and four candidates were encoded with the same LLM and pooling method, projected linearly, and compared via cosine similarity. Cross-entropy loss was applied over the similarity scores, encouraging correct pairings. For next-token prediction, we used the TinyStories [7] corpus under a standard autoregressive setting, The training set comprised 4000 batches randomly sampled from the corpus. randomly initialized language modeling head was trained to predict the next token based on preceding context. We used a held out test-set and randomly sampled tokens to predict.
GPT-2 Pretraining. To obtain a checkpoint for L2-GPT-2 (a GPT-2 model with L2-MHA), we followed the standard pretraining procedure described in [3 modifying the attention mechanism to use the L2 kernel with tied query and key matrices. This change slightly reduced the number of parameters (from 123M to 1 166M). The model was pretrained on the OpenWebTeer [T] corpus for 60 000 iterations using a batch size of 12, block size of 1024, and 40 gradient accumulation steps. Training was conducted on NXNVIDIA 14 GPUs and took about 960 GPU hours.
For fair comparison, we also pretrained baseline GPT-2 checkpoint using identical settings, differing only in the use of the standard dot-product attention mechanism.
D.3
Time Series
For time series analysis, we used MOMENT [12), family of Transformer-based foundation models for time series. We evaluate three pretrained checkpoints (AutonLab/ooENTT.L. (small, base, largel) trained on the Time Series Pile dataset [12].
During training. we kept the model backbone frozen and fine-tuned only the linear head and pooling operator (Weighted Average and Attention- -baddd) for classification, forecasting. and imputation tasks. Optimization was performed using Adam [16] with learning rate of 1 x 10-3 with batch size of 64 on single NVIDIA L4 GPU.
For classification, we run optimization for 20 epochs across six datasets, six pooling methods, three model sizes, and five random seeds, yielding 540 experiment trials and a total of approximately 22 GPU hours. For forecasting, we trained the prediction head for 10 epochs using forecasting horizon 96 future time steps across seven datasets, six pooling methods, three model sizes, and five random seeds, resulting in 630 experiment trials and approximately 90 GPU hours. For imputation, we trained the prediction head for 10 epochs across seven datasets, six pooling methods, three model sizes, and five random seeds, adding another 630 experiment trials and approximately 96 GPU hours.
E
Additional Results
Computer Vision. In addition ViT-small model, we extend our analysis evaluate whether the theoretical insights hold in larger architectures with more attention heads and blocks. Table reports the results for the same pooling benchmarks using ViT-base the backbone.
Consistent with previous findings, Weighted Average pooling maintains strong performance across tasks, reflecting its ability to adapt to context and produce stable, generalizable representations. Similar patterns emerge for Attention based pooling, which performs best in the inpainting task but does not outperform Weighted Average pooling in other settings. This suggests that Attention- based pooling may require more data and computational resources to reach its full potential.
Among flat pooling strategies, CLS pooling continues to yield the best results for classification tasks. Notably, the performance gap between CLS and Average pooling narrows, indicating that larger models can offieet suboptimal pooling through increased representational capacity.
8

---

Table 4: Mean and standard deviation of test metrics for computer vision tasks. Best performance per dataset and model is indicated in bold. Best performance among non-le non-learnable pooling methods is underlined.
Meddl
Classlication (Acomaag)
Iepointing (MSE)
Seguecttaioon (Aecarag))
Poaling
CIRAK-10
CHAAX-100
Dregshhe- 102
CLB
218-2011
Minitfaon
Caten
Orderst 102
Ostred-lll III
PaesiVUC-C3e
PsesiVOC-Ue Det
Liai (CLS))
12.14
90.07
T6.87
0.200
2.314 0,0S
0.003
72.49 0.00
28.01
Ang.
0.84
7167 0:40
90.50
73.36 0.50
08:81 00:1-00.0
0.587 0.000
0:300 = 0.05
6,266 0.035
72:00
20.00 1.45
78.44
06
72.00
0.312
0.001
Mis
11.00 0.55
74800 0.64
90
59.31 0.80
36,61 0.68
0.503 0.001
04001 = 0.45
6.251 0.000
70.21 +: 1.29
$2.88 0,67
Wi-rg
92.65 0.17
00.02:0.18
00.45:0077
01
0.596
0:000
6,270
2602 0:220
01.81 4.0.88
20.84
00:36-00:11
65.09 ::
07.68:10:97
0.162--4 0.003
0.303 000
6.383
0.015
71.30:00.00
16:14:111
NLP. Beyond the theoretical expressivity bounds shown in Figure we further examine how these bounds manifest empirically in NLP settings. To this end, we construct set of sentence variants: a base sentence (the original), two semantically similar versions created by replacing adjectives. and two dissimilar versions using unrelated words. Figure shows the resulting changes in pooled representations across different pooling strategies.
[3] CV: VIT on CIFARLO0 Dataset
similar-1
11P
Pooled Outpet Distone Cog aoddo
103
similar-2
Sum.
Max
non-similar-1
Puoeed Distanze (log sole)
non-similar-2
17.0
0.2
0.4
10.0
0.8
(2)
NLP: GPT-2 on Simple Texi
110"
10-9
Stn
Ma
10
0.0
0.2
0.4
1.0
smalll (3) TS. MOMENT small on Electricity Dataset
Avg
Max
Last
Sum
Pooling Metlod
10-1
original: The
saw significant fluctuotions last sinilar-1: The stook market encountered notable swings last week.
Stm
Mas
similar-2 22: The equity market experienced marked volatility last week.
0.0
L'0
non con-nimilar--. The stock murket encountered notable castle last week. The stock market encountered notable car last week.
Nogghborooo Tresholi
bon-einilar-a 2:
Figure 4: Empirical analysis of the expressivity power across modalities and pooling strategies. Left: Mean pooled-output distance y versus perturbation across modalities highlighting the behavior of various methods. Right: pooled-output distances for similar and dissimilar inputs, exemplifying expressivity of different strategies.
Time Series. In addition to the results presented Table we report extended empirical evaluations of pooling operators on broader range of time-series datasets. Tables [5] and provide results for classification, forecasting, and imputation tasks, respectively, Overall, the findings on these additional datasets consistent with the trends discussed in Section 5.2 further supporting
analysis.
we presented an analysis of the learned weights in the Weighted Average pooling method using the Mistral-7B model. 3 this analysis to additional models and datasets. We observe that the learned weight distributions for a given dataset remain consistent across models, with smaller models (e.g., GPT-2 family) placing more emphasis on later tokens, while larger models exhibit more uniform weighting. The average cosine similarity between Weighted Average pooling and non-learnable pooling methods follows similar trend: as model size increases, similarity to Max pooling decreases, suggesting reduced reliance on token-level extremes in larger architectures.
Weighted Average Pooling, In Section 5.3

---

Table 5: Mean and standard deviation of test accuracy for time series classification tasks. Best performance per dataset and model in bold. Best performance among non- learnabbeelle pooling methods underlined.
Model
Pooling Operator
BCG200
Electnc Devices
FoedA
FordB
SmallKitchen
Appliannes
SwedishLeeuf
Last
72.29
0.59
60.45
76.39 0.15
62.07
64.06 4: 0.75
69.2 0.74
MOMENT-
Aug
65.19 : 0.00
61.40 + 0.54
88.12 + 0.14
67.05 = 0.13
62.40 + 1.01
55.63 9.55
Sum
80.35
+: 1.82
60.62 +: 2.85
92.93 + 0.43
78.67 :: 0.72
67.13 + 2.95
79.42 :: 1.33
Max
65.19 + 0.00
61.78 + 1.24
90.42 + 0.43
72:28 + 1.00
63.94 + 1.76
58.59
6.53
W-Arg
65.19
0.00
62.04 + 0.67
87.02 + 0.64
67.36 + 0.44
02.40 + 1.01
08.81 8.41
Attn
78.84 +: 3.04
62.95 =: 2.04
93.60 +: 0.37
79.61 + 0.86
67.63 +: 2.95
75.93 :: 2.07
Lost
71.01 + 1.94
63.05 + 0.62
83.46 +: 0.40
64.60 + 0.37
03.25 + 0.54
72.03 1.09
MOMENT-bac
Avg
69.
0.00
63.71 1: 0.40
89.75 0.37
70.63 63 0.54
63.05 1: 1.60
62.92
5.93
Sum
83.30
1.51
64.30 +
92.51 0.13
0.98
+ 1.48
84.05
62.31 4:
90.23 0.80
70.91 91
64.77 :: 2.12
444
W-Arg
65.19
0,00
64. 48 + 0.69
89,93 0.36
70.98 0.62
63.20 +: 1.74
84.87
5.16
64.15 1.37
92.74 0.35
79.94
66. 76 1.39
78.66
2.25
1.st
72.67 :: 0.95
61.10 + 0.53
79:61 0.31
63.63 : 0.64
65.45 + 1.71
76.75
1.47
MOMENTTuege
Avg
66.19 + 0.00
61.72 +: 0.98
86,98 +: 0.53
67.56 #: 1.18
60.42 + 0.91
59.389
6.89
Sum
75.97
+: 2.82
63 45 : 0.76
92.85 +: 0.42
78.08 1: 0.48
64.92 +: 7.21
82:23:2 0. 88
Max
65.19 + 0.00
60. + 0.64
87.14 +: 0.39
69.52 + 0.74
02.27 + 0.64
644.00 3.75
W-Avg
65.19 + 0.00
61.44 +: 0.90
86.07 + 0.3G
67.65 : 1.09
GD.54 + 0.53
60.84 + 6.34
Autn
78.73 1: 3.09
62, 02..83.4 1: 1.38
93.02 1: 0.26
79.95 1: 1.02
65.64 1: 3.54
80.66 1.04
Table 6: Mean and standard deviation of test MSE for time series forecasting tasks. Best performance per dataset and model in bold. Best performance among non-learnable pooling methods is underlined.
Model
Operatie
Pooling
ETTh
ETT1A
ETTml
ET7a2
electisity
tralic
waster
0. 020082222 0.000
0.1815 1: 000
0,040 1: 0 000
0.107 1: 0.000
0.400 1: 0,001
0,273 :: 0.001
0.002 1 0.000
MOMENT email
0.105
0.000
0.305 + 0.001
0.079 +: 0.000
0.284 + 0.000
0.790 + 0.000
1.724 = 0.001
0.002 0.000
Sum
0.103
= 0.001
0.500 + 0.005
0.053 :: 0.000
0.141 + 0.000
0.826 + 0 0.014
1.422 = 0.050
0.005 : 0.001
Mix
0.106
= 0.001
0.311 1: 0.002
0.082 1: 0.001
0.230 1: 0.001
0.800 1: 0 001
1.750 :: 0.002
0.002 1: 0 0.000
W. Avg
0.105
0.000
0,287 à: 0.002
0.058 d: 0.005
0.184 d: 0.000
0.539 1: 0.008
0.954 :: 0,014
0.002 :: 0.000
Arm
0.100
0.001
0.313 + 0.003
0.041 : 0.001
0.097 : 0 0.000
0.475 + 0. 059
0.258 = 0.002
0.005 +: 0.000
Last
0.081
0.000
0.1933 1: 0.000
0.040 1: o 000
0.105 1: 0.000
0 39997 1: U. 000
0.265 :: 0. 0.000
0,002 1: 0 0.000
MOMENT-bO
Aug
0.105
0.000
0.304 + 0.002
0.070 +: 0.000
0.226 + 0.000
0.786 + 0.
1.719 = 0..001
0. + 0.000
0.002
0.002
0.009
0.004 ::
0.106
0.000
0.309
0.002
0.075 0.001
0.225 0.000
0.796 0.001
1.747
0.003
0.002 0.000
W.Aag
0.105
0.000
0.280 + 0.002
0.058 +: 0.000
0.178 + 0.001
0.511 + 0.006
0.857 = 0.014
0.002 + 0.000
0.009
0.012
0.043
0.097 0.001
0.507
:: 0.000
Last
0.0000
0.000
0.155
0.000
0.039 1: 0.000
0.103 1: 0 000
0.379 1: 0,000
0.272 :: 0.001
0.002 1: 0.000
MOMENT-AagE
0.105
0.000
0.306 + 0.000
0.073 + 0.000
0.207 : 0.000
0.778 + 0.000
1.711 = 0.001
0.002 + 0.000
Sum
0.101
= 0.001
0.269 :: 0.002
0.049 +: 0.000
0.126 : 0.001
0.688 : 0. 007
0.782 = 0.003
0.009 :: 0.000
Mix
0.104 = 0.001
0.306 + 0.002
0.073 + 0.000
0.206 + 0.001
0.786 + 0. 0.001
1.740 + 0.000
0.002 +: 0.000
W.Aag
0.105
0.001
0.283 + 0.000
0.053 + 0.000
0.170 + 0.003
0.584 + 0.003
0.951 = 0.005
0.002 :: 0.000
Anm
0.097 0.004
0.506 +: 0.000
0.039 + 0. 000
0. 10106622 : 0.003
0.681 +: 0. 008
0.423 = 0 0.028
0.009 2: 0.000
E.1 Additional relatel Results on Larger Models
Table E. below reports results across tasks for larger models under consistent experimental setup to our previous experiments. With larger models, trends across pooling strategies remain visible, but the absolute differences between pooling methods diminish, aligning with the theoretical interpretation.
E.2 Additional Results on Sequence Length Changes for NLP
To further examine pooling behavior, we conducted experiments with the Mistral-7B model on HellaSwag, varying the maximum input sequence length from 16 to 128 tokens. Inputs were truncated or padded as required, with padding tokens excluded from pooling operations to maintain consistency with our setup.
10

---

Table 7: Mean and standard deviation of test MSE for time series imputation tasks. Best performance per dataset and model in bold. Best performance among non-learnable pooling methods is underlined.
Model
Operince
Pooling
BThh
ETTh2
ETTml
BT7m2
electiiity
trathic
Weather
Last
0.081
0.002
0.241 + 0.002
0.051 + 0.000
0.181 # 0.001
0.753 : 0.010
1.700 = 0.016
0.002 :: 0.000
MOMENT-mal Famall
Aug
0.060
= 0.002
0:253 +: 0.002
0,050 + 0. 000
0.178 # 0.001
0.771 0.015
1,583 = 0.016
0.002 +: 0.000
Sum
0.106
0.008
0.309 + 0.041
0.054 + 0.001
0.156 + 0.010
1.072 + 0.145
2.130 + 0.238
0.087 +: 0.038
0.082
0.002
0.220 + 0.004
0.051 +: 0.000
0.174 0.003
0.720 0.013
1.211 + 0.051
0.008 +: 0.000
W-AAg
0.080 = 0.002
0:253 +: 0.002
0.050 +: 0.000
0.178 + 0.001
0.771 +: 0.015
1.582 = 0.016
0.002 :: 0.000
0.076 = 0.003
0.088 1: 0.007
0.051 1: 0.001
0.152 1: 0.001
0.379 1: 0.028
0.265 = 0.015
0.003 4: 0.001
0.082
0.001
0.243 4: 0.005
0.061 4: 0.000
0.181 :: 0.002
0.760 : 0.011
1.608 = 0.020
0.002 #: 0.000
MOMENT-POR DoN
Aug
0.005
0.002
0.769 0.016
0.002 0.000
0.103
0.002
0.255 0.022
0.054 +: 0.001
0.156 0.008
1.006 + 0.161
1.301 0.089
0.015 +: 0,009
Mux
0.002 0.082
0.219 + 0.004
0.051 + 0.000
0.173 + 0.002
0.707 + 0.014
1,000 + 0.118
0.002 +: 0.000
W-Aag
0.082
0.001
0.253 0.005
0.051 +: 0.000
0.178 0.002
0.769 :: 0.015
1.425
0.000
0.002 =: 0.000
0.072
0.003
0,082 0.002
0,050
000000
0.161
0.002
0.370 0.013
0.273
0.004
0.004 0.003
Last
0.082
0.001
0.242 + 0.005
0.051 + 0.001
0.181 + 0.001
0.762 + 0.014
1.609 = 0.017
0.002 #: 0.000
MOMENT-lage
Aug
0.081
= 0.002
0.258 + 0.007
0.050 + 0.000
0.177 + 0.001
0.753 + 0.018
1.508 = 0.011
0.002 + 0.000
Sun
0.096
0.008
0,254 H: 0.020
0.053 +: 0.001
0.177 + 0.005
0.887 + 0.075
1.150 + 0.063
0.014 #: 0.002
0.081
0.001
0.230 + 0.008
0.050 +: 0.001
0.170 + 0.002
0.706 + 008
1.158 0.025
0.003 + 0.000
W-Aag
0.081
0.002
0.258 + 0.007
0.050 + 0.000
0.176 + 0.001
0.758 + 0.018
1.508 = 0.0011
0.002 : 0.000
Ann 0.072 0.003
0.091 +: 0.006
0.060 + O. 000
0.147 +: 0. 00.003
0. 0.295-2 + 0.006
0.231 =: 0.000
0.004 +: 0.001
Table 8: Mean and standard deviation of test metrics for NLP tasks. STSB and HellaSwag are grouped under Sentiment Analysis. Values are mean + std.
Sentiment (STSH)
Seatiment (HellaSwag)
Banking (Accumay)
Twoot (Accuracy)
Next Token (Accuracy)
Model
Pooling
Last
0,288 = 0.003
0.692 + 0.011
31.497
0.178
57.005 + 0.00-1
91.039 + 0.002
Avp
0.581 + 0.000
0.773
0.000
85.390 = 0.001
69.930 + 0.008
33.893 : 0.250
5-148
Sum
0,579
= 0.002
0.776 + 0.002
82.711
0.008
61.218 + 1.19
47,053 + 0.311
0.005
0.727 0.00T
77.825
0.006
55.390
0.728
25.578
0.332
W-Avg
0.589 + 0.001
0.798 0.002
86.867
0.002
69.770 + 0.007
90.653 1 0.020
Alta
0.259 1 0.011
0.712 0.013
66.387
1.939
56.410 1 3.001
40.573 1 0.892
Last
0.303 = 0.001
0.723 0.009
34.383 = 0.268
88.886 + 0.006
89. 886 + 0.005
Avg
0.603 = 0.000
0.781 + 0.001
77. = 0.002
38.328 + 0.012
530 = 0.33T
Qwen2.5-3IB
0.604
0.005
25.227 0.013
63.665 + 0.865
Max
0.458 + 0.005
0.729 0.008
83.929 0.008
65.268 + 0.596
31,273 + 0.586
W-Avg
627 = 0. 0.003
0. 818 0. 0.003
89 903 0. 0.004
69 464 + 0.009
89.022 + 0.019
Atta
0.355 :: 0.016
0.730 t 0.011
1.054 68.929
$8.828 t 2.695
29.750 + 1.028
Last
0.503 + 0.002
0.745 +: 0.008
75.487 - 0.087
53.701 1: 0.005
88.972 1. 0.007
Avp.
0.631-
0.001
0.784 0.001
87.403
- 0.009
1 51.723
Mistra3.1-2488
Sum
0.622
0.000
0.783 0.004
87.597 0.015
62.791 + 0.976
42.306 + 0.732
Max
0.458 + 0.003
0.733 + 0.007
79.675 = 0.010
61.655 + 0.173
27,801 + 0.923
W-Avg
0.682 + 0.002
0.816 + 0.003
88.711 + 0.017
66 200 + 0.005
87,849 : 0.024
Alta
0.392 0.048
0.697 0.009
72.922
1.012
31.291 + 4.452
19.911 + 2.023
The results, summarized in Table 9 compare pooling strategies across different sequence lengths. Naturally, shorter contexts lead to performance drops due to truncation of semantically important content. To isolate the contribution of pooling itself, comparisons should be made column-wise (i.e., at fixed sequence lengths).
We find that pooling sensitivity is most pronounced at shorter lengths, especially for Last-token and Attention-based pooling. For longer contexts (64 or 128 tokens), performance stabilizes across pooling methods, and the relative differences align more closely with theoretical expectations.
E.3 Pooling and Adversarial Robustness
Beyond shaping model expressivity for downstream tasks, the choice of pooling operation also impacts the model's adversarial robustness. Our theoretical framework provides insights in this direction by interpreting neighborhood changes, introduced in Section as adversarial rather than semantic perturbations intentionally crafted to mislead the model. The analysis suggests that certain pooling operations, such Average, may naturally smooth out adversarial noise, while others like
11

---

Table Mean and standard deviation of metrics for Mistral-7F 7B on HellaSwag across different input sequence lengths. Values are mean + std. Best performance per column is in bold.
Pooling
64
Last
0.454 0.003
0.621 0.002
0.770 = 0.002
0.781 0.001
Avg
0.503 = 0.001
0.702 + 0.000
0.769 = 0.000
0.771 + 0.001
Sum
0.504
0.001
0.702 0.001
0.769 = 0.000
0.771 0.001
Max
0.435
0.003
0.616 + 0.003
0.709 = 0.001
0.700 + 0.008
W-Avg
0.523 = 0.002
0.724 1: 0.000
0.801 + 0.000
0.802 + 0.003
Attn
0.220 0.169
0.278 0.251
0.737 0.018
0.764 0.025
Table 10: Attack Success rate for different considered Pooling strategies using the FGSM adversarial attack the CIFAR-10 and CIFAR-100.
Dataset
Attack Budget
CLS
Avg
Sum
Max
W-Avg
Attention-Basel
3/255
18.56
18.94
10.92
20.32
16.93
13.9
CIFAR-I0
c 8/255
29.38
28.16
10.92
25.97
25.59
14.44
e = 3/255
34.29
32.44
20.67
31.84
31.57
29.54
CIFAR-100
8/255
43.99
41.42
20.68
38.07
40.49
32.94
Max can either amplify or ignore the perturbation depending on whether the adversarial signal falls within the selected region.
To empirically validate these insights, we apply the Fast Gradient Sign Method (FGSM) to a pre- trained ViT model evaluated CIFAR-10 and CIFAR-100. We the same attack budget the one usually used in the literature (e = 3/255 and e 2/255) and we used the same number of epochs and the same initialization for all the poolings to ensure fairness the comparison [8]. Table (10 the attack success rates for various pooling strategies under FGSM. As expected, the success rates vary across pooling methods, confirming that pooling choices can meaningfully influence robustness. Therefore, in domains where robustness is critical, such as healthcare or finance, the pooling strategy should be selected not only for its expressivity but also for its impact on adversarial resilience.
12

---

In
(a) GPT-2
(b) GPT-2
(c) L2-GPT22
(d) L2-GPT-2
2010
(c) Qwen2.5
(f) Qwen2.5
(g) Mistral-7B
(h) Mistral-7B
Son Tica
AM22
i) Llama
Llama
Figure 5: Left: Cosine similarity of weighted average pooling with other pooling methods, Right: Leamned weight distributions.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and the introduction are consistent with the provided theoretical results in Section 4 and additionally empirically validated in experimental results in Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The limitations of our work with respect to both the theoretical and empirical results are discussed in Section 6
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: All the theoretical claims are explained in the corresponding proofs (Appen- dices A B C) In addition, clearly state our problem setup detailing the assumptions
in Section which are referenced in each proof and theorem.
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: In addition to providing a detailed experimental setup in Appendix D we provide the source code reproduce our results in the supplementary materials section.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: In all experiments used publicly available datasets that be found on, e.g.. HugginggFce or other open source platforms. We provide the source code in the supplementary material for reproducibility.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
haagli///guides//eeeSbbb.m...aa htt.e/eouubbmssssooooooooo.....e for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-sourge benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
guwies/ CodeeSumiiiseeeeee onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Appendix D contains the experimental setup for all modalities, including training parameters and data splits. The code contains all hyperparameters for every model and experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report standard deviation of all results in the respective tables in Section 5.2 obtained via repeated experiments with 5 random seeds. Further details in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eπror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide the type of GPUs and number of GPU hours used for experi- ments in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our work adheres to the NeurlPS Code of Ethics. We use public datasets and publicly available models and report the limitations of our work.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
17

---

Justification: Our work has no further societal impacts apart from known impacts of LLMs and other Transformer- based models. We are not aware of any upplications of our insights that would result in negative societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not release any data models that have high risk for misuse. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
All datasets, models and repositories were cited appropriately.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
18

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwwpsssipppsseeeee thcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The source code released along with our paper is properly documented and contains the license terms.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work did not involve crowdsourcing research with human subjects. All experiments are performed on publicly available datasets.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
supplemental material is fine, but if the main contribu-
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
19

---

Justification: Our work did not involve crowdsourcing nor research with human subjects. All experiments are performed on publicly available datasets.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
Declaration LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: No LLM was used this work for the core methods.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
wwtperippscc.ooooeeeeee 2c/Conneerenccs/2225/////// MLMM
20


Reviews:
[{"review_id": "PIk1MuRPy0", "reviewer": "Reviewer_JC6b", "summary": "The paper studies the pooling operation used to combine the output token embeddings into a single embedding. They provide a theoretical framework to analyze how pooling affect the performance and offer guidelines regarding how to select the pooling function. The theoretical analysis shows that the pooling strategy affects the classifier's expressivity bound. Average pooling enhance stability by smoothing variations. Last-token and Sum pooling increases expressivity and more sensitive to minor perturbations. The paper suggests that task that requires global context would benefit from average pooling while tasks that biased toward certain local spans (such as next-token prediction) can benefit from Last-token or Sum pooling strategies. The hypothesis are verified by experiments measuring expressiveness and downstream performance.", "questions": "- (Figure 4) The experiments of replacing words by semantic similar or dissimilar words is interesting. I'm wondering if there are similar experiments of removing stop words or removing the negation word (like not). Are there any proposal on an pooling mask that masked out certain tokens?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n  - A theoretical framework to bound the expressiveness of various pooling techniques.\n  - Empirical results to back the hypothesis derived from the theoretical works.\n  - Clear guidelines of how to choose the pooling strategy.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and constructive feedback. We especially appreciate their recognition of the novelty and value of our proposed theoretical and empirical insights concerning the pooling operation in the Transformer-based Models. \n\n## Q1: What happens when stopwords or negation are removed? Could a pooling mask be used to discard such tokens?\n\nWe appreciate this suggestion and agree that extending the analysis in Figure 4 to include manipulations such as stopword removal and negation edits provides a useful perspective on how pooling functions respond to subtle semantic changes, especially in the context of token-level relevance.\n\nTo explore this, we conducted an additional analysis using the same base sentence as in Figure 4:\n\n- *Original*: \"The stock market saw significant fluctuations last week.\"\n\n- *Stopword removed*: \"Stock market saw significant fluctuations last week.\"\n\n- *Negation introduced*: \"The stock market saw *insignificant* fluctuations last week.\"\n\nWe then measured the L2 distance between the pooled output of the original and perturbed sentences for each pooling method. The results are shown below:\n| Pooling Method       \t|   Avg |   Last |   Max |    Sum |\n|:----------------------------|------:|-------:|------:|-------:|\n| Stopword removed | 8.37 |  18.93 | 19.18 | 265.00 |\n| Negation introduced |  2.38 |  16.46 |  9.78 |  21.41 |\n\nThese results suggest that Average pooling is relatively stable under small textual changes, such as those affecting global structure (e.g., stopword removal), while Last, Max, and Sum pooling are more sensitive to local token edits, such as negation. This supports the broader insight from the paper: contractive pooling methods smooth token-level effects, whereas expansive methods emphasize localized changes, making them more responsive in tasks where fine-grained distinctions (e.g., presence of negation) are crucial.\n\nWe will include this extended analysis and visualization in the updated version of the manuscript, and we thank the reviewer for motivating this additional experiment. The notion of a learnable pooling mask that can selectively attend to or discard tokens is indeed a promising direction, and we will briefly note it as future work.\n\n\n## Q2: Are there any proposals on a pooling mask that masked out certain tokens?\n\nIncorporating such a mechanism could indeed be beneficial, especially when certain tokens carry outsized semantic weight or when low-information tokens may dilute global context. However, a key challenge lies in determining which tokens to mask, especially in a generalizable way. While heuristic approaches (e.g., stop word lists) are available in NLP, such strategies do not naturally extend to other modalities like vision or time series, where \"irrelevant tokens\" are less clearly defined.\n\nIn this regard, we view weighted average pooling as a flexible alternative. Rather than applying hard masking, it learns to dynamically attenuate or amplify token contributions based on context, achieving similar outcomes without requiring explicit prior knowledge or handcrafted rules.\n\nWe will mention this connection in the revised manuscript and thank the reviewer for prompting this valuable direction for future work."}, {"review_id": "UMzclsyLAY", "reviewer": "Reviewer_duRA", "summary": "This paper presents a comprehensive theoretical and empirical analysis of pooling strategies in Transformer-based models, establishing pooling as a critical architectural component that significantly influences model expressivity and downstream task performance. This paper introduces a formal expressivity framework to quantify how different pooling methodsâ€”Average, Sum, Max, Last-token, and learnable variants (Weighted Average and Attention-based)â€”affect a modelâ€™s ability to distinguish semantically similar versus dissimilar inputs.", "questions": "This paper does not assess the robustness of different pooling strategies under adversarial or noisy input conditions. For example, it remains unclear how various pooling methods perform when input embeddings are perturbed by adversarial attacks or corrupted by noise. This omission is particularly important for real-world scenarios, where input reliability cannot always be ensured.\n\nPlease refer to the Weaknesses section for additional questions.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nBeyond theoretical contributions, the work provides clear, empirically supported recommendations for pooling selection in real-world applications. For instance, it highlights that learnable pooling (e.g., Weighted Average) can adaptively approximate optimal strategies for tasks with mixed global-local requirements, offering a practical solution for scenarios where fixed pooling methods may be suboptimal.\n\nWeaknesses:\n\n- While the paper observes that pooling strategy differences diminish in larger models (e.g., Mistral-7B, Llama3-8B), it lacks a detailed investigation into why this occurs. The absence of experiments on even larger models (e.g., 14B, 32B parameters) leaves a gap in understanding how pooling choices scale with model size. \n\n- The experiments primarily focus on short sequences and do not explore how varying input lengths affect pooling performance, a critical consideration for real-world NLP applications. Additionally, the paper neglects to analyze the computational trade-offs between simple and learnable pooling strategies, including their impact on inference efficiency and cost.  \n\n- While the study provides valuable insights into shallow models, it does not systematically investigate how pooling strategies interact with deeper Transformer architectures, which are increasingly prevalent in modern applications. Moreover, although the paper offers general recommendations, it lacks concrete, task-specific guidelines for selecting among particular pooling methods (e.g., Average, Sum, Max, Last), limiting its practical applicability. Additionally, the study overlooks the exploration of hybrid or ensemble pooling strategies (e.g., combining Average and Max pooling), which could better balance global and local feature retention. Such strategies may yield superior performance in multi-objective tasks but remain unexplored in this work.\n\n- This paper notes that Average pooling may discard local features while Max pooling can overlook global context, but it fails to propose a quantitative framework for evaluating semantic preservation in tasks requiring both fine-grained detail and efficient compression (e.g., long-document understanding).", "was_revised": true, "final_justification": "This work is interesting, and can be accepted. My concerns can be addressed as stated in the author's response.", "rebuttal": "We thank the reviewer for their thoughtful comments, constructive feedback, and detailed suggestions for improving the manuscript. Below, we address each point and outline corresponding revisions.\n## Q1: On the effect of pooling on adversarial robustness\n\nWe appreciate raising this important point. As noted in Line 344, we briefly mention robustness under perturbations as a direction for future work, but we agree that this aspect was not explored in detail in the current version. We appreciate the connection the reviewer draws between adversarial robustness and our expressivity framework (Equation 4). The measure $\\mathcal{E}_\\epsilon[f]$ captures the probability that small input changes lead to disproportionately large changes in the output – a concept closely related to Lipschitz continuity, which underlies many robustness analyses.\n\nImportantly, our expressivity formulation is defined over the full $\\epsilon$-neighborhood of an input and can be readily adapted to reflect worst-case deviation within that neighborhood. By considering the maximum rather than expected output difference, the upper bounds in Theorem 4.2 can be directly extended to the worst-case setting commonly used in adversarial robustness. Since our proofs already assume bounded input spaces and operate over the entire neighborhood, they naturally include both semantically similar perturbations and adversarial examples.\n\nTo complement this theoretical link with empirical evidence, we conducted a preliminary experiment using the Fast Gradient Sign Method (FGSM) attack on CIFAR-10 and CIFAR-100, using a ViT backbone. We report the attack success rate for two commonly used perturbation budgets ($\\epsilon = 3/255$ and $\\epsilon = 8/255$). The results, shown below, indicate that pooling choice meaningfully impacts adversarial vulnerability, in addition to clean accuracy:\n\n|Dataset|Attack Budget $\\epsilon$|Last (CLS)|Avg|Sum|Max|W-Avg|Attn|\n|--|:--:|:--:|:--:|:--:|:--:|:--:|:--:|\n|CIFAR-10|$3/255$|18.56|18.94|10.92|20.32|16.93|13.9|\n|CIFAR-10|$8/255$|29.38|28.16|10.92|25.97|25.59|14.44|\n|CIFAR-100|$3/255$|34.29|32.44|20.67|31.84|31.57|29.54|\n|CIFAR-100|$8/255$|43.99|41.42|20.68|38.07|40.49|32.94|\n\nWe plan to expand this analysis further in the final version, both theoretically and empirically, including experiments with stronger attacks (e.g., Projected Gradient Descent) and additional datasets. We thank the reviewer for motivating this valuable extension.\n\n## W1 & W3: On the effect of model depth\nWhile our theoretical analysis focuses on a $1$-layer Transformer-based model, the framework naturally extends to deeper models. Specifically, for a model $f$ with $L$ layers, we can decompose it as a composition of $L$ functions $f^{(i)}$ with $i \\leq L$, such that:\n\n$f(x) = f^{(L)}(f^{(L-1)})(\\dots f^{(1)}(x))))$\n\nFollowing standard results in Lipschitz continuity, the overall expressivity bound $\\gamma$ of the model $f(x)$ becomes the product of the expressivity of each $1$-layer function. Since each individual layer operates on a bounded input (as we assume $\\mathcal{X} \\in [0,1]^{n \\times d}$), this composition holds, and the theoretical structure remains intact.\n\nThis formulation provides insight into the diminishing influence of pooling in deeper models. As the number of layers increases, the overall bound $\\gamma$ grows, making the relative effect of the final pooling operation smaller, especially when its impact scales as, for example, $1/\\sqrt{n}$ for Average pooling. This aligns with our empirical observation in Lines 282-283, where the performance gap between pooling strategies narrows in larger models.\nTo further validate this effect in practice, and in response to the reviewer's suggestion, we extended our experiments to larger-scale models:\n- Qwen2.5 14B and Qwen2.5 32B\n- Mistral3.1 24B\n\nThe tables below report results across tasks for each model using the same experimental setup as in Section 5.2. As shown, while pooling trends are generally preserved, the absolute differences between pooling methods decrease as model size increases, consistent with the theoretical interpretation.\n\nTable 2: Test metrics for Qwen2.5 14B\n|  |STSB|Hellaswag|Banking|Tweet|Next Token|\n|--|--|--|--|--|--|\n|Last|0.288±0.003|0.692±0.011|34.497±0.178|57.008±0.004|91.039±0.002|\n|Avg|0.581±0.000|0.773±0.000|85.390±0.001|69.930±0.008|53.893±0.250|\n|Sum|0.579±0.002|0.776±0.002|82.711±0.008|61.218±1.194|47.053±0.341|\n|Max|0.488±0.005|0.727±0.001|77.825±0.006|55.390±0.728|25.578±0.332|\n|W-Avg|0.589±0.001|0.798±0.002|86.867±0.002|69.770±0.007|90.653±0.020|\n|Attn|0.259±0.011|0.712±0.013|66.387±1.939|56.410±3.001|40.573±0.892|\n\nTable 3: Test metrics for Qwen2.5 32B\n|  |STSB|Hellaswag|Banking|Tweet|Next Token|\n|--|--|--|--|--|--|\n|Last|0.303±0.001|0.723±0.009|34.383±0.268|55.886±0.006|89.886±0.005|\n|Avg|0.603±0.000|0.781±0.001|87.760±0.002|68.328±0.012|48.536±0.334|\n|Sum|0.604±0.001|0.782±0.005|85.227±0.013|63.665±0.865|39.204±0.627|\n|Max|0.488±0.005|0.729±0.008|83.929±0.008|65.268±0.596|31.273±0.586|\n|W-Avg|0.627±0.003|0.812±0.003|89.903±0.004|69.464±0.009|89.022±0.019|\n|Attn|0.355±0.016|0.730±0.011|68.929±1.054|58.828±2.695|29.750±1.028|\n\nTable 4: Test metrics for Mistral3.1 24B model. \n|  |STSB|Hellaswag|Banking|Tweet|Next Token|\n|--|--|--|--|--|--|\n|Last|0.503±0.002|0.745±0.008|75.487±0.087|53.701±0.005|88.972±0.007|\n|Avg|0.631±0.001|0.784±0.001|87.403±0.006|66.871±0.009|51.723±0.812|\n|Sum|0.622±0.004|0.783±0.004|87.597±0.015|62.791±0.976|42.306±0.732|\n|Max|0.488±0.003|0.733±0.007|79.675±0.010|61.655±0.473|27.804±0.923|\n|W-Avg|0.682±0.002|0.816±0.003|88.711±0.017|66.200±0.005|87.849±0.024|\n|Attn|0.392±0.043|0.697±0.009|72.922±1.012|31.294±4.452|19.911±2.023|\n\nWe will integrate this theoretical extension and the new empirical results into the revised version of the manuscript, and we thank the reviewer again for prompting this valuable addition.\n\n## W2: On the effect of sequence length in pooling operators\n\nTo investigate this, we conducted an additional experiment using the Mistral-7B model on the HellaSwag dataset, varying the maximum input sequence length from 16 to 128 tokens. Inputs were either truncated or padded as needed, and padding tokens were excluded from pooling operations, consistent with our prior setup.\nThe results, shown in Table 6, compare the performance of different pooling strategies across sequence lengths. Naturally, shorter inputs can lead to performance degradation due to the truncation of semantically important content. To isolate the effect of pooling itself, we recommend comparing methods column-wise (i.e., at fixed sequence lengths)\n\nTable 6: Mean and standard deviation of metrics for Mistral 7B on HellaSwag for different sequence lengths.\n| |16|32|64|128|\n|--|--|--|--|--|\n|Last|0.454±0.003|0.621±0.002|0.770±0.002|0.781±0.001|\n|Avg|0.503±0.001|0.702±0.000|0.769±0.000|0.771±0.001|\n|Sum|0.504±0.001|0.702±0.001|0.769±0.000|0.771±0.001|\n|Max|0.435±0.003|0.616±0.003|0.709±0.001|0.700±0.008|\n|W-Avg|0.523±0.002|0.724±0.000|0.801±0.000|0.802±0.003|\n|Attn|0.220±0.169|0.278±0.251|0.737±0.018|0.764±0.025|\n\nWe observe that while pooling sensitivity is more pronounced at shorter input lengths, particularly for methods like Last-token or Attention-based pooling, the relative trends between pooling strategies remain consistent. For longer sequences (64 or 128 tokens), performance stabilizes across pooling types, aligning with our theoretical bounds, which become more predictive in this regime.\n\nAdditionally, in a revised version of the manuscript, we will include a dedicated analysis of inference-time costs. As a preview, *inference latency (in ms)* for pooling on Mistral-7B with Banking77 is: Last: 0.14 ± 0.01; Avg: 0.16 ± 0.01; Sum: 0.12 ± 0.01; Max: 0.23 ± 0.01; W-Avg: 0.33 ± 0.02; Attn: 0.53 ± 0.03. These show that while learnable methods add modest overhead, it may be justified in tasks requiring higher expressivity. In contrast, flat poolings are cheaper and often competitive, highlighting a classic trade-off between inductive bias and learnable capacity. \n\nWe will include these experiments and discussion in the revised manuscript and thank the reviewer for suggesting this addition.\n\n## W3: Regarding hybrid pooling strategies\n\nWe agree with the reviewer that hybrid pooling, which combines multiple strategies to capture both global and local context, is a promising direction.\nThis idea is reflected in Figure 3 (and Figure 5), where Weighted Average pooling adaptively attends to both prominent and supporting tokens, thus effectively blending pooling behaviors. Importantly, any hybrid pooling strategy is ultimately built from flat pooling operations which are the exact components analyzed in our expressivity framework. Understanding their individual properties, as we do in this work, is a necessary step toward designing principled hybrid or adaptive approaches.\n\nWe will clarify this connection in the revised manuscript and thank the reviewer for highlighting this valuable perspective.\n\n## W4: Considering the tasks requiring both local and global context\n\nWe agree with the reviewer that certain downstream tasks may require preserving both local detail and global structure. As noted in our response to W3, such mixed-context tasks present a challenge for flat pooling strategies, which are typically biased toward one or the other. In these cases, we believe that either learned pooling methods, such as Weighted Average or Attention-based, or hybrid approaches that combine multiple fixed strategies offer a more suitable solution. This is consistent with our findings in Sections 5.2 and 5.3, where learned methods adaptively emphasize tokens based on task demands and outperform flat pooling on tasks with diverse contextual requirements.\n\nWe will clarify this point in the revised manuscript and explicitly highlight use cases where hybrid or learned pooling strategies are likely preferable. We appreciate the reviewer's suggestion in guiding this improvement."}, {"review_id": "qYuXqDm3Mx", "reviewer": "Reviewer_u5rN", "summary": "The author studied the impact of various token embedding pooling approaches both theoretically and experimentally. They first analyzed the impact of pooling the expressivity of Transformer-based models, deriving closed-form bounds on the model’s ability to distinguish similar inputs. Then, they experimentally evaluated pooling strategies across computer vision, natural language processing, and time-series tasks. Their findings suggest that tasks requiring global contextual understanding, such as text classification, benefit from contractive methods like average pooling. Conversely, tasks emphasizing local details, such as next-token prediction, benefit from expansive methods like last-token and sum pooling.", "questions": "* Can the authors elaborate on the two aspects of contractive vs. expansive and global vs. final? How would you revise your manuscript to address my comment?\n\n* The statement in lines 293-295 that “Sum pooling outperforms Average pooling in several cases, that can be attributed to the larger norm of summed representations, which results in stronger gradients and faster learning under fixed training hyperparameters”, while logical, lacks supporting evidence. If that’s the case, why does average pooling outperform sum pooling in most cases?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "This paper makes a timely contribution as pooling in Transformer-based models is an essential design choice that is underexplored. The experimental setup is modern, leveraging state-of-the-art LLMs such as Llama3, Mistral 7B, and Qwen 2.5. Note that I believe it would be beneficial to explicitly highlight the choice of these models in the main text.\n\nThe main weakness lies in the shift between the theoretical and empirical analyses. The theoretical section focuses on contractive against expansive pooling, whereas the empirical results discuss global against final pooling. For instance, average pooling is contractive and sum pooling is expansive, and the authors suggest that contractive methods benefit global tasks while expansive benefit local tasks. Yet, experiments indicate that tasks requiring global context perform better with global pooling methods like average or sum. \n* Lines 228-230: “Therefore, we posit that pooling should be selected based on task requirements: tasks emphasizing global context (e.g., image inpainting or text classification) benefit from contractive pooling, while those relying on local detail (e.g., next-token prediction) may perform better with expansive alternatives.”\n* Lines 278-281: “For tasks requiring global context, such as classification and semantic similarity, global pooling methods like Average or Sum significantly outperform Last-token pooling. Conversely, Last-token pooling yields superior performance in next-token prediction tasks.”\nTo improve the paper's narrative, both the theoretical analysis and experimental results should consistently address both contractive vs. expansive and global vs. final aspects. \n\nNote that I did not thoroughly review the supplemental material, particularly the proofs.", "was_revised": true, "final_justification": "The authors addressed my main concern, which was the connection between the theoretical analysis (expansive/contractive) and experimental (global/local). I have updated my score from 3 to 4.", "rebuttal": "We thank the reviewer for their thoughtful comments and acknowledging the value and timely contribution of our study. Below, we respond point by point to the reviewer's questions and comments, and outline the corresponding revisions we will incorporate to improve the manuscript.\n\n## W1: On the choice of LLMs used\n\nWe appreciate highlighting the importance of model selection when evaluating pooling strategies. Our aim was to ensure that the observed trends generalize across a diverse range of model capacities and architectures. To this end, we selected a representative set of widely used state-of-the-art language models, including Qwen2.5, Mistral-7B, and LLaMA3.1-8B. These models span different size classes and training pipelines, and are broadly recognized in the community. Our results across these models consistently support the theoretical insights introduced in the paper.\n\nTo further strengthen our findings, we extended our study to include larger-scale models, such as Mistral 3.1 24B, and report the results in the table below. As seen, pooling behavior is consistent even at higher parameter scales, in line with additional observations on Qwen2.5 14B and 32B (see our response to Reviewer duRA W1 & W3). These additions reinforce our claim that the conclusions drawn in this work are not model-specific, but rather reflect generalizable properties of pooling strategies.\n\nTable 2: Mean and Standard deviation of test metrics for NLP tasks for the Mistral3.1 24B model. \n| Metric  | STSB              | Hellaswag         | Banking           | Tweet             | Next Token        |\n|---------|-------------------|-------------------|-------------------|-------------------|-------------------|\n| Last    | 0.503 ± 0.002     | 0.745 ± 0.008     | 75.487 ± 0.087    | 53.701 ± 0.005    | 88.972 ± 0.007    |\n| Avg     | 0.631 ± 0.001     | 0.784 ± 0.001     | 87.403 ± 0.006    | 66.871 ± 0.009    | 51.723 ± 0.812    |\n| Sum     | 0.622 ± 0.004     | 0.783 ± 0.004     | 87.597 ± 0.015    | 62.791 ± 0.976    | 42.306 ± 0.732    |\n| Max     | 0.488 ± 0.003     | 0.733 ± 0.007     | 79.675 ± 0.010    | 61.655 ± 0.473    | 27.804 ± 0.923    |\n| W-Avg   | 0.682 ± 0.002     | 0.816 ± 0.003     | 88.711 ± 0.017    | 66.200 ± 0.005    | 87.849 ± 0.024    |\n| Attn    | 0.392 ± 0.043     | 0.697 ± 0.009     | 72.922 ± 1.012    | 31.294 ± 4.452    | 19.911 ± 2.023    |\n\nWe will include this rationale and the extended results in the revised manuscript.\n\n\n## Q1: On the local/global and expensive/contractive aspects\n\nWe thank the reviewer for this thoughtful question. We believe the distinction raised is important, and we appreciate the opportunity to clarify the relationship between these axes.\n\nIn our framework, contractive vs. expansive refers to the mathematical behavior of pooling functions, specifically how they scale model sensitivity to input perturbations (Section 4.2, Theorem 4.2). For example, Average pooling is contractive (scaling as $1/\\sqrt{n}$), while Sum and Last-token pooling are expansive or neutral. These differences affect how much variation in the input is retained in the final pooled representation.\n\nOn the other hand, the global vs. local distinction (we believe this may be what was meant by \"final\") refers to task requirements, i.e., whether the output depends on the overall input (global context) or specific positions (local context). For example:\n- Global-context tasks (e.g., classification, semantic similarity) benefit from pooling methods that integrate information across all tokens.\n- Local-context tasks (e.g., next-token prediction, forecasting) benefit from methods that emphasize recent or specific tokens, such as Last-token pooling.\n\nThese two perspectives are connected: expansive pooling tends to preserve token-level variation, making it better suited to local-context tasks; contractive pooling smooths across the input, aligning with global-context needs.\n\nFigure 2 empirically validates this theoretical framing across modalities. It shows that expansive pooling methods (e.g., Sum, Last) are more sensitive to perturbations, while Average pooling smooths input variation, closely aligning with our derived bounds. This supports our claim that expressivity behavior maps onto different task needs. To illustrate the semantic effects more concretely, Figure 4 presents examples where minor edits to input text (e.g., replacing \"fluctuations\" with \"swings\" or \"volatility\") yield small changes under Average pooling but larger differences under Last-token pooling. This demonstrates that expansive pooling captures fine-grained changes that might be critical in local-context tasks.\n\nWe finally note that while related empirical trends have been observed in prior work [1-4], to our knowledge, this is the first work to formally characterize and prove these relationships from a theoretical standpoint.\n\nWe will revise the manuscript to:\n\n- Clarify the relationship between the theoretical (contractive/expansive) and empirical (global/local) axes in the transition between Sections 4 and 5.\n- Add a brief explanation that \"final-context\" tasks, as described by the reviewer, are more precisely described as local-context tasks, and adjust our terminology accordingly for clarity.\n- Highlight Figure 2, which empirically shows how pooling methods behave under perturbations, consistent with the scaling predicted by our theoretical framework.\n- Expand on Figure 4, which illustrates how pooling methods respond to subtle semantic differences, depending on whether token-level or global meaning is more relevant.\n\nWe appreciate the reviewer’s feedback and believe these revisions will strengthen the conceptual clarity of the manuscript.\n\n[1] Tang & Yang. Pooling and attention: What are effective designs for llm-based embedding models?\n\n[2] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.\n\n[3] Lee et al. NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models.\n\n[4] Xing et al. Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.\n\n## Q2: On the performance of Sum pooling\n\nWe thank the reviewer for the question and the opportunity to clarify this point. The observation that Sum pooling occasionally outperforms Average pooling was noted specifically in the time-series modality, for all majority of the considered sizes of the MOMENT model and the different classification datasets.\n\nIn this setting, we noted that Sum pooling results in larger gradient norms compared to Average pooling, which may lead to faster early-stage optimization. However, this effect is not observed in NLP or vision tasks, suggesting that it may be modality-specific.\n\nWe hypothesize that this behavior could be influenced by normalization schemes or architectural details unique to time-series models, such as the absence of token-level layer norm or differences in input scaling. While this offers a plausible explanation, we currently do not have direct evidence to confirm the cause. We will clarify this distinction in the revised manuscript and flag this as a modality-specific observation that warrants further study. We thank the reviewer for encouraging us to articulate this more precisely."}, {"review_id": "O3cqycz9gw", "reviewer": "Reviewer_k3Rb", "summary": "This paper conducts theoretical analysis on the pooling functions of Transformer and derives expressivity bounds for various pooling strategies, including sum, max, last, and average. Through experiments, it verifies that different tasks require distinct pooling functions to capture global and local contextual understanding, respectively. The main contribution lies in establishing an expressivity bound for pooling functions. However, this bound aligns closely with human intuition, which may limit the significance of the proposed method.\n\n\nAs a researcher in the field of machine learning with a greater focus on applications, I place more emphasis on the practical significance of theoretical research.", "questions": "1. For the multi-token prediction task, which pooling function is more suitable?\n\n2. How does the proposed theoretical framework demonstrate its guiding significance for practical tasks?", "limitations": "1. Evaluation is limited to frozen backbones, leaving the effect of jointly adapting pooling and the backbone under end-to-end training less explored. \n\n\n2  \"In future work, we aim to develop hybrid pooling methods that dynamically balance global smoothing with token-level sensitivity while adapting to the TBM’s inherent smoothing behavior. We also plan to investigate how pooling impacts robustness under perturbations, derive scaling laws that govern pooling performance as datasets grow, and refine our theoretical framework with tighter bounds that  identify when specific strategies are provably optimal.\" \nAs author discussed, I believe that further work will enhance the research significance of this article and improve its chances of publication.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1.This paper conducts a novel theoretical analysis of the pooling functions in transformers. By delving into the mathematical properties and underlying mechanisms of pooling operations, it offers a fresh perspective on how these functions interact with the overall architecture, thereby enriching the theoretical foundation of transformer-based models.\n\n2.The paper clearly illustrates that different tasks require distinct strategies for capturing global and local contextual understanding. Through in-depth analysis and examples, it effectively demonstrates the importance of tailoring pooling functions to task-specific requirements, providing valuable guidance for researchers aiming to optimize model performance for various natural language processing and computer vision tasks.\n\n3.The experimental results effectively verify the effectiveness of the proposed expressivity bound. Rigorous testing across a range of datasets and scenarios provides substantial empirical evidence supporting the validity of the theoretical framework, showcasing the practical utility of the research findings.\nWeaknesses:\n1. No new model is proposed to construct a more expressive pooling layer. Merely analyzing the existing pooling functions without introducing innovative architectures limits the potential for significant practical improvements. A novel model could have better demonstrated the practical implications of the theoretical expressivity bound and potentially led to breakthroughs in model performance.\n\n2.The results regarding the expressivity bound of the pooling layer may be intuitive, and it is challenging to draw inspiration for related research areas. The lack of deep exploration into the broader implications and potential extensions of the findings restricts the impact of the research. Without clear connections to emerging trends or open problems in the field, it becomes difficult for other researchers to build upon this work, reducing its long-term value.", "was_revised": true, "final_justification": "My concern is that no new model is proposed to construct a more expressive pooling layer, and it also has not been solved, so i think it not enough for acceptance.", "rebuttal": "We thank the reviewer for their thoughtful comments and constructive feedback. We additionally are grateful to the reviewer for acknowledging the value of our theoretical insights in guiding practical aspects of machine learning applications in various domains.\n\n## W1: No new model is proposed to construct a more expressive pooling layer\n\nThank you for the observation. While it is true that we do not propose a new pooling mechanism in this work, our primary goal is to provide a theoretical and empirical foundation for understanding existing pooling strategies -- a space that has so far lacked rigorous analysis.\n\nAs discussed in the Conclusion (lines 326-330) and Limitations (lines 342-343), our expressivity framework is intended to guide the development of future pooling methods, including hybrid or adaptive strategies. By formally quantifying how pooling affects model expressivity (via the $\\gamma$ bound), and validating these findings across tasks and domains, we aim to lay the groundwork for principled design of new pooling operators.\n\nAdditionally, the paper offers practical guidance for selecting pooling methods based on task requirements (global vs. local context), data regime, and resource constraints. We believe this makes our contribution actionable, even without introducing a new architecture, and positions pooling as a meaningful design choice rather than a fixed component.\n\n## W2 & Q2: On the expressivity results may appear intuitive, broader research impact, and significance for practical tasks\n\nWe appreciate the reviewer's perspective. While the behavior of pooling functions may appear intuitive, our contribution lies in formalizing these observations into a principled framework that explains how pooling affects model expressivity and sensitivity. Whereas prior work [1-4] has lacked theoretical grounding, our aim is to provide a unified explanation of these patterns. We see this as a necessary step toward more informed architectural choices and the development of new pooling strategies guided by theory.\n\nThis framework enables a principled understanding of when and why different pooling strategies succeed, depending on task requirements. A key takeaway of our work is that no pooling method is universally optimal. Rather, the best choice depends on whether the task requires global or local context.\n\nWe also highlight a practical implication: if resources permit, adaptive pooling strategies (e.g., attention-based or weighted average) can learn to fit task-specific needs, but they come at a higher computational cost. When training time or data is limited, our framework helps guide the selection of fixed pooling strategies that match task characteristics.\n\nWe will clarify these points in the revision and thank the reviewer for prompting us to better frame the broader value of our work.\n\n\n[1] Tang & Yang. Pooling and attention: What are effective designs for llm-based embedding models?\n\n[2] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.\n\n[3] Lee et al. NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models.\n\n[4] Xing et al. Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.\n\n## Q1: For the multi-token prediction task, which pooling function is more suitable?\n\nMulti-token prediction, where several future tokens (e.g., the next 10) are predicted simultaneously via separate heads, can be viewed as a natural extension of the classical next-token prediction task. In such settings, we would expect Last-token pooling to perform well among the fixed pooling strategies, as it retains recent contextual information, which is often most relevant for autoregressive forecasting.\n\nThat said, when long-range dependencies play a more significant role. For example, when predicting tokens farther ahead (e.g., the 8th or 10th future token), fixed strategies may be less effective. In these cases, learnable pooling methods such as Weighted Average or Attention-based pooling are better suited, as they can dynamically attend to the most informative parts of the input sequence for each prediction target.\n\nWe will clarify this distinction in the revised manuscript and thank the reviewer for prompting this useful elaboration.\n\n## L1: Evaluation is limited to frozen backbones\n\nWe agree that jointly adapting the pooling mechanism and the backbone through end-to-end training is an important direction. In this work, we focused on the frozen-backbone setting to isolate the effect of pooling strategies and analyze their theoretical properties independently of backbone optimization. We will clarify the scope in the revised manuscript and highlight joint adaptation as a potential avenue for future work.\n\n## L2: On the potential for future work and extensions\n\nWe thank the reviewer for their thoughtful support. We agree that extending this work could enhance its broader impact. We will ensure that these directions are more clearly positioned in the revised manuscript as concrete next steps."}]

Author rebuttal:

Meta-review: The paper provides a comprehensive theoretical and empirical analysis of various pooling strategies in transformer models.

Following the rebuttal, three reviewers indicate a positive opinion of the paper (one accept, two borderline accepts) and note that all concerns they had have been addressed, and that the paper makes a valuable contribution. The remaining reviewer rates the paper a borderline reject, saying that the paper only analyzes existing pooling methods rather than propose a new one. However, other reviewers have highlighted the importance of the theoretical and empirical analysis and its utility in making practical decisions for real life applications. 

Accordingly, AC concurs with the majority opinion and recommends acceptance.
Decision: Accept (poster)