Token-Weighted Averaging

• Token-weighted averaging is defined as a method that assigns explicit, learnable or statically determined weights to each token, improving the aggregation in neural architectures.
• It enhances multi-vector retrieval by modifying the Chamfer-style distance and addresses attention dispersion via arithmetic and homogeneous mixing techniques.
• The technique combines local adaptive attention with global uniform averaging, demonstrating measurable gains in benchmarks like BEIR and ImageNet.

Token-weighted averaging is an approach for aggregating information from token-level representations in neural architectures, particularly within multi-vector retrieval systems and attention-based models. It modifies classic aggregation rules by assigning explicit, learnable or statically determined weights to each token, thereby permitting the model to prioritize more semantically or informationally salient tokens during scoring, retrieval, or aggregation. The technique has seen successful application in both dense retrieval frameworks and attention mechanisms, addressing inherent limitations such as dispersion and uniform weighting in large token sets.

1. Formalization in Multi-Vector Retrieval

ColBERT introduced a "late interaction" mechanism for bi-encoder dense retrieval, where queries and documents are encoded into sets of token-level embeddings and scored with a Chamfer-style distance: for each query token, the nearest document token is selected, and the sum across query tokens forms the overall score. The classic Chamfer-style distance is an unweighted sum:

$$D(q, d) = \sum_{i=1}^m \min_{j=1,\ldots,n} d(q_i, d_j),$$

where $q_i$ and $d_j$ denote token vectors for the query and document, respectively.

Token-weighted averaging generalizes this by weighting each query token's contribution:

$$D(q, d) = \sum_{i=1}^m w_i \min_{j=1,\ldots,n} d(q_i, d_j),$$

where the weights $w_i$ reflect token importance. These weights can be statically set (e.g., Inverse Document Frequency, IDF) or learned through end-to-end training. Empirical evaluations show that incorporating token weights, even while keeping token vector encodings fixed, enhances the expressiveness of late interaction mechanisms, yielding notable improvements in retrieval benchmarks such as BEIR—e.g., a +1.28% improvement in Recall@10 using IDF-based weights and +3.66% with few-shot fine-tuning (S et al., 20 Nov 2025).

2. Token-Weighted Averaging in Generalized Attention

Within the framework of generalized attention, let query, key, and value matrices $Q, K, V \in \mathbb{R}^{n \times d}$. Define logits $e_{ij} = q_i \cdot k_j^\top$ and use a transformation $\phi$ to compute normalized attention weights:

$$\alpha_{ij} = \frac{\phi(e_{ij})}{\sum_{\ell=1}^n \phi(e_{i\ell})}.$$

The output at token $i$ is a token-weighted average:

$q_i$0

Softmax attention is recovered for $q_i$1; in linearized attention, $q_i$2 decomposes multiplicatively via positive feature maps.

3. The Dispersion Property and Its Implications

An established limitation of normalized attentions over large token sets is the dispersion phenomenon: as the number of keys increases, attention weights tend uniformly towards $q_i$3. Lemma A.1 (Tran et al., 10 Jun 2025) guarantee that, under mild continuity and boundedness conditions:

$q_i$4

for all tokens $q_i$5 and positive constants $q_i$6 independent of $q_i$7. Theorem A.2 extends this to deep stacks with compactness assumptions, showing that all attention weights vanish as $q_i$8 for large $q_i$9. This undermines the ability to focus or differentiate token contributions in deep or wide attention layers.

4. Remediation: Arithmetic Averaging and Homogeneous Mixing

To address dispersion, SEMA incorporates an explicit arithmetic averaging term ("homogeneous mixing") at each Transformer block, supplementing local attention with a global uniform average. Given $d_j$0, with $d_j$1 a learned projection, the block output at token $d_j$2 is:

$d_j$3

where

$d_j$4

and $d_j$5 is a learned scalar or vector. The LocalAttn operation considers only a neighborhood $d_j$6 for efficiency and locality. In matrix form:

$d_j$7

This explicit averaging ensures each block's output retains a representation of the whole token set, mitigating focus loss in long sequences (Tran et al., 10 Jun 2025).

5. Comparative Perspective: Softmax, Linear, and Token-Weighted Aggregations

Standard attention mechanisms aggregate values as weighted averages, but differ in their definition of weights:

Mechanism	Weighting Formula	Aggregation Scope
Softmax Attention	$d_j$8	Global
Linearized Attention	$d_j$9	Global
SEMA Arithmetic Averaging	Localized $$D(q, d) = \sum_{i=1}^m w_i \min_{j=1,\ldots,n} d(q_i, d_j),$$0, plus uniform $$D(q, d) = \sum_{i=1}^m w_i \min_{j=1,\ldots,n} d(q_i, d_j),$$1 for all	Local + Global
Token-Weighted Chamfer	$$D(q, d) = \sum_{i=1}^m w_i \min_{j=1,\ldots,n} d(q_i, d_j),$$2 determined by IDF/learning	Per-token (retrieval)

SEMA splits aggregation into a locally normalized, token-adaptive attention and a global, uniform mixing term, differentiating it from pure softmax or linear schemes.

6. Empirical Impact and Observed Benefits

Token-weighted averaging demonstrates measurable benefits across tasks. In multi-vector retrieval, weighting tokens according to importance enhances retrieval effectiveness with minimal overhead; in benchmarks such as BEIR, this yields consistent gains in Recall@10 (S et al., 20 Nov 2025). In vision transformers, arithmetic averaging in SEMA improves classification accuracy: ImageNet-1K top-1 accuracy rises from 83.5% (without averaging) to 83.7% (with averaging), indicating that homogeneous mixing is an effective, lightweight correction for attention dispersion (Tran et al., 10 Jun 2025).

7. Broader Implications and Research Trajectory

Token-weighted averaging formalizes a flexible template for token aggregation, uniting retrieval architectures with attention mechanisms via explicit weighting schemes. Its demonstrated efficacy in both text and vision domains suggests broad applicability, particularly as token counts scale and models demand mechanisms to circumvent uniformity induced by deep, broad attention or retrieval layers. A plausible implication is the emergence of further hybrid models leveraging token-adaptive and uniform global mixing for greater expressiveness and robustness in token sequence modeling.