# Peering Beyond the Gradient Veil with Distributed Auto Differentiation

Bradley T. Baker  Aashis Khanal  Barak A. Pearlmutter  Vince D. Calhoun
Sergey M. Plis

###### Abstract

Although distributed machine learning has opened up many new and exciting
research frontiers, fragmentation of models and data across different
machines, nodes, and sites still results in considerable communication
overhead, impeding reliable training in real-world contexts. The focus on
gradients as the primary shared statistic during training has spawned a number
of intuitive algorithms for distributed deep learning; however, gradient-
centric training of large deep neural networks (DNNs) tends to be
communication-heavy, often requiring additional adaptations such as sparsity
constraints, compression, and quantization, to curtail bandwidth. We introduce
an innovative, communication-friendly approach for training distributed DNNs,
which capitalizes on the outer-product structure of the gradient as revealed
by the mechanics of auto-differentiation. The exposed structure of the
gradient evokes a new class of distributed learning algorithm, which is
naturally more communication-efficient than full gradient sharing. Our
approach, called distributed auto-differentiation (dAD), builds off a marriage
of rank-based compression and the innate structure of the gradient as an
outer-product. We demonstrate that dAD trains more efficiently than other
state of the art distributed methods on modern architectures, such as
transformers, when applied to large-scale text and imaging datasets. The
future of distributed learning, we determine, need not be dominated by
gradient-centric algorithms.

Machine Learning, ICML

##  1 Introduction

The distributed deep learning community has long gravitated towards methods
which share gradients during training (Bottou, 2010; Verbraeken et al., 2020).
Owing in part to the linearity of the gradient, methods like distributed
stochastic gradient descent (dSGD) have served as the backbone for large-scale
frameworks such as horovod (Sergeev & Del Balso, 2018), PyTorch (Paszke et
al., 2019), and others. When viewed through the lens of auto-differentiation
(AD) (Speelpenning, 1980), however, we can easily observe that the gradient is
computed as the outer-product of two smaller matrices, which are accumulated
during the forward and backward passes through the network. In this work, we
develop this simple fact into an elegant new framework for distributed deep
learning. We show that methods grounded in AD naturally provide a bandwidth
reduction over standard dSGD and other state of the art methods like PowerSGD,
along with competitive performance. We aim to show that much can be gained by
turning the focus of distributed learning away from gradient-centrism and
toward auto-differentiation.

The many parameters at work in deep neural networks (DNNs) require significant
amounts of data to train, with over-fitting becoming a real possibility if not
enough data are provided, or the network is not otherwise regularized. The
need for training on large amounts of data in reasonable time has led the deep
learning community to focus on data-parallel training, where models on
different (GPU) processors are synchronously trained on their respective
subsets of data (Shallue et al., 2019) maintaining the same gradient.
Distributed deep learning can also be motivated by a desire to keep local
training samples hidden. For example, the application of deep learning to
medical problems which utilize highly personal data such as medical imaging
scans or DNA sequences can require models to be trained on samples which
cannot be transferred from one data gathering site to another due to legal or
ethical considerations. These issues motivate privacy-sensitive toolboxes for
distributed learning (Plis et al., 2016).

One of the main obstacles to the scalability of distributed deep learning is
the bottleneck introduced by the large amount of information transmitted over
the network during training. Zhang et al. (Zhang et al., 2020) showed that in
$<$ 100 Gb/s networks, training runtime is significantly worse without
compression, preventing linear scaling with workers or model size.
Svyatkovskiy et al. (Svyatkovskiy et al., 2017) also showed that runtime
increased with network size when training distributed RNNs. In extreme cases,
(Li et al., 2018) where the number of parameters is much larger than network
bandwidth, communication time dominates training time. Finally, efficient
communication of statistics during training has been a preoccupation of
algorithm designers since the advent of the field (see (Dean et al.,
2012)&(Sergeev & Del Balso, 2018) for specific examples, and (Tang et al.,
2020) for survey; also methods mentioned in related works). Even if network
architectures are ideally constructed, communication will always limit the
overall runtime of distributed algorithms.

Figure 1: A high level depiction of the state of the art approaches to
distributed SGD (A), demonstration of our observation that working with AD
process directly provides bandwidth reduction (B), and a linear reduction
algorithm that significantly reduces the bandwidth (C).

If we take a step back from gradient-centric distributed methods, however, we
notice a simple but startlingly profound observation regarding the gradient’s
structure. In reverse-mode auto-differentiation, the gradient of a layer is
computed as an outer-product of the input activations to that layer and the
partial derivative with respect to that layer’s output. These two component
matrices are themselves often smaller in size than the full gradient, and
perhaps more importantly, they represent an explicit structure at work behind
the gradient computation which deserves consideration on its own ground. In
this work, we will aim to show that distributed auto-differentiation
algorithms exploiting this inherent outer-product structure exhibit myriad
benefits, such as significant bandwidth reduction and performance
improvements, when compared to standard dSGD algorithms.

Our contributions in this work can be summarized as follows: First, we
highlight a key observation that the outer product structure of the gradient
inspires a class of inherently communication efficient distributed learning
algorithms. Next, we present a method which elegantly and naturally arises
from the gradient’s outer-product structure: rank-distributed Auto
Differentiation (rank-dAD). Rank-dAD exploits the gradient’s outer-product
structure so that rank-reduced estimates of the component matrices can be
efficiently communicated to reconstruct an estimate for the true global
gradient. Through the use of a new, linear-time algorithm for structured power
iterations (SPI), rank-dAD drastically reduces bandwidth while still
maintaining model performance. The differences between the existing approaches
and our proposed step away from the gradient compression are highlighted in
Figure 1. We illustrate our methods on benchmark deep learning problems such
as digit recognition with MNIST, continuous time-series classification with a
GRU-RNN on the UEA datasets (Bagnall et al., 2018), and large-scale image
recognition with a vision-transformer (Dosovitskiy et al., 2020) on the
cifar-10 dataset.

##  2 Related Work

Most relevant work in distributed deep learning focuses on gradients as the
primary shared statistic, and applies techniques such as dropping unnecessary
values via sparsification (Aji & Heafield, 2017; Wang et al., 2018; Stich,
2018; Wen et al., 2017; Sattler et al., 2019a, b; Shi et al., 2019), mapping
values into bins via quantization (Alistarh et al., 2016; Agarwal et al.,
2018; Bernstein et al., 2018; Horváth et al., 2019; Yu et al., 2018), or
otherwise compressing gradients (Lin et al., 2017; Koloskova et al., 2019;
Mishchenko et al., 2019). Although these methods focus on gradients, the
majority of them are applicable to any shared statistic which can be used for
learning. Since we focus on the more fundamental question of what statistics
are being shared, these methods are potentially synergistic with ours.

A second class of distributed deep learning method reduces bandwidth by
following an update schedule, so gradients are not shared for every batch or
epoch. These methods, such as bursty aggregation (Zhang et al., 2017), lazy
aggregation (Chen et al., 2018), periodic averaging (Haddadpour et al., 2019),
and other scheduling strategies (Stich, 2018; Yu et al., 2019; Assran et al.,
2019), are again agnostic to the actual statistic shared—as long as the
statistic can be used for learning. With the exception of methods which
average statistics during training, the methods proposed here should be
entirely compatible with any update schedule desired, since the auto-
differentiation statistics we share can be used to reconstruct gradients at
any point during training.

The method most closely related to that proposed here is PowerSGD (Vogels et
al., 2019), which uses a QR decomposition of local gradients to estimate two
low-rank matrices which can be used to reconstruct a low-rank approximation of
the gradient. Although PowerSGD is a gradient compression method on its face,
the sharing of low-rank Q and R matrices does represent a shift away from
sharing raw gradients. Indeed, for a chosen low rank $r$, PowerSGD is able to
achieve a bandwidth per layer of $\Theta(r(h_{i}+h_{i+1}))$ for hidden layer
sizes $h_{i}$ and $h_{i+1}$. PowerSGD, which achieves a bandwidth reduction
via the path of QR decomposition rather than auto-differentiation, thus
represents a conceptually closest alternative to our method. It will be our
goal in this work to illustrate the benefits we receive by taking the path of
auto-differentiation both in terms of mathematical intuition, model
performance, and bandwidth reduction.

##  3 Methods

Let $\mathcal{N}(W)$ be a deep neural network with $L$ hidden layers. Let
$\mathbf{X}\in\mathbb{R}^{N\times M}$ be the input batch of data with $N$
samples and $M$ dimensions. Let $\mathbf{Y}\in\mathbb{R}^{N\times C}$ be the
target variable of dimension with $N$ samples and $C$ dimensions. let $h_{i}$
be the size of the $i$th hidden layer. For a feed-forward network, the weight
matrices are thus $\mathbf{W}_{i}\in\mathbb{R}^{h_{i}\times h_{i+1}}$, with
$h_{0}=M$ and $h_{L+1}=C$. Let $\varphi_{i}(x)$ be the activation used at
layer $i$ in the network, and let $\mathcal{L}$ be a given cost-function.

For a feed-forward network, the activations at layer $i$ are thus computed as

$\displaystyle\mathbf{Z}_{i}$
$\displaystyle=\mathbf{A}_{i-1}\mathbf{W}_{i}+\mathbf{B}_{i}$
$\displaystyle\mathbf{A}_{i}$ $\displaystyle=\varphi_{i}(\mathbf{Z}_{i})$ (1)

###  3.1 Reverse-Mode Auto-Differentiation

AD is a class of methods through which derivatives of functions may be
calculated during the execution of a code which evaluates that function. The
backpropagation algorithm for training deep neural networks is specific case
of reverse-mode AD, in which derivatives are propagated backwards along the
data flow graph. This is a two stage process: first a forward pass through the
function is computed, during which intermediate outputs from expressions
within the function are saved and relationships between variables are
recorded. After the forward pass, a _backward pass_ evaluates the contribution
of each intermediate variable to the derivative of the output, retracing and
combining intermediate variables and expressions until returning to the input.
(For a survey, see Baydin et al., 2018).

Following reverse mode AD, we use the chain rule to compute the derivative at
each layer. At the output layer, we first compute the gradient of the loss
w.r.t. the output activations $\nabla_{\mathbf{A}_{L}}\mathcal{L}$, and take
the Hadamard product with the derivative at the output activation
$\varphi_{L}^{\prime}(\mathbf{Z}_{L})$ to compute the gradient w.r.t. the
output, i.e., $\mathbf{\Delta}_{L}$

$\displaystyle\mathbf{\Delta}_{L}=\nabla_{\mathbf{A}_{L}}\mathcal{L}\odot\varphi_{L}^{\prime}(\mathbf{Z}_{L})$
(2)

At higher layers (where $i<L$), we can continue to compute these errors as

$\displaystyle\begin{split}\mathbf{\Delta}_{i}&=(\mathbf{\Delta}_{i+1}\mathbf{W}_{i}\odot\varphi_{i}^{\prime}(\mathbf{Z}_{i}))\end{split}$
(3)

At layer $i$, the gradient of the weights $\nabla_{W_{i}}\mathcal{L}$ at layer
$i$ can thus be computed exactly as

$\displaystyle\begin{split}\nabla_{\mathbf{W}_{i}}\mathcal{L}&=\mathbf{A}_{i-1}^{\top}\mathbf{\Delta}_{i}\end{split}$
(4)

The key insight offered to us by the reverse-mode AD perspective into deep
learning is that in many cases the dimensionality of the intermediate
variables accumulated during the forward and backwards passes of AD will be
less than that of the gradient. In other words. the gradient is a low rank
matrix in most practically relevant cases, and this rank is limited from above
by the batch size. For example, consider a matrix-vector product $y=Wx$. If
Alice is transmitting the gradient $\nabla_{W}\mathcal{L}$ to Bob, where only
Alice knows $x$, we can note that $\nabla_{W}\mathcal{L}$ is the outer product
of $x$ and $\nabla_{y}\mathcal{L}$. If $W\in\mathbb{R}^{m\times n}$, then
these two vectors together have dimensionality $m+n$, and we are often in the
regime where $m+n\ll m\times n$.

This observation suggests a novel algorithm for backpropagation via
distributed auto-differentiation (dAD). We present this algorithm and some
more communication-efficient variants in the following sections.

###  3.2 Exploiting Outer-Product Structure for Distributed Learning

The core of our work relies on a rather simple observation: standard auto-
differentiation computes the gradient via an outer-product of two smaller
matrices (see equation 4). This elementary fact encourages a new way of
looking at distributed learning in which the gradient’s component matrices are
shared instead of the full gradient itself, as is standard practice. Indeed,
in cases where the batch size is significantly less than the hidden dimension
( i.e. $N<<h_{i}$ ) if we were to simply transfer the two component matrices
of the gradient, we would already significantly reduce communication overhead.
After communicating component matrices, sites can compute exact gradients by
merely concatenating in the batch dimension, and computing the product as
normal.

Although the naïve circulation of component matrices seems an attractive
approach on its own, certain difficulties emerge motivating further work to
create a more practically useful distributed learning algorithm. One problem,
for example, with fully communicating component matrices is that bandwidth
usage is directly tied to the chosen batch size, which can be undesirable in
applications where large batch-size is desired, or where we accumulate results
over an additional sequence dimension, such as when utilizing recurrent
architectures. Our fully-realized method achieves efficient communication by
capitalizing on the outer-product structure to perform a structured rank-
reduction of the component matrices in linear time. In the next sections, we
present our complete algorithm for rank distributed auto-differentiation
(rank-dAD), the first in our newly-revealed class of efficient algorithms for
distributed auto-differentiation.

Algorithm 1 rank distributed auto-differentiation (rank-dAD)

0:
$\\{\mathcal{N}_{s}\\}_{s}^{S},\\{\mathbf{X}_{s}\\}_{s}^{S},\\{\mathbf{Y}_{s}\\}_{s}^{S}$

1: for site $s$ do

2: $\\{\mathbf{A}_{i}^{(s)}\\}_{i=0}^{L}=$ forward($\mathcal{N}_{s}$,
$\mathbf{X}_{s}$)

3: end for

4: for hidden layer $i=L$, $0<i\leq L$ do

5: for site $s$ do

6: if i == L then

7:
$\mathbf{\Delta}_{L}^{(s)}=\nabla_{\mathbf{A}_{L}}\mathcal{L}\odot\varphi^{\prime}(\mathbf{Z}_{L}^{(s)})$

8: else

9:
$\mathbf{\Delta}_{i}^{(s)}=\mathbf{\Delta}_{i+1}^{(s)}\mathbf{W}_{i}^{(s)}\odot\varphi_{i}^{\prime}(\mathbf{Z}_{i}^{(s)})$

10: end if

11: $\mathbf{Q}_{i}^{(s)},\mathbf{G}_{i}^{(s)}=$
SPI($\mathbf{\Delta}_{i}^{(s)},\mathbf{A}_{i-1}^{(s)}$)

12: end for

13: $\mathbf{Q}_{i}=$ vertcat($\\{\mathbf{Q}_{i}^{(s)}\\}_{s}^{S}$) {At
Aggregator}

14: $\mathbf{G}_{i}=$ vertcat($\\{\mathbf{G}_{i}^{(s)}\\}_{s}^{S}$) {At
Aggregator}

15:
$\hat{\mathbf{Q}}_{i},\hat{\mathbf{G}}_{i}=$SPI$(\mathbf{Q}_{i},\mathbf{G}_{i})$
{see Algorithm 2}

16: broadcastToSites($\hat{\mathbf{Q}}_{i},\hat{\mathbf{G}}_{i}$)

17: for site $s$ do

18:
$\nabla_{\mathbf{W}_{i}}^{(s)}=\hat{\mathbf{P}}_{i}^{\top}\hat{\mathbf{Q}}_{i}$

19: end for

20: end for

###  3.3 Rank distributed Auto Differentiation

The outer-product structure of the gradient invites myriad improvements for
distributed learning. In this section, we elaborate on how iterative rank-
reduction methods can capitalize off of the outer-product structure, fostering
an algorithm in which rank-reduced component matrices can be computed
efficiently and communicated instead of the full gradient. These component
matrices provide a method for estimating the true global gradient while
reducing the overall bandwidth usage from quadratic to linear with respect to
the layer size. We call this method rank distributed auto-differentiation
(rank-dAD), as it combines iterative rank-reduction with our core insight of
distributed auto-differentiation into a single, efficient algorithm.

First, we assume that all sites coordinate the initialization of local copies
of the chosen architecture - for example, they can share a choice of random
seed and probability distribution when generating initial weights. Each site
will maintain a local copy of the model weights in memory, and these models
will ultimately have equal weights. For simplicity, each site also shares a
set of hyper-parameters, such as learning rate, momentum, number of epochs,
etc. In principal, hyper-parameters and even certain architectural elements
could be allowed to vary between sites based on local needs; however, such
circumstances will not affect the overall performance of the model in terms of
its communication and computational benefits compared to standard distributed
algorithms, and so we leave these as future work.

For a given batch, rank-dAD first performs the typical forward and backward
passes from reverse-mode auto-differentiation, accumulating local activations
and partial derivatives as would normally be used for local gradient
computation. Next, these component matrices are rank-reduced along the batch
dimension. Formally, we begin with $\mathbf{A}_{i-1}\in\mathbb{R}^{N\times
h_{i-1}}$ and $\mathbf{\Delta}_{i}\in\mathbb{R}^{N\times h_{i}}$, and reduce
these to matrices $\mathbf{G}\in\mathbb{R}^{k\times h_{i-1}}$ and
$\mathbf{Q}\in\mathbb{R}^{r\times h_{i}}$ where $r$ is a chosen natural number
designating our maximum target rank. Following rank reduction, we transmit the
rank-reduced matrices to a single aggregator, or set of aggregators which will
perform the next further reduction step. At aggregator nodes, we then
concatenate the received $\mathbf{Q}$ and $\mathbf{G}$ matrices, and perform a
final rank-reduction to obtain $\hat{\mathbf{Q}}$ and $\hat{\mathbf{G}}$. Once
the reduction has obtained data from all sites in the network, we broadcast
the final reduced component matrices, and can locally estimate the gradient
from these components. For simplicity, we will only designate one node in the
network as the aggregator, as this models the architecture supported by
COINSTAC (Plis et al., 2016), our target platform; however, in principal,
multiple aggregation stages could be used, as in ring-reduce, or hierarchical
communication frameworks for example. We leave the detailed investigation of
these alternate communication setups as future work.

It remains for us to show exactly how we can compute the rank-reduced
component matrices in an efficient manner. It turns out that the structure of
the gradient as revealed by AD provides us with a unique approach to rank-
reduction which is computationally efficient, and beneficial for integration
in our full rank-dAD framework. In the next section, we will present the
details of this algorithm, which we call Structured Power Iterations (SPI).

Algorithm 2 Structured Power Iterations (SPI)

0: $\mathbf{\Delta}_{i}$ $\in$ $\mathbb{R}^{N\times h_{i}}$,
$\mathbf{A}_{i-1}$ $\in$ $\mathbb{R}^{N\times h_{i-1}}$, $r$ $\in$
$\mathbb{N}_{+}$, $n,\theta=10^{-3}$

1: $\mathbf{Q}=[],\mathbf{G}=[]$

2: $\mathbf{C}=\mathbf{A}_{i-1}\mathbf{A}_{i-1}^{T}$

3: $\mathbf{B}=\mathbf{\Delta}_{i}^{\top}\mathbf{C}$

4: for $j=0,j\leq r$ do

5: $\mathbf{g}_{0}^{j}\sim\mathcal{N}(0,1)$

6: for $k=1,k<n$ do

7:
$\mathbf{g}_{k+1}^{j}=\mathbf{B}\Delta_{i}\mathbf{g}_{k}^{j}-\mathbf{Q}(\mathbf{G}^{\top}\mathbf{g}_{k}^{j})$

8: end for

9: $\mathbf{v}=\mathbf{\Delta}_{i}\mathbf{g}_{k}^{j}$

10: $\sigma^{j}=\sqrt{\mathbf{v}^{\top}\mathbf{C}\mathbf{v}}$ {May be
avoided1}

11: $\mathbf{q}^{j}=\mathbf{A}_{i-1}^{\top}\mathbf{v}/\sigma^{j}$

12: if
$||\mathbf{g}^{j-1}-\mathbf{g}^{j}||_{2}/||\mathbf{g}^{j-1}||_{2}\leq\theta$
then

13: break

14: end if

15: $\mathbf{Q}=$concat($\mathbf{Q},\mathbf{q}^{j}$)

16: $\mathbf{G}=$concat($\mathbf{G},\mathbf{g}^{j})$

17: end for

18: return$(\mathbf{Q},\mathbf{G})$

####  3.3.1 Structured Power Iterations

Observe that we can compute the singular vector corresponding to the dominant
singular value of $\nabla_{W_{i}}\mathcal{L}$ by iterating the following
recurrence:

$\displaystyle\mathbf{g}_{k+1}^{1}$
$\displaystyle=(\nabla_{\mathbf{W}_{i}}\mathcal{L})^{\top}(\nabla_{\mathbf{W}_{i}}\mathcal{L})\mathbf{g}_{k}^{1}$
(5)

Relying on (4), and pre-computing
$\mathbf{C}=\mathbf{A}_{i-1}\mathbf{A}_{i-1}^{\top}$,
$\mathbf{B}=\mathbf{\Delta}_{i}^{\top}\mathbf{C}$, we can instead iterate:

$\displaystyle\mathbf{g}_{k+1}^{1}$
$\displaystyle=\mathbf{B}\left(\mathbf{\Delta}_{i}\mathbf{g}_{k}^{1}\right)$
(6)

Compared to the $O(h^{2})$ complexity of the iteration (5), the complexity of
the structured power iteration is just $O(h\times N)$ and since $N\ll h$ for
all practical models, it is linear in $h$. The corresponding singular value
$\sigma^{1}=\sqrt{\mathbf{v}^{T}\mathbf{C}\mathbf{v}}$ (where
$\mathbf{v}=\mathbf{\Delta}_{i}\mathbf{g}$) is computed in $O(h\times N)$ once
per singular vector, while computation of the left singular vector
$\mathbf{q}^{1}$ is just $O(h\times N)$ as
$\mathbf{q}^{1}=\mathbf{A}_{i-1}^{\top}\mathbf{v}/\sigma^{1}$.

We successively collect $(\sigma^{j}\mathbf{g}^{j},\mathbf{q}^{j})$, absorbing
singular values into one of the vectors111In practice we bypass computing
$\sigma^{j}$ and gain additional speed up because it cancels out in the outer
product since $\mathbf{q}^{j}$ contains $1/\sigma^{j}$ factor.  constructing
respective $\mathbf{G}_{j}$ and $\mathbf{Q}_{j}$ by concatenating
$\mathbf{g}$s and $\mathbf{q}$s as columns, and proceed to computing the next
singular vectors set by peeling the previously computed least-squares optimal
low-rank representation:

$\displaystyle\mathbf{g}_{k+1}^{j}$
$\displaystyle=\left((\nabla_{\mathbf{W}_{i}}\mathcal{L})^{\top}(\nabla_{\mathbf{W}_{i}}\mathcal{L})-\mathbf{Q}_{j-1}\mathbf{G}_{j-1}^{\top}\right)\mathbf{g}_{k}^{j}$
(7)
$\displaystyle=(\nabla_{\mathbf{W}_{i}}\mathcal{L})^{\top}(\nabla_{\mathbf{W}_{i}}\mathcal{L})\mathbf{g}_{k}^{1}-\mathbf{Q}_{j-1}(\mathbf{G}_{j-1}^{\top}\mathbf{g}_{k}^{j})$

We have already shown that complexity of the first component of (7) is linear
in $h$, but the second component is clearly linear in $h$ as well. Thus,
thanks to the outer-product structure of AD gradients, we can compute low rank
approximation in time linear in the layer width $h$.222To declutter notation
we have dropped the layer index on $h$, as each $h\gg N$.

We have additionally observed, that during training the true rank of
$\nabla_{\mathbf{W}_{i}}\mathcal{L}$ fluctuates and although always below $N$
may take significantly lower values than the desired $r$ we pick for the
structured power iterations. To skip computing noisy columns for our
$\mathbf{Q}$ and $\mathbf{G}$ matrices, we stop the process when
$\lVert\mathbf{g}^{j}-\mathbf{g}^{j+1}\rVert_{2}/\lVert\mathbf{g}^{j}\rVert_{2}<\theta$,
where we set the threshold $\theta$ to $10^{-3}$.

###  3.4 Back Propagation Through Time

So far, we have treated rank-dAD and SPI as applied to matrices where one
dimension corresponds to the batch-size, and the other corresponds to the
number of neurons in the hidden layer. In applications to recurrent networks
and backpropogation through time (BPTT), we require some further consideration
to clarify how we can apply rank-dAD.

In standard rank-dAD, we approximate the exact gradient by locally reducing
$N\times h_{i}$ matrices to $r\times h_{i}$ matrices. These matrices are then
stacked along the rank dimension, and further reduced before broadcasting back
to local sites. For BPTT, we can assume that parameters for a given recurrent
network are tied over the length of the sequence, and the gradient for these
weights can be computed by aggregation over this sequence. Thus, a sensible
extension to rank-dAD accomodates BPTT by stacking matrices into a joint batch
and time dimension of size $N\cdot T$ where $T$ is the length of the sequence.
We then reduce this $(N\cdot T)\times h_{i}$ matrix as before.

This approach to BPTT with rank-dAD allows us to obtain $O(k\times h_{i})$
communication, with only a factor of $T$ increase in initial local complexity.
We thus avoid expensive gradient reduction over the length of the sequence,
replacing that operation completely with structured power iterations.

##  4 Results

This section presents the results of several experiments which illustrate the
communication and computational benefits of rank-dAD. We begin with small-
scale experiments using a feed-forward architecture for digit classification
of the MNIST data set, and then show a similar small example with a GRU-based
recurrent architecture used for multivariate time-series classification on
several datasets from the UEA repository. Finally, we show how rank-dAD can be
used for tasks using modern transformer architectures for both sentiment
analysis of the IMBD dataset, and a larger scale experiment using Vision
Transformers (Dosovitskiy et al., 2020) for image recognition on Cifar-10.
Table 1 provides information on the three architectures which were used for
experiments in this paper.

Architecture | Hidden Layer Sizes | Sequence Length | Depth |  Datasets Tested
---|---|---|---|---
Feed Forward | 1024,1024 | - | 2 |  MNIST
GRU-RNN | 512,256 | 256 | 2 |  Spoken Arabic Digits, PENS-SF, NATOPS, Pen-Digits
Vision Transformer | 128 | 50 (patch size 32) | 12 |  CIFAR-10
Table 1: Data sets and architectures tested as part of the experiments for this paper. Manufacturer | Cores | Memory | GPUs
---|---|---|---
AMD | 64 | 512 GB | 1$\times$Nvidia 2080
Nvidia DGX-1 | 40 | 512 GB | 8$\times$Nvidia V100
Dell | 40 | 192 GB | 4$\times$Nvidia V100
Table 2: Specs for the SLURM cluster used to run the experiments described in
§4.

All experiments were run on a SLURM cluster which submits jobs to one of the
26 machines on the same network. Specs for these machines are provided in
Table 2. For all experiments, all networks were implemented in PyTorch 1.7.1
with Python 3.8 using an Adam optimizer with a fixed learning rate of
$10^{-4}$ and batch size of $64$ per site. We performed $k=5$-fold cross-
validation for all experiments, and plot the average results with error bars
across these folds. We use the gloo distributed backend for communication
between nodes, with all distributed communication methods implemented in
native PyTorch.

In figure 2, the top two panels compare the Area-Under the Curve (AUC) for
digit recognition on MNIST between power-SGD and our method. Table 3 shows the
average per-batch runtime for rank-dAD as a ratio of the average dSGD runtime.
Rank-dAD provides comparable performance to dSGD, regardless of the choice of
rank, and provides a speedup of between 15 and 25 times over that of vanilla
dSGD.

Mode/Sites | 4 | 6 | 8 | 10 | 12 | 14 | 16 | 18
---|---|---|---|---|---|---|---|---
rank-dAD | 0.063 | 0.055 | 0.049 | 0.046 | 0.039 | 0.040 | 0.037 | 0.038
Table 3: For a feed-forward network trained on the MNIST data set, average
per-batch runtime of rank-dAD as a ratio of the per-batch runtime of dSGD.
Even with only 4 sites, rank-dAD sees an over 15 times runtime decrease, and
with 18 sites, over 25 times.
Figure 2: The average test AUC across sites for power-SGD and rank-dAD with
increasing rank. Top panel: AUC over 50 epochs for the MNIST dataset trained
on a Feed Forward Network. Bottom panel: AUC for the Spoken Arabic Digits
dataset trained on a GRU-RNN.

Figure 3 plots the effective rank of the gradient as computed by rank-dAD
during training. We notice that as the model trains, the rank needed for
reliable estimation, and thus overall communication, decreases in all layers
of the model.

Figure 3: Effective rank for the MNIST dataset trained with a Feed Forward
network across training time, where the initial maximum rank was set to 10.

The bottom two panels of figure 2 compare the performance of rank-dAD with
power-SGD for a GRU-RNN used to classify the Spoken Arabic Digits data set for
different choices of maximum effective rank. Again rank-dAD performs
comparably to power-SGD during training, regardless of the choice of rank. In
addition, in figure 4 we show the AUC compared between PowerSGD and rank-dAD
for the GRU-RNN architecture on three additional datasets from the UEA
repository (Bagnall et al., 2018).

Figure 4: The test AUC for a GRU-RNN trained with rank-dAD compared with the
same architecture trained with PowerSGD. Each curve in the two plots provides
the AUC over training for a different maximum rank.

Finally, in figure 5, we were interested in examining how in large-scale
settings with modern architectures, rank-dAD and dSGD compare with a more
rudimentary baseline. As a baseline, we send only the top 3 columns from the
activation and delta matrices to the aggregator and use these to compute the
gradient. This amounts to effectively reducing the batch-size to 3, and
discarding the majority of each batch. We illustrate in this figure how
quickly it takes rank-dAD and dSGD to train to AUC comparable to this
baseline. Using a Vision-Transformer as our base architecture for image
recognition on CIFAR-10, our results clearly illustrate that rank-dAD provides
a clear speedup over dSGD, achieving comparable AUC in a much faster runtime.
Since our baseline represents a lower bound on the communication allowed
between sites, we also illustrate here that rank-dAD, despite having the same
communication as the baseline, is able to achieve comparble AUC faster,
illustrating the clear benefit of using the rank-reduction method to leverage
information from all samples.

Figure 5: Using a Vision Transformer, the log speedup for achieving matching
baseline AUC on CIFAR-10 trained with rank-dAD and dSGD compared to the
baseline of sharing the top 3 columns from the activation and delta matrices.

##  5 Discussion

This section presents an analysis of the theoretical and empirical results
provided above, taking note of how each result contributes to support our
claims.

###  5.1 Performance

Because dAD involves the transmission of full activations and deltas to all
sites, the gradients computed by this method exactly matches those which would
be computed in the pooled case, or in distributed SGD. Thus, dAD is well-
suited to applications where the exact gradients are required, and its
bandwidth improvements over vanilla dSGD make it the favorable choice in such
cases.

In many applications, however, low-rank approximations may be sufficient, and
methods like rank-dAD may be applied. For different initial ranks, rank-dAD
performs on par with PowerSGD on MNIST, and often better on the UEA datasets
(see figure 2). We attribute our improved performance over PowerSGD to a
stronger robustness of our low-rank approximation of the gradient in the
$L_{2}$ sense. For the application to the transformer, rank-dAD sees a small
performance hit when compared to the pooled case, which we believe is
attribute to layer norms being computed locally for the distributed
transformers.

Figure 6: Bandwidth per layer as a function of the layer size for the
considered methods.

###  5.2 Limitations and Future Work

The key insight of our work is that auto-differentiation provides a unique and
useful perspective into distributed deep learning, which can be utilized both
to improve the bandwidth of distributed algorithms and to examine the dynamics
involved in learning. Figure 6 puts potential bandwidths saving in
perspective.

##### Extension to General-Purpose AD

AD has applications to machine learning beyond deep learning, such as generic
gradient-based optimization with the Hessian (Pearlmutter, 1994; Agarwal et
al., 2017), or Bayesian posterior inference in MCMC (Meyer et al., 2003).
Further work is required to look into auto-differentiation as a potentially
distributable process in and of itself; however, such work would open up a
much wider domain of machine learning to the insights provided here.

##### The Problem with Convolutions

In their current form, dAD and rank-dAD share input activations for the given
layers in the network. For feed-forward and recurrent networks (as well as
transformers), the bandwidth improvements provided over dSGD are obvious.
Convolutional layers, however, present a bandwidth problem for these methods,
because the size of the resulting output activations from a convolutional
layer tend to be much larger in size than the number of parameters within that
layer. Thus, further work is needed in examining AD applied to convolutional
layers to see if bandwidth reduction is available without the addition of
heuristics.

##### Effective Gradient Rank for Introspection and DNN Dynamics

Although our initial approach was to use the unique structure provided to us
by AD to compute an accurate low-rank approximation, the apparent dynamics
this approach reveals beg for further empirical and theoretical analysis. The
Singular Value Decomposition has been used to study the dynamics of training
in networks with linear activations (Saxe et al., 2013), and it is possible
that distributed models may be provided with such an analysis for free when
using our method. By opening the black box of AD, we may get a peek into the
black box of deep learning for free.

##  6 Conclusions

In this work, we took a step back from standard gradient-based methods for
distributed deep learning, and presented a novel algorithm for distributed
auto-differentiation (dAD). The insight that the intermediate outer product
factors gathered by AD can be transferred instead of the full gradient
provides a significant bandwidth reduction over full-gradient methods like
dSGD without a loss in performance. Furthermore, we are able to show that the
structure of AD can be further capitalized on to reduce bandwidth again by
half, since for standard backpropagation, the global delta values can be back-
propagated through the network as long as the activations are still shared.
Finally, we push dAD even further by exploiting the explicit outer-product
structure to obtain low rank approximations for the gradient in terms of low-
rank versions of the intermediate statistics. With this, rank-dAD provides an
intriguing method for adaptively reducing bandwidth where the chosen rank is
an upper-limit on communication. We are also able to analyze how the effective
rank of the gradient changes during training, and thus obtain introspective
information about the learning dynamics. It has been our goal to illustrate
that auto-differentiation provides a rich landscape for further exploration
into distributed deep learning. The reduction in bandwidth, intuitive
algorithms, and competitive performance we have demonstrated here represent
the first of potentially many benefits available to distributed machine
learning practice and theory.

## Acknowledgements

Removed for anonymity

## References

  * Agarwal et al. (2017) Agarwal, N., Bullins, B., and Hazan, E.  Second-order stochastic optimization in linear time.  _Journal of Machine Learning Research_ , 18(116):1–40, 2017.  URL https://jmlr.org/papers/volume18/16-491/16-491.pdf.  Also arXiv:1602.03943.
  * Agarwal et al. (2018) Agarwal, N., Suresh, A. T., Yu, F. X. X., Kumar, S., and McMahan, B.  CPSGD: Communication-efficient and differentially-private distributed SGD.  In _Advances in Neural Information Processing Systems_ , pp. 7564–7575, 2018.
  * Aji & Heafield (2017) Aji, A. F. and Heafield, K.  Sparse communication for distributed gradient descent.  _arXiv preprint arXiv:1704.05021_ , 2017.
  * Alistarh et al. (2016) Alistarh, D., Grubic, D., Li, J., Tomioka, R., and Vojnovic, M.  QSGD: Communication-efficient SGD via gradient quantization and encoding.  _arXiv preprint arXiv:1610.02132_ , 2016.
  * Assran et al. (2019) Assran, M., Loizou, N., Ballas, N., and Rabbat, M.  Stochastic gradient push for distributed deep learning.  In _International Conference on Machine Learning_ , pp. 344–353. PMLR, 2019.
  * Bagnall et al. (2018) Bagnall, A., Dau, H. A., Lines, J., Flynn, M., Large, J., Bostrom, A., Southam, P., and Keogh, E.  The UEA multivariate time series classification archive, 2018.  _arXiv preprint arXiv:1811.00075_ , 2018.
  * Baydin et al. (2018) Baydin, A. G., Pearlmutter, B. A., Radul, A. A., and Siskind, J. M.  Automatic differentiation in machine learning: A survey.  _Journal of Machine Learning Research_ , 18, 2018.
  * Bernstein et al. (2018) Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anandkumar, A.  signSGD: Compressed optimisation for non-convex problems.  _arXiv preprint arXiv:1802.04434_ , 2018.
  * Bottou (2010) Bottou, L.  Large-scale machine learning with stochastic gradient descent.  In _Proceedings of COMPSTAT’2010_ , pp. 177–186. Springer, 2010\.
  * Chen et al. (2018) Chen, T., Giannakis, G., Sun, T., and Yin, W.  LAG: Lazily aggregated gradient for communication-efficient distributed learning.  In _Advances in Neural Information Processing Systems_ , pp. 5050–5060, 2018.
  * Dean et al. (2012) Dean, J., Corrado, G., Monga, R., Chen, K., Devin, M., Mao, M., Ranzato, M., Senior, A., Tucker, P., Yang, K., et al.  Large scale distributed deep networks.  In _Advances in neural information processing systems_ , pp. 1223–1231, 2012.
  * Dosovitskiy et al. (2020) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Haddadpour et al. (2019) Haddadpour, F., Kamani, M. M., Mahdavi, M., and Cadambe, V.  Local SGD with periodic averaging: Tighter analysis and adaptive synchronization.  In _Advances in Neural Information Processing Systems_ , pp. 11082–11094, 2019.
  * Horváth et al. (2019) Horváth, S., Kovalev, D., Mishchenko, K., Stich, S., and Richtárik, P.  Stochastic distributed learning with gradient quantization and variance reduction.  _arXiv preprint arXiv:1904.05115_ , 2019.
  * Koloskova et al. (2019) Koloskova, A., Lin, T., Stich, S. U., and Jaggi, M.  Decentralized deep learning with arbitrary communication compression.  _arXiv preprint arXiv:1907.09356_ , 2019.
  * Li et al. (2018) Li, Y., Park, J., Alian, M., Yuan, Y., Qu, Z., Pan, P., Wang, R., Schwing, A., Esmaeilzadeh, H., and Kim, N. S.  A network-centric hardware/algorithm co-design to accelerate distributed training of deep neural networks.  In _2018 51st Annual IEEE/ACM International Symposium on Microarchitecture (MICRO)_ , pp. 175–188. IEEE, 2018.
  * Lin et al. (2017) Lin, Y., Han, S., Mao, H., Wang, Y., and Dally, W. J.  Deep gradient compression: Reducing the communication bandwidth for distributed training.  _arXiv preprint arXiv:1712.01887_ , 2017.
  * Meyer et al. (2003) Meyer, R., Fournier, D. A., and Berg, A.  Stochastic volatility: Bayesian computation using automatic differentiation and the extended Kalman filter.  _The Econometrics Journal_ , 6(2):408–420, 2003\.
  * Mishchenko et al. (2019) Mishchenko, K., Gorbunov, E., Takáč, M., and Richtárik, P.  Distributed learning with compressed gradient differences.  _arXiv preprint arXiv:1901.09269_ , 2019.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.  Pytorch: An imperative style, high-performance deep learning library.  _arXiv preprint arXiv:1912.01703_ , 2019.
  * Pearlmutter (1994) Pearlmutter, B. A.  Fast exact multiplication by the Hessian.  _Neural Computation_ , 6(1):147–160, 1994.
  * Plis et al. (2016) Plis, S. M., Sarwate, A. D., Wood, D., Dieringer, C., Landis, D., Reed, C., Panta, S. R., Turner, J. A., Shoemaker, J. M., Carter, K. W., et al.  COINSTAC: a privacy enabled model and prototype for leveraging and processing decentralized brain imaging data.  _Frontiers in Neuroscience_ , 10:365, 2016.
  * Sattler et al. (2019a) Sattler, F., Wiedemann, S., Müller, K.-R., and Samek, W.  Sparse binary compression: Towards distributed deep learning with minimal communication.  In _2019 International Joint Conference on Neural Networks (IJCNN)_ , pp. 1–8. IEEE, 2019a.
  * Sattler et al. (2019b) Sattler, F., Wiedemann, S., Müller, K.-R., and Samek, W.  Robust and communication-efficient federated learning from non-iid data.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2019b.
  * Saxe et al. (2013) Saxe, A. M., McClelland, J. L., and Ganguli, S.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  _arXiv preprint arXiv:1312.6120_ , 2013.
  * Sergeev & Del Balso (2018) Sergeev, A. and Del Balso, M.  Horovod: fast and easy distributed deep learning in TensorFlow.  _arXiv preprint arXiv:1802.05799_ , 2018.
  * Shallue et al. (2019) Shallue, C. J., Lee, J., Antognini, J., Sohl-Dickstein, J., Frostig, R., and Dahl, G. E.  Measuring the effects of data parallelism on neural network training.  _Journal of Machine Learning Research_ , 20:1–49, 2019\.  URL https://www.jmlr.org/papers/volume20/18-789/18-789.pdf.  Also arXiv:1811.03600v3.
  * Shi et al. (2019) Shi, S., Wang, Q., Zhao, K., Tang, Z., Wang, Y., Huang, X., and Chu, X.  A distributed synchronous SGD algorithm with global Top-$k$ sparsification for low bandwidth networks.  In _2019 IEEE 39th International Conference on Distributed Computing Systems (ICDCS)_ , pp. 2238–2247. IEEE, 2019.
  * Speelpenning (1980) Speelpenning, B.  _Compiling Fast Partial Derivatives of Functions given by Algorithms_.  PhD thesis, Department of Computer Science, University of Illinois at Urbana-Champaign, USA, 1980.  AAI8017989.
  * Stich (2018) Stich, S. U.  Local SGD converges fast and communicates little.  _arXiv preprint arXiv:1805.09767_ , 2018.
  * Svyatkovskiy et al. (2017) Svyatkovskiy, A., Kates-Harbeck, J., and Tang, W.  Training distributed deep recurrent neural networks with mixed precision on gpu clusters.  In _Proceedings of the Machine Learning on HPC Environments_ , pp. 1–8. 2017.
  * Tang et al. (2020) Tang, Z., Shi, S., Chu, X., Wang, W., and Li, B.  Communication-efficient distributed deep learning: A comprehensive survey.  _arXiv preprint arXiv:2003.06307_ , 2020.
  * Verbraeken et al. (2020) Verbraeken, J., Wolting, M., Katzy, J., Kloppenburg, J., Verbelen, T., and Rellermeyer, J. S.  A survey on distributed machine learning.  _ACM Computing Surveys (CSUR)_ , 53(2):1–33, 2020\.
  * Vogels et al. (2019) Vogels, T., Karimireddy, S. P., and Jaggi, M.  PowerSGD: Practical low-rank gradient compression for distributed optimization.  _arXiv preprint arXiv:1905.13727_ , 2019.
  * Wang et al. (2018) Wang, H., Sievert, S., Liu, S., Charles, Z., Papailiopoulos, D., and Wright, S.  Atomo: Communication-efficient learning via atomic sparsification.  In _Advances in Neural Information Processing Systems_ , pp. 9850–9861, 2018.
  * Wen et al. (2017) Wen, W., Xu, C., Yan, F., Wu, C., Wang, Y., Chen, Y., and Li, H.  Terngrad: Ternary gradients to reduce communication in distributed deep learning.  In _Advances in neural information processing systems_ , pp. 1509–1519, 2017.
  * Yu et al. (2019) Yu, H., Yang, S., and Zhu, S.  Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pp. 5693–5700, 2019.
  * Yu et al. (2018) Yu, M., Lin, Z., Narra, K., Li, S., Li, Y., Kim, N. S., Schwing, A., Annavaram, M., and Avestimehr, S.  Gradiveq: Vector quantization for bandwidth-efficient gradient aggregation in distributed cnn training.  In _Advances in Neural Information Processing Systems_ , pp. 5123–5133, 2018.
  * Zhang et al. (2017) Zhang, H., Zheng, Z., Xu, S., Dai, W., Ho, Q., Liang, X., Hu, Z., Wei, J., Xie, P., and Xing, E. P.  Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters.  In _2017 USENIX Annual Technical Conference (USENIX ATC 17)_ , pp. 181–193, 2017.
  * Zhang et al. (2020) Zhang, Z., Chang, C., Lin, H., Wang, Y., Arora, R., and Jin, X.  Is network the bottleneck of distributed training?  In _Proceedings of the Workshop on Network Meets AI & ML_, pp. 8–13, 2020.