Backprop with Approximate Activations for Memory-efficient Network Training

1 INTRODUCTION . Deeper neural network models are able to express more complex functions , and recent results have shown that with the use of residual ( He et al. , 2016a ) and skip ( Huang et al. , 2017 ) connections to address vanishing gradients , such networks can be trained effectively to leverage this additional capacity . As a result , the use of deeper network architectures has become prevalent , especially for visual inference tasks ( He et al. , 2016b ) . The shift to larger architectures has delivered significant improvements in performance , but also increased demand on computational resources . In particular , deeper network architectures require significantly more on-device memory during training—much more so than for inference . This is because training requires retaining the computed activations of all intermediate layers since they are needed to compute gradients during the backward pass . The increased memory footprint means fewer training samples can fit in memory and be processed as a batch on a single GPU . This is inefficient : smaller batches are not able to saturate all available parallel cores , especially because computation in “ deeper ” architectures is distributed to be more sequential . Moreover , smaller batches also complicate the use of batch-normalization ( Ioffe & Szegedy , 2015 ) , since batch statistics are now computed over fewer samples making training less stable . These considerations often force the choice of architecture to be based not just on optimality for inference , but also practical feasibility for training—for instance , deep residual networks for large images drop resolution early , so that most layers have smaller sized outputs . While prior work to address this has traded-off memory for computation ( Martens & Sutskever , 2012 ; Chen et al. , 2016 ; Gruslys et al. , 2016 ; Gomez et al. , 2017 ) , their focus has been on enabling exact gradient computation . However , since stochastic gradient descent ( SGD ) inherently works with noisy gradients at each iteration , we propose an algorithm that computes reasonably approximate gradients , while significantly reducing a network ’ s memory footprint and with virtually no additional computational cost . Our work is motivated by distributed training algorithms that succeed despite working with approximate and noisy gradients aggregated across multiple devices ( Recht et al. , 2011 ; Dean et al. , 2012 ; Seide et al. , 2014 ; Wen et al. , 2017 ) . We propose using low-precision approximate activations—that require less memory—to compute approximate gradients during back-propagation ( backprop ) on a single device . Note that training with a lowerprecision of 16- instead of 32-bit floating-point representations is not un-common . But this lower precision is used for all computation , and thus allows only for a modest lowering of precision , since the approximation error builds up across the forward and then backward pass through all layers . In this work , we propose a new backprop implementation that performs the forward pass through the network at full-precision , and incurs limited approximation error during the backward pass . We use the full-precision version of a layer ’ s activations to compute the activations of subsequent layers . However , once these activations have been used in the forward pass , our method discards them and stores a low-precision approximation instead . During the backward pass , gradients are propagated back through all the layers at full precision , but instead of using the original activations , we use their low-precision approximations . As a result , we incur an approximation error at each layer when computing the gradients to the weights from multiplying the incoming gradient with the approximate activations , but ensure the error in gradients going back to the previous layer is minimal . Our experimental results show that even using only 4-bit fixed-point approximations , for the original 32-bit floating-point activations , causes only minor degradation in training quality . This significantly lowers the memory required for training , which comes essentially for “ free ” —incurring only the negligible additional computational cost of converting activations to and from low precision representations . Our memory-efficient version of backprop is thus able to use larger batch sizes at each iteration—to fully use available parallelism and compute stable batch statistics—and makes it practical for researchers to explore the use of much larger and deeper architectures than before . 2 RELATED WORK . A number of works focus on reducing the memory footprint of a model during inference , e.g. , by compression Han et al . ( 2015 ) and quantization Hubara et al . ( 2017 ) , to ensure that it can be deployed on resource-limited mobile devices , while still delivering reasonable accuracy . These methods still require storing full versions of model weights and activations during network training , and assume there is sufficient memory to do so . However , training requires significantly more memory than inference because of the need to store all intermediate activations . And so , memory can be a bottleneck during training , especially with the growing preference for larger and deeper network architectures . A common recourse to this has been to simply use multiple GPUs during training . But , this is inefficient due to the overhead of intra-device communication , and often under-utilizes the available parallelism on each device—computation in deeper architectures is distributed more sequentially and , without sufficient data parallelism , often does not saturate all GPU cores . A popular strategy to reduce training memory requirements is “ checkpointing ” . Activations for only a subset of layers are stored at a time , and the rest recovered by repeating forward computations ( Martens & Sutskever , 2012 ; Chen et al. , 2016 ; Gruslys et al. , 2016 ) . This affords memory savings with the trade-off of additional computational cost—e.g. , Chen et al . ( 2016 ) propose a strategy that requires memory proportional to the square-root of the number of layers , while requiring up to the computational cost of an additional forward pass . In a similar vein , Gomez et al . ( 2017 ) considered network architectures with “ reversible ” or invertible layers . This allows re-computing intermediate input activations of reversible layers from their outputs during the backward pass . These methods likely represent the best possible solutions if the goal is restricted to computing exact gradients . But SGD is fundamentally a noisy process , and the exact gradients computed over a batch at each iteration are already an approximation—of gradients of the model over the entire training set ( Robbins & Monro , 1985 ) . Researchers have posited that further approximations are possible without degrading training ability , and used this to realize gains in efficiency . For distributed training , asynchronous methods ( Recht et al. , 2011 ; Dean et al. , 2012 ) delay synchronizing models across devices to mitigate communication latency . Despite each device now working with stale models , there is no major degradation in training performance . Other methods quantize gradients to two ( Seide et al. , 2014 ) or three levels ( Wen et al. , 2017 ) so as to reduce communication overhead , and again find that training remains robust to such approximation . Our work also adopts an approximation strategy to gradient computation , but targets the problem of memory usage on a each device . We approximate activations , rather than gradients , with lower-precision representations , and by doing so , we are able to achieve considerable reductions in a model ’ s memory footprint during training . Note that since our method achieves a constant factor saving in memory for back-propagation across any group of layers , it can also be employed within checkpointing to further improve memory cost . It is worth differentiating our work from those that carry out all training computations at lowerprecision Micikevicius et al . ( 2017 ) ; Gupta et al . ( 2015 ) . This strategy allows for a modest loss in precision : from 32- to 16-bit representations . In contrast , our approach allows for much greater reduction in precision . This is because we carry out the forward pass in full-precision , and approximate activations only after they have been used by subsequent layers . Our strategy limits accumulation of errors across layers , and we are able to replace 32-bit floats with 8- and even 4-bit fixed-point approximations , with little to no effect on training performance . Note that performing all computation at lower-precision also has a computational advantage : due to reduction in-device memory bandwidth usage ( transferring data from global device memory to registers ) in Micikevicius et al . ( 2017 ) , and due to the use of specialized hardware in Gupta et al . ( 2015 ) . While the goal of our work is different , our strategy can be easily combined with these ideas : compressing intermediate activations to a greater degree , while also using 16- instead of 32-bit precision for computation . 3 PROPOSED METHOD . 3.1 BACKGROUND . A neural network is composition of linear and non-linear functions that map the input to the final desired output . These functions are often organized into “ layers ” , where each layer consists of a single linear transformation—typically a convolution or a matrix multiply—and a sequence of nonlinearities . We use the “ pre-activation ” definition of a layer , where we group the linear operation with the non-linearities that immediately preceed it . Consider a typical network whose lth layer applies batch-normalization and ReLU activation to its input Al : i followed by a linear transform : [ Batch Norm . ] Al:1 = ( σ2l + ) −1/2 ◦ ( Al : i − µl ) , µl = Mean ( Al : i ) , σ2 = Var ( Al : i ) ; ( 1 ) [ Scale & Bias ] Al:2 = γl ◦Al:1 + βl ; ( 2 ) [ ReLU ] Al:3 = max ( 0 , Al:2 ) ; ( 3 ) [ Linear ] Al : o = Al:3 ×Wl ; ( 4 ) to yield the output activations Al : o that are fed into subsequent layers . Here , each activation is a tensor with two or four dimensions : the first indexing different training examples , the last corresponding to “ channels ” , and others to spatial location . Mean ( · ) and Var ( · ) aggregate statistics over batch and spatial dimensions , to yield vectors µl and σ2l with per-channel means and variances . Element-wise addition and multiplication ( denoted by ◦ ) are carried out by “ broadcasting ” when the tensors are not of the same size . The final operation represents the linear transformation , with × denoting matrix multiplication . This linear transform can also correspond to a convolution . Note that ( 1 ) - ( 4 ) are defined with respect to learnable parameters γl , βl , and Wl , where γl , βl are both vectors of the same length as the number of channels in Al , and Wl denotes a matrix ( for fullyconnected layers ) or elements of a convolution kernel . These parameters are learned iteratively using SGD , where at each iteration , they are updated based on gradients—∇γl , ∇βl , and ∇Wl—of some loss function with respect to these parameters , computed on a batch of training samples . To compute gradients with respect to all parameters for all layers in the network , the training algorithm first computes activations for all layers in sequence , ordered such that each layer in the sequence takes as input the output from a previous layer . The loss is computed with respect to activations of the final layer , and then the training algorithm goes through all layers again in reverse sequence , using the chain rule to back-propagate gradients of this loss . For the lth layer , given the gradients ∇Al : o of the loss with respect to the output , this involves computing gradients ∇γl , ∇βl , and∇Wl with respect to the layer ’ s learnable parameters , as well as gradients∇Al : i with respect to its input for further propagation . These gradients are given by : [ Linear ] ∇W = ATl:3 × ( ∇Al : o ) , ∇Al:3 = ( ∇Al : o ) ×WTl ; ( 5 ) [ ReLU ] ∇Al:2 = δ ( Al:2 > 0 ) ◦ ( ∇Al:3 ) ; ( 6 ) [ Scale & Bias ] ∇βl = Sum ( ∇Al:2 ) , ∇γl = Sum ( Al:1 ◦ ( ∇Al:2 ) ) , ∇Al:1 = γl ◦ ∇Al:2 ; ( 7 ) [ Batch Norm . ] ∇Al : i = ( σ2l + ) −1/2 ◦ [ ∇Al:1 −Mean ( ∇Al:1 ) −Al:1 ◦Mean ( Al:1 ◦ ∇Al:1 ) ] ; ( 8 ) where Sum ( · ) and Mean ( · ) involve aggregation over all but the last dimension , and δ ( A > 0 ) a tensor the same size as A that is one where the values in A are positive , and zero otherwise . When the goal is to just compute the final output of the network , the activations of an intermediate layer can be discarded during the forward pass as soon as we finish processing the subsequent layer or layers that use it as input . However , we need to store all these intermediate activations during training because they are needed to compute gradients during back-propagation : ( 5 ) - ( 8 ) involve not just the values of the incoming gradient , but also the values of the activations themselves . Thus , training requires enough available memory to hold the activations of all layers in the network .

Summary: This paper proposes to use 8/4-bit approximation of activations to save the memory cost during gradient computation.  The proposed technique is simple and straightforward. On the other hand, the proposed method only saves up to a constant cost of the storage. With the constant factor (4x, 8x) depending on whether fp16 or fp32 is used during computation. Notably, there is a small but noticeable accuracy drop in the final trained model using this mechanism.