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: Overview:  This paper shows that the forward pass of a fully-connected layer (generalized to convolutions) followed by a nonlinearity in a neural network is equivalent to an iteration of a prox algorithm, where different regularizers in the objective of the related prox problem correspond to different nonlinearities such as ReLu. This connection is quite interesting. They further relate different stochastic prox algorithms to different dropout  layers and show results of improved performance on CIFAR-10 and CIFAR-100 on several architectures. The paper is well-written.