Revisiting Locally Supervised Learning: an Alternative to End-to-end Training

1 INTRODUCTION . End-to-end ( E2E ) back-propagation has become a standard paradigm to train deep networks ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; Szegedy et al. , 2015 ; He et al. , 2016 ; Huang et al. , 2019 ) . Typically , a training loss is computed at the final layer , and then the gradients are propagated backward layer-by-layer to update the weights . Although being effective , this procedure may suffer from memory and computation inefficiencies . First , the entire computational graph as well as the activations of most , if not all , layers need to be stored , resulting in intensive memory consumption . The GPU memory constraint is usually a bottleneck that inhibits the training of state-of-the-art models with high-resolution inputs and sufficient batch sizes , which arises in many realistic scenarios , such as 2D/3D semantic segmentation/object detection in autonomous driving , tissue segmentation in medical imaging and object recognition from remote sensing data . Most existing works address this issue via the gradient checkpointing technique ( Chen et al. , 2016 ) or the reversible architecture design ( Gomez et al. , 2017 ) , while they both come at the cost of significantly increased computation . Second , E2E training is a sequential process that impedes model parallelization ( Belilovsky et al. , 2020 ; Löwe et al. , 2019 ) , as earlier layers need to wait for their successors for error signals . As an alternative to E2E training , the locally supervised learning paradigm ( Hinton et al. , 2006 ; Bengio et al. , 2007 ; Nøkland & Eidnes , 2019 ; Belilovsky et al. , 2019 ; 2020 ) by design enjoys higher memory efficiency and allows for model parallelization . In specific , it divides a deep network into several gradient-isolated modules and trains them separately under local supervision ( see Figure 1 ( b ) ) . Since back-propagation is performed only within local modules , one does not need to store all ∗Corresponding author . intermediate activations at the same time . Consequently , the memory footprint during training is reduced without involving significant computational overhead . Moreover , by removing the demands for obtaining error signals from later layers , different local modules can potentially be trained in parallel . This approach is also considered more biologically plausible , given that brains are highly modular and predominantly learn from local signals ( Crick , 1989 ; Dan & Poo , 2004 ; Bengio et al. , 2015 ) . However , a major drawback of local learning is that they usually lead to inferior performance compared to E2E training ( Mostafa et al. , 2018 ; Belilovsky et al. , 2019 ; 2020 ) . In this paper , we revisit locally supervised training and analyse its drawbacks from the informationtheoretic perspective . We find that directly adopting an E2E loss function ( i.e. , cross-entropy ) to train local modules produces more discriminative intermediate features at earlier layers , while it collapses task-relevant information from the inputs and leads to inferior final performance . In other words , local learning tends to be short sighted , and learns features that only benefit local modules , while ignoring the demands of the rest layers . Once task-relevant information is washed out in earlier modules , later layers can not take full advantage of their capacity to learn more powerful representations . Based on the above observations , we hypothesize that a less greedy training procedure that preserves more information about the inputs might be a rescue for locally supervised training . Therefore , we propose a less greedy information propagation ( InfoPro ) loss that aims to encourage local modules to propagate forward as much information from the inputs as possible , while progressively abandon task-irrelevant parts ( formulated by an additional random variable named nuisance ) , as shown in Figure 1 ( c ) . The proposed method differentiates itself from existing algorithms ( Nøkland & Eidnes , 2019 ; Belilovsky et al. , 2019 ; 2020 ) on that it allows intermediate features to retain a certain amount of information which may hurt the short-term performance , but can potentially be leveraged by later modules . In practice , as the InfoPro loss is difficult to estimate in its exact form , we derive a tractable upper bound , leading to surrogate losses , e.g. , cross-entropy loss and contrastive loss . Empirically , we show that InfoPro loss effectively prevents collapsing task-relevant information at local modules , and yields favorable results on five widely used benchmarks ( i.e. , CIFAR , SVHN , STL-10 , ImageNet and Cityscapes ) . For instance , it achieves comparable accuracy as E2E training using 40 % or less GPU memory , while allows using a 50 % larger batch size or a 50 % larger input resolution with the same memory constraints . Additionally , our method enables training different local modules asynchronously ( even in parallel ) . 2 WHY LOCALLY SUPERVISED LEARNING UNDERPERFORMS E2E TRAINING ? . We start by considering a local learning setting where a deep network is split into multiple successively stacked modules , each with the same depth . The inputs are fed forward in an ordinary way , while the gradients are produced at the end of every module and back-propagated until reaching an earlier module . To generate supervision signals , a straightforward solution is to train all the local modules as independent networks , e.g. , in classification tasks , attaching a classifier to each module , and computing the local classification loss such as cross-entropy . However , such a greedy version of the standard supervised learning ( greedy SL ) algorithm leads to inferior performance of the whole network . For instance , in Table 1 , we present the test error of a ResNet-32 ( He et al. , 2016 ) on CIFAR-10 ( Krizhevsky et al. , 2009 ) when it is greedily trained with K modules . One can observe a severe degradation ( even more than 15 % ) with K growing larger . Plausible as this phenomenon seems , it remains unclear whether it is inherent for local learning and how to alleviate this problem . In this section , we investigate the performance degradation issue of the greedy local training from an information-theoretic perspective , laying the basis for the proposed algorithm . Linear separability of intermediate features . In the common case that greedy SL operates directly on the features output by internal layers , a natural intuition is to investigate how these locally learned features differ from their E2E learned counterparts in task-relevant behaviors . To this end , we fix the networks in Table 1 , and train a linear classifier using the features from each layer . The test errors of these classifiers are presented in the left plot of Figure 2 , where the horizontal axis denotes the indices of layers . The plot shows an intriguing trend : greedy SL contributes to dramatically more discriminative features with the first one ( or few ) local module , but is only able to slightly improve the performance with all the consequent modules . In contrast , the E2E learned network progressively boosts the linear separability of features throughout the whole network with even more significant effects in the later layers , surpassing greedy SL eventually . This raises an interesting question : why does the full network achieve inferior performance in greedy SL compared to the E2E counterpart , even though the former is based on more discriminative earlier features ? This observation appears incompatible with prior works like deeply-supervised nets ( Lee et al. , 2015a ) . Information in features . Since we use the same training configuration for both greedy SL and E2E learning , we conjecture that the answer to the above question might lie in the differences of features apart from merely separability . To test that , we look into the information captured by the intermediate features . In specific , given intermediate feature h corresponding to the input data x and the label y ( all of them are treated as random variables ) , we use the mutual information I ( h , x ) and I ( h , y ) to measure the amount of all retained information and task-relevant information in h , respectively . As these metrics can not be directly computed , we estimate the former by training a decoder with binary cross-entropy loss to reconstruct x from h. For the latter , we train a CNN using h as inputs to correctly classify x , and estimate I ( h , y ) with its performance . Details are deferred to Appendix G. The estimates of I ( h , x ) and I ( h , y ) at different layers are shown in the middle and right plots of Figure 2 . We note that in E2E learned networks , I ( h , y ) remains unchanged when the features pass through all the layers , while I ( h , x ) reduces gradually , revealing that the models progressively discard task-irrelevant information . However , greedily trained networks collapse both I ( h , x ) and I ( h , y ) in their first few modules . We attribute this to the short sighted optimization objective of earlier modules , which have relatively small capacity compared with full networks and are not capable of extracting and leveraging all the task-relevant information in x , as the E2E learned networks do . As a consequence , later modules , even though introducing additional parameters and increased capacity , lack necessary information about the target y to construct more discriminative features . Information collapse hypothesis . The above observations suggest that greedy SL induces local modules to collapse some of the task-relevant information that may be useless for short-term performance . However , the information is useful for the full model . In addition , we postulate that , although E2E training is incapable of extracting all task-relevant information at earlier layers as well , it alleviates this phenomenon by allowing a larger amount of task-irrelevant information to be kept , even though it may not be ideal for short-term performance . More empirical validation of our hypothesis is provided in Appendix A . 3 INFORMATION PROPAGATION ( INFOPRO ) LOSS . In this section , we propose an information propagation ( InfoPro ) loss to address the issue of information collapse in locally supervised training . The key idea is to enforce local modules to retain as much information about the input as possible , while progressively discard task-irrelevant parts . As it is difficult to estimate InfoPro loss in its exact form , we derive an easy-to-compute upper bound as the surrogate loss , and analyze its tightness . 3.1 LEARNING TO DISCARD USELESS INFORMATION . Nuisance . We first model the task-irrelevant information in the input data x by introducing the concept of nuisance . A nuisance is defined as an arbitrary random variable that affects x but provides no helpful information for the task of interest ( Achille & Soatto , 2018 ) . Take recognizing a car in the wild for example . The random variables determining the weather and the background are both nuisances . Formally , given a nuisance r , we have I ( r , x ) > 0 and I ( r , y ) = 0 , where y is the label . Without loss of generality , we suppose that y , r , x and h form the Markov chain ( y , r ) → x→ h , namely p ( h|x , y , r ) = p ( h|x ) . As a consequence , for the intermediate feature h from any layer , we obviously have I ( r , x ) ≥ I ( r , h ) . Nevertheless , we postulate that maxr I ( r , h ) > 0 . This assumption is mild since it does not hold only when h strictly contains no task-irrelevant information . Information Propagation ( InfoPro ) Loss . Thus far , we have been ready to introduce the proposed InfoPro loss . Instead of overly emphasizing on learning highly discriminative features at local modules , we also pay attention to preventing collapsing useful information in the feed-forward process . A simple solution to achieve this is maximizing the mutual information I ( h , x ) . Ideally , if there is no information loss , all useful information will be retained . However , it goes to another extreme case where the local modules do not learn any task-relevant feature , and is obviously dispensable . By contrast , in E2E training , intermediate layers progressively discard useless ( taskirrelevant ) information as well as shown above . Therefore , to model both effects simultaneously , we propose the following combined loss function : LInfoPro ( h ) = α [ −I ( h , x ) + βI ( r∗ , h ) ] , α , β ≥ 0 , s.t . r∗ = argmax r , I ( r , x ) > 0 , I ( r , y ) =0 I ( r , h ) , ( 1 ) where the nuisance r∗ is formulated to capture as much task-irrelevant information in h as possible , and the coefficient β controls the amount of information that is propagated forward ( first term ) and task-irrelevant information that is discarded ( second term ) . Notably , we assume that the final module is always trained using the normal E2E loss ( e.g. , softmax cross-entropy loss for classification ) weighted by the constant 1 , such that α is essential to balance the intermediate loss and the final one . In addition , LInfoPro ( h ) is used to train the local module outputting h , whose inputs are not required to be x . The module may stack over another local module trained with the same form of LInfoPro ( h ) but ( possibly ) different α and β . Our method differs from existing works ( Nøkland & Eidnes , 2019 ; Belilovsky et al. , 2019 ; 2020 ) in that it is a non-greedy approach . The major effect of minimizing LInfoPro ( h ) can be described as maximally discarding the task-irrelevant information under the goal of retaining as much information of the input as possible . Obtaining high short-term performance is not necessarily required . As we explicitly facilitate information propagation , we refer to LInfoPro ( h ) as the InfoPro loss . 3.2 UPPER BOUND OF LInfoPro The objective function in Eq . ( 1 ) is difficult to be directly optimized , since it is usually intractable to estimate r∗ , which is equivalent to disentangling all task-irrelevant information from intermediate features . Therefore , we derive an easy-to-compute upper bound of LInfoPro as an surrogate loss . Our result is summarized in Proposition 1 , with the proof in Appendix B . Proposition 1 . Suppose that the Markov chain ( y , r ) → x → h holds . Then an upper bound of LInfoPro is given by LInfoPro ≤ −λ1I ( h , x ) − λ2I ( h , y ) , L̄InfoPro , ( 2 ) where λ1 = α ( 1− β ) , λ2 = αβ . For simplicity , we integrate α and β into two mutually independent hyper-parameters , λ1 and λ2 . Although we do not explicitly restrict λ1 ≥ 0 , we find in experiments that the performance of networks is significantly degraded with λ1 → 0+ ( see Figure 4 ) , or say , β → 1− , where models tend to reach local minima by trivially minimizing I ( r∗ , h ) in Eq . ( 1 ) . Thus , we assume λ1 , λ2 ≥ 0 . With Proposition 1 , we can optimize the upper bound L̄InfoPro as an approximation , circumventing dealing with the intractable term I ( r∗ , h ) in LInfoPro . To ensure that the approximation is accurate , the gap between the two should be reasonably small . Below we present an analysis of the tightness of L̄InfoPro in Proposition 2 , ( proof given in Appendix C ) . We also empirically check it in Appendix H. Proposition 2 provides a useful tool to examine the discrepancy between LInfoPro and its upper bound . Proposition 2 . Given that r∗ = argmaxr , I ( r , x ) > 0 , I ( r , y ) =0 I ( r , h ) and that y is a deterministic function with respect to x , the gap = L̄InfoPro − LInfoPro is upper bounded by ≤ λ2 [ I ( x , y ) − I ( h , y ) ] . ( 3 )

Summary: The paper analyzes the pitfalls of locally supervised learning from the point of view of information propagation and proposes a new auxiliary loss that can facilitate locally supervised learning. The proposed loss, "infopro loss", is then relaxed to a tractable upper bound, which is then used instead. To implement the loss, mutual information is approximated with a decoder, as well as a classifier. The authors further introduce now contrastive learning fits in the framework as a lower bound maximization process regarding mutual information. The experimental results on standard datasets demonstrate the efficacy of the proposed method.