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 contributes two techniques to improve generalization of semantic segmentation networks. The first technique is an adaptation of the test-time behaviour of batch normalization, where the statistics of the sample under consideration are blended into the training-time statistics of the batch-norm layer. The second technique is to use test time augmentations to first derive a pseudo-labeling from the sample and then to update parts of the weights to maximize the probability of the pseudo-labels. The paper additionally contributes a multi-dataset evaluation procedure and shows favorable performance of the proposed techniques on this evaluation protocol.