PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sunghwan Hong, Seokju Cho, Jisu Nam, Stephen Lin, Seungryong Kim

arXiv:2207.10866v1cs.CV

TL;DR

Few-shot segmentation must aggregate sparse support-query matching information despite challenges from correlation-map tokenization and limited local or long-range context. VAT addresses this with a 4D Convolutional Swin Transformer, pyramidal guidance, and an appearance-aware decoder. It achieves state-of-the-art performance on standard few-shot segmentation benchmarks and semantic correspondence, with higher memory and runtime costs and substantial correlation-map storage requirements.

  • Problem

    Few-shot segmentation needs reliable cost aggregation from support-query correlation maps, but tokenization can reduce local context and inductive bias while prior aggregators have limited contextual coverage.

  • Method

    VAT uses a 4D Convolutional Swin Transformer preceded by overlapping small-kernel convolutions, coarse-to-fine pyramidal guidance, and a query-appearance-aware decoder.

  • Results

    VAT attains state-of-the-art performance on all standard few-shot segmentation benchmarks and also on semantic correspondence.

  • Takeaways & Limitations

    The results support cost aggregation as a central component for both few-shot segmentation and semantic correspondence, with potential for general matching.

  • Takeaways & Limitations

    VAT is relatively slower and uses more memory; explicitly storing global correlation maps becomes memory-expensive as spatial resolution increases and can hinder training with sufficient batch size.

Abstract

from arXiv · show

This paper presents a novel cost aggregation network, called Volumetric Aggregation with Transformers (VAT), for few-shot segmentation. The use of transformers can benefit correlation map aggregation through self-attention over a global receptive field. However, the tokenization of a correlation map for transformer processing can be detrimental, because the discontinuity at token boundaries reduces the local context available near the token edges and decreases inductive bias. To address this problem, we propose a 4D Convolutional Swin Transformer, where a high-dimensional Swin Transformer is preceded by a series of small-kernel convolutions that impart local context to all pixels and introduce convolutional inductive bias. We additionally boost aggregation performance by applying transformers within a pyramidal structure, where aggregation at a coarser level guides aggregation at a finer level. Noise in the transformer output is then filtered in the subsequent decoder with the help of the query's appearance embedding. With this model, a new state-of-the-art is set for all the standard benchmarks in few-shot segmentation. It is shown that VAT attains state-of-the-art performance for semantic correspondence as well, where cost aggregation also plays a central role.

1 Introduction

Few-shot segmentation reduces reliance on pixel-wise labels by predicting query masks from a handful of support images and masks. VAT reformulates the task as semantic correspondence and introduces transformer-based cost aggregation to address limitations of prototype, raw-correlation, and local-convolution approaches.

  • Motivation: Few-shot segmentation predicts a query image’s pixel labels using only a handful of support images and associated masks.This setting reduces reliance on manually labeled segmentation maps when adding new classes.
  • Motivation: Prototype-based methods can disregard pixel-level support-query relationships and feature spatial structure, producing sub-optimal results.
  • Motivation: Semantic correspondence provides a reformulation that seeks pixel-level matches across semantically similar images despite appearance and geometric variation.Its classical pipeline includes feature extraction, cost aggregation, and flow estimation.
  • Motivation: Raw correlation maps can remain ambiguous under repetitive patterns or background clutter, while HSNet’s 4D convolutions cannot aggregate long-range context.
  • Method: VAT introduces a 4D Convolutional Swin Transformer that combines high-dimensional self-attention with overlapping small-kernel convolutions.The convolutions provide local context at every pixel and convolutional inductive bias before transformer processing.
  • Results: VAT attains state-of-the-art performance across standard few-shot segmentation benchmarks and also for semantic correspondence.The authors identify cost aggregation as important for both tasks and potentially useful for general matching.

2 Related Work

Prior few-shot segmentation methods use prototypes, correlation maps, graph attention, or feature-refining transformers, while semantic correspondence emphasizes reliable matching under substantial variation. VAT differs by applying a 4D Convolutional Swin Transformer directly to cost aggregation.

  • Few-shot Segmentation: Prototype-based few-shot segmentation extracts support-derived prototypes to identify foreground features in query images.
  • Few-shot Segmentation: Correlation-map methods preserve pairwise support-query relationships, using graph attention, 4D convolutions, or optimal transport for feature matching.
  • Few-shot Segmentation: Some transformer-based approaches refine feature maps with attention or enhancement modules without aggregating matching scores.
  • Few-shot Segmentation: VAT instead proposes a 4D Convolutional Swin Transformer for cost aggregation.
  • Semantic Correspondence: Semantic correspondence seeks matches between semantically similar images despite large intra-class appearance and geometric variations.Its latest methods focus on cost aggregation to obtain reliable correspondences.
  • Vision Transformer: Vision transformers are widely used in computer vision, but their quadratic sequence complexity creates substantial computational burden.Swin Transformer is presented among efficient transformer approaches intended to reduce this load.

3 Methodology

VAT formulates few-shot segmentation around pixel-level support–query matching and aggregates the resulting 4D correlation maps with a pyramidal transformer encoder. Its 4D convolutional embedding, shifted-window attention, guided coarse-to-fine processing, and affinity-aware decoder address local-context, long-range, and upsampling challenges.

  • 3.1 Problem Formulation: Few-shot segmentation predicts masks for unseen query classes from only a few annotated support examples, using episodic training to mitigate overfitting.The model learns mappings from support sets and query images to query masks.
  • 3.1 Problem Formulation: Prototype-based methods can disregard pixel-level support–query relationships and feature spatial structure, producing sub-optimal segmentation.Pairwise matching is therefore used to represent correspondence more explicitly.
  • 3.3 Feature Extraction and Cost Computation: Query and masked-support features are extracted at multiple levels, then normalized inner products form stacked 4D correlation maps across spatially matched feature levels.Support masks filter background information before correlation computation.
  • 3.4 Pyramidal Transformer Encoder: The Volumetric Transformer Module extends Swin self-attention to 4D correlation maps using partitioned windows and shifted windows for inter-window interactions.This reduces token-processing demands while retaining local-window attention and cross-window communication.
  • 3.4 Pyramidal Transformer Encoder: The Volumetric Convolution Module replaces non-overlapping embedding with overlapping small-kernel 4D convolutions, adding local context and convolutional inductive bias before transformer aggregation.The design targets boundary discontinuities, limited local context, and weak inductive bias from non-overlapping operations.
  • 3.4 Pyramidal Transformer Encoder: Guided Pyramidal Processing upsamples each coarser aggregated correlation map and adds it to the next finer level as guidance.The process continues to the finest level before decoding.
  • 3.5 Affinity-Aware Transformer Decoder: The Affinity-aware Transformer Decoder averages support dimensions, concatenates the result with query appearance embeddings, and refines upsampled predictions.Appearance affinity filters matching-score noise, while higher-resolution structure helps preserve fine details.

4 Experiments

VAT is evaluated on standard few-shot segmentation and semantic correspondence benchmarks, with ablations examining its components, aggregators, convergence, and efficiency. It achieves state-of-the-art results across the reported benchmarks while incurring modest memory and runtime costs.

  • Few-shot Segmentation Results: VAT sets a new state-of-the-art on PASCAL-5i, COCO-20i, and FSS-1000 few-shot segmentation benchmarks.On PASCAL-5i it outperforms competing methods on almost all folds for mIoU and FB-IoU and also leads in mBA.
  • Few-shot Segmentation Results: 4.6% higher mIoU than HSNet on FSS-1000 with ResNet50 demonstrates VAT’s large reported margin on that dataset.VAT also outperforms FSOT on FSS-1000.
  • Ablation Study: Each VAT component improves performance, while Swin Transformer aggregation and VCM produce especially substantial gains over the baseline.Adding ATD yields a large mBA gain by using higher-resolution appearance structure to refine fine details.
  • Ablation Study: VAT quickly exceeds HSNet during convergence despite starting at a lower mIoU.The ablation attributes this convergence behavior to the early convolutions.
  • Ablation Study: VAT outperforms alternative aggregators by a large margin, while Linear Transformer and Fastformer achieve similar performance.Standard Transformer requires approximately 84 GB per batch, making it infeasible in the reported setting; MLP-Mixer requires substantially more memory.
  • Efficiency: VAT uses 0.3 GB more memory and runs 5 ms slower than alternatives, which the authors describe as a minor sacrifice for better performance.The comparison uses one NVIDIA GeForce RTX 3090 GPU and an Intel Core i7-10700 CPU.
  • Semantic Correspondence Results: On semantic correspondence benchmarks, VAT sets a new state-of-the-art or achieves the second-highest PCK, supporting cost aggregation across both tasks.PF-PASCAL is described as nearly saturated, whereas VAT has a large margin over other methods on PF-WILLOW.

5 Conclusion

The paper presents VAT, a cost aggregation network built around a 4D Convolutional Swin Transformer, pyramidal processing, and appearance-guided decoding. It reports state-of-the-art performance on standard benchmarks for both few-shot segmentation and semantic correspondence.

  • Conclusion: VAT addresses correlation-map tokenization issues with a 4D Convolutional Swin Transformer preceded by small-kernel convolutions.The convolutions provide local context and convolutional inductive bias.
  • Conclusion: Pyramidal transformer processing and appearance-guided decoding are added to boost aggregation and filter transformer outputs.Coarser aggregation guides finer aggregation in the pyramidal structure.
  • Conclusion: VAT attains state-of-the-art performance across standard benchmarks for few-shot segmentation and semantic correspondence.The conclusion identifies cost aggregation as central to both tasks.

Appendix

The appendix provides additional experimental and qualitative material for the benchmarks used in the paper.

  • Appendix: The appendix includes experimental settings, further ablations, additional semantic-correspondence results, and qualitative results across all benchmarks.The listed semantic-correspondence benchmarks include SPair-71k, PF-PASCAL, and PF-WILLOW.

Appendix A. Experimental Setting for Semantic Correspondence

Semantic correspondence experiments use standard benchmarks and evaluate VAT with a task-specific objective while retaining the model components. The evaluation uses PCK and AEPE-related flow supervision under provided keypoints.

  • Datasets: The semantic-correspondence evaluation covers SPair-71k, PF-PASCAL, and PF-WILLOW under a common protocol.SPair-71k contains 70,958 image pairs, while PF-PASCAL and PF-WILLOW are smaller-scale datasets.
  • Evaluation Metric: PCK measures the ratio of predicted keypoints lying within a thresholded distance of their ground-truth locations.The threshold is α·max(H, W), where α is the scale parameter and H and W are image dimensions.
  • Implementation Details: For semantic correspondence, VAT retains its components but uses a different objective function with provided ground-truth keypoints.The implementation computes AEPE by averaging Euclidean distances between ground-truth and estimated flow, using an L2 loss.

Appendix B. Additional Ablation Study

The appendix evaluates backbone, augmentation, decoder, and convolutional embedding choices. Results favor transformer-based cost aggregation with overlapping convolutions, while data augmentation can substantially reduce few-shot segmentation performance.

  • Feature Backbone: Convolution- and transformer-based backbones attain similar performance in the reported pyramidal-backbone comparison.The authors conjecture that pretraining on the same dataset produces nearly alike representations, while noting other backbone structures may differ.
  • Data Augmentation: Data augmentation severely affects overall few-shot segmentation performance, with CATs augmentation performing worst on PASCAL-5i and second worst on FSS-1000.The authors caution that the small 0.3% gap to the best FSS-1000 result could change on retrial.
  • Data Augmentation: Augmentation effects may differ between semantic correspondence and few-shot segmentation because VAT aggregates correlation maps whose matching distributions can be altered by image transformations.The authors also conjecture that combining augmentations increases learning difficulty and can reduce accuracy.
  • ATD: A 3% mIoU gap separates the transformer-equipped ATD from the convolution-only AD, with similar differences for FB-IoU and mBA.The comparison supports more effective aggregation, filtering, and integration of correlation maps with appearance embeddings when transformers are used.
  • VCM: Swin Transformer achieves the best VCM results, while Linear Transformer and Fastformer show similar results.Replacing VCM with VEM gives Swin Transformer 89.9% mIoU, a 0.4% drop; the authors relate its smaller decline to relative position bias.

Appendix C. Limitations

The principal limitation is the memory cost of explicitly computing and storing global correlation maps. This cost grows with spatial resolution and can prevent training with sufficient batch size.

  • Memory Cost: Global correlation maps require substantial memory because a 128×128 feature resolution produces correlation maps of at least 128^4.Including level dimensions and other pyramidal levels makes training difficult even on NVIDIA GeForce RTX-3090 GPUs.
  • Practical Boundary: The coarse-to-fine architecture does not make training feasible at high resolutions with a sufficient batch size.The authors identify this computational burden as a limitation on the approach’s accessibility.

Appendix D. More Results

Additional results cover semantic correspondence, mBA comparisons, and qualitative evaluations across few-shot segmentation and correspondence benchmarks. The appendix reports per-class and per-fold analyses alongside visual examples.

  • Semantic Correspondence: Table 7 provides per-class quantitative evaluation on SPair-71k against a broad set of semantic correspondence methods.The comparison includes CNNGeo, WeakAlign, NC-Net, HPF, SFNet, DCC-Net, GSF, SCOT, DHPF, CHM, MMNet, PMNC, and CATs.
  • mBA Comparisons: Tables 8 and 9 provide per-fold mBA comparisons on PASCAL-5i and COCO-20i.The authors obtained HSNet and RePRI results using released pretrained weights and code, while omitting CyCTR because official code and weights were unavailable.
  • Qualitative Results: Qualitative results are reported for PASCAL-5i, COCO-20i, FSS-1000, PF-PASCAL, PF-WILLOW, and SPair-71k.The visual evaluations correspond to Figures 2 through 6.
Loading 2207.10866v1…