Stable Test-Time Training for Semantic Segmentation with Output Contrastive Loss

Abstract

Deep learning-based models have achieved impressive performance on public segmentation benchmarks, yet generalizing to unseen environments remains challenging. Test-time training (TTT) addresses this by adapting source-pretrained models during evaluation. While existing TTT methods have shown promise in image classification, they often exhibit instability with small test batches and class imbalance—challenges that intensify in semantic segmentation tasks. To tackle this issue, we present Output Contrastive Loss (OCL) to improve the stability of contrastive loss when applied to TTT for segmentation. OCL applies contrastive loss directly to the output space, avoiding the need for extra regularization, and employs a high temperature to prevent model collapse. To further stabilize the TTT process, we integrate BN statistics Modulation and Stochastic Restoration techniques. Extensive experiments across diverse datasets, settings, architectures, and pretrained methods demonstrate consistent performance improvements, achieving a 7.5 mIoU gain on the GTA→CS benchmark and showing effectiveness even with domain adaptation pretraining. Code is available at https://github.com/dazhangyul23/OCL.