---
title: 'Why Mean Pooling Works: Quantifying Second-Order Collapse in Text Embeddings'
created: 2026-05-05
updated: 2026-05-05
type: concept
tags:
- embedding
- architecture
- retrieval
- training
sources:
- raw/papers/2026/04/2604.27398.md
---

# Why Mean Pooling Works: Second-Order Collapse

## Core Problem Definition

**Formal Statement.** Given a text $t_i$ with token embeddings $\mathbf{X}_i = [\mathbf{x}_{i,1}, \ldots, \mathbf{x}_{i,n_i}] \in \mathbb{R}^{d \times n_i}$, mean pooling produces:

$$\bm{\mu}(\mathbf{X}_i) = \frac{1}{n_i} \sum_{j=1}^{n_i} \mathbf{x}_{i,j}$$

This is the **first-order statistic** (mean) of the token embedding distribution. The **second-order statistic** (covariance) is:

$$\bm{\Sigma}(\mathbf{X}_i) = \frac{1}{n_i} \sum_{j=1}^{n_i} (\mathbf{x}_{i,j} - \bm{\mu}(\mathbf{X}_i))(\mathbf{x}_{i,j} - \bm{\mu}(\mathbf{X}_i))^\top$$

**The concern:** Mean pooling discards all information beyond the first-order statistic. Two texts with *identical means but different covariances* would be mapped to the same embedding — a "second-order collapse." This raises the question: does this collapse actually happen in practice, and if so, does it matter?

**The paper's thesis:** Mean pooling works *not because* it preserves all information, but because modern text encoders (especially contrastive fine-tuned ones) naturally concentrate token embeddings, making the second-order information less discriminative.

## Method

### Second-Order Collapse Metric (SOCM)

**When collapse arises:** First-order statistics are similar but second-order statistics differ:
$$\bm{\mu}(\mathbf{X}_1) \approx \bm{\mu}(\mathbf{X}_2) \land \bm{\Sigma}(\mathbf{X}_1) \neq \bm{\Sigma}(\mathbf{X}_2)$$

**Distance definitions** (assuming unit-norm means $\|\bm{\mu}(\mathbf{X}_i)\|_2 = 1$):

First-order distance (scaled squared Euclidean):
$$d_\mu := \|\bm{\mu}(\mathbf{X}_1) - \bm{\mu}(\mathbf{X}_2)\|_2^2 / 4 \in [0, 1]$$

Second-order distance (scaled Bures-Wasserstein distance):
$$d_\Sigma := \text{tr}\left(\bm{\Sigma}(\mathbf{X}_1) + \bm{\Sigma}(\mathbf{X}_2) - 2(\bm{\Sigma}(\mathbf{X}_1)^{1/2} \bm{\Sigma}(\mathbf{X}_2) \bm{\Sigma}(\mathbf{X}_1)^{1/2})^{1/2}\right) / 4 \in [0, 1]$$

**SOCM formula:**
$$\text{SOCM} = (1 - d_\mu) \cdot d_\Sigma \in [0, 1]$$

**Interpretation:**
- SOCM = 1: Maximum collapse (identical means, completely different covariances)
- SOCM = 0: No collapse (either means differ or covariances match)

**Theoretical justification:** SOCM corresponds to the interaction term in the decomposition of the $L_2$-Wasserstein distance:
$$W_2^2(\mathcal{N}(\bm{\mu}_1, \bm{\Sigma}_1), \mathcal{N}(\bm{\mu}_2, \bm{\Sigma}_2)) = \|\bm{\mu}_1 - \bm{\mu}_2\|^2 + \text{tr}(\bm{\Sigma}_1 + \bm{\Sigma}_2 - 2(\bm{\Sigma}_1^{1/2} \bm{\Sigma}_2 \bm{\Sigma}_1^{1/2})^{1/2})$$

SOCM captures the *unexplained* part of this distance after accounting for mean differences.

### SOCM Properties

The metric satisfies five desirable properties:
1. **(a) Extreme collapse:** SOCM = 1 $\Leftrightarrow$ $d_\mu = 0 \land d_\Sigma = 1$
2. **(b) No collapse:** SOCM = 0 $\Leftrightarrow$ $d_\mu = 1 \lor d_\Sigma = 0$
3. **(c) Monotonicity in $d_\mu$:** $\partial \text{SOCM} / \partial d_\mu \leq 0$ (decreasing means distance reduces collapse)
4. **(d) Monotonicity in $d_\Sigma$:** $\partial \text{SOCM} / \partial d_\Sigma \geq 0$ (increasing covariance distance increases collapse)
5. **(e) Interaction:** $\partial^2 \text{SOCM} / \partial d_\mu \partial d_\Sigma \leq 0$ (the effect of $d_\Sigma$ diminishes as $d_\mu$ grows)

### Experimental Procedure

1. Prepare a text pair dataset $D = \{(t_1, t_2)\}$
2. For each text, obtain token embeddings $\mathbf{X}_i$ from model $f$
3. Normalize: $\mathbf{X}_i^{\text{norm}} = [\mathbf{x}_{i,1}/\|\bm{\mu}(\mathbf{X}_i)\|, \ldots, \mathbf{x}_{i,n_i}/\|\bm{\mu}(\mathbf{X}_i)\|]$
4. Compute SOCM for each pair $(\mathbf{X}_1, \mathbf{X}_2)$

## Models Evaluated

The paper evaluates multiple model families:
- **BERT** (pretrained backbone)
- **GTE-base** (contrastive fine-tuned)
- Additional contrastive fine-tuned encoders vs. their pretrained backbones

## Experimental Results

### Finding 1: Contrastive Fine-Tuning Reduces Collapse

Contrastive fine-tuned text encoders show **significantly lower SOCM** than their pretrained backbones. This means:
- Fine-tuned models produce token embeddings that are more concentrated within each text
- The spatial structure (covariance) varies less between different texts
- Mean pooling is less lossy for fine-tuned models

**Visualization (Figure 1):** PCA projection shows BERT token embeddings for two different texts with overlapping means but different spreads (high SOCM). GTE-base shows tighter clusters with more distinct means (low SOCM).

### Finding 2: Token Embedding Concentration Explains Robustness

The robustness of mean pooling comes from **token embedding concentration** within each text, which arises through Transformer layers. As token embeddings become more concentrated:
- The covariance $\bm{\Sigma}(\mathbf{X}_i)$ becomes more uniform across texts
- $d_\Sigma$ decreases
- SOCM decreases

This concentration is a natural consequence of the Transformer's self-attention mechanism, which progressively aggregates information.

### Finding 3: SOCM Correlates with Downstream Performance

Models with lower SOCM (less collapse) tend to perform better on downstream tasks:
- Information retrieval
- Semantic similarity
- Classification tasks

This suggests that robustness to mean pooling collapse is one factor in the success of modern text encoders.

### Model Comparison

| Model Type | Typical SOCM | Mean Pooling Effectiveness |
|------------|-------------|---------------------------|
| Pretrained backbone (e.g., BERT) | Higher | More information lost |
| Contrastive fine-tuned (e.g., GTE) | Lower | Less information lost |

## Comparison with Alternative Approaches

| Approach | Representation | Preserves 2nd-Order | Computational Cost |
|----------|---------------|---------------------|-------------------|
| Mean pooling | Single vector | No | $O(nd)$ |
| ColBERT (late interaction) | Token list | Yes | $O(n_q \cdot n_d)$ |
| Optimal transport (WMD) | Token list | Yes | $O(n^3 \log n)$ |
| GaussCSE | Gaussian (mean + covariance) | Yes | $O(d^2)$ |
| BERTScore | Token list | Partial | $O(n_q \cdot n_d)$ |

The paper argues that mean pooling remains dominant not despite losing second-order information, but because modern encoders minimize the *impact* of that loss through concentration.

## Critical Analysis

### Strengths
1. **Rigorous formalization**: SOCM is well-motivated theoretically with clean properties and a connection to Wasserstein distance
2. **Explains an empirical mystery**: Why does such a simple aggregation method work so well? This paper provides a satisfying answer
3. **Actionable insight**: If you want mean pooling to work better, train your encoder to produce more concentrated token embeddings (e.g., via contrastive learning)
4. **Broad applicability**: The metric can be applied to any embedding model to assess pooling robustness

### Weaknesses and Open Questions
1. **Correlation vs. causation**: SOCM correlates with downstream performance, but is lower SOCM *causing* better performance, or is it a side effect of better training?
2. **Unit-norm assumption**: The SOCM formula assumes $\|\bm{\mu}(\mathbf{X}_i)\|_2 = 1$. While common in practice, this limits applicability to models that don't normalize
3. **Scope limited to second-order**: The paper acknowledges that even higher-order statistics could collapse. Why is second-order the right level of analysis?
4. **Text pair construction**: The results depend on how text pairs are selected. Adversarial pairs might reveal different collapse patterns
5. **Causal mechanism**: The paper shows that Transformer layers cause concentration, but doesn't fully explain *why* self-attention produces this effect
6. **Practical guidance limited**: The paper explains why mean pooling works but doesn't provide concrete guidance on when to use alternatives

### Implications for Our Work
- **For [[embedding]] systems**: When choosing pooling strategies, measure SOCM to assess whether mean pooling is sufficient for your use case
- **For [[retrieval]]**: If using mean pooling, prefer contrastive fine-tuned encoders (lower SOCM = less information loss)
- **For [[tachiom]]**: Multivector retrieval methods like ColBERT avoid the mean pooling limitation entirely by preserving token-level representations
- **For model design**: If you need richer text representations than mean pooling provides, consider GaussCSE-style approaches that explicitly model covariance
- **For [[agent]] memory**: When encoding agent experiences/token sequences, be aware that mean pooling may lose important structural information unless the encoder produces concentrated embeddings

## Related
- [[embedding]] — Text embedding methods and representations
- [[retrieval]] — Information retrieval paradigms
- [[tachiom]] — Multivector retrieval with token-level representations
- [[inference]] — Inference efficiency for embedding models