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: The authors propose a new framework for constructing teacher models in student-teacher model training, which involves a combination of spatial smoothing (piecewise copying of students to teacher) and the usual temporal smoothing (trajectory-wise parameter averaging).   They situate spatial smoothing inside a more general framework of spatial-temporal smoothing, and recover temporal smoothing and the newly proposed spatial smoothing as special cases of spatial-temporal smoothing.  They go on to show its utility in two learning scenarios (semi-supervised learning and self-supervised learning).