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: The paper proposes to approximate matrix multiplications (and convolutions) by randomly sampling column-row pairs and only multiplying submatrices. This is used to compute stochastic gradients which subsequently speed up training. Different sampling modes are presented and compared in experiments. Theoretical and practical training speed-ups are reported on standard architectures and datasets.