CAT-Seg: Cost Aggregation for Open-Vocabulary Semantic Segmentation

In plain words

CLIP can tell whether a whole picture matches a sentence. Labelling every pixel with a word needs more. This method starts from a simple grid: for each pixel and each candidate class name, how well do their embeddings agree? That grid is then treated like a small image and cleaned up. Neighbouring pixels are encouraged to agree, and the candidate classes are made to compete with each other, so a label that fits the whole scene beats one that only looks similar. The cleaned-up grid is enlarged back to full resolution to give the final mask.

The method also adapts CLIP itself, but lightly. Only a few projection matrices inside CLIP's attention layers are trained; the rest of CLIP stays frozen, which makes training cheaper than tuning everything. The aim is to label classes that the model did not see during training.

TL;DR

• CAT-Seg builds a cosine-similarity cost volume between CLIP pixel embeddings and class-name text embeddings, refines it with spatial and class aggregation, and decodes it with an upsampling decoder §3.1 §3.2 §3.3 §3.4.
• Only the query and value projections of the attention layers in both CLIP encoders are trained; the paper reports this as the best trade-off among the fine-tuning variants it tested §3.4 tab. 6.
• Trained only on COCO-Stuff, the author-reported results with ViT-L/14 are 16.0 mIoU on ADE-847 and 23.8 on PC-459 tab. 1.
• The released code is MIT-licensed and differs from the paper in several places (decoder feature pairing, a 256-class cap, pooling at evaluation, the loss on ignored pixels). See Paper vs code before re-implementing.

Problem

The task labels each pixel with a class that is given as free text, so the label set can change at test time and can include classes unseen in training §3. Region-to-text methods depend on a region generator trained on a limited amount of annotated data §2. Decoders placed on top of CLIP image embeddings can lose CLIP's text alignment for unseen classes, and fine-tuning the encoders directly tends to overfit to seen classes, so most methods freeze CLIP §1.

Method

Dense image and text embeddings come from CLIP. Following an earlier method, the last attention layer of the image encoder is modified to remove its pooling effect §3.1. Each class name is turned into a text embedding §3.1 app. B.

The cost volume holds, for each of the H x W image positions i and each of the N_C classes n, the cosine similarity between the image and text embeddings §3.1 eq. 1. A single convolution is applied to each class slice on its own, giving a cost embedding with d_F channels §3.1.

Spatial aggregation runs separately for each class, to exploit spatial smoothness of the image, with two Swin transformer blocks: one attends within local windows, the second within shifted windows §3.2 eq. 2. Class aggregation runs at each position over the class axis, so that classes are reasoned about together. It uses a transformer layer without positional embedding and linear attention. Its output does not depend on class order, and it handles a varying number of classes §3.3 eq. 3. The two aggregations alternate §3.

The upsampling decoder takes the aggregated volume, enlarges it, concatenates a matching feature map from intermediate CLIP ViT layers, and applies a 3x3 convolution. This repeats N_U = 2 times, followed by a prediction head §3.4 app. B fig. 3.

Embedding guidance feeds linearly projected CLIP image and text embeddings into the query and key inputs of both aggregators. The value inputs do not receive guidance §3.4 eq. 4.

Efficient fine-tuning trains only the query and value projections of CLIP's attention layers, in both encoders. The other CLIP weights stay frozen. The paper compares this with prompt tuning, full fine-tuning, attention-only tuning, query-key-only and key-value-only tuning, and single-encoder variants §3.4 §4.4 tab. 6.

Training uses a per-pixel binary cross-entropy loss §4.2. Inference uses sliding windows on 640 x 640 inputs (see Hyperparameters) app. B.

# Forward pass for one batch of images  [§3.1] [§3.2] [§3.3] [§3.4]
# Input:  images (B, 3, 384, 384); class_names (N_C strings)
# Output: per-class logits (B, N_C, 96, 96)  [app. B]

# Dense CLIP image embeddings, one per patch, on a 24 x 24 grid  [app. B]
# The last attention layer is modified to remove pooling  [§3.1]
image_embeds = clip_dense_image(images)              # (B, 24, 24, d)

# Text embedding per class, from the prompt "A photo of a {class}"  [app. B]
# The code uses another template; see Paper vs code
text_embeds = clip_text(class_names)                 # (N_C, d)

# Cost volume: cosine similarity per position and class  [§3.1] [eq. 1]
cost = cosine_similarity(image_embeds, text_embeds)  # (B, H*W, N_C)

# Cost embedding: one convolution applied to each class slice on its own
# d_F = 128  [§3.1] [§4.2]
cost_embed = cost_conv(cost)                         # (B, H*W, N_C, 128)

# Alternate spatial and class aggregation, N_B = 2 times  [§3] [§4.2]
for block in range(N_B):
    # Spatial: per class, a Swin pair with CLIP guidance on Q and K  [§3.2]
    for n in range(N_C):
        cost_embed[:, :, n] = swin_pair(cost_embed[:, :, n], guide_v)
    # Class: per position, linear attention over classes  [§3.3] [eq. 3]
    for i in range(H * W):
        cost_embed[:, i] = class_attention(cost_embed[:, i], guide_t)

# Upsampling decoder, applied per class  [§3.4] [app. B]
# Stage 1 target 48 x 48, stage 2 target 96 x 96  [app. B]
for stage in range(N_U):
    # Enlarge the aggregated volume  [§3.4]
    volume = upsample_2x(volume)
    # Concatenate the matching CLIP feature map, then convolve  [§3.4]
    volume = conv3x3(concat(volume, clip_map[stage]))
logits = prediction_head(volume)                     # (B, N_C, 96, 96)

# Per-pixel binary cross-entropy against one-hot class targets  [§4.2]
# Handling of unlabelled pixels is not stated in the paper
loss = binary_cross_entropy(logits, one_hot(targets))

Key equations (paper numbering):

C(i,n) = \frac{D^V(i)\cdot D^L(n)}{\|D^V(i)\|\,\|D^L(n)\|} eq. 1

Eq. 2 applies the spatial block per class, eq. 3 the class block per position, and eq. 4 adds the projected CLIP guidance to the inputs of both blocks §3.2 §3.3 §3.4.

Figure 1 (introduction) compares feature aggregation with cost aggregation. Figure 2 shows the raw cost volume from frozen CLIP, the fine-tuned cost volume and the aggregated cost volume for a seen and an unseen class. Figure 3 is the overall architecture. Figure 4 compares with SAN qualitatively on PC-459 and on MESS. Figure 5 shows feature versus cost aggregation on an unseen class. Figure 6 shows t-SNE plots of CLIP dense embeddings before and after adaptation §1 §3 §4.3 §4.4 (figure numbers follow source order).

Hyperparameters

Results

Author-reported. Setup: models trained on COCO-Stuff, mIoU on the standard benchmarks, backbones ViT-B/16 and ViT-L/14@336px §4.1 §4.3 tab. 1.

Further results. On MESS the authors note weaker scores in Medical Sciences: 28.09 (ViT-B) and 24.70 (ViT-L), below the best listed baseline in that column. On Engineering (ViT-B), 20.34 is below SAN's 23.58 tab. 2. Inference with the training resolution instead of patch inference gives 14.6 ADE-847 and 22.1 PC-459 with ViT-L app. C. With other VLMs, EVA-02-CLIP-L/14 gives 16.4 ADE-847 and SigLIP-ViT-L/16 gives 18.0 app. C.

Paper vs code

Code: official repository cvlab-kaist/CAT-Seg (MIT), commit 6d3a188af95165147fe2f34a8237fa7d2633e784. The paper links no code. The project page links the code (its Code button points to KU-CVLAB/CAT-Seg, which I could not confirm as the same repository from the two URLs alone). Ownership is by the authors: the project page lists the six authors and the code header says "Written by Seokju Cho and Heeseong Shin" model.py. The LICENSE copyright line reads KU-CVLAB LICENSE.

Differences found when comparing paper text and code:

• Upsampling operator. The paper uses bilinear upsampling followed by a single 3x3 convolution §3.4. The code uses a learned 2x2 transposed convolution followed by two 3x3 convolutions, each with GroupNorm and ReLU [code: cat_seg/modeling/transformer/model.py#L540-L555] [code: cat_seg/modeling/transformer/model.py#L520-L537].
• CLIP feature maps for the decoder. The paper pairs the 48 x 48 map with the 8th CLIP layer and the 96 x 96 map with the 4th layer (ViT-B/16), and with the 16th and 8th layers for ViT-L/14 app. B. The code hooks blocks 4 and 8 (0-based 3 and 7) for ViT-B, and 8 and 16 for ViT-L [code: cat_seg/cat_seg_model.py#L84-L87]. The earlier block feeds the 2x transposed convolution (48 x 48) and the later block feeds the 4x transposed convolution (96 x 96) cat_seg_model.py [code: cat_seg/cat_seg_model.py#L81-L82]. The paper's pairing is reversed.
• Guidance projections. The paper specifies linear projections P^V and P^L §3.4. The code uses a 3x3 convolution with ReLU for the image guidance [code: cat_seg/modeling/transformer/model.py#L615-L618] and a linear layer with ReLU for the text guidance, applied to the prompt-averaged, normalised text embedding model.py [code: cat_seg/modeling/transformer/model.py#L711-L714]. The Swin wrapper also adds a LayerNorm on the guidance model.py model.py, which the paper does not mention.
• Class-axis pooling. The paper describes class aggregation per position §3.3. The code average-pools the 24 x 24 grid 2x2 before class attention and bilinearly upsamples afterwards model.py model.py. The training configs set this pooling to 2 vitb_384.yaml. The eval script sets it to 1 eval.sh.
• Class padding and cap. The paper does not mention either. The code pads the class axis to 256 with learnable tokens, which join the attention [code: cat_seg/modeling/transformer/model.py#L397-L402] [code: cat_seg/modeling/transformer/model.py#L419-L421]. When there are more than 256 classes it keeps the 256 with the highest maximum cosine score over positions, and sets the others to logit -100 [code: cat_seg/modeling/transformer/model.py#L693-L701] [code: cat_seg/modeling/transformer/model.py#L720-L723]. ADE-847 (847 classes) and PC-459 (459 classes) are affected at test time.
• Ignored pixels in the loss. The paper says per-pixel binary cross-entropy §4.2. The code computes the loss over all pixels. Pixels with label 255 get all-zero targets, so they count as negatives [code: cat_seg/cat_seg_model.py#L165-L172] vitb_384.yaml.
• Prompt template. The paper gives "A photo of a {class}" as an example and says other prompts were not explored app. B. The code uses the single template "A photo of a {} in the scene" [code: cat_seg/modeling/transformer/cat_seg_predictor.py#L84-L85], selected by vitb_384.yaml.
• Input size to CLIP. The paper gives a training resolution of 384 app. B. The code resizes the crop to 336 for ViT-L/14@336px before the image encoder, and to 384 for ViT-B/16 cat_seg_model.py [code: cat_seg/cat_seg_model.py#L146].
• Dense last layer. The paper follows an earlier method and removes the pooling effect in the last attention layer §3.1. The code replaces that layer with a value-only path: query and key outputs are discarded, and the input of the class token is added to every position model_vpt.py (split at model_vpt.py, residual at [code: cat_seg/third_party/model_vpt.py#L225]).
• Window attention. The paper uses Swin blocks §3.2. The code's window attention has no relative-position bias table [code: cat_seg/modeling/transformer/model.py#L65-L114].
• Sliding-window inference. The paper averages overlapping patch predictions and uses the image resized to 384 for the global view app. B. The code resizes the whole image to 640 x 640, averages the patch windows, then averages that 1:1 with the global-view output cat_seg_model.py cat_seg_model.py. The paper does not give this weight.
• Evaluation settings. Config default has sliding-window testing off [code: configs/config.yaml#L65]. The eval script enables it and sets class-axis pooling to 1 for all six test sets eval.sh [code: eval.sh#L49-L50] eval.sh eval.sh [code: eval.sh#L88-L89] eval.sh. The paper does not mention this.
• Learning-rate schedule, warmup and gradient clipping. The paper states none of these. The code uses WarmupCosineLR with zero warmup iterations [code: configs/config.yaml#L37-L38] config.yaml and clips the gradient norm of all trainable parameters to 0.01 [code: configs/config.yaml#L44-L48] train_net.py.
• Consistent with the paper: query and value projections of both encoders are trainable, and the CLIP learning rate is 0.01 times the base rate [code: cat_seg/cat_seg_model.py#L58-L75] train_net.py. Class aggregation uses linear attention model.py, and the aggregation dimension (128), layer count (2) and decoder count (2) match §4.2.

Pitfalls

• The 256-class cap changes what is evaluated. ADE-847 and PC-459 have more than 256 classes. The released model drops the rest, so a re-implementation that follows the paper only will not match those two benchmarks [code: cat_seg/modeling/transformer/model.py#L693-L701].
• Padding tokens take part in attention. Learnable tokens fill the class axis to 256 and join the linear attention, so outputs depend on how many classes are padded model.py.
• Train and test pooling differ in the released scripts. The config trains with class-axis pooling 2; the eval script runs with 1 vitb_384.yaml eval.sh. Reproduce evaluation with the script, not only the config, and expect the reported numbers to depend on this.
• The loss includes ignored pixels as negatives, which changes the training signal wherever unlabelled pixels exist cat_seg_model.py.
• Gradient clipping at 0.01 on the total norm is not mentioned in the paper and is not ablated config.yaml. Treat it as part of the recipe.
• Table backbones are unclear. Tables 3, 4, 6 and 8 do not state the backbone. The final rows of Tables 4 and 6 match the ViT-L/14 row of Table 1, so the setting is probably ViT-L/14, but that is inferred from the numbers, not stated tab. 4 tab. 6. Table 8's totals (433.7 M parameters) are also not labelled with a backbone tab. 8.
• Table numbers follow the order of the source; check against the PDF before citing them.
• Efficiency numbers were measured on one RTX 3090 and are not normalised for other hardware tab. 8.
• The authors list ambiguous ground truth as a limitation for the evaluation sets app. E.
• The ViT-L/14 model sees the 384 crop resized to 336 [code: cat_seg/cat_seg_model.py#L146], which the paper does not state.
• The fine-tuning table ranks key-value tuning slightly above query-value tuning on ADE-847 (16.1 vs 16.0), while the text picks query-value as the best trade-off tab. 6 §4.4.

Lineage

• Builds on: CLIP [radford2021learning] for the encoders and text prompts; dense CLIP features from an earlier method [zhou2022extract] §3.1; cost aggregation from correspondence work [kendall2017end] and CATs [cho2021cats] [cho2022cats++] [hong2022cost] §3; Swin transformer [liu2021swin] and linear transformer [katharopoulos2020transformers] §3.2 §3.3; FPN-style decoding [lin2017feature] §3.4.
• Compared against: SPNet, ZS3Net, LSeg, LSeg+, OpenSeg, ZegFormer, ZSSeg, OVSeg, ZegCLIP, SAN and ODISE tab. 1; on MESS also X-Decoder and OpenSeeD tab. 2. FC-CLIP is discussed in related work only §2.