Faster Neural Network Training with Approximate Tensor Operations

1 Introduction . Approximation techniques for faster inference and training of deep neural networks ( DNNs ) have received considerable attention . Examples include quantization [ 1–5 ] , low-rank and structured-sparse models [ 6–9 ] , weight extrapolations [ 10 ] , and partial/asynchronous gradient updates in the context of distributed training [ 11 , 12 ] . Sampling-based approximations were used to accelerate inference [ 13 , 14 ] , but using them in training [ 15–17 ] has not been systematically studied nor demonstrated end-to-end GPU performance benefits in practice . We propose a novel approach to accelerating DNN training by systematically approximating tensor operations via sampling . At a high level , the original matrix products and convolutions are replaced with their faster approximate versions . The approximation is applied separately to each tensor operation , keeping the network architecture and tensor dimensions intact , thereby facilitating the adoption of this technique in existing DNN training frameworks , potentially in combination with other approximation techniques . Furthermore , when combined with distributed training , our technique allows for seamless reduction in the communication bandwidth and increased performance gains . We begin by reviewing the plethora of existing methods for approximating matrix multiplication . We compare several known algorithms [ 18–25 ] , and find column-row sampling ( CRS ) [ 20 ] to be the most suitable for approximating matrix multiplications in training . In order to compute the product of two matrices A > B , the CRS algorithm samples the columns of A > and the corresponding rows of B thus constructing smaller matrices which are then multiplied as usual . This method incurs low sampling overheads and lends itself to an efficient implementation using existing dense matrix product routines . CRS minimizes the approximation error for the Frobenius norm of the resulting matrix while keeping the approximation unbiased . ∗A Viterbi fellow 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Sampling-based approximations can be interpreted as a form of Dropout [ 26 ] , and we discuss the similarities and differences between the two . While Dropout aims to prevent overfitting , we focus on approximations as means to accelerate training by reducing the amount of computation . In this work we aim to answer two main questions . First , can neural networks be trained while using approximate tensor operations ? Second , what are the relations between using exact or approximate operations during training ? We start by analyzing the simpler case of linear regression , where we can derive the effects of approximations in closed form . We define a new loss function that takes the sampling into account , and observe that the resulting gradients differ from the exact training due to the dependency between sampled features . To this end , we propose a new Bernoulli-CRS variant which achieves statistical independence of samples , study its properties , and show that in linear regression it is equivalent to dynamic L2 weight regularization of the original , non-approximate loss . We then turn to the more general case of non-linear deep neural networks . We show that using sampling-based approximations in the backward pass provides the same convergence guarantees as the exact SGD for bounded weights . The convergence result holds for unbounded weights as well if approximation is applied only to the weight gradients and if the activation functions are bounded . We also study a new TopK-CRS algorithm which deterministically selects the top-k column-row pairs with the highest norms . We show that this algorithm is equivalent to the minimal mean square error estimator ( MMSE ) in case column-row pairs are pairwise independent with zero mean . Last , we generalize matrix product approximation to convolutions and analyze the approximation error to derive the optimal sampling policy . This allows us to apply approximations to training of convolutional neural networks ( CNNs ) . We implement our techniques in PyTorch [ 27 ] 2 and evaluate them on several DNN topologies , including MLP and CNN networks on MNIST [ 28 ] , Wide ResNet 28-10 [ 29 ] on CIFAR-10 [ 30 ] , and ResNet-50 and ResNet-152 [ 31 ] on ImageNet [ 32 ] . We demonstrate up to 66 % reduction in the number of FLOPs and up to 1.33x faster training time with little or no degradation in model accuracy . We develop another flavor of TopK-CRS which samples according to the weight norms only . When sampling the same subset of weights for different workers in a data-parallel setting , our sampling technique enables reducing the amount of gradient communication between workers . Notably , our algorithm is compatible with the standard AllReduce approach used in distributed deep learning . We implement an AllReduce scheme that takes advantage of the smaller gradient footprint and demonstrate 1.09x-1.37x speedup in multi-node training . Our contributions are as follows : • We derive general convergence guarantees for training with approximate tensor operations . • We develop novel sampling algorithms and analyze their theoretical properties . • We extend sampling-based algorithms to fast approximation of multi-channel convolutions . • We show that our approach can reduce the computation , communication and total training time on several popular neural network architectures with little or no accuracy degradation . 2 Related work . To the best of our knowledge , we are the first to study the application of sample-based approximations of tensor operations to speed up DNN training . However , there have been several prior efforts to accelerate DNN computations via approximation which we survey below . Several works accelerate inference through model compression [ 33–39 ] . A large body of work is devoted to quantization and low-precision datatypes ( see for example [ 1 , 2 , 5 ] ) . Approximation was used to extrapolate weight values [ 10 ] . Another approach enforces low-rank or structured-sparse structure on the layers , resulting in lower computational cost both for training and inference [ 6–9 ] . Other works accelerate inference by approximating large matrices as products of lower-ranked ones [ 40 , 41 ] or through locality-sensitive hashing [ 42 ] . 2https : //github.com/acsl-technion/approx In the context of distributed training , several works targeted communication bottlenecks by gradient quantization [ 3 , 4 ] , delayed weight updates [ 11 , 12 ] and low-rank approximation of the gradient matrix [ 43 ] . These methods are complementary and compatible with ours . Sampling-based approximations were used to accelerate inference [ 13 , 14 ] , but using them for training [ 15–17 ] has not been systematically studied nor shown to speed up training on GPUs without accuracy degradation . Sub-sampling whole layers was shown to enable training of very deep CNNs [ 44 ] . 2.1 Approximate matrix multiplication . There are several known algorithms for approximating matrix product . However , only those that meet the following requirements will be effective for DNN training . First , the algorithm should apply to dense matrices of arbitrary dimensions . Second , to reduce training time , the overall multiplication including input transformation should be faster than the original product . Last , the algorithm should be amenable to efficient implementation on commodity hardware . Using these criteria , we consider the following algorithms : Random walk [ 18 ] This algorithm performs random walks on a graph representation of the input matrices , but is applicable to non-negative matrices only . Random projections [ 21–23 ] The two matrices to be multiplied are first projected into a lowerdimensional subspace by a scaled random sign matrix . These algorithms require both input matrices to be roughly square , otherwise the cost of projection will be similar to the cost of original product . In DNNs , however , it is common for one dimension to be smaller than the other . FFT [ 24 , 25 ] These algorithms represent each column-row outer product as a polynomial multiplication and then calculate it using Fast Fourier Transform . The complexity depends on the sparsity of the input matrices , decreasing as the sparsity increases . Therefore , these algorithms might not be effective for dense matrices . SVD [ 19 , 33 , 36 ] Several algorithms replace one input matrix with its low-rank approximation using truncated SVD . These algorithms are suitable for inference where the weight matrix factorization can be pre-computed offline , but are not applicable to training since the high cost of factorization is incurred in every matrix product . Column-row sampling ( CRS ) [ 19 , 20 ] The sampling algorithm approximates matrix product A > B by sampling k columns of A > and respective rows of B to form smaller matrices , which are then multiplied as usual . We choose CRS as the basis for our current work because it meets all the criteria above : It is applicable to fully-connected layers of any size , its effectiveness does not depend on the matrix contents , its sampling is computationally lightweight , and may use regular matrix multiplication algorithms since the sampled sub-matrices remain dense . 2.2 CRS . Let A ∈ Rn×m , B ∈ Rn×p . Their product A > B is approximated as a weighted sum of outer products between sampled columns of A > and corresponding rows of B : A > B ≈ k∑ t=1 1 kpit A > ( it ) B ( it ) ( 1 ) where A > ( i ) , B ( i ) denote the matrix i ’ th column and row respectively , k is the number of samples ( satisfying 1 ≤ k ≤ n ) , { pi } ni=1 is a probability distribution over the column-row pairs of A > , B and it ∈ { 1 , ... , n } . This algorithm allows linear reduction in complexity from O ( mnp ) to O ( mkp ) . ( 1 ) can also be expressed as A > DS > SDB , where D ∈ Rn×n is a diagonal scaling matrix with : ( D ) j , j = 1√ kpj ( 2 ) and S ∈ Rk×n is a sampling matrix that selects k features , possibly with replacement . S is a random matrix , where each row has 1 in one entry and zeros in others . In each row , the probability of having the non-zero entry in column j is pj . Drineas et al . [ 20 ] show that CRS is unbiased : E [ A > DS > SDB ] = A > B ( 3 ) and also derive upper bounds for the expected Frobenius and spectral norms of the error matrix∣∣∣∣A > B −A > DS > SDB∣∣∣∣ . They show that the error is minimized when the sampling probabilities are proportional to the product of the column-row Euclidean norms : pi = |A ( i ) ||B ( i ) |∑n j=1|A ( j ) ||B ( j ) | ( 4 ) In which case the expected Frobenius error is : 1 k ( k∑ t=1 |A ( it ) ||B ( it ) | ) 2 − 1 k ∣∣∣∣A > B∣∣∣∣2 F ( 5 ) 2.3 Approximate Tensor Operations and Dropout . Sampling-based approximations can be interpreted as a flavor of Dropout [ 26 ] , a popular technique to prevent overfitting by randomly zeroing individual activations during training . However , the sparsity pattern resulting from Dropout is unstructured and therefore can not be exploited efficiently by GPUs despite recent advances in structured sparsity support [ 8 ] . Prior works on fast Dropout training [ 45 , 46 ] are different than ours and do not demonstrate acceleration of large networks while maintaining accuracy . Some works proposed non-uniform Dropout probabilities for individual activations [ 47 ] or full channels [ 48 ] . Their sampling heuristics are different from ours which are derived from optimal approximation . Furthermore , they use Dropout only for preventing overfitting and do not leverage it to speed up training . In our experiments we demonstrate the utility of sampling-based approximations for DNNs with and without Dropout . Conversely , we did not observe improved accuracy from approximations which could have been attributed to reduced overfitting .

Summary: This paper adapts the CRS method for approximate matrix multiplication into a new method, Bernoulli-CRS, that yields unbiased gradient estimates. The paper then analyzes this method on linear regression and shows the new method and giving unbiased estimates of a gradient for the original objective with a regularization term. They then devise a new routine for approximate backpropagation that has unbiased estimates with bounded second moments. The paper then introduces a deterministic variant called top-k sampling and show the performance of the 3 methods (CRS, Bernoulli-CRS, top-k) in experiments in terms of performance, compute, and speed.