PaperlayerView exact revision on arXiv

Source-linked AI summary

MaskMerge: Structure-aware Token Merging for Efficient Instance Segmentation

Bavesh Balaji, Letian Jiang, Alexander Wong, David Clausi, Sirisha Rambhatla

arXiv:2610.22921v1cs.CV

TL;DR

Vision Transformer token merging reduces computation, but feature similarity alone may not preserve instance identity when adjacent objects look alike. MaskMerge constrains feature-based merging with coarse predicted masks, improving accuracy–efficiency results on PanNuke and iSAID while matching ALGM on Cityscapes at lower computation.

  • Problem

    High-resolution ViT segmentation is computationally costly, and adjacent same-class instances can have similar features despite requiring distinct instance masks.

    Source: §1 · pp. 1–2

  • Method

    MaskMerge uses predicted intermediate mask structure to constrain feature-based token aggregation to tokens assigned to the same predicted instance.

    Source: §1 · p. 2

  • Results

    MaskMerge raises mAP from 30.1 to 32.1 on PanNuke and from 33.1 to 36.2 on iSAID, reducing computation from 229 to 153 and 198 to 182 GFLOPs, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Results support using task-level instance structure, not appearance similarity alone, to guide token reduction for instance segmentation.

    Source: §5 · p. 8

Abstract

from arXiv · show

Token merging improves the efficiency of Vision Transformers by combining redundant tokens, but existing methods largely rely on feature similarity to determine which tokens to merge. This criterion is effective for semantic segmentation, yet can fail in instance segmentation, where adjacent objects of the same class often have highly similar features. We identify this failure mode as cross-instance token fusion. We introduce MaskMerge, an instance-aware token merging method that uses coarse intermediate mask predictions to constrain feature-based merging and preserve instance boundaries. We evaluate MaskMerge on two diverse instance segmentation domains characterized by high inter-instance visual similarity: PanNuke and iSAID. Against the SOTA token merging method ALGM, MaskMerge improves mAP from 30.1 to 32.1 on PanNuke while reducing computation from 229 to 153 GFLOPs and running 1.18x faster, and from 33.1 to 36.2 on iSAID while reducing computation from 198 to 182 GFLOPs and running 1.23x faster. Additionally, we also test it on Cityscapes and match ALGM's performance while using fewer FLOPs. The gains also transfer to an MAE pretrained backbone, demonstrating that instance-aware merging enables more aggressive compression when appearance similarity alone cannot reliably preserve object identity.

1. Introduction

Vision Transformers offer strong dense recognition but incur substantial computation, while feature-similarity merging can confuse adjacent, visually similar instances. MaskMerge uses predicted instance structure to guide merging, improving the accuracy–efficiency trade-off on PanNuke and iSAID.

1. Introduction

  • The approach targets costly ViT dense prediction, where high-resolution inputs preserve spatial detail and dense representations contain redundancy that motivates token reduction.Token merging combines redundant tokens while retaining correspondence to their original spatial support.Source: §1 · pp. 1–2
  • Adjacent same-class instances can have highly similar features despite belonging to different objects, making similarity-only merging prone to cross-instance token fusion.At the same merge budget, ALGM produces 5.7× and 2.9× more cross-instance merges than MaskMerge on PanNuke and iSAID, respectively.Source: §1 · p. 2, Figure 2 · p. 3
  • MaskMerge uses coarse intermediate masks to restrict candidate matches to tokens assigned to the same predicted instance before feature-similarity selection.The method uses structure emerging within the segmentation model rather than introducing an independent grouping mechanism.Source: §1 · p. 2
  • 32.1 mAP versus ALGM’s 30.1 on PanNuke and 36.2 versus 33.1 on iSAID, with computation reduced from 229 to 153 and 198 to 182 GFLOPs, respectively.MaskMerge also matches ALGM at 19.8 AP on Cityscapes using 298 rather than 347 GFLOPs.Source: §1 · p. 2
  • Cross-instance fusion is concentrated in crowded regions with adjacent, visually similar instances, and can force similarity-based methods to retain more tokens as merge rates rise.A safer criterion can support more aggressive compression while preserving instance boundaries.Source: §1 · p. 2

2. Related Works

ViT segmentation models use object-specific representations and mask-prediction heads to support predictions for multiple instances. Prior token-reduction work spans general vision and semantic segmentation, where dense-output recovery and spatial fidelity constrain efficiency methods.

2.1. ViTs for Instance Segmentation

  • Query-based models represent objects with learned queries, while MaskFormer, Mask2Former, and EoMT develop mask-based or generic ViT segmentation frameworks.MaskMerge builds on this line by enhancing efficiency while preserving instance structure.Source: §2.1 · p. 3

2.1.1. Token reduction in general.

  • General token reduction includes pruning, merging, hybrid strategies, compressed representations, quantization, and alternative attention mechanisms.Approaches differ in reduction timing, recoverability of token identities, routing overhead, and optimization complexity.Source: §2.1.1 · pp. 3–4

2.1.2. Token reduction for semantic segmentation.

  • Semantic segmentation requires per-location output recovery and boundary preservation, while task-specific methods address recoverability through halting, clustering, token sharing, or local and global merging.The cited discussion notes overhead, model complexity, or adaptability constraints for several approaches, and that prior methods had not been evaluated for instance segmentation.Source: §2.1.2 · p. 4

3. Methodology

MaskMerge progressively converts fixed patch tokens into content-aware tokens, first reducing local redundancy and then merging tokens guided by instance structure. It restores the reduced sequence to a dense spatial grid for segmentation.

3. Methodology

  • MaskMerge uses object-level mask proposals as proxies for token merging, jointly refining tokens and performing segmentation within a transformer architecture.Its two-stage strategy first merges tokens in local neighborhoods, then consolidates tokens belonging to the same instance.Source: §3 · p. 4, §3.2 · p. 4

3.1. Vision Transformers with Patch-based Tokenization

  • Fixed-size patch tokenization gives all image regions equal token counts, so simple regions incur redundant computation while complex boundaries remain under-represented.The method addresses this limitation by dynamically merging tokens while preserving object-level structure.Source: §3.1 · p. 4

3.2. Mask-based Token Merging

  • Mask-based merging progressively transforms fixed patch tokens into content-aware tokens through local and same-instance merging stages.The stages remove redundancy in small spatial neighborhoods, then consolidate tokens belonging to the same instance.Source: §3.2 · p. 4

3.2.1. Local Token Merging

  • Tokens are grouped into non-overlapping spatial windows, where average pairwise cosine similarity determines whether tokens are merged.A similarity above threshold τ triggers a size-aware average; local merging removes redundancy while preserving spatial locality.Source: §3.2.1 · pp. 4–5

3.2.2. Object-level Adaptive Token Merging (OLAM)

  • OLAM constrains feature-similarity merging with coarse instance masks, reducing the risk of merging tokens from different objects.It computes bipartite matching pairs from the masked similarity matrix and merges each pair by weighted averaging.Source: §3.2.2 · pp. 5–6

3.2.3. Token Unmerging

  • A compact index list records each merged token’s original contributors, allowing its embedding to be copied back to the original spatial positions.This sparse bookkeeping reduces memory overhead from O(N^2) to O(N) and restores a dense sequence for the decoder.Source: §3.2.3 · p. 6

4. Experiments

Across PanNuke, iSAID, and Cityscapes, the experiments test MaskMerge’s accuracy–efficiency trade-off, backbone generalization, and transfer to Mask2Former. Results favor instance-aware merging on the visually homogeneous datasets, retain gains with MAE pretraining, and show benefits when mask prediction is decoupled from the backbone.

4.1. Experimental Setup

  • The evaluation covers PanNuke, iSAID, and Cityscapes, compares ToMe and ALGM, and tests MAE-pretrained ViT-B and Mask2Former transfer.The primary setup uses EoMT with a DINOv2 ViT-B backbone; merging follows ALGM’s local merging in the first encoder block and applies instance-aware merging at layer 5.Source: §4.1 · p. 6, §4 · p. 6

4.2. Main Results

  • 2.0 mAP higher than ALGM on PanNuke, MaskMerge uses 1.50× fewer FLOPs and achieves 1.18× higher throughput; on iSAID it gains 3.1 mAP, uses 1.08× fewer FLOPs, and runs 1.23× faster.On Cityscapes, it matches ALGM at 19.8 mAP with fewer FLOPs; throughput is similar because coarse-mask prediction adds runtime overhead.Source: §4.2 · p. 7, Table 1 · p. 7
  • On PanNuke, MaskMerge reaches 32.1 mAP at 153 GFLOPs, versus ToMe’s 29.5 mAP at 369 GFLOPs; on iSAID, it raises mAP from 31.5 to 36.2 while using fewer FLOPs.The results indicate that more aggressive appearance-only merging does not provide the same accuracy–efficiency trade-off as incorporating instance structure.Source: §4.2 · p. 7

4.3. Generalization Across Backbone Pre-training

  • With an MAE-pretrained ViT-B on PanNuke, MaskMerge improves overall mAP from 32.0 to 34.0 and small-instance performance from 31.5 to 33.4.The advantage persists under a different self-supervised pretraining strategy and is not specific to DINOv2 features.Source: §4.3 · p. 7, Table 2 · p. 7

4.4. Extension to Mask2Former

  • For Mask2Former, a two-layer convolutional head predicts coarse instance structure from backbone features using decoder masks from the frozen model as training targets.This single-pass variant constrains merging without requiring a separate dense forward pass at inference.Source: §4.4 · p. 7, Table 3 · p. 8
  • With Mask2Former, single-pass MaskMerge retains 32.53 mAP, 1.36 points below the standard model, compared with 30.69 for ToMe and 24.44 for ALGM.The results support the benefit of instance-aware merging when mask prediction is decoupled from the backbone.Source: §4.4 · pp. 7–8, Table 3 · p. 8

4.5.1. Impact of each module

  • Replacing cosine-similarity global merging with OLAM raises mAP from 29.5 to 33.4; combined with local merging, OLAM reaches 32.1 mAP at 153 GFLOPs versus 30.1 mAP at 229 GFLOPs.In the combined setting, throughput increases from 121.5 to 142.8 images/s, indicating that OLAM supports more aggressive early token reduction.Source: §4.5.1 · p. 8, Table 5 · p. 8

4.5.2. OLAM module position

  • Applying OLAM earlier improves efficiency but can lower accuracy; delaying it raises accuracy to 33.8 mAP at layer 7 while increasing computation and reducing throughput.Layer 5 is the chosen balance, reaching 32.1 mAP at 153 GFLOPs and 142.8 images/s.Source: §4.5.2 · p. 8, Table 4 · p. 8

5. Conclusion

MaskMerge uses intermediate coarse masks to constrain token aggregation, preserving instance structure while enabling more aggressive token reduction. It improves accuracy–efficiency tradeoffs on PanNuke and iSAID and matches ALGM on Cityscapes at lower computation.

5. Conclusion

  • MaskMerge achieves consistently better accuracy–efficiency tradeoffs than cosine-based merging on PanNuke and iSAID, where visually similar instances are common.The advantage persists across ViT scales and different backbone pre-training strategies.Source: §5 · p. 8
  • MaskMerge matches ALGM on Cityscapes at lower computation and transfers to Mask2Former.Source: §5 · p. 8
  • Intermediate coarse masks constrain token aggregation, preserving instance structure while enabling more aggressive token reduction.The conclusion presents instance identity, alongside feature similarity, as a consideration for effective token reduction in instance segmentation.Source: §5 · p. 8

Related papers

Loading 2610.22921v1…