Learnable Token Merging in Transformer Models

• Learnable token merging is a method that adaptively compresses token sequences in transformers by combining input tokens using learned soft assignment rather than fixed pooling.
• It employs parameterized and differentiable techniques, such as those seen in PatchMerger and DTEM, to maintain key information while significantly reducing computational cost.
• The approach extends across various domains—including vision, video, and genomic models—while raising open challenges in adaptive compute and integration with downstream tasks.

Searching arXiv for the main paper and closely related token-merging works to ground the article. Learnable token merging is a family of sequence-compression mechanisms for transformers in which multiple input tokens are combined into a smaller token set through parameterized or differentiable assignment and aggregation rules, rather than by fixed pooling, hard pruning, or purely algorithmic matching. In vision, the topic emerged from the observation that flat-token Vision Transformers keep nearly constant sequence length across depth, so later blocks repeatedly pay the cost of long sequences; in broader settings, learnable token merging has also been used for graph-structured brain connectomes, code models, video understanding, diffusion tokenizers, and genomic sequence models (Renggli et al., 2022). Across these domains, the central technical question is how to reduce token count while preserving task-relevant information: by learned soft routing into latent slots, learned prompt-conditioned clustering, differentiable relaxations of hard matching, saliency-guided retention, or hybrid schemes in which only part of the pipeline is learned (Lee et al., 2024).

1. Origins in Vision Transformers

The canonical starting point is PatchMerger, introduced for Vision Transformers as a learnable bottleneck inserted between two consecutive encoder blocks (Renggli et al., 2022). The motivating inefficiency is that once an image is split into patch tokens, standard ViTs preserve essentially the same token count through all layers, even though self-attention scales quadratically with the number of tokens $N$ and the MLP and projection layers scale at least linearly in $N$. PatchMerger addresses this by placing a single merge module in the middle of the network, so the first half of the backbone operates on the original token sequence and the second half operates on a much shorter sequence.

The PatchMerger mechanism is a learned soft aggregation. If the incoming normalized tokens are $\mathbf{X}\in\mathbb{R}^{N\times D}$, the module learns a matrix $W\in\mathbb{R}^{D\times M}$ and produces

$$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$

Elementwise, each input token $x_n$ is scored against each learned query vector $w_m$, the softmax is applied over output slots for each input token, and each merged token is a weighted sum of the inputs. This makes PatchMerger neither hard pruning nor clustering with explicit nearest-centroid assignment. The learned parameters are only the columns of $W$, and the module is explicitly described as similar to non-scaled dot-product attention with fixed learned queries and input tokens serving as both keys and values (Renggli et al., 2022).

A defining property of this first generation is that the number of output tokens is fixed rather than adaptive per example. In the main PatchMerger experiments, models with 49, 196, or 256 input patch tokens are all reduced to 8 output tokens. This fixed bottleneck was chosen partly because it remains compatible with resolution changes at fine-tuning time: the first half of the network may see more tokens at higher resolution, but the second half still processes only the fixed output count (Renggli et al., 2022). The paper explicitly identifies adaptive per-input compute as future work.

Empirically, the strongest case in PatchMerger is at large model scale. Representative results include ViT-H/14 versus Merger ViT-H/14, where ImageNet fine-tuned accuracy is 87.97 versus 87.90 while FLOPs drop from 4275.92 ExaFLOPs to 2207.33 ExaFLOPs; the paper highlights this as 51.6% of the FLOPs and 63.5% of the runtime (Renggli et al., 2022). The same large-scale pattern appears for ViT-L/16 and for V-MoE-H/14, whereas smaller models such as Merger ViT-S/32 degrade more noticeably. This established a recurrent theme in later work: learnable merging is often most convincing when later transformer blocks dominate cost and the backbone is large enough to tolerate aggressive compression.

2. Core Mechanisms of Learnable Merging

Most learnable token merging methods can be organized by what is learned: merge targets, assignment weights, merge-specific embeddings, saliency scores, or latent grouping structure. The simplest mechanism is learned soft assignment to a fixed number of output slots. PatchMerger exemplifies this pattern, with learned queries $W$ defining output slots and a column-wise softmax routing each input token across them (Renggli et al., 2022). A related single-stage clustering formulation appears in TC-BrainTF, where $K$ learnable prompt tokens $N$0 act as community prototypes, each ROI token is softly assigned to prompts through

$N$1

in the paper’s notation, and the merged representation is formed by

$N$2

as a weighted aggregation of the original token embeddings (Yang et al., 2024). In both cases, merging is differentiable and soft, but the semantic interpretation differs: PatchMerger presents a mid-network bottleneck for efficiency, whereas TC-BrainTF interprets merged tokens as latent communities in a connectome.

A second design axis is whether the representation used for deciding merges is itself learned specifically for merging. DTEM makes this distinction explicit by introducing a decoupled embedding module $N$3 so that merge similarity is computed from $N$4 rather than directly from the backbone’s intermediate features (Lee et al., 2024). Similarity is

$N$5

and the paper argues that backbone features are optimized for contextual encoding and downstream tasks, not specifically for redundancy detection. DTEM then trains the merge embedding through a continuous relaxation of top-$N$6 bipartite matching, producing a soft adjacency matrix and a soft weighted merge operator. This separates “features for the task” from “features for deciding what to merge,” and the empirical ablations show that the decoupled embedding, more than differentiability alone, drives the gains over ToMe (Lee et al., 2024).

A third mechanism is learned saliency-guided retention combined with heuristic assignment. In long-form video, the learnable VTM method predicts saliency scores from key features as

$N$7

uses these scores to sample target tokens, and then still performs similarity-based assignment and average-pooling merge afterward (Lee et al., 2024). This partially learnable structure separates two decisions: which tokens should remain as merge targets, and which source tokens should be assigned to those targets. The merge operator itself remains non-learned, but the retention policy is learned.

A fourth mechanism appears in domain-specific settings where token groups are predefined or locally constrained. In code models, learnable token merging can mean attention-style pooling over subtokens that are already known to belong to the same semantic unit via tokenizer alignment. If $N$8 denotes the subtoken positions of one code token, the learnable merge is

$N$9

with only a single learned vector $\mathbf{X}\in\mathbb{R}^{N\times D}$0 added to the model (Saad et al., 19 Jul 2025). In MergeDNA, by contrast, the grouping itself is learned through stacked local token-merging blocks over DNA bases, with repeated local merges producing variable-length tokens and a source matrix preserving span provenance (Li et al., 17 Nov 2025). These examples broaden the concept beyond vision: learnable token merging need not always mean discovering groups globally from scratch; it can also mean learning how to pool within externally supplied or locally constrained groups.

3. Training Regimes and Differentiable Relaxations

A central technical difficulty is that useful merge policies are often discrete, while end-to-end learning requires gradients. Different methods resolve this tension differently. PatchMerger avoids discrete grouping altogether: the routing weights are continuous softmax probabilities, so the module is trained jointly with the whole network from scratch during upstream pretraining, with no special auxiliary loss, no distillation loss, and no custom regularization specific to the merger (Renggli et al., 2022). Its simplicity is one reason it remains a reference point.

DTEM is designed precisely around the discreteness problem. It takes ToMe-style Bipartite Soft Matching as the target hard operator, then constructs a continuous relaxation by repeatedly applying a global softmax over bipartite edges,

$\mathbf{X}\in\mathbb{R}^{N\times D}$1

updating similarities to suppress reuse of the same source node, and accumulating these soft selections into a final soft adjacency $\mathbf{X}\in\mathbb{R}^{N\times D}$2 (Lee et al., 2024). The merge operator is also relaxed: tokens in the destination set absorb weighted contributions from softly assigned source tokens, while effective token sizes are updated consistently. At inference, the method reverts to hard ToMe-style matching, but with similarities computed from the learned decoupled embeddings rather than directly from backbone keys. This is one of the clearest examples of learnable token merging in the strong sense: the merge representation is specialized, the grouping is differentiably approximated, and the training signal flows directly through the merge path.

LTM-Transformer takes a different route. Rather than relaxing a known hard matching algorithm, it formulates merging through a learned mask $\mathbf{X}\in\mathbb{R}^{N\times D}$3 that maps $\mathbf{X}\in\mathbb{R}^{N\times D}$4 tokens to $\mathbf{X}\in\mathbb{R}^{N\times D}$5 merged tokens by

$\mathbf{X}\in\mathbb{R}^{N\times D}$6

The paper motivates the mask by an Information Bottleneck objective

$\mathbf{X}\in\mathbb{R}^{N\times D}$7

derives a separable variational upper bound for the IB loss, and designs the mask module to mimic a gradient-descent step on that upper bound across layers (Wang et al., 2024). Here the learning problem is not expressed as discrete top-$\mathbf{X}\in\mathbb{R}^{N\times D}$8 grouping, but as soft matrix-valued aggregation optimized through the backbone’s standard task loss after a warm-up period without merging.

Some later methods are only partially learnable. QuickMerge++ uses differentiable Gumbel-Softmax selection,

$\mathbf{X}\in\mathbb{R}^{N\times D}$9

followed by weighted averaging within heuristic clusters, but its salience is mainly derived from attention entropy and its clustering rule is described procedurally rather than as a learned parametric module (Liu et al., 16 Aug 2025). The paper’s own analysis makes clear that the merge policy is only partially learned: the autoregressive prior is clearly learned, whereas salience estimation, clustering, and budgeting are mostly heuristic.

The same caution applies to methods such as MergeTok and MergeVQ. MergeTok uses ToMe-derived grouping internally as a structural prior rather than learning merge assignments directly; the encoder features determining similarity are learned end-to-end, but the source map is discrete and inherited from ToMe-style matching (Zhang et al., 29 May 2026). MergeVQ likewise uses ToMe-style merging after self-attention blocks in the encoder to separate top-level semantics from local detail before quantization, but the merge rule itself remains heuristic and the novelty lies in how the merge ancestry is exploited by source recovery and generation (Li et al., 1 Apr 2025). These cases suggest a useful taxonomy: some systems learn the merge operator itself, while others learn around a fixed merger and turn its output into a supervisory or architectural signal.

4. Adaptivity, Hierarchy, and Variable-Length Tokenization

A major line of development concerns adaptivity: whether token count or merge structure should vary across inputs, depths, or operating points. PatchMerger is explicitly non-adaptive in this sense: it always emits a fixed number of output tokens, and the paper identifies adaptive per-input compute as future work (Renggli et al., 2022). By contrast, video-oriented and autoregressive methods often make adaptivity central.

In long-form video, learnable VTM argues that token similarity alone is insufficient and that saliency should govern what remains unmerged. The method processes spatiotemporal token grids $W\in\mathbb{R}^{D\times M}$0, learns saliency scores, samples targets accordingly, and then merges sources into them (Lee et al., 2024). The paper reports that this learned saliency-based retention yields better or comparable performances on LVU, COIN, and Breakfast, while reducing memory costs by 84% and boosting throughput by approximately 6.89 times compared to baseline algorithms. This suggests that adaptivity is especially valuable when redundancy is heterogeneous across both space and time.

QuickMerge++ extends the adaptive theme to autoregressive generation. It defines a compressed sequence $W\in\mathbb{R}^{D\times M}$1 with $W\in\mathbb{R}^{D\times M}$2, uses attention-entropy-derived salience, Gumbel-Softmax selection, and cluster-wise weighted averaging,

$W\in\mathbb{R}^{D\times M}$3

then trains a lightweight autoregressive prior over the merged-token sequence (Liu et al., 16 Aug 2025). The paper emphasizes dynamic budgets and reports mean token count dropping from 128 to 54, latency from $W\in\mathbb{R}^{D\times M}$4 to $W\in\mathbb{R}^{D\times M}$5 ms, and KV memory from $W\in\mathbb{R}^{D\times M}$6 to $W\in\mathbb{R}^{D\times M}$7 MB on A100, alongside improved PPL, accuracy, and FVD relative to listed baselines. At the same time, the paper contains explicit inconsistencies about whether salience is based primarily on attention norm magnitude or entropy, and it does not fully formalize the budget estimator. It is therefore best treated as a hybrid dynamic token merging framework rather than a fully specified end-to-end learnable merger.

Variable-length tokenization for diffusion introduces a different kind of adaptivity: changing token length while preserving cross-length representational alignment. The 2026 learnable global merging paper argues that truncation-based variable-length tokenizers induce a cross-length shift in latent geometry, and proposes reducing token count by merging instead: $W\in\mathbb{R}^{D\times M}$8 Here the merge pattern $W\in\mathbb{R}^{D\times M}$9 is not data-dependent but learned globally from position embeddings $$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$0, so the same hierarchy of agglomerative merges is available at generation time (Lee et al., 18 Jun 2026). The paper derives the distortion identity

$$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$1

showing that cross-length deviation is minimized when merged clusters contain similar tokens. It then combines this with merged positional embeddings, proportional attention, and optional length-specific LoRA adapters. On ImageNet $$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$2, the method improves average cross-length CKNNA from 0.34 for nested dropout to 0.76, and markedly narrows the joint-versus-independent DiT gap (Lee et al., 18 Jun 2026). This is a notable shift in emphasis: token merging becomes not just an efficiency device, but a way to make variable-length latent representations geometrically compatible across operating points.

Hierarchical merging also appears outside vision. MergeDNA treats tokenization itself as a learnable in-network process, stacking local-window token-merging blocks over DNA bases and propagating a source matrix $$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$3 so that merged tokens can later be unmerged by

$$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$4

Its local encoder uses top-$$\mathbf{Y}=\mathrm{softmax}\!\left((\mathbf{X}W)^T\right)\mathbf{X} =\mathrm{softmax}\!\left(W^T\mathbf{X}^T\right)\mathbf{X}, \qquad \mathbf{Y}\in\mathbb{R}^{M\times D}.$$5 pair selection within windows, repeated across layers, so segmentation is learned implicitly through hierarchical local merges rather than predefined vocabulary units (Li et al., 17 Nov 2025). This suggests that “learnable token merging” can denote not only adaptive reduction inside a model, but also adaptive token formation itself.

5. Relation to Training-Free Merging and Common Misconceptions

A persistent misconception is that token merging is synonymous with learnable token merging. The literature makes the opposite point. ToMe is explicitly nonparametric: it partitions tokens into two sets, computes cosine similarity on normalized keys, selects top cross-partition matches, and merges by weighted averaging without any learnable merge parameters (Bolya et al., 2022). ATC likewise presents a parameter-free alternative based on agglomerative clustering of self-attention keys with cosine distance, no learned routing weights, no learned queries, and no cluster-center parameters (Haurum et al., 2024). Both methods are strong baselines precisely because they show that much of token-merging performance can come from the grouping structure itself rather than from learning an explicit merge module.

This matters because several later works argue against conflating “content-dependent” with “learnable.” ToMe is content-dependent because it uses learned backbone features, especially attention keys, but the merge rule is training-free (Bolya et al., 2022). ATC is also content-dependent and in some settings outperforms prior merging and pruning methods, including at low keep rates, while remaining parameter-free (Haurum et al., 2024). HoliTom shows the same phenomenon in video LLMs: it combines global redundancy-aware temporal segmentation, attention-guided spatial reduction, density-based clustering, and inner-LLM similarity merging, all without any trained merge module, and reports reducing computational costs to 6.9% of FLOPs while maintaining 99.1% of the original performance on LLaVA-OneVision-7B (Shao et al., 27 May 2025). These systems are not learnable mergers in the strict sense, yet they are sophisticated adaptive merging frameworks.

Another misconception is that merging is just pruning with averaging. Multiple papers stress the distinction. PatchMerger is “not token pruning, clustering, or hard assignment” but a learned soft aggregation mechanism (Renggli et al., 2022). ToMe frames merging as preferable to pruning because it preserves information in aggregated form rather than discarding it (Bolya et al., 2022). ALGM argues that dense prediction is particularly hostile to pruning because every token corresponds to a region that must eventually be labeled, whereas merging allows information to remain active in later layers (Norouzi et al., 2024). HoliTom similarly positions inner-LLM merging as safer than hard pruning after outer compression has already concentrated information (Shao et al., 27 May 2025). This suggests a general principle: pruning optimizes sparsity by deletion, whereas merging optimizes compression by controlled aggregation.

A third misconception is that learnability automatically implies better deployment properties. Many learnable methods require either full pretraining integration, differentiable relaxations with training-time overhead, or downstream fine-tuning. PatchMerger is trained from scratch during upstream pretraining (Renggli et al., 2022); LTM-Transformer is trained from scratch and adds training complexity due to cluster updates and variational distribution estimation (Wang et al., 2024); DTEM supports modular training on frozen backbones but still needs optimization of the decoupled embedding modules (Lee et al., 2024). By contrast, ToMe, ATC, ALGM in direct mode, and the adaptive thresholding method for semantic communication can all be inserted without retraining, albeit with different trade-offs (Bolya et al., 2022). This does not make training-free methods universally preferable, but it does mean that learnable token merging should be assessed against deployment constraints, not only against idealized accuracy–FLOPs frontiers.

6. Empirical Patterns, Open Problems, and Research Directions

Across the literature, several empirical regularities recur. First, larger backbones tolerate merging better. PatchMerger is strongest at ViT-L, ViT-H, and corresponding V-MoE scales, whereas smaller models suffer more under aggressive compression (Renggli et al., 2022). DTEM also shows larger gains over ToMe at higher compression on DeiT-B than on DeiT-S (Lee et al., 2024). This suggests that redundancy exploitable by learned merging increases with scale, or at least that larger models have more capacity to reorganize information after compression.

Second, where merging occurs matters as much as how it is parameterized. PatchMerger’s ablation found that placing the merger too early hurts because useful representations have not yet formed, while the middle of the network gives comparable performance to the original backbone (Renggli et al., 2022). ALGM’s segmentation results similarly show that local merging is safe in the first layer but global merging becomes reliable only later, motivating its local-then-global schedule (Norouzi et al., 2024). The video study “How Do Your Tokens Merge?” reaches an analogous conclusion for training-free video merging: late-layer merging is far safer than early-layer merging, especially when temporal semantics are delicate (Pollard et al., 4 Jun 2025). A plausible implication is that future learnable systems should make depth-dependent merging a first-class design variable rather than a fixed schedule.

Third, learnability can target different subproblems. Some methods learn assignment weights directly, as in PatchMerger and TC-BrainTF (Renggli et al., 2022, Yang et al., 2024). Some learn merge-specific embeddings, as in DTEM (Lee et al., 2024). Some learn saliency or retention policies while leaving assignment heuristic, as in VTM (Lee et al., 2024). Some learn around a fixed merger, exploiting merge structure for alignment, quantization regularization, or source recovery, as in MergeTok and MergeVQ (Zhang et al., 29 May 2026, Li et al., 1 Apr 2025). This diversity suggests that “learnable token merging” is not a single technique but a design space spanning routing, grouping, pooling, compression scheduling, and representation alignment.

Fourth, interoperability with downstream architectures remains a major issue. QuickMerge++ is motivated by autoregressive compatibility and therefore adds a learned prior over merged-token sequences (Liu et al., 16 Aug 2025). The diffusion VLT work requires globally known merge patterns so that merged positional embeddings and proportional attention are available during generation (Lee et al., 18 Jun 2026). MergeVQ and MergeTok turn source maps into structural signals for reconstruction or codebook regularization rather than merely using merging to save compute (Li et al., 1 Apr 2025, Zhang et al., 29 May 2026). These examples indicate that once merging is placed inside generative or reconstruction-sensitive systems, positional semantics, latent alignment, and reversibility become as important as FLOPs.

Several open problems are explicit in the cited work. PatchMerger calls out adaptive per-input compute as future work (Renggli et al., 2022). ATC suggests that a stronger learnable merger might preserve its local bottom-up inductive bias while learning the similarity metric, linkage rule, stopping criterion, or aggregation weights (Haurum et al., 2024). The video saliency-merging paper notes that its actual merge decisions remain discrete and non-differentiable even though saliency is learned indirectly (Lee et al., 2024). MergeDNA leaves the exact similarity-score parameterization, merge formula, and differentiable relaxation under-specified in the main text (Li et al., 17 Nov 2025). The variable-length diffusion tokenizer shows that globally shared merge patterns improve generative compatibility, but at the cost of per-sample optimality (Lee et al., 18 Jun 2026). These tensions point toward a broader research agenda: combining strong structural priors from training-free merging with learned, task-aware, and length-aware adaptation.

In that sense, learnable token merging has evolved from a mid-network bottleneck for ViTs into a general framework for adaptive coarsening of token sequences. Its strongest formulations are differentiable and merge-specific, as in DTEM or LTM-Transformer (Lee et al., 2024, Wang et al., 2024); its most widely deployable baselines are still training-free, as in ToMe and ATC (Bolya et al., 2022, Haurum et al., 2024); and its newest applications increasingly treat merging not only as acceleration, but as a representational primitive for variable-length generation, tokenizer design, and structured latent organization (Lee et al., 18 Jun 2026).