AC-GC: Lossy Activation Compression with Guaranteed Convergence

1 Introduction . Stochastic Gradient Descent ( SGD ) has proven efficient and effective for optimizing Deep and Convolutional Neural Networks ( DNNs and CNNs ) . However , due to deeper and automatically generated networks [ 20 , 22 , 44 , 52 ] , improvement of accuracy has caused a rapid increase in training memory requirements , which are dominated by the temporary storage of activations between the forward and backward pass of the back-propagation algorithm [ 45 , 48 ] . Reducing memory consumption leads to faster training and , thus , more effective research of DNN models and applications . However , doing this by decreasing the batch size has many drawbacks . On parallel processors , such as GPUs , a small batch size can lead to poor compute saturation , and reduced training throughput [ 47 ] . Smaller batch sizes also introduce errors that impact convergence and accuracy [ 18 ] . Over 50GB of memory is required to train some networks , e.g. , GPIPE [ 22 ] . Many works have examined reducing activation storage overheads . Lossy compression of activations in memory can reduce memory footprint without network modifications [ 6 , 14 , 25 , 27 ] . Error bounded lossy compression ( EBC ) [ 27 ] has bounded activation error , however , it uses an empirical study to select an error target . Activations can also be offloaded to an external memory ( e.g . CPU DRAM ) , using either an uncompressed link [ 31 , 45 ] or compressed link [ 14 , 46 ] . Activation compression and offloading have performance overheads from 5 % to 60 % [ 7 , 14 , 25 , 27 , 45 ] . Reduced precision 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . training has the side effect of reducing activation size [ 9 , 51 , 57 ] . Finally , restructuring networks to be reversible [ 17 ] or efficient scheduling of network layers [ 8 ] can reduce memory use . Prior lossy and reduced precision approaches [ 6 , 9 , 14 , 25 , 51 , 57 ] utilize automated searches or hand-tuning to determine compression rates , which increase training time and have the potential to select poor compression/accuracy trade-offs . Ideally , an activation compression method has high compression and minimal decrease in trained accuracy . Tuning to achieve this is prohibitively expensive . For instance , selecting a fixed-point integer ( fixpoint ) compression rate for ImageNet/ResNet50 using Grid Search uses 16 training runs ( Figure 1 ) . Using SuccessiveHalving [ 26 ] can decrease training time , however even with aggressive resource allocations ( e.g . 24 GPU-days , Figure 1 ) total training time is still high . With low resource settings , methods such as SuccessiveHalving [ 26 ] and Hyperband [ 36 ] allocate little time to some configurations , increasing the likelihood that compression artifacts [ 6 , 25 ] can be missed , resulting in poor accuracy . Additionally , when tuning hyperparameters , lossy compression makes it difficult to determine the cause of degraded accuracy . Finally , prior compression methods have an unknown impact on convergence behavior . In this work , we present a framework for lossy Activation Compression with Guaranteed Convergence ( AC-GC ) . To our knowledge , our work is the first to prove convergence bounds on SGD with activation compression . AC-GC involves allowing an increase in the bound on the expected loss , which we trade-off for increased compression . We formulate this as a constrained optimization problem : maximizing compression subject to a bounded increase in loss . Doing so allows using a single hyperparameter to correlate convergence bounds with the activation error , which creates compression methods that are iteration , network , and dataset agnostic . Having convergence bounds known a priori allows a user to set a tolerable error rate before training , avoiding compression rate search cost entirely . Our contributions : • We prove convergence bounds on SGD under error bounded activation compression with weak assumptions on convexity . • We express activation compression and convergence as a constrained optimization problem and analyze the activation error tolerance of common DNN layers within this framework . • We combine these error bounds with fixpoint , image , and error bounded compression , to create methods with guaranteed convergence and a compression/accuracy trade-off known prior to training . 2 Preliminaries . DNNs are commonly used on problems involving a sum , for instance , minimize the total error on a set of training images . The loss L for such a problem takes the form L ( θ ) = ∑ n f ( θ , Xn ) ( 1 ) where f represents the loss of one example input Xn with weights θ. Stochastic Gradient Descent ( SGD ) is typically used to optimize these finite sums , using the iteration θ ( t+1 ) = θ ( t ) − α∇θf ( θ ( t ) , Xnt ) ( 2 ) where α is the learning rate , t is the iteration , ∇θf represents the gradient of f with respect to θ , and nt is a randomly chosen training example index from the distribution over n such that E [ ∇θf ( θ ( t ) , Xnt ) ] = ∇θL ( θ ( t ) ) . Figure 2a shows the computation graph for the back-propagation [ 48 ] algorithm for a single DNN layer without compression . Back-propagation is often used as it allows efficient calculation of parameter gradients at the expense of storing activations [ 48 ] . A DNN layer is any linear or non-linear function , e.g. , a convolution or ReLU activation . The functions fwd , bwd_param , bwd_act are algorithmic implementations of the layer function and the gradients w.r.t . θ and X . In the forward pass , each layer calculates an output activation Y = fwd ( X ) which is fed to subsequent layers . These activations are temporarily stored after use to avoid a performance penalty from recalculating them in the backward pass . To our knowledge , all frameworks opt to store activations [ 42 , 53 ] . In the backward pass , parameter gradients and activation gradients are calculated using bwd_param and bwd_act . Parameter gradients are used to update the parameters ( Eqn . ( 2 ) ) , and activation gradients are sent downward to the next layer . Depending on the layer type and its derivatives , the bwd_param and bwd_act functions may require activations to be stored . For example , it is computationally efficient to store the input X for convolution layers [ 42 , 53 ] . The many layer types place a diverse set of constraints on the activation storage . Activation compression addresses one of the most significant contributors to memory consumption in DNNs . In the forward pass , an activation can be lazily compressed after its last usage ( C , Figure 2b ) . Eagerly compressing activations would require storing both a compressed and uncompressed copy until its last use . In the backward pass , activations are decompressed before their first use ( D , Figure 2b ) . The backward pass begins only after the forward pass is completed for all layers , resulting in a large reuse distance for stored activations . Compression can thus be performed off the critical path in parallel with compute , with low performance overheads from 4 % -30 % [ 7 , 25 ] , provided that sufficient resources are available . We denote the uncompressed activation as X = ( xnchw ) ∈ RN×C×H×W , where N , C , H , and W represent the batch size , channel , height and width , respectively . In the uncompressed backward pass , gradients are calculated from the saved activations , parameters , and gradients from the upward layer ( Figure 2 with a ) ; we write this as∇θf ( θ , X ) = bwd_param ( X , θ , ∇Y f ( θ , X ) ) . Lossy compression involves discarding some information of the activation to increase the compression rate . In information theory , this is referred to as the rate-distortion trade-off . In our model , the rate refers to the activation error , and the distortion refers to resulting impacts on gradient error and thus accuracy after training . We model lossy compression between the forward and backward pass as an independent perturbation on each value in the activation , ∆X ∈ RN×C×H×W . Thus , in the compressed backward pass , the perturbed activation is X + ∆X . We denote the approximate gradient resulting from lossy error as ∇̂θf ( θ , X ) : = bwd_param ( X + ∆X , θ , ∇Y f ( θ , X ) ) , and the corresponding gradient error as ∆∇θf ( θ , X ) : = ∇̂θf ( θ , X ) −∇θf ( θ , X ) . 3 Guaranteed Convergence . This section details how gradient error ∆∇θf ( θ , X ) impacts convergence of SGD . Following this , the lossy compression error ∆X can be expressed in terms of the gradient error , and the compression rate for many methods can be determined . For example , the bitwidth b of fixpoint compression of an activation with range ( −1 , 1 ) is b ≥ − log2 |∆xnchw|+ ... ( 3 ) where the remaining terms are constants determined by the rounding mode , sign , etc . Any compression method with bounded activation error for a given rate can be combined with the error bounds from this work to guarantee convergence ( Sections 3.2 and 4 ) . 3.1 SGD Convergence . We will briefly summarize uncompressed convergence of SGD from Karimi et al . [ 28 ] . Consider training using t iterations of SGD with loss L ( θ ( t ) ) , with a constant learning rate α and initial point θ ( 0 ) . We assume that L has an optimal point θ ( ∗ ) and satisfies E [ ‖∇θf ( θ , Xnt ) ‖2 ] ≤ V 2 for all θ and some V 2 . We refer to V 2 as the variance . With some assumptions and problem-defined constants C1 and C2 ( Appendix A ) , Karimi et al . [ 28 ] demonstrate that the expected error at iteration t is E [ L ( θ ( t ) ) − L ( θ ( ∗ ) ) ] ≤ ( 1− C1α ) t ( L ( θ ( 0 ) ) − L ( θ ( ∗ ) ) ) + C2αV 2 ( 4 ) Initially when training , fast convergence occurs as ( 1− C1α ) t approaches zero ( Figure 3a ) . Later in training , C2αV 2 dominates , resulting in an approximately constant expected error ( Figure 3b ) . The key observation of this result is that the final loss scales with gradient variance , E [ L ( θ ( ∞ ) ) ] ∝ V 2 . Many DNN classes fall under this progress bound as it uses relatively weak assumptions and does not require a convex f . DNNs using ReLU activations are Lipschitz continuous [ 55 ] . Furthermore , those with an L2 loss are piecewise strongly convex , which implies that the required Polyak-Łojasiewic condition is satisfied locally [ 40 ] . Networks this does not apply to could use another progress bound [ 3 , 28 , 41 ] .

Summary: The authors bound the increase in loss in the case that activations used for backpropagation have errors. By doing so, they can set a fixed acceptable level of increase in the loss, and solve for the error level that can be tolerated for each activation used in backpropagation (and therefore deduce a level of compression). This analysis is relatively agnostic to the precise implementation of the compression. They evaluate their approach on a variety of datasets, including relatively large datasets such as ImageNet, demonstrating that their method can achieve high compression levels with relatively little tuning.