Spatial Ensemble: a Novel Model Smoothing Mechanism for Student-Teacher Framework

1 Introduction . Recent years have witnessed the success of the student-teacher framework in self-supervised [ 1–5 ] and semi-supervised [ 6–8 ] learning tasks . In this framework , the teacher is responsible for generating surrogate supervision signals to guide the learning of the student on unlabeled data , for example , category IDs for classification , feature patterns for contrastive learning , etc . Many important works under this framework share a common model smoothing technique , i.e. , the Temporal Moving Average ( TMA ) , to generate the teacher from the student model . For example , Mean Teacher [ 6 ] for semi-supervision averages historical student weights as the teacher . MoCo [ 1 ] for self-supervision employs a slowly progressing encoder ( teacher ) , driven by a momentum update with the query encoder ( student ) . Another state-of-the-art method for self-supervision BYOL [ 2 ] employs a similar strategy to update the target ( teacher ) network with a slow-moving average of the online ( student ) network . We thus recognize that model smoothing with TMA is of central importance for obtaining a reliable teacher for these works . In this paper , we rethink the mechanism of TMA and summarize its importance as two-fold . First , TMA enables the teacher to absorb parameters from different student editions and thus benefits from ∗corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the model ensemble . Second , TMA constrains the variance of the teachers to be small to avoid inconsistent labels produced during two adjacent updates . An abrupt change to the pseudo labels will cripple training and possibly leads to optimization divergence . This is particularly important for self-supervised learning where annotations are unavailable . We propose a novel model smoothing mechanism named “ Spatial Ensemble ” . Compared with TMA , Spatial Ensemble facilitates similar model smoothing benefits with a different mechanism , as illustrated in Figure 1 . In TMA , the teacher absorbs each historical student as a whole . A temporally near student is absorbed into the teacher with a larger weighting factor due to the Exponential Moving Average effect . In contrast , in Spatial Ensemble , the teacher randomly picks up one or more small fragments of the student to replace the corresponding fragments of the teacher during each update . All the fragments are weighted equally in the teacher model . In this way , Spatial Ensemble stitches different fragments of previous student models into a unity , yielding the “ Spatial Ensemble ” effect . In spite of their fundamental differences , Spatial Ensemble facilitates similar benefits of model ensemble and smoothed update . The smoothed update is maintained by replacing only a small portion of parameters at each iteration . On the challenging self-supervised learning task , which is very sensitive to the quality of surrogate supervision signals , we observe that Spatial Ensemble achieves comparable performance with temporal moving average . It shows that our Spatial Ensemble can serve as a basic model smoothing technique for the student-teacher framework . An important advantage of Spatial Ensemble lays in its complementarity to Temporal Moving Average . To explore their complementary benefits , we integrate both the temporal smoothing and spatial smoothing mechanism into a unified method named Spatial-Temporal Smoothing ( STS ) . STS does not directly absorb a whole student , or simply replace a fragment . Instead , it updates a fragment of the teacher model with a temporal moving average of the corresponding fragment of the student . This “ mix-up ” model smoothing benefits from the mutual complementarity of temporal moving average and Spatial Ensemble . Experiments on both self-supervised and semi-supervised tasks show that using STS for model smoothing obtains general ( sometimes significant ) improvement over a majority of state-of-the-art methods . We summarize our contribution as follows : 1 . We propose a novel model smoothing mechanism named Spatial Ensemble . Compared with the popular temporal moving average , Spatial Ensemble achieves a comparable model smoothing effect with a different mechanism . 2 . We integrate the spatial and temporal smoothing mechanism into a unified method named SpatialTemporal Smoothing . STS benefits from the mutual complementarity of two mechanisms and consistently improves student-teacher methods on self- and semi-supervised learning tasks . 3 . We notice two recognizable advantages of the model learned through STS , including strong robustness against various data corruptions , and superior transferring capability to the object detection task . 2 Related Work . In this section , we present a brief review of the development on semi- and self-supervised learning , as well as the student-teacher framework . We mainly focus on the recent deep learning literature related to our method . Student Teacher Layer 1 Layer 2 … Layer N-1 Layer N Layer 1 Layer 2 … Layer N-1 Layer N Student Teacher ( c ) SE between models Layer 1 Layer 2 … Layer N-1 Layer N Layer 1 Layer 2 … Layer N-1 Layer N ( d ) STS between models Student Teacher Layer 1 Layer 2 … Layer N-1 Layer N Layer 1 Layer 2 … Layer N-1 Layer N ( b ) TMA between models Semi-supervised learning aims at achieving comparable performance to fully supervised methods , usually with a small portion of labeled data and the rest as unlabeled data . It is of great importance for real scenarios where annotations are difficult/expensive to acquire , such as medical data . Pseudo labels [ 9 , 10 ] and consistency regularization [ 6 , 7 , 11 ] are two essential techniques in many recent semi-supervised approaches . Pseudo-Labeling [ 9 ] generates soft pseudo labels for the unlabeled data using the logits predicted by the model from the latest epoch . Laine et al . [ 10 ] improve the pseudo label generation process by leveraging a temporal ensembling prediction of different epochs . Mean-teacher [ 6 ] further simplifies this process by introducing a teacher model , which is a temporal ensembling of the student model through the exponential moving average , to generate pseudo labels directly . Besides , mean-teacher employs a consistency regularization to encourage similar predictions from the teacher and the student model . FixMatch [ 7 ] assumes that the consistency should also hold between images applied with augmentations of different degrees , and applies both a weak and a strong set of data augmentations . Self-supervised learning largely narrows down its performance gap with supervised approaches and is receiving increasing attention in the community . Recent self-supervised methods can be roughly categorized into two genres , i.e. , approaches based on pretext task and approaches based on contrastive learning . There are many kinds of insightful pretext task designs . One common method is generating surrogate supervision signals by feature clustering [ 12–14 ] , image augmentation ( usually rotation [ 15 ] ) , relative order of clipped image patches [ 16–18 ] , etc . Another mainstream approach is to design a generative pretext task , e.g. , image colorization [ 19 , 20 ] , image inpainting [ 21 ] and image denoising [ 22 ] , and usually employ auto-encoders or generative adversarial networks [ 23 ] to transfer images . Instead of designing a pretext task , contrastive learning approaches [ 1 , 4 , 24–29 ] aim at maximizing the feature similarities with augmented views produced from the identical image sample , while minimizing those from different image samples . Student-teacher framework is a popular architecture that has been extensively used in many recent state-of-the-art methods , especially for semi- and self-supervised learning tasks . Usually , temporal moving average is employed to generate the teacher from historical student models . The optimization goal is to ensure the prediction consistency between the student and the teacher . Besides pulling different views of the same image closer , MoCo [ 1 ] employs a memory bank to store negative samples and pushes views from different images apart . BYOL [ 2 ] removes negative samples and leverages an asymmetric student-teacher architecture to prevent collapsing . An extra predictor is appended to the online ( student ) network to regress the output of the target ( teacher ) network . 3 Method . Preliminaries . Given a same network architecture , ΘT and ΘS are the parameters of the teacher and student model . Both ΘT and ΘS contains n units of parameters , i.e. , ΘT = { θT1 , θT2 , · · · , θTn } and ΘS = { θS1 , θS2 , · · · , θSn } . A unit may correspond to a layer , a neuron or a channel , depending on the desired spatial granularity . To analyze the model update in consecutive iterations , we further define them as functions of the iteration t , i.e. , ΘT ( t ) = { θTi ( t ) } and ΘS ( t ) = { θSi ( t ) } . 3.1 Spatial Ensemble . We present a new model smoothing method named Spatial Ensemble ( SE ) , as illustrated in Figure 2 ( b ) . During each iteration , SE randomly picks up some units of the student model to replace the corresponding units of the teacher model , leaving the remaining parts of the teacher unchanged . We formulate this process as : ΘT ( t ) = { PiθTi ( t− 1 ) + ( 1− Pi ) θSi ( t− 1 ) } , ( 1 ) where Pi ∼ Ber ( p ) ( i.e. , the Bernoulli distribution with success probability p ) is a binary variable . We note that all Pi ∼ Ber ( p ) ( i = 1 , 2 , · · · , n ) are mutually independent to each other and have the same success probability p. It determines whether the i-th unit of the teacher parameters is to be preserved . Specifically , if Pi = 1 , the i-th unit parameters is to be preserved . If Pi = 0 , the i-th unit is to be replaced . Since Pi ∼ Ber ( p ) , larger p results in lower replacing frequency and higher proportion of preserved units . Therefore , we name p as the “ preserving probability ” . For clarity , we omit the iteration indicator in all the equations in the following part of this paper and rewrite Equation 1 as : ΘT ⇐= { PiθTi + ( 1− Pi ) θSi } . ( 2 ) To investigate the mechanism of SE , we quantitatively analyze the smoothing effect on both the model parameters and the generated surrogate supervision signals . Specifically , we calculate the Mean-Square Errors ( MSE ) between teachers at two adjacent epochs , as well as the MSE between their outputs . We visualize the MSE values at the early training stages ( epoch 1 to 20 ) in Fig . 3 , from which we make two observations : Remark 1 : Spatial Ensemble smooths the teacher update . It is observed that without model smoothing , the MSE of both model parameters and supervision signals maintain relatively high values , indicating strong vibration during the teacher update . In another word , the teacher model and its output signals undergo substantial instability . In contrast , when using the TMA and the proposed SE to update the teacher , the MSE of both model parameters and supervision signals are much smaller ( than Non-smoothing ) , yielding a smoothed teacher update . Remark 2 : Spatial Ensemble promotes convergence . Using TMA and SE , the MSE values gradually decrease as the training epochs increases . Such result indicates that the up-to-date teacher model gradually become closer to its prior edition , therefore approaching convergence . Without TMA or SE , the Non-smoothing mode maintains a roughly constant MSE value . We thus infer that Spatial Ensemble ( as well as the TMA ) promotes convergence of student-teacher framework . We provide the pseudocode of SE in Algorithm 1 , which is easy to implement with only a few lines of code . In Section 4.3 , we further analyze the effect of p on student-teacher learning performance . Algorithm 1 Pseudocode of SE . # net_s , net_t : student/teacher network # m : momentum for TMA # p : probability of a student layer belonging to the unpicked fragment . params_s = net_s.parameters ( ) params_t = net_t.parameters ( ) for s , t in zip ( params_s , params_t ) : if random.rand ( ) < p : continue else : t = s Algorithm 2 Pseudo code of STS . # net_s , net_t : student/teacher network # m : momentum for TMA # p : probability of a student layer belonging to the unpicked fragment . params_s = net_s.parameters ( ) params_t = net_t.parameters ( ) for s , t in zip ( params_s , params_t ) : if random.rand ( ) < p : continue else : t = m * t + ( 1-m ) * s

Summary: Different from the popular EMA mechanism in self-supervised and semi-supervised learning, this paper proposes a new mechanism for model smoothing in spatial ensemble perspective as complementary to the popular temporal moving average.  By combining the proposed Spatial Ensemble and the popular Temporal Ensemble, this work proposes to integrate temporal moving average and spatial ensemble into a unified model smoothing method, named as Spatial-Temporal Smoothing (STS). Experiments on self-supervised and semi-supervised tasks show that STS brings improvement to some methods.