📄 Semantically Aware Optimal Transport for Dense Label Transfer

Summary

The paper addresses dense label transfer for semantic segmentation: projecting per-pixel class labels from a labeled source domain onto an unlabeled target domain while the vision transformer backbone stays frozen. Standard optimal-transport (OT) baselines match patches with a fixed similarity measure. Under domain shift, this confuses lookalike but semantically different regions, ignores class-frequency differences and scales poorly to thousands of tokens. The authors propose SAOT, which learns the transport cost itself through bilevel training with sparse target point labels. The cost mixes a trainable appearance metric with a penalty based on source class prototypes. It is plugged into an entropic, unbalanced OT problem solved over k-nearest-neighbour blocks. The resulting pseudo-labels train a residual MLP decoder, which is then refined by EMA-teacher self-training. On GTA5→Cityscapes with frozen DINOv2 ViT-L/14, the full pipeline reaches 25.7% mIoU, about 3.8× nearest-neighbour transfer. The authors report this as the best result among frozen-backbone methods. This shows that learning the matching cost, not the transport machinery alone, drives the gains.

Key points

• Setup: GTA5→Cityscapes at 504×1008 resolution (2,592 tokens per image) with frozen DINOv2 ViT-L/14. The block-sparse solver uses k = 128 neighbours, cutting coupling size from O(nm) to O(nk).
• Main results: Fixed-cost OT, UOT and RWOT score only 6.5–6.8% mIoU, while balanced OT collapses to 0%. Raw SAOT reaches 12.6%, rising to 15.8% when aggregating 3 source images. SAOT+Decoder reaches 25.7% and labels 500 images in 1.2 s.
• Per-class results: Road scores 90.3%, car 76.2% and building 71.5%. Six small or rare classes stay at 0% IoU, including traffic sign, rider and bicycle.
• Refinement and transport behaviour: EMA refinement with hard labels and label smoothing lifts mIoU from 23.7% to 25.7%, whereas soft KL targets drop it to 9.9%. The learned cost cuts mass sent to wrong classes from 68% to 41%.
• Second benchmark and context: On PASCAL VOC 2012 train→val with ViT-B/16, the full pipeline scores 47.5% versus KNN-Global at 45.6%. Fully fine-tuned UDA methods (DAFormer, HRDA, MIC) reach 68–76% and Rein reaches 78.4%. SAOT uses a 5.3M-parameter head and trains in about 6 GPU-hours.