Statistically Calibrated Scaling for Token Merging in Transformers

Token merging accelerates Transformer inference by clustering similar tokens to reduce sequence length (retention ratio $r$), but distorts attention outputs, inducing covariate shift in residual streams and performance collapse under high compression. Existing heuristics, such as proportional attention, mitigate mild compression effectively but degrade sharply at aggressive ratios due to unaddressed energy drift and biased attention distributions. We reframe token merging as a statistical reconstruction problem in high dimensions and introduce an asymptotic radial-angular decomposition of the reconstruction error, an analytical framework decoupling magnitude and distributional distortions. Minimizing this decomposed risk under minimal assumptions of finite second moments and variance stationarity yields closed-form optimal corrections governed by a single scaling factor $\sqrt{r}$: scaling merged values and shrinking merged logits toward the cluster-size prior. This calibrates both energy balance and distributional fidelity. Extensive experiments on vision Transformers demonstrate superior accuracy and robustness across compression levels.

Modern AI models often process images or other inputs by breaking them into many small pieces, called tokens. To make these models faster, a common approach is to merge similar tokens so the model has fewer pieces to process. However, when too many tokens are merged, the model can become less reliable because the merged information no longer matches what the original model was trained to expect. This paper studies why this happens and proposes a simple fix. We show that token merging can distort both the strength of the information and how the model decides which tokens to focus on. Our method applies a small automatic adjustment based only on the amount of token reduction. It does not require retraining or extra tuning. Experiments on image recognition and vision-language tasks show that this adjustment helps models stay accurate and stable even when many tokens are merged.