SEDONA: Search for Decoupled Neural Networks toward Greedy Block-wise Learning

1 INTRODUCTION . Backpropagation ( Rumelhart et al. , 1986 ) has made a significant contribution to the success of deep learning as the core learning algorithm for SGD-based optimization . However , backpropagation is sequential in nature and supports only synchronous weight updates . Specifically , the limited concurrency in backpropagation breaks down into two locking problems ( Jaderberg et al. , 2017 ) . First , update locking – a forward pass must complete first before any weight update . Second , backward locking – gradient computation of upper layers must precede that of lower layers . Also , backpropagation may be biologically implausible since the human brain prefers local learning rules without the global movement of error signals ( Crick ( 1989 ) ; Marblestone et al . ( 2016 ) ; Lillicrap et al . ( 2020 ) ) . Greedy block-wise learning is a competitive alternative to backpropagation that overcomes these limitations . It splits layers into a stack of gradient-isolated blocks , each of which is trained with local error signals . Therefore , it is possible to simultaneously compute the gradients for different network components with more fine-grained locks . Limiting the depth of error propagation graphs also reduces the vanishing gradient and increases memory efficiency . Recently , Belilovsky et al . ( 2019 ) , Nøkland & Eidnes ( 2019 ) , Belilovsky et al . ( 2020 ) , and Löwe et al . ( 2019 ) empirically demonstrated that greedy block-wise learning could yield competitive performance to end-to-end backpropagation . However , greedy block-wise learning introduces a group of new architectural decisions . Let us consider a case where we want to decouple an L-layer network into K blocks for a given K ∈ { 1 , . . . , L } . Then , the number of all possible groupings is ( L−1 K−1 ) . If we want to choose one of M candidates of auxiliary networks to generate local error gradients , we would have to consider ( L−1 K−1 ) MK−1 different configurations . If local signals are not representative of the global goal , then the final performance would be damaged significantly . In this work , we introduce a novel search method named SEDONA ( SEarching for DecOupled Neural Architectures ) , which allows efficient search of decoupled neural architectures toward greedy block-wise learning . Given a base neural network , SEDONA optimizes the validation loss by grouping layers into blocks and selecting the best auxiliary network for each block . Inspired by DARTS ( Liu et al. , 2019 ) , we first relax the decision variables representing error propagation graphs and auxiliary networks to continuous domains . We then formulate a bilevel optimization problem for the decision variables , which is solved via gradient descent . Our key contributions are summarized as follows . 1 . To the best of our knowledge , this work is the first attempt to automate the discovery of decoupling neural networks for greedy block-wise learning . We propose an efficient search method named SEDONA , which finds decoupled error propagation graphs and auxiliary heads suitable for successful greedy training . 2 . Through extensive experiments on image classification tasks , we show that our locally optimized networks outperform not only the ones trained with end-to-end backpropagation but also two state-of-the-art greedy-learning methods of DGL ( Belilovsky et al. , 2020 ) and PredSim ( Nøkland & Eidnes , 2019 ) in CIFAR-10 , Tiny-ImageNet and ImageNet . 3 . SEDONA discovers decoupled architectures for VGG ( Simonyan & Zisserman , 2015 ) and ResNet ( He et al. , 2016 ) variants with only 0.25× width on CIFAR-10 . The discovered networks are transferable to Tiny-ImageNet and ImageNet in which they are high-performing enough to beat backpropagation and other greedy-leaning methods . It means that no search is required for every pair of a network and a dataset . 4 . Finally , based on experimental results , we analyze the common characteristics among the architectures that are favorable to greedy-learning , such as avoiding the shallow first block and using deeper auxiliary networks for lower blocks . 2 PROBLEM STATEMENT AND MOTIVATION . In typical neural network training , backpropagation computes the gradients of weights with respect to the global loss by the chain rule ( Rumelhart et al. , 1986 ) . On the other hand , in greedy block-wise learning ( Löwe et al. , 2019 ; Belilovsky et al. , 2020 ) , the network is split into several subnetworks ( i.e . blocks ) , each of which consists of one or more consecutive layers . Next , each block is attached to a small neural network called the auxiliary network that computes its own objective ( i.e . local loss ) , from which layer weights are optimized by propagating error signals within the block . Naturally , each block can independently perform parameter updates even while other blocks process forward passes . Figure 1 illustrates the high-level overview of greedy block-wise learning . For successful greedy block-wise learning , one must make two design decisions beforehand : ( i ) how to split the original network into a set of subnetworks , and ( ii ) which auxiliary network to use for each subnetwork . Finding the best configuration to both problems requires significant time and effort from human experts . We empirically show that the performance of greedy block-wise learning is critically sensitive to these two design choices in Appendix A . This sensitivity introduces a paradox of replacing backpropagation with greedy block-wise learning . If one has to put significant cost and time through a series of experiments to discover a workable configuration , then the benefit of greedy learning ( e.g . reduction of training time ) is diluted . Unfortunately , there has not been a generally acceptable practice to answer these two design choices . Therefore , this work aims to propose an automated search method for the discovery of the best configuration , which has not been discussed so far . Although there have been several works on modifying backward computation graphs ( Bello et al. , 2017 ; Alber et al. , 2018 ; Xu et al. , 2018 ) , they still rely on global end-to-end learning and focus on finding new optimizers , weight update formulas , or error propagation rules , assuming that the backward computation graphs are never discontinuous . In this work , we instead concern ourselves with making backward computation graphs discontinuous , i.e . finding optimal points where we stop gradient flow and use local gradients instead . 3 SEDONA : SEARCHING FOR DECOUPLED NEURAL ARCHITECTURES . We begin by describing the search space ( Section 3.1 ) and present a continuous relaxation strategy that allows searching for the decoupled architecture using stochastic gradient descent on the validation loss ( Section 3.2 ) . Finally , we propose a bilevel optimization that makes SEDONA computationally efficient ( Section 3.3 ) and the discretization on relaxed decision variables ( Section 3.4 ) . 3.1 THE SEARCH SPACE . Our goal is to vertically split an input network into a set of blocks to enable greedy block-wise training . This task is equivalent to solving the following two problems ( Figure 2 ) : ( i ) finding layers at which we discontinue the gradient flow and ( ii ) assigning the most appropriate auxiliary network to each of those layers for generating local error signals . We assume that an input network consists of L convolutional layers , each of which may contain normalization and pooling layers . If the network has residual connections , we regard each residual block as a single layer to simplify our notation . We focus on splitting the convolutional layers only and treat the last fully connected ( FC ) layers as the last block ’ s auxiliary network . To facilitate our exposition , we introduce two binary variables for every layer . First , the signal variable α ( l ) ∈ { 0 , 1 } denotes whether layer l should utilize local gradients ( i.e . the last layer of a block ) or backpropagated gradients ( i.e . the inside layer ) . In other words , if α ( l ) = 1 , the layer l is trained using local gradients , hence becoming the last layer in a gradient-isolated block . Otherwise , the layer is trained with backpropagated gradients , therefore becoming the inside layer in a block . Second , the auxiliary type variable β ( l ) ∈ { 0 , 1 } M denotes a one-hot encoding of which auxiliary is chosen out of M possible candidates for layer l. We here assume that every layer selects its own auxiliary network ; later , we leave it only for the last layer of each block . Note that the last layer of the whole network does not need these variables because the final classifier acts as its auxiliary head . Now , we can formulate the loss for layer l as follows : L ( l ) train ( θ , φ , β ) = 1 |D| ∑ ( x , y ) ∈D ` ( l ) ( y , o ( l ) ) where o ( l ) ( x ; θ ) = M∑ m=1 β ( l ) m · f ( l ) m ( a ( l ) ( x ; θ ) ; φ ( l ) m ) . ( 1 ) D is the training data , ` ( l ) is the loss function , { f ( l ) m } are M candidates for auxiliary networks , and a ( l ) is the activations , θ is the network layer weights , and φ ( l ) m is them-th auxiliary network weights . The error gradients at each non-last layer l are calculated as δ ( l ) = α ( l ) · δ ( l ) local + ( 1− α ( l ) ) · δ ( l ) bp , where δ ( l ) local = ∇a ( l ) L ( l ) train , δ ( l ) bp = δ ( l+1 ) · ∂a ( l+1 ) ∂a ( l ) , ( 2 ) which is equivalent to choosing between local and backpropagated error gradients . Algorithm 1 : SEDONA – Searching for Decoupled Neural Architectures Initialize signal variables α and auxiliary type variables β as 0 . Pretrain layer and auxiliary network weights ( θ , φ ) with Eq . ( 5 ) – ( 6 ) and store N sets of weights with the best validation accuracies in memoryM . while ( α , β ) not converged do Sample one set of layer and auxiliary network weights ( θ , φ ) fromM Calculate ( θT , φT ) by performing T gradient steps on ( θ , φ ) with Eq. ( 5 ) – ( 6 ) . Update ( α , β ) by descending ∇ ( α , β ) L ( L ) val ( θT , φT , α , β ) with Eq. ( 7 ) . Calculate ( θ1 , φ1 ) by performing single gradient step on ( θ , φ ) with Eq. ( 5 ) – ( 6 ) . Evict the oldest weights inM and save the updated weights ( θ1 , φ1 ) intoM . end Obtain decoupling by discretizing α and β 3.2 CONTINUOUS RELAXATION . We optimize over α and β after relaxing them to a continuous domain , in a similar way to DARTS ( Liu et al. , 2019 ) . It enables a gradient-based optimization , which is significantly faster than directly optimizing over the discrete domain using reinforcement learning or evolutionary algorithms . This benefit accelerates especially when the search space is large , i.e . when the input network is deep . Henceforth , we let α ( l ) ∈ R2 and β ( l ) ∈ RM . As a result , we relax the categorical choice of an auxiliary head to a weighted mixture of all possible auxiliary networks by replacing Eq . ( 1 ) with ō ( l ) = M∑ m=1 softmax ( β ( l ) ) m · f ( l ) m ( a ( l ) ; φ ( l ) m ) . ( 3 ) Also , replacing Eq . ( 2 ) , the gradients for layer l become a mix of local and backpropagated gradients : δ̄ ( l ) = ᾱ ( l ) 1 · δ ( l ) local + ᾱ ( l ) 2 · δ ( l ) bp , ( 4 ) which is the affine combination between δ ( l ) local = ∇a ( l ) L ( l ) train and δ ( l ) bp = δ̄ ( l+1 ) · ∂a ( l+1 ) ∂a ( l ) . We use ᾱ ( l ) = softmax ( α ( l ) ) and δ̄ ( L ) = ∇a ( L ) L ( L ) .

Summary: The paper proposes a method for decoupled training of neural networks called SEDONA. In the spirit of recent trends in greedy layer-wise and indirect training, SEDONA allows gradient information to flow either from the next layer as in backpropagation or from an auxiliary head, trying to make a prediction using the current layer's output. Since a direct search for the best decoupled configuration results into probing a combinatorial number of splits, authors propose a continuously relaxed formulation which they later discretize. Transferring the found decoupled configurations to datasets different from the "pretraining" ones shows improvements in terms of validation accuracy and training time.