Adaptive Generalization for Semantic Segmentation

1 INTRODUCTION . The current state of the art for semantic segmentation ( Long et al. , 2015 ; Chen et al. , 2018b ) lacks direly in out-of-distribution robustness , i. e. when the training and testing distributions are different . Numerous studies have investigated this issue , with a primary focus on image classification ( Arjovsky et al. , 2019 ; Bickel et al. , 2009 ; Li et al. , 2017a ; Torralba and Efros , 2011 ; Volpi et al. , 2018 ) . However , the conclusion from a recent study of domain generalization methods ( Gulrajani and Lopez-Paz , 2020 ) , spanning more than three years of research , is disillusioning : Empirical Risk Minimization ( ERM ) , which is based on the i. i. d. assumption of the training and testing distributions , is still highly competitive . This is in stark contrast to the evident advances in the area of domain adaptation , both for image classification ( Ben-David et al. , 2010 ; Ganin et al. , 2016 ; Long et al. , 2016 ; Xie et al. , 2018 ) and semantic segmentation ( Araslanov and Roth , 2021 ; Yang and Soatto , 2020 ; Vu et al. , 2019 ) . This setup , however , assumes access to an unlabelled test distribution at training time , whereas in the generalization setting considered here , only one test sample is accessible at inference time , and no knowledge between the subsequent test samples must be shared . In this work , we study the generalization problem of semantic segmentation from synthetic data ( Richter et al. , 2016 ; Ros et al. , 2016 ) through the lens of adaptation . Instead of modifying the model architecture ( Pan et al. , 2018 ) or the training process ( Chen et al. , 2020 ; 2021 ; Yue et al. , 2019 ) , we enhance the inference procedure with two orthogonal techniques inspired by domain adaptation methods ( Araslanov and Roth , 2021 ; Li et al. , 2017b ) . The first mechanism leverages normalization layers ( Ioffe and Szegedy , 2015 ) in modern convolutional neural networks ( CNNs ) ( He et al. , 2016 ) by integrating feature statistics of the test sample . Expanding upon previously attained conclusions for the image classification task ( Schneider et al. , 2020 ) , we find that this strategy not only improves the segmentation accuracy , but also the calibration quality of the prediction confidence . Our second contribution is a self-supervised loss , allowing the model to adapt to a single test sample with a few parameter updates . Thirdly , we set up an extensive empirical study following a rigorous evaluation protocol , allowing us to establish that the two techniques are complementary . Our study of this one-sample adaptation process reveals a consistent improvement in out-of-distribution robustness over the baseline in all benchmark scenarios , and yields a new state of the art in segmentation accuracy , substantially surpassing that of previous work in virtually all considered settings . 2 RELATED WORK . Our work contributes to recent research on generalization of semantic segmentation models , and relates to the studies on feature normalization ( Pan et al. , 2018 ; Schneider et al. , 2020 ) and test-time training ( Sun et al. , 2020 ) . While the focus in previous investigations was the training strategy ( Yue et al. , 2019 ) and model design ( Pan et al. , 2018 ) , we exclusively study the test-time inference process here . Yue et al . ( 2019 ) augmented the synthetic training data by transferring style from real images . Assuming access to a classification model trained on real images , Chen et al . ( 2020 ) regularize the training on synthetic data by ensuring feature proximity of the two models via distillation , and seek layer-specific learning rates for improved generalization . Advancing the distillation technique , Chen et al . ( 2021 ) devise a contrastive loss that facilitates model invariance to standard image augmentations . Pan et al . ( 2018 ) heuristically add instance normalization ( IN ) layers to the network . We remark on two limitations of these works that we address here : First , these methods assume access to a distribution of real images during training ( Chen et al. , 2020 ; 2021 ; Yue et al. , 2019 ) ( as opposed to only for pre-training of the backbone ) , or require a modification of the network architecture ( Pan et al. , 2018 ) . Our work requires neither , hence the presented techniques apply even post-hoc to the already ( pre- ) trained models to improve their generalization . Second , as we discuss and address in Sec . 5 , the evaluation protocol used by previous studies exhibits a number of shortcomings . Normalization . Batch Normalization ( BN ; Ioffe and Szegedy , 2015 ) and other normalization techniques have been increasingly linked to model robustness ( Deecke et al. , 2019 ; Huang et al. , 2019 ; Schneider et al. , 2020 ; Wang et al. , 2018 ; Wu and He , 2018 ) . The most commonly used BN , Layer Normalization ( LN ; Ba et al. , 2016 ) , and Instance Normalization ( IN ; Ulyanov et al. , 2016 ) also affect the model ’ s expressive power , which can be further enhanced by their arrangements ( Nam and Kim , 2018 ; Luo et al. , 2019 ) . In a domain adaptation setting , Li et al . ( 2017b ) use source-domain statistics during training while replacing them with target-domain statistics during inference . More recently , Schneider et al . ( 2020 ) combine the source and target statistics during inference , but the statistics are weighted depending on the number of samples that these statistics aggregate . Nado et al . ( 2020 ) propose using batch statistics during inference from the target domain instead of the training statistics acquired from the source domain . Our comprehensive empirical study complements these results by demonstrating improved generalization of semantic segmentation models . Test-time training . Test-time training ( TTT ) refers to updating the model parameters also at inference time with a self-supervised loss incurred on a single unlabelled test sample . This technique has been recently applied with success for improving the robustness of image classification models ( Sun et al. , 2020 ; Wang et al. , 2021a ) . The design of the self-supervised task is crucial , and the techniques developed for image classification are unsuitable for dense prediction tasks , such as semantic segmentation . Nevertheless , recent work explored such losses in domain adaptation scenarios ( Araslanov and Roth , 2021 ) , and a number of other works exploit domain-specific knowledge from medical imaging ( Varsavsky et al. , 2020 ) or first-person vision ( Cai et al. , 2020 ) . Setup comparison . Most of these technically related works ( Schneider et al. , 2020 ; Sun et al. , 2020 ; Wang et al. , 2021a ) focus on the problem of domain adaptation in the context of image classification . They typically assume access to a number of samples ( or even all test images ) from the target distribution at training time . Our work instead addresses semantic segmentation in the domain generalization setting , which is fundamentally different as it only necessitates a single datum from the test set . In this scenario , simple objectives , such as entropy minimization employed by Tent ( Wang et al. , 2021a ) , improve the baseline accuracy only moderately . By contrast , our one-sample adaptation with pseudo labels accounts for the inherent uncertainty in the predictions , which proves substantially more effective , as the comparison to Tent in Sec . 5.3 reveals . Our task is also different from few-shot learning ( e. g. , ( Finn et al. , 2017 ) ) , where the model may adapt at test time using a small annotated set of image samples . No such annotation is available in our setup ; our model adjusts to the test sample in a completely unsupervised fashion , requires neither proxy tasks to update the parameters ( Sun et al. , 2020 ) nor any knowledge of the test domain . 3 INSTANCE-ADAPTIVE BATCH NORMALIZATION . Batch Normalization ( BN ; Ioffe and Szegedy , 2015 ) has become an inextricable component of modern CNNs ( He et al. , 2016 ) . Although BN was originally designed for improving training convergence , there is now substantial evidence that it plays an important role in model robustness ( Nado et al. , 2020 ) , including domain generalization ( Pan et al. , 2018 ) . Let z ∈ RB , C , H , W denote a feature tensor , with C channels and batch size B , produced by a convolutional layer at resolution H ×W . We omit the layer indexing , as the following presentation applies to all layers in which BN is applied . At training time , BN first computes the mean and the standard deviation for each of the C channels batch-wise , i. e. µc = 1 BHW ∑ i , j , k zi , c , j , k , σ 2 c = 1 BHW ∑ i , j , k ( zi , c , j , k − µc ) 2 . ( 1 ) The normalized features ẑ follow from applying these statistics , ẑi , c , j , k = ( zi , c , j , k−µc ) / √ σ2c+ϵ . ( 2 ) Notably , this process differs from the normalization used at inference time . At training time , every BN layer maintains a running estimate of µc and σc across the training batches , which we denote here as µ̂c and σ̂c . At test time , it is an established practice to normalize w. r. t. µ̂c and σ̂c in Eq . ( 2 ) , instead of the test-batch statistics , a scheme we refer to as train BN ( t-BN ) . BN and generalization . In the context of out-of-distribution generalization , the running statistics µ̂c and σ̂c derive from the source data and can be substantially different had they been computed using the target images . This discrepancy is generally known as the covariate shift problem . Domain adaptation methods , which assume access to the ( unlabelled ) target distribution , often alleviate this issue with a technique typically referred to as Adaptive Batch Normalization ( AdaBN ; Li et al. , 2017b ) . The key idea behind the method is to simply replace the source running statistics with those of the target . Instead of alternating BN layers between the training and testing modes , recent work ( Nado et al. , 2020 ) studies prediction-time BN ( p-BN ) , which maintains BN layers in ªtraining modeº also at inference time . That is , p-BN replaces the running statistics from the training time , µ̂c and σ̂c , with the statistics µc and σc of the current test batch . Such a seemingly innocuous change is shown to benefit model robustness for image classification ( Nado et al. , 2020 ; Schneider et al. , 2020 ) . In contrast to AdaBN and p-BN , which utilize either the whole target distribution or a number of samples , our study of model generalization assumes only a single target example to be available Ð the one that our model receives as the input at inference time . A viable alternative to AdaBN and p-BN is to compute the statistics per sample , which amounts to replacing BN layers with Instance Normalization ( IN ) layers ( Ulyanov et al. , 2016 ) after model training . However , this may cause another extreme scenario for covariate shift , since such replacement may significantly interfere with the statistics of the activations in the intermediate layers with which the network was trained . Moreover , previous work showed that IN layers hurt the discriminative power of the model and can improve model robustness only in combination with BN ( Pan et al. , 2018 ) , thereby requiring changes in the model architecture . Instance-adaptive batch normalization . Following Schneider et al . ( 2020 ) , we compute a weighted average of the running statistics from training on source data and the target statistics , which in contrast to Schneider et al . ( 2020 ) always derives from a single sample in our work . Let µ̂ ( s ) c and σ̂ ( s ) c denote the running estimate of the mean and the standard deviation at training time on the source data ( for feature channel c ) . Similarly , define µ ( t ) c and σ ( t ) c as the mean and the standard deviation computed from a feature tensor z ( t ) ∈ RC , H , W , coming from a single target sample : µ ( t ) c = 1 HW ∑ j , k z ( t ) c , j , k , σ ( t ) 2 c = 1 HW ∑ j , k ( z ( t ) c , j , k − µ ( t ) c ) 2 . ( 3 ) At inference time , we compute the new mean and standard deviation , µ ( ∗ ) c and σ ( ∗ ) c , as follows : µ ( ∗ ) c = αµ ( t ) c + ( 1− α ) µ̂ ( s ) c , σ ( ∗ ) 2 c = ασ ( t ) 2 c + ( 1− α ) σ̂ ( s ) c . ( 4 ) We then use µ ( ∗ ) c and σ ( ∗ ) c in place of µc and σc in Eq . ( 2 ) to normalize the features . Note that this does not affect the behaviour of the BN layers at training time and applies only at test time . Since this approach combines the inductive bias coming in the form of the running statistics from the source domain with statistics extracted from a single test instance , we refer to this technique as Instance-adaptive Batch Normalization ( IaBN ) . Discussion . Setting α = 0 in Eq . ( 4 ) defaults to the established procedure , t-BN , and uses only the running statistics from training on the source domain at inference time . Conversely , α = 1 corresponds to Instance Normalization ( Ulyanov et al. , 2016 ) or , equivalently , to the p-BN strategy ( Nado et al. , 2020 ) with a batch size of 1 . This is in contrast to ( Nado et al. , 2020 ) , where µ ( t ) c and σ ( t ) c are averaged over a batch of target images , whereas we rely on a single test sample in this work . Our experiments in Sec . 5.1 expand upon previous analyses ( Schneider et al. , 2020 ) with an extensive emperical study of this normalization strategy for semantic segmentation . Note that Batch Instance Normalization ( Nam and Kim , 2018 ) used a similar weighting approach for adaptive stylization . However , α was a training parameter , whereas we use α only for model selection using a validation set . In Sec . 5 we empirically verify that IaBN yields a consistent boost of the segmentation accuracy in the out-of-distribution scenario . Perhaps more surprisingly , we also find significant improvements in model calibration ( in terms of the expected calibration error , ECE ( Naeini et al. , 2015 ) ) . We leverage this further in devising a self-supervised approach for test-time training of the segmentation model .

Summary: This paper presents a method on generalization of segmentation from synthetic data to real street scene data. To adapt the model pre-trained with synthetic source domain data such as GTA and SYNTHIA to target domain such Cityscapes, BDD and IDD, the authors proposes Instance-adaptive Batch Normalization (IaBN). In addition, as a method to learn from a single sample, TT-SEG was proposed in which a pseudo GT mask is estimated from augmented images with the initial model and the last parts of the network are finetuned with the pseudo GT mask. In the experiments, the proposed method which employs both IaBN and TT-SEG successfully outperformed all the baselines, which indicated the effectiveness of the proposed method.