Deep Layers as Stochastic Solvers

1 INTRODUCTION . Deep learning has revolutionized computer vision and natural language processing and is increasingly applied throughout science and engineering ( LeCun et al. , 2015 ) . This has motivated the mathematical analysis of various aspects of deep networks , such as the capacity and uniqueness of their representations ( Soatto & Chiuso , 2014 ; Papyan et al. , 2018 ) and their global training convergence properties ( Haeffele & Vidal , 2017 ) . However , a complete characterization of deep networks remains elusive . For example , Bernoulli dropout layers are known to improve generalization ( Srivastava et al. , 2014 ) , but a thorough theoretical understanding of their behavior remains an open problem . While basic dropout layers have proven to be effective , there are many other types of dropout with various desirable properties ( Molchanov et al. , 2017 ) . This raises many questions . Can the fundamental block of layers that consists of a dropout layer followed by a linear transformation and a non-linear activation be further improved for better generalization ? Can the choice of dropout layer be made independently from the linear transformation and non-linear activation ? Are there systematic ways to propose new types of dropout ? We attempt to address some of these questions by establishing a strong connection between the forward pass through a block of layers in a deep network and the solution of convex optimization problems of the following form : minimize x∈Rd F ( x ) + g ( x ) , F ( x ) def = 1 n n∑ i fi ( a > i x ) . ( 1 ) Note that when fi ( a > i x ) = 1 2 ( a > i x − yi ) 2 and g ( x ) = ‖x‖22 , Eq . ( 1 ) is standard ridge regression . When g ( x ) = ‖x‖1 , Eq . ( 1 ) has the form of LASSO regression . We show that a block of layers that consists of dropout followed by a linear transformation ( fullyconnected or convolutional ) and a non-linear activation has close connections to applying stochastic solvers to ( 1 ) . Interestingly , the choice of the stochastic optimization algorithm gives rise to commonly used dropout layers , such as Bernoulli and additive dropout , and to a family of other types of dropout layers that have not been explored before . As a special case , when the block in question does not include dropout , the stochastic algorithm reduces to a deterministic one . Our contributions can be summarized as follows . ( i ) We show that a forward pass through a block that consists of Bernoulli dropout followed by a linear transformation and a non-linear activation ∗The work was done during an internship at Intel Labs . is equivalent to a single iteration of τ -nice Proximal Stochastic Gradient , Prox-SG ( Xiao & Zhang , 2014 ) when it is applied to an instance of ( 1 ) . We provide various conditions on g that recover ( either exactly or approximately ) common non-linearities used in practice . ( ii ) We show that the same block with an additive dropout instead of Bernoulli dropout is equivalent to a single iteration of mS2GD ( Konečnỳ et al. , 2016 ) – a mini-batching form of variance-reduced SGD ( Johnson & Zhang , 2013 ) – applied to an instance of ( 1 ) . ( iii ) By expressing both fully-connected and convolutional layers ( referred to as linear throughout ) as special cases of a high-order tensor product ( Bibi & Ghanem , 2017 ) , we derive a formula for the Lipschitz constant L of ∇F ( x ) . As a consequence , we can compute the optimal step size for the stochastic solvers that correspond to blocks of layers . We note that concurrent work ( Sedghi et al. , 2019 ) used a different analysis strategy to derive an equivalent result for computing the singular values of convolutional layers . ( iv ) We validate our theoretical analysis experimentally by replacing blocks of layers in standard image classification networks with corresponding solvers and show that this improves the accuracy of the models . 2 RELATED WORK . Optimization algorithms can provide insight and guidance in the design of deep network architectures ( Vogel & Pock , 2017 ; Kobler et al. , 2017 ; Yang et al. , 2016 ; Zhang & Ghanem , 2018 ) . For example , Yang et al . ( 2016 ) have proposed a deep network architecture for compressed sensing . Their network , dubbed ADMM-Net , is inspired by ADMM updates ( Boyd et al. , 2011 ) on the compressed sensing objective . Similarly , Zhang & Ghanem ( 2018 ) demonstrated that unrolling a proximal gradient descent solver ( Beck & Teboulle , 2009 ) on the same problem can further improve performance . The work of Kobler et al . ( 2017 ) demonstrated a relation between incremental proximal methods and ResNet blocks ; based on this observation , they proposed a new architecture ( variational networks ) for the task of image reconstruction . Amos & Kolter ( 2017 ) proposed to embed optimization problems , in particular linearly-constrained quadratic programs , as structured layers in deep networks . Meinhardt et al . ( 2017 ) replaced proximal operators in optimization algorithms by neural networks . Huang & Van Gool ( 2017 ) proposed a new matrix layer , dubbed ReEig , that applies a thresholding operation to the eigenvalues of intermediate feature representations that are stacked in matrix form . ReEig can be tightly connected to a proximal operator of the set of positive semi-definite matrices . Sulam et al . ( 2018 ) proposed a new architecture based on a sparse representation construct , Multi-Layer Convolutional Sparse Coding ( ML-CSC ) , initially introduced by Papyan et al . ( 2017 ) . Sparsity on the intermediate representations was enforced by a multi-layer form of basis pursuit . This body of work has demonstrated the merits of connecting the design of deep networks with optimization algorithms in the form of structured layers . Yet , with few exceptions ( Amos & Kolter , 2017 ; Sulam et al. , 2018 ) , previous works propose specialized architectures for specific tasks . Our work aims to contribute to a unified framework that relates optimization algorithms to deep layers . A line of work aims to provide rigorous interpretation for dropout layers . For example , Wager et al . ( 2013 ) showed that dropout is linked to an adaptively balanced ` 2-regularized loss . Wang & Manning ( 2013 ) showed that approximating the loss with a normal distribution leads to a faster form of dropout . Gal & Ghahramani ( 2016a ; b ) developed a framework that connects dropout with approximate variational inference in Bayesian models . We provide a complementary perspective , in which dropout layers arise naturally in an optimization-driven framework for network design . 3 UNIFIED FRAMEWORK . This section is organized as follows . We introduce our notation and preliminaries in Section 3.1 . In Section 3.2 , we present a motivational example relating a single iteration of proximal gradient descent ( Prox-GD ) on ( 1 ) to the forward pass through a fully-connected layer followed by a nonlinear activation . We will show that several commonly used non-linear activations can be exactly or approximately represented as proximal operators of g ( x ) . In Section 3.3 , we unify fully-connected and convolutional layers as special cases of a high-order tensor product . We propose a generic instance of ( 1 ) in a tensor setting , where we provide a formula for the Lipschitz constant L of the finite sum structure of ( 1 ) . In Section 3.4 , we derive an intimate relation between stochastic solvers , namely τ -nice Prox-SG and mS2GD , and two types of dropout layers . Figure 1 shows an overview of the connections that will be developed . 3.1 NOTATION AND PRELIMINARIES . As we will be working with tensors , we will follow the tensor notation of Kolda & Bader ( 2009 ) . The order of a tensor is the number of its dimensions . In particular , scalars are tensors of order zero , vectors are tensors of order one , and matrices are tensors of order two . We denote scalars by lowercase letters a , vectors by bold lowercase letters a , and matrices by bold capital letters A . We use subscripts ai to refer to individual elements in a vector . Tensors of order three or more will be denoted by cursive capital lettersA ∈ RJ1×J2×···×Jn . Throughout the paper , we will handle tensors that are of at most order four . High-order tensors with a second dimension of size equal to one are traditionally called vector tensors and denoted ~A ∈ RJ1×1×J3×J4 . We use A ( i , j , k , z ) to refer to an element in a tensor and A ( i , j , k , : ) to refer to a slice of a tensor . The inner product between tensors of the same size is denoted 〈A , B〉 = ∑ i1 , ... , iN A ( i1 , . . . , iN ) B ( i1 , . . . , iN ) . The squared Frobenius norm of a tensor A is defined as ‖A‖2F = 〈A , A〉 . Lastly , the superscripts > and H are used to denote the transpose and the Hermitian transpose , respectively . 3.2 MOTIVATIONAL INSIGHT : NON-LINEAR ACTIVATIONS AS PROXIMAL OPERATORS . As a motivating example , we consider the lth linear layer in a deep network that is followed by a non-linear activation ρ , i.e . xl = ρ ( Axl−1 + b ) , where A ∈ Rn2×n1 and b ∈ Rn2 are the weights and biases of the layer and xl−1 and xl are the input and output activations , respectively . Now consider an instance of ( 1 ) with a convex function g ( x ) and F ( xl ) = 1 2 ‖A > xl − xl−1‖2 − b > xl = 1 2 n1∑ i ( A > ( i , : ) xl − xl−1i ) 2 − b > xl , ( 2 ) where A > ( i , : ) is the ith row of A > . Such an objective can be optimized iteratively in xl using Prox-GD with the following update equation : xl ← Prox 1 L g ( ( I− 1 L AA > ) xl + 1 L ( Axl−1 + b ) ) , ( 3 ) where the Lipschitz constant L = λmax ( AA > ) and λmax ( . ) denotes the maximum eigenvalue . By initializing the iterative optimization at xl = 0 , it becomes clear that a single iteration of ( 3 ) is equivalent to a fully-connected layer followed by a non-linearity that is implemented by the proximal operator ( Fawzi et al. , 2015 ) . The choice of g ( x ) determines the specific form of the non-linearity ρ . Several popular activation functions can be traced back to their corresponding g ( x ) . The ReLU , which enforces non-negative output activations , corresponds to the indicator function g ( x ) = 1x≥0 ; the corresponding instance of problem ( 1 ) is a non-negative quadratic program . Similar observations for the ReLU have been made in other contexts ( Amos & Kolter , 2017 ; Papyan et al. , 2017 ) . We g ( x ) 1x≥0 γ 2 ∑ i max 2 ( −xi − λ , 0 ) −γ ∑ i log ( xi ) −γ ∑ i log ( 1− x2i ) Proxg ( η ) max ( 0 , η ) { η−γλ 1+γ if η ≤ −λ η if η ≥ −λ 1 2η + √ 1 4η 2 + γ Root of cubic polynomial Shape Activation = ReLU ( η ) = LeakyReLU ( η ) ≈ Softplus ( η ) ≈ Tanh ( η ) Table 1 : Different choices of g ( x ) , their corresponding proximal operators , and their relation to common activation functions . Squared hinge loss regularization of the activations yields a generalized Leaky ReLU . Log-barriers recover smooth activations , such as SoftPlus , Tanh , or Sigmoid . Derivations can be found in supplementary material . observe that many other activation functions fit this framework . For example , when g ( x ) is a squared hinge loss , i.e . γ2 ∑ i max 2 ( −xi − λ , 0 ) , a single update of ( 3 ) is equivalent to a linear layer followed by a Leaky ReLU . Table 1 lists some other choices of g ( x ) and their induced activations . Note that g ( x ) is not required to exhibit a simple , coordinate-wise separable structure . More complex functions can be used , as long as the proximal operator is easy to evaluate . Interesting examples arise when the output activations have matrix structure . For instance , one can impose nuclear norm regularization g ( X ) = ‖X‖∗ to encourage X to be low rank . Alternatively , one can enforce positive semi-definite structure on the matrix X by defining g ( X ) = 1X 0 . A similar activation has been used for higher-order pooling ( Huang & Van Gool , 2017 ) . In what follows , we will show that this connection can be further extended to explain dropout layers . Interestingly , specific forms of dropout do not arise from particular forms of objective ( 1 ) , but from different stochastic optimization algorithms that are applied to it .

Summary: This paper theoretically verifies an equivalence between stochastic solvers on a particular class of convex optimization problems and a forward pass through a dropout layer followed by a linear layer and a non-linear activation. Experiments show that replacing a block of layers with multiple iterations of the corresponding solver improves classification accuracy. My detailed comments are as follows.