Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models

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.

W1: No new model is proposed to construct a more expressive pooling layer

Thank 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.

As 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.

Additionally, 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.

W2 & Q2: On the expressivity results may appear intuitive, broader research impact, and significance for practical tasks

We 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.

This 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.

We 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.

We will clarify these points in the revision and thank the reviewer for prompting us to better frame the broader value of our work.

[1] Tang & Yang. Pooling and attention: What are effective designs for llm-based embedding models?

[2] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.

[3] Lee et al. NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models.

[4] Xing et al. Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.

Q1: For the multi-token prediction task, which pooling function is more suitable?

Multi-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.

That 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.

We will clarify this distinction in the revised manuscript and thank the reviewer for prompting this useful elaboration.

L1: Evaluation is limited to frozen backbones

We 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.

L2: On the potential for future work and extensions

We 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.

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.

Q1: On the effect of pooling on adversarial robustness

We 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.

Importantly, 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.

To 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:

We 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.

W1 & W3: On the effect of model depth

While 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:

$f(x) = f^{(L)}(f^{(L-1)})(\dots f^{(1)}(x))))$

Following 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.

This 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. To further validate this effect in practice, and in response to the reviewer's suggestion, we extended our experiments to larger-scale models:

• Qwen2.5 14B and Qwen2.5 32B
• Mistral3.1 24B

The 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.

Table 2: Test metrics for Qwen2.5 14B

Table 3: Test metrics for Qwen2.5 32B

Table 4: Test metrics for Mistral3.1 24B model.

We 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.

W2: On the effect of sequence length in pooling operators

To 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. The 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)

Table 6: Mean and standard deviation of metrics for Mistral 7B on HellaSwag for different sequence lengths.

We 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.

Additionally, 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.

We will include these experiments and discussion in the revised manuscript and thank the reviewer for suggesting this addition.

W3: Regarding hybrid pooling strategies

We agree with the reviewer that hybrid pooling, which combines multiple strategies to capture both global and local context, is a promising direction. This 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.

We will clarify this connection in the revised manuscript and thank the reviewer for highlighting this valuable perspective.

W4: Considering the tasks requiring both local and global context

We 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.

We 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.

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.

W1: On the choice of LLMs used

We 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.

To 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.

Table 2: Mean and Standard deviation of test metrics for NLP tasks for the Mistral3.1 24B model.

We will include this rationale and the extended results in the revised manuscript.

Q1: On the local/global and expensive/contractive aspects

We 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.

In 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.

On 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:

• Global-context tasks (e.g., classification, semantic similarity) benefit from pooling methods that integrate information across all tokens.
• Local-context tasks (e.g., next-token prediction, forecasting) benefit from methods that emphasize recent or specific tokens, such as Last-token pooling.

These 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.

Figure 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.

We 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.

We will revise the manuscript to:

• Clarify the relationship between the theoretical (contractive/expansive) and empirical (global/local) axes in the transition between Sections 4 and 5.
• 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.
• Highlight Figure 2, which empirically shows how pooling methods behave under perturbations, consistent with the scaling predicted by our theoretical framework.
• 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.

We appreciate the reviewer’s feedback and believe these revisions will strengthen the conceptual clarity of the manuscript.

[1] Tang & Yang. Pooling and attention: What are effective designs for llm-based embedding models?

[2] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.

[3] Lee et al. NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models.

[4] Xing et al. Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.

Q2: On the performance of Sum pooling

We 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.

In 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.

We 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.

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.

Q1: What happens when stopwords or negation are removed? Could a pooling mask be used to discard such tokens?

We 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.

To explore this, we conducted an additional analysis using the same base sentence as in Figure 4:

• Original: "The stock market saw significant fluctuations last week."

• Stopword removed: "Stock market saw significant fluctuations last week."

• Negation introduced: "The stock market saw insignificant fluctuations last week."

We then measured the L2 distance between the pooled output of the original and perturbed sentences for each pooling method. The results are shown below:

These 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.

We 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.

Q2: Are there any proposals on a pooling mask that masked out certain tokens?

Incorporating 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.

In 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.

We will mention this connection in the revised manuscript and thank the reviewer for prompting this valuable direction for future work.