Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation

In plain words

Few-shot segmentation asks a model to outline an object in a new picture after seeing only a few example pictures with that object already outlined. A common shortcut is to squash the examples into one summary vector and compare the new picture against that. This paper instead compares every spot in the new picture with every spot in the examples. That produces a large table of similarity scores, and the table is noisy: repeated textures and cluttered backgrounds create false matches.

The method cleans up that table with a transformer, the same kind of attention model used in language and vision. The table is too big to attend over directly, so the attention is run in small windows that slide across both images at once. Before the windows, a few small overlapping convolutions mix each spot with its neighbours, so that the window edges do not cut off local context. The cleanup runs from coarse to fine, and the coarse result guides the finer one. A last decoder looks at the new picture's own appearance to sharpen the outline.

The authors report state-of-the-art numbers on the three standard benchmarks. They also show that the same kind of cleanup helps with finding matching points between different photos of similar objects, which is a neighbouring task.

TL;DR

• VAT aggregates the 4D correlation volume (every query position against every support position) with a 4D convolutional Swin Transformer, inside a three-level pyramid, then decodes a mask with an appearance-aware transformer.
• Backbones are ImageNet ResNet-50/101, frozen; training is AdamW at 5e-4 with no stated weight decay in the paper.
• Author-reported: PASCAL-5i 1-shot mIoU 67.9 (ResNet-101) vs HSNet 66.2; COCO-20i 41.3 vs 39.2 (ResNet-50); FSS-1000 1-shot 90.3 vs 86.5 (ResNet-101) tab. 1 tab. 2 tab. 3.
• Costs are high: the full correlation volume is materialised, and the authors say it limits training resolution app. C. A transformer on the unsplit volume needed about 84 GB per batch tab. 5.
• The released code differs from the paper on data augmentation, K-shot voting and the appearance-affinity layers; see Paper vs code before reproducing numbers.

Problem

Few-shot segmentation must segment unseen classes from a few annotated support images, with no retraining for the new class §1 §3.1. Prototype methods compress the support set and lose pixel-level pair relations and spatial structure §1. Pairwise methods built on a 4D correlation volume are limited by the receptive field of fixed 4D convolutions and by the lack of adaptive, content-dependent aggregation §1 §3.2. Transformers can aggregate pixel pairs globally, but a standard transformer over all query-support positions is computationally infeasible, and tokenising the volume by collapsing one image's spatial axes discards the spatial structure that matching needs §3.4.1. Non-overlapping patch embeddings also cut local context at token boundaries and give weak inductive bias §3.4.2.

Method

The task is episodic: an episode gives a support set of K image-mask pairs and a query image, and the model predicts the query mask §3.1. The method has five parts.

1. Feature extraction and correlation §3.3. Query and support images pass through a CNN. The support features are multiplied by the resized support mask, so background is zeroed before matching §3.3. Correlation is the ReLU of the cosine similarity between each query feature and each support feature eq. 1. The paper stacks the correlation maps of all layers with the same spatial size into one volume for each pyramid level p, with layers conv3_x (p=3), conv4_x (p=4) and conv5_x (p=5) used §3.3 §4.1.

2. Volumetric convolution module (VCM) §3.4.2 fig. 4. The paper first considers an embedding that reshapes the 4D volume into non-overlapping 4D windows, similar to a ViT patch embedding (VEM) §3.4.1. It replaces this with a stack of overlapping 4D convolutions. Each stage applies spatial 4D max-pooling, an overlapping 4D convolution, ReLU and group normalisation. The output is projected to a fixed embedding dimension D = 128 §4.1. The paper argues that this adds convolutional inductive bias and handles window boundaries better §3.4.2 app. B.

3. Volumetric transformer module (VTM) §3.4.1 fig. 3. Swin Transformer is extended from 2D to 4D. The volume is partitioned into non-overlapping 4D windows of size n along each axis, self-attention runs inside each window, then the windows are shifted by floor(n/2) along each axis and attention runs again §3.4.1 fig. 3. The window size is 4 §4.1. Relative position bias is a 4D table expanded from the 2D Swin design §3.4.1. Layer norm and MLP blocks are unchanged §3.4.1. The aggregation is written as A^p = VTM(M^p) eq. 2.

4. Residual connections and pyramidal guidance §3.4.2 §3.4.3. Each transformer output is added back to its input as a residual, because random initialisation makes early matching scores noisy §3.4.2. The pyramid runs from the coarsest correlation level to the finest. The coarser aggregated map is bilinearly upsampled and added to the next level's VCM output as guidance §3.4.3. The encoder is written A^p = VTM(VCM(C^p) + up(A^{p+1})) eq. 4.

5. Affinity-aware transformer decoder (ATD) §3.5 eq. 5. The aggregated volume is averaged over the support spatial dimensions to leave a query-side map §3.5. This is concatenated with a linear projection of the query's appearance embedding, passed through a transformer, and bilinearly upsampled. The mask is m_q = ATD([A'^p, P(F_q)]) eq. 5. The appearance embedding comes from conv2_x, conv3_x and conv4_x, projected to 16, 32 and 64 channels respectively, and it is used for FSS-1000 and, with fewer layers, for other benchmarks §4.1.

K-shot. The model runs once per support pair. Per-pixel foreground votes are summed, and a pixel is foreground when the sum divided by K exceeds the threshold tau = 0.5 §3.6 §4.1.

Loss. The text does not write out the segmentation loss. The released code uses a two-class cross-entropy on logits upsampled to input resolution vat.py. Treat this as code-only (see Paper vs code).

Pseudocode. The paper has no algorithm box; "Algorithm 1" is cited in app. B but is not in the source. The sketch below follows §3.

# VAT forward pass for one query and one support pair  [§3.2] [§3.3]
# Input:  query_img (B, 3, H, W); support_img (B, 3, H, W)
# Input:  support_mask (B, H, W)
# Output: logit_mask (B, 2, H, W)

def vat_forward(query_img, support_img, support_mask):
    # Frozen ResNet features, conv2_x to conv5_x  [§4.1]
    query_feats = backbone_features(query_img)
    support_feats = backbone_features(support_img)
    # Zero the support features outside the support mask  [§3.3]
    support_feats = mask_features(support_feats, support_mask)
    # Stacked correlation per pyramid level, coarse to fine  [§3.3] [eq. 1]
    corr_levels = stacked_correlation(query_feats, support_feats)
    # Pyramidal encoder: each level is guided by the coarser one  [§3.4.3]
    aggregated = pyramid_encoder(corr_levels)
    # Decoder fuses the query appearance embedding  [§3.5] [eq. 5]
    return decoder(aggregated, query_feats)

# Cosine correlation between query and support feature maps  [§3.3] [eq. 1]
# query_feat: (C, h_q, w_q)   support_feat: (C, h_s, w_s)

def correlation(query_feat, support_feat):
    # Scale each spatial feature vector to unit length  [§3.3]
    query_unit = unit_length(query_feat, dim=channels)
    support_unit = unit_length(support_feat, dim=channels)
    # Inner product for every query-support pair, clipped at zero  [eq. 1]
    corr = relu(inner_product(query_unit, support_unit))
    # Output shape (h_q, w_q, h_s, w_s)
    return corr

# Volumetric convolution module (VCM), replaces VEM  [§3.4.2] [fig. 4]
# Input:  corr (B, L, h_q, w_q, h_s, w_s)
# Output: (B, D, h_q', w_q', h_s', w_s')

def volumetric_conv(corr, stages):
    # Each stage: 4D max-pool, overlapping 4D conv, group norm, ReLU  [§3.4.2]
    for stage in stages:
        # Kernel 3 on every axis; stride differs per axis
        # [code: model/ours.py#L18-L73]
        corr = conv4d(corr, kernel=3, stride=stage.stride, padding=1)
        corr = group_norm(corr, groups=stage.groups)  # group counts: not stated
        corr = relu(corr)
    return corr

# Pyramidal transformer encoder, coarse (conv5) to fine (conv3)  [§3.4.3]
# Residual placement follows the code and [§3.4.2]; [eq. 4] omits it

def pyramid_encoder(corr_levels):
    aggregated = None
    for corr in corr_levels:
        embedded = volumetric_conv(corr, stages)  # [§3.4.2]
        if aggregated is not None:
            # Coarser result, upsampled bilinearly, added as guidance  [eq. 4]
            embedded = embedded + upsample_bilinear(aggregated)
        # Shifted-window 4D self-attention, added back as a residual  [§3.4.2]
        aggregated = embedded + swin4d(embedded)
    return aggregated

# Affinity-aware transformer decoder  [§3.5] [eq. 5]
# appearance_feats: query embeddings with 16, 32 and 64 channels  [§4.1]

def decoder(aggregated, appearance_feats):
    # Average over support positions, keep only the query grid  [§3.5]
    x = mean(aggregated, over=support_dims)  # (B, D, h_q', w_q')
    for stage, appearance in zip(decoder_stages, appearance_feats):
        # Concatenate query appearance; off for some layers  [§4.1]
        x = concat(x, appearance)
        # 2D shifted-window attention, window size 8
        # [code: model/ours.py#L112-L130]
        x = swin2d(x)
        # Bilinear upsampling before the next stage
        # [code: model/ours.py#L193-L197]
        x = upsample_bilinear(x)
    # Two 3x3 convs to two logits per pixel  [code: model/ours.py#L132-L136]
    return conv_head(x)

# K-shot inference by voting over per-shot foreground masks  [§3.6]
# Input:  support_pairs, a list of K (support_img, support_mask) pairs

def predict_k_shot(query_img, support_pairs, tau):
    # One forward pass per support pair; count foreground votes  [§3.6]
    votes = zeros(query_img_size)
    for support_img, support_mask in support_pairs:
        logit_mask = vat_forward(query_img, support_img, support_mask)
        votes = votes + argmax(logit_mask, over=channels)
    # Paper: sum divided by K exceeds tau; code differs, see Paper vs code
    return votes / len(support_pairs) > tau

Hyperparameters

Values in the first block are stated in the paper. Values marked "code only" are taken from the released repository and are not in the paper text.

Results

Author-reported throughout. Setup: ResNet-50 or ResNet-101 backbones, frozen, evaluated with mIoU and FB-IoU on PASCAL-5i, COCO-20i and FSS-1000 §4.2. The tables below reproduce the rows that report numbers. Per-fold COCO and FSS values are not split out in the source. Per-class SPair-71k numbers app. D and per-fold mBA tables app. D are not reproduced here.

Table 1, PASCAL-5i tab. 1. Per-fold IoU, mean mIoU, FB-IoU and mBA. Dashes mean not reported in the source.

Table 2, COCO-20i tab. 2. ResNet-50 only. Columns are per-fold IoU, mean, FB-IoU and mBA.

Table 3, FSS-1000 tab. 3. Mean mIoU, FB-IoU and mBA.

Component ablation, FSS-1000 tab. 4. The backbone is not stated for this table. The full-model row matches the ResNet-101 FSS result in tab. 3.

Aggregator ablation, FSS-1000 tab. 5. Only the VTM aggregator is replaced. The shot setting is not stated. Timings are on one RTX 3090 and an i7-10700 §4.4. OOM means out of memory.

Semantic correspondence tab. 6 §4.4. Trained on the PF-PASCAL training split for PF-PASCAL and PF-WILLOW, and on SPair-71k for SPair-71k §4.4. PCK at the listed alpha thresholds. Rows marked * use pretrained weights from the authors or numbers from the papers tab. 6. F.T. = backbone fine-tuned. Data augmentation as in CATs app. A. Columns for SPair-71k are alpha_bbox 0.03, 0.05, 0.10, 0.15; for PF-PASCAL alpha_img; for PF-WILLOW alpha_bbox 0.05, 0.10, 0.15.

The text describes PF-PASCAL as nearly saturated and VAT as on par there §4.4.

Appendix results (author-reported).

• Backbone ablation, FSS-1000 mIoU, 1-shot / 5-shot app. B: ResNet-50 90.1 / 90.7; ResNet-101 90.3 / 90.8; PVT 90.0 / 90.6; Swin 89.8 / 90.2.
• Augmentation, PASCAL-5i mIoU per fold and FSS-1000 mIoU, 1-shot app. B:

• Decoder ablation, FSS-1000 app. B: convolutional decoder mIoU 87.3 / 88.8, FB-IoU 92.2 / 93.2, mBA 66.8 / 67.2; transformer decoder (ATD) 90.3 / 90.8, 94.0 / 94.4, 68.0 / 68.6 (1-shot / 5-shot).
• VCM vs VEM, FSS-1000 app. B: VEM + Linear Transformer 87.0 / 87.4 mIoU; VEM + Fastformer 87.1 / 87.6; VEM + Swin 89.9 / 90.5; VCM + Linear 87.7 / 88.3; VCM + Fastformer 87.8 / 88.2; VCM + Swin 90.3 / 90.8 (1-shot / 5-shot).

Paper vs code

Compared against the official repository, Seokju-Cho/Volumetric-Aggregation-Transformer, at commit 5534a15ca3808e5d85c742a075debf6d0eac5a45 (default branch, shallow clone). The project page linked from the abstract links the repository, and its owner is an author README.md.

• Training augmentation contradicts the paper. The paper says no augmentation is used §4.1. Every released PASCAL, COCO and FSS config sets CATS_AUGMENTATIONS: true, which applies seven photometric transforms with p=0.2 each config.yaml config.yaml config.yaml dataset.py dataset.py. The appendix says the same augmentation lowers PASCAL mIoU app. B. Which setting produced the reported numbers is not stated; treat both as unresolved.
• K-shot voting uses a different rule. The paper divides the vote sum by K and thresholds at 0.5 §3.6. The code divides by the per-image maximum vote count, clamped to at least 1, then thresholds at 0.5 vat.py. A pixel is therefore foreground when its votes reach half of the image's most-voted count, not half of K.
• Appearance affinity for COCO. The paper excludes conv4_x for PASCAL and COCO §4.1. The code gives COCO (False, False, True), which also drops conv3_x and keeps only conv2_x learner.py. PASCAL matches the paper: conv4 is off, conv3 and conv2 are on learner.py.
• Code-only regularisation. Dropout2d(p=0.5) is applied to the projected conv4 and conv3 appearance maps during training ours.py ours.py. The paper does not mention dropout.
• Convolution implementation. The paper describes overlapping 4D convolutions §3.4.2. The code implements each "4D" conv as a 2D convolution over query axes plus a 2D convolution over support axes, summed our_conv4d.py. This is a factorised form, not a full 4D kernel. Downsampling uses max-pool with kernel equal to stride and ceil_mode=True our_conv4d.py.
• Residuals. Equation 4 has no residual path, but the text adds one eq. 4 §3.4.2. The code adds the residual after the transformer and after adding the upsampled guidance ours.py.
• Model selection uses the evaluation fold. The training script keeps the checkpoint with the best mIoU on the validation split train.py. For PASCAL, both val and test load the same val split of the current fold pascal.py. Reported numbers therefore use a fold that selected the checkpoint. The paper does not state this.
• Not in the paper: batch size 16, 300 or 400 epochs, no weight decay, constant learning rate config.yaml train.py. The decoder uses Swin window size 8 and 3x3 convolutions ours.py.
• Semantic correspondence is not in this branch. The README says that semantic-matching code is on a separate semantic-matching branch README.md. It was not checked here, so the tab. 6 numbers are unverified against code.
• Licence text. The repository's LICENSE is MIT but names a different copyright holder ("Hamacojr") LICENSE. code_license records MIT as stated in the file.
• Paper reference. The README cites arXiv 2112.11685 for the paper and 2207.10866 for the camera-ready version README.md. Only 2207.10866 is used here.

Pitfalls

• Augmentation is ambiguous. The paper says none; the released configs use one photometric set. Check which one you need before comparing with the reported numbers §4.1 config.yaml.
• K-shot rule. Follow the paper's rule or the code's? They differ for K > 1. Choose the one your comparison uses and say so. The code's normalisation depends on the whole image, not only on K vat.py.
• Memory grows with the full volume. The correlation volume is materialised for every query-support pair. A standard transformer needed about 84 GB per batch tab. 5. VAT's memory is 3.8 GB in that table, whose batch size is not stated tab. 5. Training resolution is limited by the volume's 4D size: at 128x128 feature maps the volume is at least 128^4 app. C.
• Window and shift details. Swin windows are 4 on every axis §4.1. Shifts are window/2 on every second block swin4d.py. The number of layers per correlation group depends on the backbone: 3, 6 or 23, and 4 vat.py.
• Interpolation. Masks and feature maps use bilinear resizing with align_corners=True throughout vat.py ours.py. Other choices can shift results by small amounts.
• Epsilon in normalisation. The cosine correlation adds 1e-5 to the norm correlation.py. Low-norm features are affected.
• Ablation backbone not stated. The component table tab. 4 gives no backbone. Do not assume ResNet-101 for other tables without checking.
• Author-reported results. Every number here is from the paper. The per-fold and mBA numbers depend on the selection issue above.
• Figure and table numbers. The numbers come from the source's order of \begin{table} and \begin{figure} environments, not from a compiled PDF. The aggregator table is numbered tab. 5 on that basis.
• Appendix cross-reference. app. B refers to an "Algorithm 1" for the VEM that is not present in the source.

Lineage

• Builds on: Swin Transformer, arXiv:2103.14030 (shifted windows, relative position bias) §3.4.1. The 4D extension and the overlapping early convolutions are this paper's changes §3.4.1 §3.4.2.
• Builds on: HSNet, arXiv:2104.01538 (4D convolutions over correlation, ResNet features, pyramid) §1 §3.3 §3.4.3 §4.1. VAT replaces the fixed 4D kernels with transformer aggregation.
• Builds on: CATs (Cho et al., NeurIPS 2021, cited as cho2021semantic), a transformer for cost aggregation in semantic correspondence, which flattens the volume and is the main comparison for the semantic task §1 §3.4.1 tab. 6.
• Compared against: PFENet, ASGNet, RePRI, CyCTR, CWT, DAN, FWB, PANet, PMM, RPMM, FSOT tab. 1 tab. 2 tab. 3; for semantic correspondence, NC-Net, SCOT, CHM, MMNet, PMNC, DHPF, CATs tab. 6.
• Aggregator variants compared: standard transformer, MLP-Mixer, centre-pivot 4D convolution (HSNet), Linear Transformer, and Fastformer, arXiv:2108.09084 tab. 5. The bibliography gives an arXiv id only for Fastformer.
• Superseded by: not known from content in this wiki.