Conditional Networks

1 INTRODUCTION . Deep learning has achieved great success in many core artificial intelligence ( AI ) tasks ( Hinton et al. , 2012 ; Krizhevsky et al. , 2012 ; Brown et al. , 2020 ) over the past decade . This is often attributed to better computational resources ( Brock et al. , 2018 ) and large-scale datasets ( Deng et al. , 2009 ) . Collecting and annotating datasets which represent a sufficient diversity of real-world test scenarios for every task or domain is extremely expensive and time-consuming . Hence , sufficient training data may not always be available . Due to many factors of variation ( e.g. , weather , season , daytime , illumination , view angle , sensor , and image quality ) , there is often a distributional change or domain shift that can degrade performance in real-world applications ( Shimodaira , 2000 ; Wang & Schneider , 2014 ; Chung et al. , 2018 ) . Applications in remote sensing , medical imaging , and Earth observation commonly suffer from distributional shifts resulting from atmospheric changes , seasonality , weather , use of different scanning sensors , different calibration and other variations which translate to unexpected behavior at test time ( Zhu et al. , 2017 ; Robinson et al. , 2019 ; Ortiz et al. , 2018 ) . In this work , we present a novel neural network architecture to increase robustness to distributional changes ( See Figure 1 ) . Our framework combines conditional computation ( Dumoulin et al. , 2018 ; 2016 ; De Vries et al. , 2017 ; Perez et al. , 2018 ) with a task specific neural architecture for better domain shift generalization . One key feature of this architecture is the ability to exploit extra information , often available but seldom used by current models , through a conditioning network . This results in models with better generalization , better performance in both independent and identically distributed ( i.i.d . ) and non- i.i.d . settings , and in some cases faster convergence . We demonstrate these methodological innovations on an aerial building segmentation task , where test images are from different geographic areas than the ones seen during training ( Maggiori et al. , 2017 ) and on the task of Tumor Infiltrating Lymphocytes ( TIL ) classification ( Saltz et al. , 2018 ) . We summarize our main contributions as follows : • We propose a novel architecture to effectively incorporate conditioning information , such as metadata . • We show empirically that our conditional network improves performance in the task of semantic segmentation and image classification . • We study how conditional networks improve generalization in the presence of distributional shift . 2 BACKGROUND AND RELATED WORK . Self-supervised learning . Self-supervised learning extracts and uses available relevant context and embedded metadata as supervisory signals . It is a representation learning approach that exploits a variety of labels that come with the data for free . To leverage large amounts of unlabeled data , it is possible to set the learning objectives such that supervision is generated from the data itself . The selfsupervised task , also known as pretext task , guides us to a supervised loss function ( Gidaris et al. , 2018 ; Oord et al. , 2018 ; He et al. , 2019 ; Chen et al. , 2020 ) . However , in self-supervised learning we usually do not emphasize performance on this auxiliary task . Rather we focus on the learned intermediate representation with the expectation that this representation can carry good semantic or structural meanings and can be beneficial to a variety of practical downstream tasks . Conditional networks can be seen as a self-supervision approach in which the pretext task is jointly learned with the downstream task . Our proposed modulation of a network architecture based on an auxiliary network ’ s intermediate representation can also be seen as an instance of knowledge transfer ( Hinton et al. , 2015 ; Urban et al. , 2016 ; Buciluǎ et al. , 2006 ) . Because the auxiliary network has an additional task signal – metadata prediction – information about this task can be transferred to the main task network . Conditional Computation . Ioffe and Szegedy designed Batch Normalization ( BN ) as a technique to accelerate the training of deep neural networks ( Ioffe & Szegedy , 2015 ) . BN normalizes a given mini-batch B = { Fn , . , . , . } Nn=1 of N feature maps Fn , . , . , . as described by the following Equation : BN ( Fn , c , h , w|γc , βc ) = γc Fn , c , h , w − EB [ F. , c , . , . ] √ VarB [ F. , c , . , . ] + + βc , ( 1 ) where c , h and w are indexing the channel , height and width axis , respectively , γc and βc are trainable scale and shift parameters , introduced to keep the representational power of the original network , and is a constant factor for numerical stability . For convolutional layers the mean and variance are computed over both the batch and spatial dimensions , implying that each location in the feature map is normalized in the same way . De Vries et al . ( 2017 ) ; Perez et al . ( 2018 ) introduced Conditional Batch Normalization ( CBN ) as a method for language-vision tasks . Instead of setting γc and βc in Equation 1 directly , CBN defines them as learned functions βn , c = βc ( qn ) and γn , c = γc ( qn ) of a conditioning input qn . Note that this results in a different scale and shift for each sample in a mini-batch . Scale ( γn , c ) and shift ( βn , c ) parameters for each convolutional feature are generated and applied to each feature via an affine transformation . Feature-wise transformations frequently have enough capacity to model complex phenomena in various settings ( Dumoulin et al. , 2018 ) . For instance , they have been successfully applied to neural style transfer ( Dumoulin et al. , 2016 ) and visual question answering ( Perez et al. , 2018 ) . This kind of conditional computation scheme is not tied to the type of normalization used . Wu & He ( 2018 ) recently proposed Group Normalization ( GN ) . GN divides feature maps into groups and normalizes the features within each group . GN only uses the layer dimension , hence its computation is independent of batch size . Ortiz et al . ( 2020 ) proposed Local Context Normalization ( LCN ) to encourage local contrast enhancement by normalizing features based on a spatial window around it and the filters in its group . Recently , Michalski et al . ( 2019 ) showed that Conditional Group Normalization ( CGN ) offer performance similar to CBN . In this work , we show results using CBN and CGN . Conditional normalization methods have been applied to tasks related to generalization , such as few-shot learning ( Jiang et al. , 2018 ; Tseng et al. , 2020 ) and domain adaption ( Su et al. , 2020 ) . Su et al . propose to use conditional normalization and an adversarial loss for domain adaption in object detection . In contrast to this work , we propose a method for implicit conditioning on an auxiliary task to leverage available metadata . 3 FORMULATION AND NETWORK ARCHITECTURE . 3.1 PROBLEM ABSTRACTION . We first establish notation . Let x be an image input , associated with a ground-truth target y for the main task ( e.g . a segmentation mask ) . Let available extra annotation for x be denoted by t. The main network is trained to predict y given x and contextual information from an auxiliary network . The auxiliary network learns to predict t , also given x . Features z of an intermediate layer of the auxiliary network are used to transform the main task network ’ s layers using conditional normalization parameterized by a learned function of z . The motivation for this method of implicit conditioning is the following : 1 . Since t ’ s are only used as training targets , auxiliary annotation is not required at test time . 2 . During training , the auxiliary network learns ( via backpropagation ) to visually capture in- formation predictive of t. At test time , the auxiliary network reasons about the unavailable t in terms of visual patterns that correlate with auxiliary annotations of training data . Note that this allows the distribution of auxiliary information at test time to differ from the training data ( see for example our experiments on out-of-distribution generalization in remote sensing in Section 4.1 ) . While the first statement is true for any multi-task architecture , the second statement describes the flexibility of the proposed method in leveraging auxiliary information of varying degrees of relevance . Obviously , the modulation will help most , if the auxiliary information is maximally relevant for the main task . Since the mapping from z to the modulation parameters is trained with the main task ’ s training signal , the network can learn to discard components of z that are not useful for the main task . It is also possible for the network to learn a constant identity transformation of the main network ’ s features in case no correlation is found . This reduces the potential of negative transfer learning between unrelated tasks common in multi-task learning ( Ruder , 2017 ) . To provide an example of how this method can help to exploit inexpensive metadata . Consider the task of segmenting satellite imagery of different regions on the globe . We can use the prediction of geographic coordinates , which are often logged by default when building satellite imagery datasets , as the auxiliary task . In this case , the auxiliary network may learn to capture visual characteristics that are distinctive for each region in the training set , such as a predominance for smaller buildings . This would provide a useful inductive bias for the segmentation network , even for regions with very different coordinates . By using feature modulation to integrate this contextual information , we hypothesize that the main network can learn more general purpose features , which can be attended to based on the context . 3.2 NETWORK ARCHITECTURE . Our proposed architecture modification transforms any standard neural network with normalization layers into one that incorporates conditioning information t. In each convolutional block of the neural network we substitute the normalization layer by it ’ s conditional counterpart . We refer to this family of networks as Conditional Networks . Figure 3 shows this extension applied to the popular U-Net Ronneberger et al . ( 2015 ) architecture . U-Net is an encoder-decoder network architecture with skip connections . Figure 3 shows the auxiliary network on the left modulating the modified U-Net on the right . The conditioning network is a convolutional architecture ( LeCun et al. , 1998 ) followed by a fully-connected layer that predicts metadata tn as a function of the input image xn . The pre-activation features before the output layer are used as z ( xn ) . The functions βn , c ( z ( xn ) ) and γn , c ( z ( xn ) ) mapping z ( xn ) to the scale and shift parameters are implemented with a multi- layer perceptron ( MLP ) . Using the latent representations instead of directly using tn allows us to leverage combinations of features that were useful in localizing images from previously seen data , potentially improving generalization . Because all its parts are differentiable , conditional networks can be trained end-to-end using gradient-based optimization . Our full objective is described in Equation 2 , where α is a hyperparameter balancing the main and auxiliary losses . Lmain task represents a standard main task loss . Lmain task depends on the task , such as Jaccard , cross-entropy , and dice for semantic segmentation . Lconditioning ensures the conditioning networks correctly predicts tn . Lcond . net = Lmain task + α · Lconditioning ( 2 )

Summary: This paper aims to tackle the out-of-distribution generalization problem where a model needs to generalize to new distributions at test time. The authors propose to utilize some extra information like the additional annotations as the conditional input and output the affine transformation parameters for the batch normalization stage. This extra information helps the backbone network get a more general representation from the training set thus the model is robust to the distribution shift when testing. Experiments are conducted on the Aerial Image Labeling and the Tumor-Infiltrating Lymphocytes datasets which correspond to the image segmentation and classification task respectively.