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: This paper proposes a differentiable architecture search approach for splitting a deep network into locally-trained blocks to achieve training speedup. The approach achieves better performance than using backprop on small datasets (CIFAR10 and TinyImageNet), and comparable or slightly improved performance on ImageNet with 2x claimed training speedup. Learned network architecture choices seem to transfer between datasets.