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: This paper studies activation compression to reduce the memory consumption when training models with stochastic gradient descent. The goal is to establish an activation compression framework with convergence guarantees, while eliminate the tuning overhead from conventional search based method for compression precision selection. The contributions of this paper come in three folds:  1. The authors first established a connection between the activation compression error and the convergence bound (i.e. the upper bound on training loss).  2. Built on the above theory, the authors formulate the activation compression precision selection problem as a constrained optimization problem, which can accommodate typical DNN layers such as linear layers, convolution layers and normalization layers.  3. Empirically, the authors combine this optimization-based activation precision selection method with various compression methods (with bounded errors at each compression precision). They show that it could attain matching accuracy (minimal acc degradation compared to no activation compression) compared to conventional grid search method for activation precision without the grid search (i.e. hyperparameter tuning) overhead.