💻 computer science

ClustViT: Clustering-based Token Merging for Semantic Segmentation

ClustViT addresses the quadratic complexity of Vision Transformers in semantic segmentation by introducing a trainable Cluster module that merges similar tokens guided by pseudo-clusters and a Regenerator module that restores fine details, achieving significant computational efficiency and faster inference without compromising accuracy.

Original authors: Fabio Montello, Ronja Güldenring, Lazaros Nalpantidis

Published 2026-05-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Fabio Montello, Ronja Güldenring, Lazaros Nalpantidis

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 describe a complex painting to a friend over the phone. The painting has a huge, boring blue sky, a few distinct trees, and a tiny, detailed bird.

A standard "Vision Transformer" (the AI model this paper improves) tries to describe every single square inch of that painting with equal detail. It spends just as much time and brainpower describing the empty blue sky as it does the tiny bird. This is incredibly accurate, but it's also slow and uses a lot of energy—like trying to carry a heavy backpack full of rocks just to walk to the mailbox.

The authors of this paper, ClustViT, asked: "Why describe every single pixel if we already know which parts are the same?"

Here is how their solution works, broken down into simple steps:

1. The Problem: Too Much Noise

Current AI models for "semantic segmentation" (which is just a fancy way of saying "labeling every pixel in an image") are great at recognizing things. But they are slow because they treat every part of the image as equally important. If you are a robot trying to drive through a field, you don't need to analyze every single blade of grass individually; you just need to know "this is grass."

2. The Solution: The "Grouping" Strategy

The authors created a new system called ClustViT. Think of it as a smart editor that looks at the image and says, "Okay, these 100 pixels are all just 'sky,' so let's group them together and treat them as one single piece of information."

They do this using two special tools inside the AI's brain:

The Cluster Module (The Grouping Tool):
Imagine you have a pile of mixed Lego bricks. Instead of looking at every single brick, you quickly sort them into buckets: "Reds," "Blues," and "Special Pieces."
In the AI, this module looks at the image and uses a "cheat sheet" (learned from the ground truth labels) to find pixels that belong to the same semantic category (like "water" or "grass"). It merges all those similar pixels into one representative token.
- The Result: The AI now has to process 100 pixels as if they were just 1. This makes the math much faster.
The Regenerator Module (The Detail Restorer):
Here is the tricky part: If you just merge the pixels, you lose the fine details needed to draw the final picture perfectly.
So, after the AI does its fast, grouped processing, the Regenerator steps in. It takes that single "grouped" piece of information and says, "Okay, I know this represents the sky, so I will expand it back out to fill the original space."
- The Result: The AI gets the speed of the group, but the final output still looks like it has all the original fine details.

3. The "Cheat Sheet" (Pseudo-Clusters)

How does the AI know which pixels to group? It uses a clever training trick. During training, the AI is shown the "correct answer" (the perfect map of the image). It uses this map to create a "pseudo-cluster" guide. It learns: "Oh, I see that all these pixels are labeled 'water,' so I should merge them." It learns to group things based on meaning, not just random similarity.

4. The Results: Faster, Lighter, Still Smart

The authors tested this on three different types of images:

ADE20K: A mix of indoor and outdoor scenes (very complex).
SUIM: Underwater scenes (mostly water, some objects).
RumexWeeds: Agricultural fields (mostly grass/weeds).

What happened?

Speed: The new model was up to 1.64 times faster than the standard model.
Efficiency: It used up to 2.18 times less computing power (energy).
Accuracy: It stayed almost as accurate as the slow, heavy model.

The Sweet Spot:
The paper notes that this works best in "robotic" scenarios where there is a lot of background (like a robot looking at a field or underwater). In these cases, the AI can merge huge chunks of "background" into single tokens, saving massive amounts of energy. In very chaotic, complex images, the savings are smaller, but the model still works well.

Summary

ClustViT is like a smart assistant who realizes that when describing a room, you don't need to list every single grain of dust on the floor. You can say "the floor" (grouping the dust) and then zoom in only on the important furniture. This makes the description much faster to generate, but the listener still gets a clear picture of the room.

This allows robots to "see" and understand their world much faster and with less battery power, without losing the ability to spot important details.

1. Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance in computer vision but suffer from quadratic computational complexity relative to input size. While token pruning and merging techniques (token skimming) have successfully reduced costs for image classification, they struggle in semantic segmentation (a dense prediction task).

The Conflict: Classification can discard "irrelevant" tokens, but segmentation requires preserving spatial details for every pixel to ensure compatibility with standard decoder heads.
The Gap: Existing methods often rely on low-level token similarity or structured, fixed-size compression (e.g., merging only $2\times2$ patches), which fails to adapt to the semantic content of the image or generalizes poorly to dense prediction tasks.

2. Methodology: ClustViT

The authors propose ClustViT, an architecture that integrates a trainable Cluster module and a Regenerator module into a standard ViT backbone. The core innovation is using semantic information derived from ground-truth segmentation masks to guide token merging, rather than simple pixel proximity.

A. Architecture Overview

Standard ViT Backbone: The process begins like a standard ViT, splitting the image into patches and projecting them into tokens.
Cluster Module (Injection Point): At a predefined layer ($ip$) within the encoder, a lightweight MLP predicts cluster assignments for each token.
- Pseudo-Clusters: During training, the model uses ground-truth masks to generate "pseudo-clusters." If all pixels in a patch belong to the same semantic class, the patch is assigned that class ID; otherwise, it is marked as mixed (0).
- Merging: Tokens assigned to the same cluster are aggregated into a single representative token (via mean pooling). Tokens in complex/mixed regions remain unclustered to preserve detail.
- Result: The sequence length is reduced, lowering the computational cost of subsequent Transformer layers.
Regenerator Module: After the final encoder layer, the compressed sequence must be restored to the original resolution for the decoder.
- The module uses the stored cluster assignments to expand representative tokens back to their original positions.
- A refining MLP reconstructs the fine-grained features, ensuring the output sequence matches the input length and is compatible with off-the-shelf segmentation heads (e.g., Segmenter, UPerNet).

B. Training Strategy

Dual-Task Loss: The model is trained end-to-end with a combined loss function:
$L = L_{segm} + \lambda L_{clust}$
- $L_{segm}$ : Standard cross-entropy loss for semantic segmentation.
- $L_{clust}$ : Cross-entropy loss forcing the cluster module to predict the pseudo-clusters derived from the ground-truth mask.
- $\lambda$ : A balancing hyperparameter (set to 0.1 in experiments).

3. Key Contributions

Semantic-Guided Token Merging: Unlike previous methods that merge based on visual similarity or fixed grids, ClustViT merges tokens based on semantic content (derived from segmentation masks), allowing for unstructured, content-adaptive compression.
End-to-End Trainable Backbone: The clustering component is integrated directly into the ViT backbone and trained jointly with the segmentation task, eliminating the need for separate policy networks or pre-training steps.
Regenerator for Dense Prediction: The introduction of a Regenerator module solves the compatibility issue, allowing compressed token representations to be fully restored for dense prediction heads without losing spatial resolution.
Efficiency-Accuracy Trade-off: The method significantly reduces computational cost (GFLOPs) and inference time while maintaining competitive accuracy, particularly in scenarios with large background areas common in robotics.

4. Experimental Results

The authors evaluated ClustViT on three datasets: ADE20K (general scenes), SUIM (underwater robotics), and RumexWeeds (agricultural robotics). They compared against standard ViT and the state-of-the-art Content-aware Token Sharing (CTS) method.

Performance Metrics:
- Computational Efficiency: Achieved up to 2.18× fewer GFLOPs and 1.64× faster inference (images per second) compared to baseline ViT.
- Accuracy: Maintained comparable segmentation accuracy (mIoU). On ADE20K, the best configuration (ClustViT-bk3,ip3) achieved a 1.47× speedup with only a ~1.66 mIoU point drop.
- Robotics Scenarios: The method excelled on datasets with large background areas (SUIM and RumexWeeds). For example, on RumexWeeds, it achieved 1.54× FPS and 2.12× GFLOP gains while slightly improving mIoU over the baseline.
Ablation Studies:
- Cluster Count ( $k$ ): Increasing $k$ (more clusters) improves speed but reduces accuracy as more details are merged. $k=3$ offered the best balance.
- Injection Point ($ip$): Compressing earlier in the network ($ip=2$) maximized speed but hurt accuracy. A middle injection point ($ip=3$ or $4$) provided the optimal trade-off.
- Model Scale: The efficiency gains were most pronounced in larger models (Base/Large), where the overhead of the clustering MLP was negligible compared to the savings in attention computation.

5. Significance and Impact

Robotics Applicability: The paper addresses a critical bottleneck for deploying ViTs on real-world robotic systems (e.g., autonomous agriculture, underwater exploration) where computational resources are limited and scenes often contain large, uniform backgrounds.
Paradigm Shift: It demonstrates that token compression in dense prediction does not require sacrificing spatial fidelity if the merging is guided by semantic understanding rather than just visual similarity.
Generalizability: By supporting standard decoder heads (Segmenter and UPerNet), ClustViT offers a drop-in replacement for existing ViT backbones, making it highly practical for the broader computer vision community.

In conclusion, ClustViT successfully bridges the gap between the high accuracy of Vision Transformers and the efficiency requirements of real-time robotic perception by introducing a semantically aware, trainable token merging mechanism.