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 proposes a strategy for training feed-forward networks in a more memory-efficient manner by employing local as opposed to end-to-end supervision. End-to-end/global (E2E) supervision as the dominant paradigm in training deep networks considers a loss function at the very end of the network for backpropagation of the resulting gradients, whereas local supervision injects supervisory signals (such as the same E2E objective, e.g. classification) at intermediate layers in the network. The benefit of such intermediate supervision is the ability to train larger networks in smaller chunks piece by piece, where each individual training is more memory efficient due to reduced need to store activations (and weights and biases) in GPU memory. As a drawback, however, it had been shown earlier that such local training is less optimal than global training in terms of the achievable generalization performance. The authors propose a new training strategy that aims at combining the memory efficiency of local supervision and piecewise training with the error performance of global training. Considering a given intermediate layer, the paper motivates to maximize the mutual information between the activations in this layer and the input signal to retain relevant information, while minimizing the mutual information of the activations and a nuisance variable, where the nuisance is defined as having no mutual information with the target variable (e.g. the classification prediction). The authors argue that this local supervision allows to train the features at the intermediate layer such that they carry relevant information from the input to the target variable without resorting to direct supervision with the target variable. Direct computation of the nuisance variable is infeasible and the authors propose a bounded approximation.