🤖 machine learning

MergeOver: Post-Training Token Merging for Recursive Vision Transformers

MergeOver is a post-training method that integrates Token Merging into recursively weight-shared Vision Transformers to significantly reduce memory usage and latency on edge devices while maintaining high accuracy, all without requiring costly retraining.

Original authors: Junseo Kim, Uraz Odyurt, Amirreza Yousefzadeh

Published 2026-08-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Junseo Kim, Uraz Odyurt, Amirreza Yousefzadeh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot to recognize cats, dogs, and cars just by looking at pictures. To do this, the robot uses a special kind of brain called a "Vision Transformer." Think of this brain as a team of tiny detectives, where every single pixel in a photo gets its own detective. The more pixels there are, the more detectives you need, and the more they have to talk to each other to figure out what they are seeing. This is great for accuracy, but it's a huge problem for small computers, like the ones in your smartwatch or a Raspberry Pi, because they don't have enough memory or battery power to handle thousands of detectives chatting at once.

To fix this, scientists have tried two main tricks. The first is "recursive weight-sharing," which is like having one super-detective who does the same job over and over again, reusing their notes instead of hiring a new team for every step. This saves space but makes the detective work slower and heavier. The second trick is "token merging," where you realize that some detectives are looking at the same thing, so you just merge them into one big detective to save time. The big question this paper tackles is: What happens if you try to use both tricks at the same time without having to retrain the whole robot from scratch? It's like trying to teach a single detective to merge with their clones while they are already working, without messing up the map they are following.

The authors of this paper, Junseo Kim and his team from the University of Twente, say that trying to mix these two methods is usually a disaster because the rules of the game get broken. They propose a new method called MergeOver to solve this puzzle. Imagine the robot's brain is a multi-story building where the detectives move up floors. Every time they move to a new floor, the building's architecture requires them to stand in a perfect square grid. But merging detectives usually messes up that square grid, leaving empty spots that break the elevator. MergeOver introduces a clever "Unmerge" trick: it temporarily splits the merged detectives back into their original grid positions just long enough to get through the elevator, and then merges them again once they are safe on the next floor. This allows the robot to keep its memory usage low without needing to be retrained.

The team tested this on a famous picture dataset called ImageNet-1K, using different types of computers, from powerful graphics cards to a tiny Raspberry Pi 5. They found that MergeOver works, but the results depend heavily on how many pictures you ask the robot to look at at once (the "batch size"). When the robot looks at just one picture at a time, the extra work of merging and unmerging actually slows things down a bit. However, when the robot looks at 16 pictures at once, the magic happens. On a powerful GPU, this setup reduced the memory needed by 38.4% and actually made the robot 21.7% faster. On the tiny Raspberry Pi, it made the robot 17.6% faster for batches of 16. The best part is that the robot didn't get much dumber; its accuracy only dropped by 1.47 percentage points, which the authors say is a small price to pay for such a big speed boost.

However, the paper is careful to note that this isn't a magic wand that fixes everything instantly. The authors explicitly argue against the idea that you can just slap this on any model and expect it to work perfectly everywhere. They found that for single-picture tasks (batch size 1), the method actually slows things down on powerful computers because the overhead of the "Unmerge" trick isn't worth it for just one image. They also point out that while their method saves a lot of memory and speeds up batch processing, it doesn't necessarily beat all other types of efficient AI models that were designed from the ground up to be fast. The results are measured and real, but they suggest that this is a promising starting point—a "baseline"—for combining these techniques, rather than a final, perfect solution. The authors suggest that future work might need to combine MergeOver with other tricks, like better software optimization, to make it even faster, especially for those single-picture tasks on small devices.

=== SUMMARY ===

Technical Summary: MergeOver

Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance in computer vision but face significant barriers to deployment on resource-constrained edge hardware due to massive parameter counts and quadratic computational complexity ( $O(N^2)$ ) relative to sequence length. While recursive weight-sharing (e.g., in the Sliced Recursive Transformer, SReT) effectively reduces parameter counts by reusing weights across multiple passes, it exacerbates computational bottlenecks and Peak Activation Memory (PAM). Conversely, Token Merging (ToMe) offers a post-training mechanism to reduce sequence length and mitigate these bottlenecks.

However, integrating these two paradigms is non-trivial. Direct integration creates severe architectural friction:

Spatial Constraints: SReT relies on hierarchical convolutional pooling layers that require rigid 2D spatial grids. Standard token merging destroys the perfect square property of the sequence length, causing dimensional clashes.
Group Constraints: SReT's Sliced Group Self-Attention (SGA) requires sequence lengths to be evenly divisible by specific group sizes. Unconstrained merging violates these boundaries.
Tracking Constraints: SReT employs spatial permutations to shift attention windows. ToMe relies on a token-mass vector to maintain proportional attention; if this vector is not permuted synchronously with the features, the attention mechanism is corrupted.

Existing efficient architectures often require retraining, distillation, or structural redesigns, leaving a gap for post-training compression methods that combine recursive weight-sharing with dynamic token reduction.

Methodology: MergeOver

The authors propose MergeOver, a post-training approach that integrates ToMe into SReT without retraining. The methodology addresses the three incompatibilities through specific mechanisms:

1. Unmerge Tracking Stack (Resolving Spatial Constraints)

To satisfy the geometric requirements of inter-stage convolutional pooling, MergeOver implements an Unmerge stack.

Mechanism: During a stage, when tokens are merged, the corresponding "Unmerge" operation (which duplicates merged features back to their original spatial coordinates) is pushed onto a Last-In-First-Out (LIFO) stack.
Restoration: Before the sequence enters the inter-stage convolutional pooling layer, the stack is popped, and Unmerge operations are applied in reverse order. This restores the sequence to its original length and 2D grid structure, allowing the convolutional kernels to operate correctly while the Transformer blocks process a compressed sequence.

2. Constraint-Safe Merge-Rate Adjustment (Resolving Group Constraints)

The merge rate ( $r$ ) cannot be arbitrary. It must satisfy:

Divisibility: The resulting sequence length ( $N-r$ ) must be divisible by the Least Common Multiple (LCM) of the group sizes ( $g$ ) used in the current and subsequent recursive iterations.
Bipartite Limits: ToMe's Bipartite Soft Matching (BSM) cannot merge more than half the tokens in a single iteration ( $r \le \lfloor N/2 \rfloor$ ).
Implementation: The target reduction is adjusted to the nearest "safe" value ( $r_{safe}$ ) that satisfies $(N-r) \mod g = 0$ and $N-r \ge g$ .

3. Synchronized Token-Mass Tracking (Resolving Tracking Constraints)

To maintain proportional attention during SReT's spatial permutations:

A token-mass tensor ( $s$ ) is initialized as ones and updated across layers to track the number of original patches each token represents.
Synchronization: The mass tensor is routed in parallel through the exact same permutation and inverse permutation functions as the feature sequence. This ensures that the weights used in the attention calculation (Equation 2) remain aligned with their corresponding tokens.
Weighted Merging: Merging is performed via a weighted average using the mass tensor to prevent "heavy" tokens from being underrepresented.

4. Token Reduction Scheduling

The paper evaluates four scheduling strategies to determine the optimal reduction trajectory:

Global Schedules: Constant and Linearly decreasing (standard ToMe baselines).
Stage-wise Schedules:
- Exponential: Progressively smaller reductions within a stage.
- Single-shot: Token reduction occurs only at the first block of each stage, maintaining a fixed sequence length for subsequent recursive iterations within that stage.

Key Contributions

MergeOver Formulation: A post-training framework enabling token merging within hierarchical recursive transformers without retraining.
Architectural Integration: Solving the specific incompatibilities between SReT's spatial/attention constraints and ToMe's merging logic via Unmerge stacks and synchronized mass tracking.
Scheduling Analysis: An evaluation of token reduction strategies, identifying a stage-wise single-shot strategy as the most compatible with SReT's architecture.
Cross-Platform Benchmarking: A comprehensive evaluation on ImageNet-1K across GPU, x86 CPU, and ARM CPU (Raspberry Pi 5) platforms.

Experimental Results

Evaluated on the ImageNet-1K validation set using the SReT-Tiny-Distill baseline:

Accuracy vs. Efficiency Trade-off

Selected Configuration: The stage-wise single-shot schedule with an initial reduction fraction of $\rho_{shot} = 0.25$ .
Accuracy: Reduces Top-1 accuracy by only 1.47 percentage points (from 77.39% to 75.92%).
FLOPs: Reduces theoretical FLOPs by 22.0%.

Hardware Performance (Batch Size 16)

GPU (RTX 4060 Ti):
- Throughput: Increases by 21.7%.
- Peak Activation Memory (PAM): Reduces by 38.4%.
- Note: At Batch Size 1, throughput decreases by 21.7% due to overhead, indicating the method favors batched execution on GPUs.
x86 CPU (Intel Core Ultra 9):
- Latency: Reduces by 30.0%.
- Throughput: Increases by 42.9% (observed in the aggressive $\rho_{shot}=0.40$ configuration, which incurs a 5.03% accuracy drop).
- Memory: Process Resident Set Size ( $\Delta$ RSS) decreases by 22.0%.
- Note: Unlike GPU, MergeOver accelerates single-stream inference (Batch Size 1) on x86, reducing latency by 17.9%.
ARM CPU (Raspberry Pi 5):
- Latency: Reduces by 17.6% at Batch Size 16 (and 2.4% at Batch Size 1).
- Throughput: Increases by 21.4% at Batch Size 16.
- Memory: $\Delta$ RSS decreases by 39.6% at Batch Size 16.

Significance and Claims

The paper claims that MergeOver successfully bridges the gap between recursive weight-sharing and token merging, a combination previously unexplored. Its primary significance lies in:

Post-Training Viability: It demonstrates that complex architectural integrations can be achieved without the cost of retraining, distillation, or structural redesign.
Recovering Efficiency: It recovers a meaningful portion of the throughput and memory costs introduced by recursive weight-sharing, making recursive models more viable for edge deployment.
Platform Dependence: The results highlight that the benefits of token reduction are highly dependent on the processing platform and batch size. While single-stream GPU inference suffers from overhead, batched execution and CPU-based edge devices (x86 and ARM) show substantial latency and memory improvements.
Baseline Establishment: MergeOver provides a foundational baseline for combining token merging with hierarchical recursive transformers, paving the way for future low-level kernel optimizations and further compression techniques (e.g., quantization).

The authors remain modest, acknowledging that the evaluation is limited to the SReT-Tiny-Distill model and that the observed performance gains are not yet isolated from specific architectural properties or low-level kernel optimizations. They position MergeOver as a starting point for further research into multi-axis compression for edge AI.