Test-Time Training for Semantic Segmentation with Output Contrastive Loss

TL;DR

This paper introduces Output Contrastive Loss (OCL), a novel test-time training method for semantic segmentation that stabilizes adaptation and improves generalization to unseen environments by adapting contrastive loss to the output space.

Contribution

It adapts contrastive loss to the output space for test-time training in semantic segmentation, enhancing stability and generalization during online adaptation.

Findings

OCL improves segmentation performance across diverse scenarios.

The method is effective even with models pre-trained via domain adaptation.

OCL demonstrates robustness and adaptability in unseen environments.

Abstract

Although deep learning-based segmentation models have achieved impressive performance on public benchmarks, generalizing well to unseen environments remains a major challenge. To improve the model's generalization ability to the new domain during evaluation, the test-time training (TTT) is a challenging paradigm that adapts the source-pretrained model in an online fashion. Early efforts on TTT mainly focus on the image classification task. Directly extending these methods to semantic segmentation easily experiences unstable adaption due to segmentation's inherent characteristics, such as extreme class imbalance and complex decision spaces. To stabilize the adaptation process, we introduce contrastive loss (CL), known for its capability to learn robust and generalized representations. Nevertheless, the traditional CL operates in the representation space and cannot directly enhance predictions. In this paper, we resolve this limitation by adapting the CL to the output space, employing a high temperature, and simplifying the formulation, resulting in a straightforward yet effective loss function called Output Contrastive Loss (OCL). Our comprehensive experiments validate the efficacy of our approach across diverse evaluation scenarios. Notably, our method excels even when applied to models initially pre-trained using domain adaptation methods on test domain data, showcasing its resilience and adaptability.\footnote{Code and more information could be found at~ \url{https://github.com/dazhangyu123/OCL}}