Memory-efficient training with streaming dimensionality reduction

As the market demand for deep neural networks grows, both data centers and deep neural network frameworks are constantly being improved to accommodate the ever changing demands of machine learning workloads. To fully exploit the hardware available within the data center, deep neural network frameworks will have to natively support methods for model and data parallelism. Ideally, this support would allow for a linear scaling of the model training speedup with the number of available computing resources. These hardware resources include traditional central processing units (CPUs), graphical/tensor processing units (GPUs/TPUs), as well as data communication buses such as peripheral connection interconnect express (PCIe) buses.

The most straightforward method, data parallelism, simply duplicates the machine learning model across all of the available resources as efficiently as possible [1, 2]. These multiple instances can then be used to process batches of the data in parallel. In between processing the data, weight and gradient information is transmitted throughout the computing platform so that the duplicate models can be updated. The synchronization of the network models is a time consuming process due to the limited communication bandwidth available and the large size of the models. In some cases, communication of the gradient and weight information can be the most time consuming step during the training process, taking up more than 60 % of the overall training time [14]. Even in hybrid approaches using model parallelism, where separate parts of the model can operate independently but synchronously or even asynchronously, transmission of gradient information and model synchronization are important bottlenecks.

To deal with this problem, different data reduction approaches have been considered. The simplest, gradient quantization, is to reduce the precision of the gradient information transmitted [2]. Other approaches, such as gradient sparsification, transmit only the largest values of the gradient, or, only sparsely communicate the weight updates when they reach a threshold of significance [2]. Much effort has been invested in these approaches, enabling reductions in the volume of gradient data transmission by several hundred times without reducing accuracy [15]. However, such approaches may have high local memory and computational overhead, requiring significant logic to be performed over the full set of estimated gradients and weights within the parallel model.

Principal component analysis provides an efficient method for compressing the gradient information, by extracting the most important components. As such, it could significantly contribute to the effectiveness of parallelization of gradient updates.

The movement of data, weight matrices, and gradients is a major source of latency and resource utilization in deep neural networks. One approach to minimizing this movement is to build nonvolatile memories that can simultaneously store the weights and compute the multiply and accumulate operations commonly implemented in GPUs and TPUs [8, 10]. Performing these operations in the nonvolatile memory itself eliminates the cycling of data in and out of the local GPU/TPU memories.

The principal advantage afforded by this scheme is the amount of information that can be stored compared to other approaches. On one hand a TPU can store approximately 65 kilobytes of weight data within its primary systolic array, which is only a small fraction of an overall deep neural network model. On the other hand, a nonvolatile memory array based on flash memories, resistive random accesses memory (ReRAM), or magnetic tunnel junctions, could potentially store terabytes of model data at much greater density [16, 8]. Different methodologies for implementing the operation of these arrays include both analog and digital approaches, such as binary neural networks [6, 10, 17]. In either case, multiply and accumulate operations are encoded in charge and current summation, natively implemented via Kirchoff’s law, allowing for massively parallel multiplication of the model data with the input activation or error. One such emerging computing architecture that potentially exploits this physics is the three dimensional system on a chip (also known as 3DSoC) architecture. The 3DSoC incorporates dense vertical ReRAM with both silicon and carbon-nanotube digital logic [18].

The principal weaknesses of such approaches include the high programming energy of nonvolatile memories, which would favor sparse updates to the network, as well as the difficulty of storing and processing the immense quantities of training and gradient data. Whereas the model can be stored in the nonvolatile memory arrays, there is at present no compatible ultra dense short term memory array with which to store the training and gradient data, though potential approaches are under investigation [19]. Existing implementations of L1 and L2 caches could be used to temporarily store this data for transfer into the nonvolatile memory arrays, but efficient compression of the training and gradient data would be required.

Batch PCA significantly reduces the number of hardware updates for each training data set. It allows for outer-product updates in which all of the components are updated at once in the most important directions. It also compresses a number of input data into a single best average update further reducing the rate of hardware updates.

Memory-efficient compression of large matrices without explicitly calculating them has a robust history of theoretical investigation [20, 21, 22]. For one approach, streaming principal component analysis (SPCA), the objective is to extract the top $k$ principal components ($\hat{U}_{PCA}$) from a stream of incoming data $\{\vec{x}_{i}\}$. Due to the high dimensionality of $\vec{x}_{i}$, calculation of the covariance matrix $\sum_{i}\vec{x}_{i}\vec{x}_{i}^{T}$ and its singular value decomposition (SVD) would be costly, particularly in terms of memory, as the covariance matrix (as in Fig. 1) might not even be able to fit within the available memory. Inspired by neural biology, Oja proposed a linearized version of stochastic power iterations that can be used to extract one, or, with deflation, multiple principal components of a random data stream [23, 24]. More modern approaches, such as history PCA, rely instead on $QR$-factorization to ensure rank orthonormality [25, 26, 27, 13]. Update and reorthonormalization of the $Q$ matrix allows the algorithm to span a greater fraction of the input data space, reducing the effective eigengap of the stochastic power iteration and improving the rate of convergence[28, 29].

As we discuss below, our proposed approach does not map perfectly to streaming PCA; we are trying to extract the singular values of a general rectangular gradient matrix, rather than a square and symmetric covariance matrix. We combine streaming PCA methods with bi-iterative power iterations to perform stochastic approximation of the top singular vectors. Such an approach was first used by Strobach to perform subspace tracking [30]. To the best of our knowledge, such an approach has not been used previously to extract the top $k$ principal components of a neural network gradient matrix during training.

In neural networks, finding a solution to a classification problem is, for a given network structure of weights and neurons, mapped to minimizing a scalar loss function, $\ell$, which correlates inputs with correct classifications. In classical gradient descent, we define the differential of the loss function $d\ell$ over all of the weights $\vec{\Theta}$ in a layer as

During training, we compute $\partial\ell/\partial\vec{\Theta}$ at every step. The negative of this gradient gives the direction of steepest descent. Choosing a step of $\Delta\vec{\Theta}=-\alpha(\partial\ell/\partial\vec{\Theta})$ changes the loss function by

where $\alpha$ is the learning rate.

Though we have presented the weights in vector form thus far, we could instead suppose that they are arrayed in a matrix $\hat{\Theta}$, as is typical in neural networks. In that case we recognize the norm taken in Eq. (2) as the Frobenius norm of $\nabla_{\hat{\Theta}}\ell$. It is well known that Eq. (2) can then be written as

where the $\sigma_{j}$ are the singular values of $\nabla_{\hat{\Theta}}\ell$ in its singular value decomposition. This leads us to the following observation. Suppose $\nabla_{\hat{\Theta}}\ell$ could be well-approximated by a truncated singular value decomposition using only the first $k$ ranks. We write this approximation as $\nabla_{\hat{\Theta}}^{(k)}\ell$. Then moving along $\nabla^{(k)}_{\hat{\Theta}}\ell$ instead of $\nabla_{\hat{\Theta}}\ell$ produces $\Delta\ell$ almost as large as the exact gradient descent, up to a small error on the order of $O(\sigma_{k+1}^{2}/\sigma_{1}^{2})$. We conjecture that using the truncated decomposition to propose weight matrix updates may still allow one to minimize the loss function, but with a smaller memory footprint than would be needed to store the full-rank gradient matrix.

In practice, $\Delta\hat{\Theta}$ is typically not computed with the true gradient of $\ell$, but a stochastic approximation of the gradient over a relatively small batch of samples of size $B$, much less than the size of the full dataset. This balances the acceleration of training available from the pipelining of $B$ inferences at a time against the desired stochasticity of the sampled gradient. For our proposed method, the equivalent stochastic algorithm would generate an approximation of $\tilde{\nabla}^{(k,B)}_{\hat{\Theta}}\ell$ which is randomly sampled from a $B$-batch of the overall training set and is also of tunable truncation rank $k$.

In order to generate a $k$-rank stochastic approximation of the matrix $\Delta\hat{\Theta}\in\mathbb{R}^{m\times n}$, we start from the truncated singular value decomposition for the approximate gradient

where $\hat{\Delta}\in\mathbb{R}^{m\times k}$ is a matrix composed of the top $k$ left singular vectors (that is, singular backpropagated errors) and $\hat{X}\in\mathbb{R}^{n\times k}$ is a matrix of the top $k$ right singular vectors (that is, singular activations). $\hat{\Sigma}$ is a square $k\times k$ matrix with singular values $\vec{\sigma}$ on the diagonal and zeroes elsewhere. Fig. 2 depicts the breakdown and dimensionality of this decomposition with respect to the weight matrix in the network layer. The memory cost of this decomposed form is $mk+k+nk=k(m+n+1)$, which, for sufficiently small $k$, is significantly less than the memory overhead of the full $m\times n$ matrix. The challenge now is to determine the correct values of $\hat{X}$, $\hat{\Delta}$, and $\vec{\sigma}$ in a computationally efficient way, given a list of input activations $\vec{x}_{i}$ and backpropogated vectors $\vec{\delta}_{i}$ of a fully connected layer in a neural network.

Performing the singular value decomposition using traditional methods will fail to deliver efficient memory and compute footprints. Suppose that for a particular batch of inputs $\{\vec{x}_{i}\}$ and errors $\{\vec{\delta}_{i}\}$ one constructs the update matrix $\Delta\hat{\Theta}=\sum_{i}\vec{\delta}_{i}\vec{x}_{i}^{T}$ as usual and then takes its full singular value decomposition. In addition to requiring the calculation [$O(Bmn)$ operations] and storage [$O(mn)$] of the full array, an additional $O(n^{2}m)$ operations (assuming $n<m$) are required to perform the singular value decomposition.

Calculations of the truncated SVD requiring only the top $k$ ranks can be dramatically more efficient, with a complexity of $O(k^{2}m)$. Like Lanczos’s method, these methods are based on the use of power iterations to extract the top eigenvalues and eigenvectors. If these power iterations are replaced with stochastic power iterations—that is, if the gradient is replaced by a stochastic approximation from the input $\{\vec{x}_{i}\}$ and $\{\vec{\delta}_{i}\}$ values—then the computational complexity can be dramatically reduced, and in fact it is possible to avoid directly calculating or storing the batch gradient. Since the batch already gives a stochastic approximation of the gradient, it is plausible that further stochastic approximation may not yield significantly reduced performance, while reducing memory and computational overheads.

The general approach for a multi-rank update of the $\hat{X}$ and $\hat{\Delta}$ matrices follows directly from conventionally used algorithms in streaming PCA, with the caveat that the target matrix is a general rectangular matrix rather than a square and symmetric covariance matrix. From iteration to iteration $\hat{X}$ can be updated as:

where $c$ is a convergence coefficient which gradually decays to zero. The matrix $\hat{d}$ is an update to $\hat{X}$ based on block-size $b$ entries of new data such that

where $\hat{x}\in\mathbb{R}^{b\times n}$ and $\hat{\delta}\in\mathbb{R}^{b\times m}$ are matrices constructed from $b$ sets of $\vec{x}$ and $\vec{\delta}$ respectively. After updating the values of $\hat{X}$, the matrix can be reorthonormalized using $QR$-factorization, where

defines the $QR$-factorization of $\hat{X}$ with the $R$ matrix discarded, and, to ensure the uniqueness of the factorization, the $R$ matrix is defined with non-negative entries on the diagonal. The matrix $\hat{\Delta}$ can be updated in a reciprocal fashion, making this algorithm fall within the family of bi-iterative stochastic power iterations. The vector $\vec{\sigma}$ can be similarly updated with a block average of the data by

which performs inner products over the singular vectors with the incoming data streams to determine the singular values of the matrix. The matrix operation evaluates as $\left(\sum_{\text{rows}}A\odot B\right)_{j}=\sum_{i}A_{ij}B_{ij}$. The rank-wise renormalization of the update $\hat{d}$ in Eq. 6 is critical as it ensures that the learning rate is consistent from epoch to epoch. As the accuracy of the network increases, the magnitude of the input activations would remain approximately constant, but the backpropogated errors’ magnitudes will decline, reducing the effective learning rate during training if they are not rescaled. In classic streaming PCA approaches, this role is filled by adaptively choosing the learning rate.

Since the $QR$-factorization is the most computationally intensive part of the update, the update can be done more sparsely over blocks of size $b<B$, where $b$ is the block size and $B$ is the batch size rather than performing a $QR$-factorization for every member of the batch. In addition to the existing hyperparameters, batch size and learning rate, new hyperparameters, rank, convergence coefficient, and block size, determine the performance of the training algorithm. In traditional PCA, the starting singular vectors are randomly initialized, but in our algorithm, the previous set of singular vectors and values is used as the starting condition for the next batch. Depending on the learning rate and batch size, then, significant prior history and the rate of change of the gradient interact.

We looked at two slightly different versions of our stochastic low rank approximation algorithm. Algorithm 1, dubbed streaming batch principal component analysis (SBPCA), uses a fixed block size $b$ and a convergence coefficient $c=1/(i+1)$, where $i$ is the block index within a batch. To increase the stochasticity of the algorithm while minimizing the number of QR factorizations, we also investigated a variable block size version dubbed SBPCA-variable (SBPCAV) where the convergence coefficient is fixed as $c=1/2$ but the block size increases as $2^{i}$. The first value ($i=0$) in each block of SBPCAV will induce a large rotation in gradient space, reducing the influence of prior batch history by more strongly biasing this single randomly chosen input. SBPCA, by contrast, does not favor the ordering of inputs. Despite these differences, our investigation did not find a large difference between the performance of the two algorithms.

To evaluate the streaming batch PCA approach, three image datasets of increasing size are analyzed: CIFAR-10, CIFAR-100 [31] and ImageNet [32]. Fig. 3 shows details of the input datasets. The CIFAR-10 dataset consists of $60\,000$ color images of size $32\times 32$ organized into 10 classes, with $6\,000$ images per class. The CIFAR-100 dataset is similar to CIFAR-10 dataset, with the exception that it has 100 classes containing 600 images each. The largest dataset used for this investigation is ImageNet, which consists of more than a million training images organized into $1\,000$ classes.

To investigate the algorithm, we use a well known and well understood network structure based on AlexNet [33]. This network, which uses 5 convolutional layers and 3 fully connected layers, was used to achieve a high (40.9 %) Top-5 classification rate on the ImageNet LSVRC-2010 dataset. We duplicate the structure exactly, including using the same rectifying linear unit (ReLU) activation function and four max pooling layers. The training penalty over the dataset is calculated using the traditional cross-entropy loss function for deep neural networks. To better control for the impact of the streaming batch PCA algorithm, however, we omit certain more complex features such as regularization and data augmentation used during training in the original AlexNet implementation. To account for the smaller data set size, the fully connected layers are scaled down proportionally with image size and class size (Fig. 3). Since only the large fully connected layers have significant training data overhead, the streaming batch PCA algorithm is only implemented on these, leaving the convolutional layers to be trained with full rank batch data.

Since the streaming batch PCA algorithm is a form of gradient compression which reduces the dimensionality of the network training data, we explicitly compare it to another form of network dimensionality reduction: random dropout of neurons during training (Fig. 4) [34]. While streaming batch PCA compresses the network training data in an abstract vector space, random dropout (in this case with a dropout fraction of 0.5), implements a real space reduction in both the network size and in the gradient. Note that, where streaming batch PCA and dropout are used concurrently, the activations and backpropogated errors are given no special treatment by streaming batch PCA: the dropped out neurons simply report values of 0 which are seamlessly integrated into the streaming batch PCA dataflow.

To reduce the number of $QR$-factorizations, the block size $b$ for the streaming batch PCA approach was set to $B/4$, with $B$ the batch size, for all cases with $B$ restricted to powers of $2^{n}$. Where the batch size is noted for the streaming batch PCA and streaming batch PCA-variable algorithms, the batch size for the streaming batch PCA-variable approach is set to $B-1$ to accommodate the fixed convergence coefficient rule. We did not perform a search to optimize the hyperparameters, but used parameters consistent with existing practice, taking batch sizes from 128 to 2,048 and learning rates of approximately $10^{-2}$. We do notice that, at low ranks, there is significant loss of gradient information and magnitude due to the approximation which can reduce the training speed. To compensate for this, we keep the convolutional layer learning rate $\alpha_{\text{conv}}$ fixed at $10^{-2}$ but investigate the impact of modifying the learning rate $\alpha_{\text{fc}}$ within the fully connected layers for both streaming batch PCA and MBGD.