Lifted Proximal Operator Machines

We propose a new optimization method for training feed-forward neural networks. By rewriting the activation function as an equivalent proximal operator, we approximate a feed-forward neural network by adding the proximal operators to the objective function as penalties, hence we call the lifted proximal operator machine (LPOM). LPOM is block multi-convex in all layer-wise weights and activations. This allows us to use block coordinate descent to update the layer-wise weights and activations in parallel. Most notably, we only use the mapping of the activation function itself, rather than its derivatives, thus avoiding the gradient vanishing or blowup issues in gradient based training methods. So our method is applicable to various non-decreasing Lipschitz continuous activation functions, which can be saturating and non-differentiable. LPOM does not require more auxiliary variables than the layer-wise activations, thus using roughly the same amount of memory as stochastic gradient descent (SGD) does. We further prove the convergence of updating the layer-wise weights and activations. Experiments on MNIST and CIFAR-10 datasets testify to the advantages of LPOM. Introduction Feed-forward deep neural networks (DNNs) are cascades of fully connected layers and there are no feedback connections. In recent years, with the advances in hardware and dataset sizes, feed-forward DNNs have become standard in many tasks, such as image recognition (Krizhevsky, Sutskever, and Hinton, 2012), speech recognition (Hinton et al., 2012), natural language understanding (Collobert et al., 2011), and as a building block of the Go game learning system (Silver et al., 2016). For several decades, training a DNN is accomplished by optimizing a highly nonconvex and nested function of the network weights. The predominant method for training DNNs is stochastic gradient descent (SGD) (Rumelhart, Hinton, and Williams, 1986), whose effectiveness has been demonstrated by the successes of DNNs in various real-world applications. Recently, many variants of SGD have been proposed, which use adaptive learning rates and momentum terms, e.g., Nesterov momentum (Sutskever et al., 2013), AdaGrad (Duchi, Copyright c © 2019, Association for the Advancement of Artificial Intelligence (www.aaai.org). All rights reserved. Hazan, and Singer, 2011), RMSProp (Dauphin, de Vries, and Bengio, 2015), and Adam (Kingma and Ba, 2014). SGD and its variants use a few training samples to estimate the full gradient, making the computational complexity of each iteration small. Moreover, the estimated gradients have noise, which is helpful for escaping saddle points (Ge et al., 2015). However, they have some drawbacks as well. One major problem is the vanishing or blow-up gradient issue, where the magnitudes of gradients decrease or increase exponentially with the number of layers. This causes slow or unstable convergence, especially in very deep networks. This flaw can be remitted by using non-saturating activation functions, such as rectified linear unit (ReLU), and modified network architectures, such as ResNet (He et al., 2016). However, the fundamental problem remains. Furthermore, they cannot deal with nondifferentiable activation functions directly (e.g., binarized neural networks (Hubara et al., 2016)) and do not allow parallel weight updates across the layers. For more discussions on the limitations of SGD, please refer to (Taylor et al., 2016). The drawbacks of SGD motivate research on alternative training methods for DNNs. Recently, training a feed-forward neural network is formulated as a constrained optimization problem, where the network activations are introduced as auxiliary variables and the network configuration is guaranteed by layer-wise constraints (Carreira-Perpinan and Wang, 2014). It breaks the dependency among the nested functions into equality constraints, so many standard optimization methods can be utilized. Some methods of this type of approach were studied and they differed in how to handle the equality constraints. Carreira-Perpinan and Wang (2014) approximated the equality constraints via quadratic penalties and alternately optimized network weights and activations. Zeng et al. (2018) introduced one more block of auxiliary variables per layer and also approximated the equality constraints via quadratic penalties. Inspired by alternating direction method of multiplier (ADMM) (Lin, Liu, and Su, 2011), Taylor et al. (2016) and Zhang, Chen, and Saligrama (2016) used the augmented Lagrangian approach to obtain exact enforcement of the equality constraints. However, the two methods involved the Lagrange multipliers and nonlinear constraints, thus were more memory demanding and more difficult in optimization. Motivated by the fact that the ReLU activation function is equivalent to a simple constrained convex minimization problem, Zhang and Brand (2017) relaxed the ar X iv :1 81 1. 01 50 1v 1 [ cs .L G ] 5 N ov 2 01 8 nonlinear constraints as penalties, which encode the network architecture and the ReLU activation function. Thus, the nonlinear constraints no longer exist. However, their approach is limited to the ReLU function and does not apply to other activation functions. Askari et al. (2018) followed this idea by considering more complex convex optimization problems and discussed several types of non-decreasing activation functions. However, their methods to update the weights and activations are still limited to the ReLU function. Their approach cannot outperform SGD and can only serve for producing good initialization for SGD. Actually, we have found that their formulation was incorrect (see Subsection “Advantages of LPOM”). This paper makes the following contributions: • We propose a new formulation to train feed-forward DNNs, which we call the lifted proximal operator machine (LPOM)1. LPOM is block multi-convex, i.e., the problem is convex w.r.t. weights or activations of each layer when the remaining weights and activations are fixed. In contrast, almost all existing DNN training methods do not have such a property. This greatly facilitates the training of DNNs. • Accordingly, we apply block coordinate descent (BCD) to solve LPOM, where the layer-wise weights and activations can be updated in parallel. Most notably, the update of the layer-wise weights or activations only utilizes the activation function itself, rather than its derivatives, thus avoiding the gradient vanishing or blow-up issues in gradient based training methods. Moreover, LPOM does not need more auxiliary variables than the layer-wise activations, thus its memory cost is close to that of SGD. We further prove that the iterations to update layer-wise weights or activations are convergent. • Since only the activation function itself is involved in computation, LPOM is able to handle general non-decreasing2 Lipschitz continuous activation functions, which can be saturating (such as sigmoid and tanh) and non-differentiable (such as ReLU and leaky ReLU). So LPOM successfully overcomes the computation difficulties when using most of existing activation functions. We implement LPOM on fully connected DNNs and test it on benchmark datasets, MNIST and CIFAR-10, and obtain satisfactory results. For convolutional neural networks (CNNs), since we have not reformulated pooling and skip-connections, we leave the implementation of LPOM on CNNs to future work. Note that the existing non-gradient based approaches also focus on fully connected DNNs first (Carreira-Perpinan and Wang, 2014; Zeng et al., 2018; Taylor et al., 2016; Zhang, Chen, and Saligrama, 2016; Zhang and Brand, 2017; Askari et al., 2018). Related Work In a standard feed-forward neural network, the optimization problem to train an n-layer neural network for classification Patent filed. Although our theories can be easily extended to non-monotone activation functions, such a case is of less interest in reality. tasks is: min {W i} ` ( φ(Wn−1φ(· · ·φ(W φ(W X)) · · · )), L ) , (1) where X ∈ Rn1×m is a batch of training samples, L ∈ Rc×m denotes the corresponding labels, n1 is the dimension of the training samples,m is the batch size, c is the number of classes, {W i}n−1 i=1 are the weights to be learned in which the biases have been omitted for simplicity, φ(·) is an elementwise activation function (e.g., sigmoid, tanh, and ReLU), and `(·, ·) is the loss function (e.g., the least-square error or the cross-entropy error). Here the neural network is defined as a nested function, where the first layer function of the neural network is φ(W X), the i-th layer (i=2, · · · , n) function has the form φ(W X), and X is the output of the (i−1)th layer function. A common approach to optimize (1) is by SGD, i.e., calculating the gradient w.r.t. all weights of the network using backpropagation and then updating the weights by gradient descent. By introducing the layer-wise activations as a block of auxiliary variables, the training of a neural network can be equivalently formulated as an equality constrained optimization problem (Carreira-Perpinan and Wang, 2014): min {W i},{Xi} `(X, L) s.t. X=φ(W i−1Xi−1), i=2, 3, · · · , n, (2) where X is the activation of the i-th layer and other notations are the same as those in (1). The constraints in (2) ensure that the auxiliary variables {X}i=2 exactly match the forward pass of the network. Compared with problem (1), problem (2) is constrained. But since the objective function is not nested, hence much simpler, such an equivalent reformulation may lead to more flexible optimization methods. Note that when using SGD to solve problem (1), it actually works on problem (2) implicitly as the activations {X}i=2 need be recorded in order to compute the gradient. Inspired by the quadratic-penalty method, CarreiraPerpinan and Wang (2014) developed the method of auxiliary coordinates (MAC) to solve problem (2). MAC uses quadratic penalties to approximately enforce equality constraints and tries to solve the fo