Dynamic Key Merging in Vision-Language Models

• Dynamic Key Merging is a token reduction paradigm that adaptively compresses visual tokens based on image complexity, preserving task performance.
• It employs Dynamic Token Merging (DToMe) to selectively compute token similarities and perform weighted merging according to per-layer budgets.
• Virtual Token Unmerging (VTU) reconstructs full-length sequences in LLMs, enabling efficient cross-modal attention with minimal performance loss.

Dynamic Key Merging is a token reduction paradigm for vision-LLMs (VLMs) that targets the inefficiencies associated with fixed-length visual token outputs in transformer architectures. Central to this approach is the DyMU framework, which consists of Dynamic Token Merging (DToMe) for content-adaptive token compression and Virtual Token Unmerging (VTU) for attention-consistent expansion within LLMs. DyMU operates entirely training-free, introducing dynamic, image-dependent token budgets and efficient attention mechanisms that preserve downstream task performance across diverse VLM and LLM backbones (Wang et al., 23 Apr 2025).

1. Dynamic Token Merging: Principles and Workflow

Dynamic Token Merging (DToMe) addresses the inefficiency of uniform token budgets per image by selectively merging redundant visual token embeddings in a vision transformer’s output sequence. Given a token sequence $x_i\in\mathbb{R}^{N_i\times D}$ at transformer layer $i$, DToMe seeks to minimize $N_i$ dynamically in response to image complexity.

Key Steps

• Token-Pair Similarity via Bipartite Matching: Tokens are split into two interleaved sets $\mathbb{A}$ and $\mathbb{B}$ (e.g., by even/odd indices). For each $t\in\mathbb{A}$, a similarity score is computed via dot product between key projections $k_i$:

$$t_B = \arg\max_{n\in\mathbb{B}} \langle k_i[t], k_i[n] \rangle,\quad S_i[t] = \langle k_i[t], k_i[t_B] \rangle$$

• Thresholding and Dynamic Merge Budget: A per-layer target merge count $r_i$ is selected. Offline, the algorithm determines a threshold $\tau_i$ such that the number of retained edges $S_i^{(b)}[t] > \tau_i$ matches $B r_i$ over a batch of $B$ images. At inference, $\tau_i$ is fixed; images with higher redundancy yield more merges.
• Merging Procedure: For each selected edge ($S_i[t]>\tau_i$), tokens are merged via a position-weighted average:

$$x_i[t_B] \leftarrow \frac{|P_i[t]|\,x_i[t]+|P_i[t_B]|\,x_i[t_B]}{|P_i[t]|+|P_i[t_B]|}$$

Position sets $\mathcal{P}_i[t]$ and $\mathcal{P}_i[t_B]$ are updated, with the merged token $t$ dropped from the active sequence.

• Size-Weighted Self-Attention: To retain semantic coverage post-merge, scaled-dot attention incorporates a logarithmic token-size bias:

$$A = \mathrm{Softmax}\left(\frac{QK^T}{\sqrt{d}} + \log[|P_i[1]|,\ldots,|P_i[N_i]|]\right)$$

• Image Complexity and Token Survival: JPEG-compression-based complexity $C(I)=\mathrm{size}_{\mathrm{JPEG}(I)}/(H\times W)$ correlates strongly with post-merge token count.
• Layer-By-Layer Adaptivity: Merge budgets $\{r_i\}$ may be fixed or decay linearly, enabling end-to-end compression schedule optimization.

Inference Pseudocode

Input: image I of size HxW x_1 <- ViT_patch_projection(I) # N1 x D P_1[t] <- {t} for t=1..N1 for i in 1..L: [x_i, k_i, Q_i, K_i, V_i] <- TransformerBlock(x_i, P_i) split tokens into A,B for each t in A: find t_B = argmax_{n in B} <k_i[t],k_i[n]> S[t] = <k_i[t],k_i[t_B]> find all edges where S[t] > tau[i] for each edge t->t_B: merge t into t_B via weighted average update P_i[t_B] = P_i[t_B] ∪ P_i[t] drop token t x_{i+1} <- x_i[ tokens with nonempty P_i ] P_{i+1} <- corresponding P_i endfor return x_{L+1}, P_{L+1} # final variable-length token set

2. Virtual Token Unmerging in LLMs

Virtual Token Unmerging (VTU) simulates the attention dynamics of a full-length sequence in LLMs when fed compressed visual tokens. This approach enables downstream models—trained on fixed-length sequences—to preserve alignment and performance without explicit architectural retraining.

Mechanism

• Mapping and Sparse Expansion: Given a compressed set $e_{\mathrm{un}}\in\mathbb{R}^{N_{un}\times D}$, VTU constructs a sparse expansion matrix $M\in\{0,1\}^{N\times N_{un}}$, assigning each position in $N$ to its representative token:

$e = M\,e_{\mathrm{un}}$

• Sequence-Independent Layers: For elementwise operations $f$ (e.g., GeLU, LayerNorm), evaluation is reduced to:

$f(e) = M\,f(e_{\mathrm{un}})$

• RoPE-Based Self-Attention: Rotary Position Embedding (RoPE) attention, sensitive to absolute positions, is rederived in the compressed domain via four sparse-matrix terms:

$$A = C\,M\,Q_{\mathrm{un}}\,K_{\mathrm{un}}^T\,M^T\,C + S\,M\,Q_{\mathrm{un}}\,K_{\mathrm{un}}^T\,M^T\,S + \ldots$$

This reduces complexity from $O(N^2D)$ to $O(N\,N_{un}\,D)$.

• Attention Output: Softmax and value lookup are realized with:

$$f(e) = \mathrm{Softmax}\Bigl(\frac{A}{\sqrt{D}}\Bigr)\,V, \quad V = M\,V_{\mathrm{un}}$$

• Re-Merging in Subsequent Layers: To restore $N_{un}$-row representations for subsequent blocks:

$$e_{\mathrm{un}}^{(\text{next})} = (M^TM)^{-1} M^T f(e)$$

This averages rows associated with the same group. Empirically, this incurs minimal degradation (<3%) in end-to-end performance.

3. Computational Efficiency and Performance Profile

DyMU introduces substantial computational savings with negligible loss of accuracy. Reported statistics include:

Configuration	Token Count (LLaVA-1.5)	FLOPs (Self-Attn)	Performance (VQA avg.)
Full	576	1.36 GFLOPs	55.8
DyMU-low	$89 \pm 27$	72 MFLOPs	54.5 (97.7%)
DyMU-mid	$195 \pm 47$	358 MFLOPs	55.3 (99.1%)
DyMU-high	$394 \pm 57$	1.27 GFLOPs	56.0 (100.4%)

For DyMU-low, token reduction reaches $84.5\%$ versus full model, with FLOPs savings up to $19\times$, while performance remains within $97.7\%$ of baseline. Compared to ToMe—a previous fixed-length, training-free technique—DyMU-low obtains higher average scores (54.5 vs. 52.6) and adaptively varies token count per image.

A plausible implication is that token-budgeted models using DyMU can flexibly allocate compute without retraining or manual thresholding, which is substantiated by its compatibility with diverse visual encoders (CLIP, SigLIP) and LLMs (Vicuna-7B, Qwen-2.5).

4. Adaptive Compression and Control Mechanisms

DyMU’s dynamic scheme confers granular control over computational expenses via per-layer and per-image merge budgets. The offline threshold calibration ensures average budgets are met across batches, while individual images—assessed by JPEG-based complexity—yield contextually appropriate token reductions. This adaptation enables finer resource allocation, which may be critical for real-time or resource-constrained inference scenarios.

Layer-by-layer budget strategies (constant or decaying) allow further fine-tuning of the compression curve, facilitating trade-offs between throughput and task-specific accuracy.

5. Integration and Applicability Across Model Architectures

DyMU is expressly designed as a training-free, plug-and-play solution. Its token merging and virtual expansion modules are compatible with mainstream transformer vision encoders (including AnyRes variants), and downstream LLMs with standard attention stacks and RoPE position embeddings. The absence of retraining requirements streamlines deployment onto existing model checkpoints, contributing to broad utility in both research prototypes and production VLM systems.

Empirical evaluation on image and video understanding benchmarks confirms near-parity in downstream performance (exceeding $96\%$ of baseline accuracy with $\sim$15\%$ token budgets), signaling robust generalization across tasks. The framework’s modularity suggests straightforward extension to other multimodal architectures using fixed-length cross-modal sequences.

6. Context, Comparison, and Implications

DyMU’s innovation lies in dynamic, content-adaptive merging and efficient re-expansion, contrasting with traditional static token selection paradigms. Compared to fixed-length methods (e.g., ToMe), DyMU achieves higher accuracy while providing per-image flexibility. This suggests a paradigm shift in VLM design, emphasizing input-dependent computational scaling.

Experimental results on VQA-style benchmarks and ablation studies demonstrate that DyMU’s approximation in VTU incurs only minor accuracy penalties, even with large reductions in token count and floating-point operations. A plausible implication is that future VLM deployments may increasingly favor dynamic key merging architectures for efficiency-critical applications.

No significant controversies are noted; however, ongoing work may investigate the limits of compression for highly complex images and the interaction of variable-length representations with specialized LLM attention mechanisms. The methodology described in DyMU is fully documented in (Wang et al., 23 Apr 2025) and remains a reference point for dynamic token reduction research in the vision-language domain.