--- title: 'Co-Me: Confidence-Guided Token Merging' url: https://www.emergentmind.com/topics/confidence-guided-token-merging-co-me type: topic --- # Co-Me: Confidence-Guided Token Merging Confidence-Guided Token Merging (Co-Me) is an acceleration mechanism for visual geometric transformers (ViTs) that reduces inference latency by selectively merging low-confidence tokens without requiring any retraining or finetuning of the base model. Co-Me employs a distilled confidence predictor to identify which tokens are least critical to downstream geometric prediction, merging these before each attention or MLP block and restoring original shapes for dense prediction. This strategy enables substantial reductions in both computational complexity and runtime, making high-resolution, long-sequence 3D perception practical in real-time systems while preserving accuracy in depth, pose, and reconstruction tasks [2511.14751]. ## 1. Motivation in Visual Geometric Transformers Modern single-pass 3D reconstruction frameworks such as VGGT [Wang et al. 2024] and MapAnything [Keetha et al. 2025] are based on Vision Transformer (ViT) architectures. These ingest sequences of image tokens (patch embeddings) originating from multi-view or video-frame stacks, and produce dense outputs including depth maps, point clouds, and camera pose estimates. The self-attention mechanism in ViTs incurs quadratic computational complexity, $\mathcal{O}(N^2d)$, where $N$ is the number of tokens and $d$ the feature dimension. As $N$ increases due to higher spatial resolutions or longer frame sequences, inference latency rapidly becomes impractical for real-time or edge deployments. Efforts to reduce $N$, the sequence length processed by attention layers, are thus central for enabling fast 3D perception. Prior approaches—token pruning (e.g., DynamicViT, A-ViT) and similarity-based merging (e.g., ToMe, FastVGGT)—have limitations: pruning methods typically incur severe spatial information loss and require retraining, while similarity-based merging only modestly reduces compute and does not align with the regions of geometric interest highlighted by transformer attention. Notably, similarity heuristics fail to capture task-relevant coverage except at extreme sequence lengths. ## 2. Confidence Predictor Architecture and Training Co-Me introduces a lightweight “confidence predictor,” $f'$, distilled from the frozen ViT encoder. The overall model is denoted $\mathcal F = f_2 \circ f_1$, with the predictor inserted after layer 15 of the encoder stack. The architecture of $f'$ is as follows: - A single-layer MLP projects patch embeddings to a low-dimensional latent space. - A single-head self-attention module aggregates spatial context across the latent tokens. - A compact Conv2D “head” produces per-patch confidences $\mathcal C' \in \mathbb{R}^{H\times W}$. This module introduces approximately 0.2% overhead relative to full ViT runtime. Training the predictor avoids any retraining of $\mathcal F$. Instead, $f'$ is supervised to reproduce the *ranking* of confidences from the original model’s output. For each patch $i$, the average downstream confidence from $\mathcal C = f_2(f_1(x))$ is $\bar c_i$; for the predictor, $c_i' = \mathrm{Avg}(\mathcal C')_i$. The set of ordered pairs is $$\mathcal{P} = \{(i, j) \mid \bar c_i > \bar c_j\}.$$ The training objective is logistic ranking loss: $$ \mathcal{L}_{\text{conf}} = \frac{1}{|\mathcal{P}|} \sum_{(i,j)\in\mathcal{P}} \log\bigl(1+\exp(c_j' - c_i')\bigr). $$ This loss emphasizes correct ordering rather than magnitude, aligning predictor outputs with the subset of patches least relied upon by the transformer’s attention. Empirically, high-confidence patch predictions localize to image regions rich in geometric information (texture, stable cues), whereas low-confidence outputs correspond to background or occlusion areas. ## 3. Token Merging Scheme Tokens are grouped into $N/n$ contiguous spatial groups $G_1,\dots,G_{N/n}$ of size $n$ (typically $n=4$). For each group, group-average confidence is computed: $$ \tilde c_i = \frac{1}{n}\sum_{k\in G_i} c_k'. $$ Given a merge ratio $p \in (0,1)$, the $p$-fraction of groups with the lowest $\tilde c_i$ are marked for merging, i.e., $\tilde c_i < \tau$ where $\tau$ is the $p$-th percentile. Prior to each attention or MLP block in $f_2$, tokens within each group are merged as follows: - For groups marked to be merged ($m_i=1$): replace $G_i$ with its mean, $G_i \gets \{\frac{1}{n}\sum_{x\in G_i} x\}$. - For unmerged groups: preserve original tokens. Concatenate the resulting tokens across groups. After the block, the merged tokens are split back to the original sequence length: merged tokens are replicated $n$ times, while unmerged groups pass unchanged. This “merge and split” paradigm preserves the full spatial coverage, ensuring all downstream heads operate on the original token count. The process is entirely differentiable and does not affect downstream dense prediction outputs. ### Inference Pseudocode ```python # x: input patches; p: merge ratio; n: group size z = f1(x) c_pred = f'(z) group_conf = [mean(c_pred[i*n:(i+1)*n]) for i in range(N//n)] tau = percentile(group_conf, p*100) m = [1 if conf