DeepPCR: Parallelizing Sequential Operations in Neural Networks
Parallelization techniques have become ubiquitous for accelerating inference
and training of deep neural networks. Despite this, several operations are
still performed in a sequential manner. For instance, the forward and backward
passes are executed layer-by-layer, and the output of diffusion models is
produced by applying a sequence of denoising steps. This sequential approach
results in a computational cost proportional to the number of steps involved,
presenting a potential bottleneck as the number of steps increases. In this
work, we introduce DeepPCR, a novel algorithm which parallelizes typically
sequential operations in order to speed up inference and training of neural
networks. DeepPCR is based on interpreting a sequence of $L$ steps as the
solution of a specific system of equations, which we recover using the Parallel
Cyclic Reduction algorithm. This reduces the complexity of computing the
sequential operations from $\mathcal{O}(L)$ to $\mathcal{O}(\log_2L)$, thus
yielding a speedup for large $L$. To verify the theoretical lower complexity of
the algorithm, and to identify regimes for speedup, we test the effectiveness
of DeepPCR in parallelizing the forward and backward pass in multi-layer
perceptrons, and reach speedups of up to $30\times$ for the forward and
$200\times$ for the backward pass. We additionally showcase the flexibility of
DeepPCR by parallelizing training of ResNets with as many as 1024 layers, and
generation in diffusion models, enabling up to $7\times$ faster training and
$11\times$ faster generation, respectively, when compared to the sequential
approach.
Federico Danieli, Miguel Sarabia, Xavier Suau, Pau Rodríguez, Luca Zappella
# DeepPCR: Parallelizing Sequential Operations in Neural Networks

###### Abstract

Parallelization techniques have become ubiquitous for accelerating inference and training of deep neural networks. Despite this, several operations are still performed in a sequential manner. For instance, the forward and backward passes are executed layer-by-layer, and the output of diffusion models is produced by applying a sequence of denoising steps. This sequential approach results in a computational cost proportional to the number of steps involved, presenting a potential bottleneck as the number of steps increases. In this work, we introduce DeepPCR, a novel algorithm which _parallelizes typically sequential operations_ in order to speed up inference and training of neural networks. DeepPCR is based on interpreting a sequence of \(L\) steps as the solution of a specific system of equations, which we recover using the _Parallel Cyclic Reduction_ algorithm. This reduces the complexity of computing the sequential operations from \(\mathcal{O}(L)\) to \(\mathcal{O}(\log_{2}L)\), thus yielding a speedup for large \(L\). To verify the theoretical lower complexity of the algorithm, and to identify regimes for speedup, we test the effectiveness of DeepPCR in parallelizing the forward and backward pass in multi-layer perceptrons, and reach speedups of up to \(30\times\) for the forward, and \(200\times\) for the backward pass. We additionally showcase the flexibility of DeepPCR by parallelizing training of ResNets with as many as 1024 layers, and generation in diffusion models, enabling up to \(7\times\) faster training and \(11\times\) faster generation, respectively, when compared to the sequential approach.

## 1 Introduction

Neural Networks (NNs) have proven very effective at solving complex tasks, such as classification [26; 14], segmentation [5; 30], and image or text generation [26]. Training NNs, however, is a computationally demanding task, often requiring wall-clock times in the order of days, or even weeks [35; 18], before attaining satisfactory results. Even inference in pre-trained models can be slow, particularly when complex architectures are involved [4]. To reduce training times, a great effort has been invested into speeding up inference, whether by developing dedicated software and hardware [7; 22; 23], or by investigating algorithmic techniques such as (early) pruning [28; 40; 20; 27; 43; 9].

Another possibility for reducing wall-clock time, and the one we focus on in this work, consists in parallelizing computations that would otherwise be performed sequentially. The most intuitive approach to parallelization involves identifying sets of operations which are (almost entirely) independent, and executing them concurrently. Two paradigms that follow this principle are _data-parallelization_, where multiple datapoints are processed simultaneously in batches; and _model-parallelization_, where the model is split among multiple computational units, which perform their evaluations in parallel [1]. Still, certain operations which are key for training and inference in NNs have a sequential structure. The forward and backward pass of a NN are examples of such operations, where activations
(or gradients) are computed sequentially, one layer at a time. Moreover, some generative models suffer from similar shortcomings: in diffusion models (DMs), for example, the output image is generated through a sequence of denoising steps [36]. Sequential operations such as these require a computational effort which grows linearly with the sequence length \(L\) (that is, with the number of layers, or denoising steps), which represents a bottleneck when \(L\) is large. Given the prevalence of these operations, any effort towards their acceleration can result in noticeable speed gains, by drastically reducing training and inference time. Further, faster computations may allow exploration of configurations which were previously unfeasible due to the excessive time required to perform these operations sequentially: for example, extremely deep NNs, or diffusion over tens of thousands of denoising steps.

In this work we introduce DeepPCR, a novel method which provides a flexible framework for turning such sequential operations into parallel ones, thus accelerating operations such as training, inference, and the denoising procedure in DMs.

The core idea behind DeepPCR lies in interpreting a sequential operation of \(L\) steps as the solution of a system of \(L\) equations, as illustrated in Sec. 2. DeepPCR assumes the output of each step only depends on that of the previous one, that is, the sequence satisfies the Markov property. If this holds, we can leverage the specific structure of the resulting system to tackle its solution in parallel, using the Parallel Cyclic Reduction algorithm (PCR) [10; 2]. This algorithm, described in Sec. 3, guarantees the recovery of the solution in \(\mathcal{O}(\log_{2}L)\) steps, rather than the \(\mathcal{O}(L)\) steps required for its sequential counterpart. In our test, this translates into inference speedups of up to \(30\times\) for the forward pass and \(200\times\) for the backward pass in certain regimes, and \(11.2\times\) speedup in image generation via diffusion, as shown in Fig. 1. The reduced computational complexity comes in exchange for higher memory and computational intensity. Therefore, in Sec. 4.1 we investigate in detail regimes for speedup, as well as the trade-off between our method and the sequential approach, considering as model problems the forward and backward passes through multi-layer perceptrons (MLPs) of various sizes. In Sec. 4.2 we then observe how this translates into speedups when training ResNet architectures. Finally, in Sec. 4.3 we showcase how DeepPCR can be applied to accelerate other types of sequential operations as well, choosing as example the denoising procedure in DMs.

Previous WorkThe idea of parallelizing forward and backward passes through a DNN was spearheaded in [13; 32; 24; 31; 41], under the concept of _layer-parallelization_. For the most part, these approaches have been limited to accelerating the training of deep ResNets [15], since they rely on the interpretation of a ResNet as the discretization of a time-evolving differential equation [6], whose solution is then recovered in a time-parallel fashion [11].

More closely resembling our approach is the work in [39], where the authors start by interpreting a sequential operation as the solution of a large system of equations, which is then targeted using parallel solvers. They too focus on accelerating forward and backward passes on ResNets, but also consider some autoregressive generative models (specifically, MADE [12] and PixelCNN++ [38]), similarly to what is done in [44]. The main difference between our approach and the one in [39] lies in the solvers used for tackling the target system in parallel. They rely on variations of Jacobi iterations [34], which are very cost-efficient, but "fall short when the computational graph [of the sequential operation considered] is closer to a Markov chain" [39]: we can expect the convergence of Jacobi to fall to \(\mathcal{O}(L)\) in that case, thus providing no speedup over the sequential approach. By contrast, our method specifically targets Markov sequences, solving them with complexity \(\mathcal{O}(\log_{2}L)\), and is in this sense complementary to theirs. We point out that a similar theoretical foundation for our method was proposed in [33], however it was not verified experimentally, nor has it been considered for applications other than forward and backward passes acceleration.

Figure 1: DeepPCR allows executing sequential operations, such as denoising in latent diffusion, in \(\mathcal{O}(\log_{2}L)\) time, as opposed to the \(\mathcal{O}(L)\) needed for the traditional approach (\(L\) being the number of steps). In our experiments, DeepPCR achieves a \(\mathbf{11.2\times}\)**speedup for image generation with latent diffusion** with respect to the sequential baseline, with comparable quality in the recovered result.
Main ContributionsThe main contributions of this work can be summarized as follows:

1. We propose DeepPCR, a novel algorithm for parallelizing sequential operations in NN training and inference, reducing the complexity of these processes from \(\mathcal{O}(L)\) to \(\mathcal{O}(\log_{2}L)\), \(L\) being the sequence length.
2. We analyze DeepPCR speedup of forward and backward passes in MLPs, to identify high-performance regimes of the method in terms of simple architecture parameters, and we discuss the trade-offs between memory consumption, accuracy of the final solution, and speedup.
3. We showcase the flexibility of DeepPCR applying it to accelerate training of deep ResNet [15] on MNIST [8], and generation in Diffusion Models trained on MNIST, CIFAR-10 [25] and CelebA [29]. Results obtained with DeepPCR are comparable to the ones obtained sequentially, but are recovered up to \(7\times\) and \(11\times\) faster, respectively.

## 2 Turning sequential operations into systems of equations

Our approach is rooted in casting the application of a sequence of \(L\) steps as the solution of a system of \(L\) equations, which we then proceed to solve all at once, in parallel. In this section, we illustrate a general framework to perform this casting and recover the target system. Specific examples for the applications considered in our work (namely forward and backward passes, and generation in diffusion models) are described in appendix A. The algorithm for the parallel solution of the recovered system is outlined in Sec. 3.

Consider a generic sequence of steps in the form \(\mathbf{z}_{l}=f_{l}(\mathbf{z}_{l-1})\), for \(l=1,\ldots,L\), starting from \(\mathbf{z}_{0}=f_{0}(\mathbf{x})\). The various \(f_{l}\) could represent, for example, the application of layer \(l\) to the activations \(\mathbf{z}_{l-1}\) (if we are considering a forward pass), or the application of the \(l\)-th denoising step to the partially recovered image \(\mathbf{z}_{l-1}\) (if we are considering a diffusion mechanism). Notice we are assuming that the output of each step \(\mathbf{z}_{l}\) depends only on that of the previous step \(\mathbf{z}_{l-1}\) and no past ones: that is, we are considering sequences that satisfy the _Markov_ property (a discussion on the limitations related to this assumption, and possible workarounds to relax it, is provided in appendix B). We can collate this sequence of operations into a system of equations for the collated variable \(\mathbf{z}=[\mathbf{z}_{0}^{T},\ldots,\mathbf{z}_{L}^{T},]^{T}\), and obtain:

\[\mathcal{F}(\mathbf{z})=\left[\begin{array}{cccc}\mathbf{z}_{0}-f_{0}(\mathbf{x})\\ \mathbf{z}_{1}-f_{1}(\mathbf{z}_{0})\\ \vdots\\ \mathbf{z}_{L}-f_{L}(\mathbf{z}_{L-1})\end{array}\right]=\left[\begin{array}{cccc}I &&&&\\ -f_{1}(\cdot)&I&&&\\ &\ddots&\ddots&\\ &&&-f_{L}(\cdot)&I\end{array}\right]\left[\begin{array}{cccc}\mathbf{z}_{0}\\ \mathbf{z}_{1}\\ \vdots\\ \mathbf{z}_{L}\end{array}\right]-\left[\begin{array}{cccc}f_{0}(\mathbf{x})\\ \mathbf{0}\\ \vdots\\ \mathbf{0}\end{array}\right]=\mathbf{0}. \tag{1}\]

Notice that, to better highlight the structure of the operator involved, we are abusing matrix notation and considering that the "multiplication" of \(f_{l}(\cdot)\) with \(z_{l-1}\) results in its application \(f_{l}(z_{l-1})\), although \(f_{l}\) is generally a nonlinear operator. To tackle the nonlinearity (when present), we use Newton's method [34]. In more detail, denoting with a superscript \(k\) the Newton iteration, we start from an initial guess for iteration \(k=0\), namely \(\mathbf{z}=\mathbf{z}^{0}\), and iteratively update the solution \(\mathbf{z}^{k+1}=\mathbf{z}^{k}+\delta\mathbf{z}^{k}\) by solving the linearized system

\[J_{\mathcal{F}}|_{\mathbf{z}^{k}}\,\delta\mathbf{z}^{k}=-\mathcal{F}(\mathbf{z}^{k}), \tag{2}\]

until we reach convergence. \(\left.J_{\mathcal{F}}\right|_{\mathbf{z}^{k}}\) denotes the Jacobian of the global sequential operation \(\mathcal{F}(\mathbf{z})\) evaluated at the current iteration \(\mathbf{z}^{k}\). This Jacobian defines the target system we need to solve, and obeys a very specific structure: taking the derivative of (1) with respect to \(\mathbf{z}\), and expanding (2), we see that

\[(2)\Longleftrightarrow\left[\begin{array}{cccc}I&&&&\\ -\left.J_{f_{1}}\right|_{\mathbf{z}_{0}^{k}}&I&&\\ &&\ddots&\ddots&\\ &&&-\left.J_{f_{L}}\right|_{\mathbf{z}_{L-1}^{k}}&I\end{array}\right]\left[\begin{array} []{cccc}\delta\mathbf{z}_{0}^{k}\\ \delta\mathbf{z}_{1}^{k}\\ \vdots\\ \delta\mathbf{z}_{L}^{k}\end{array}\right]=\left[\begin{array}{cccc}f_{0}(\mathbf{x })-\mathbf{z}_{0}^{k}\\ f_{1}(\mathbf{z}_{0}^{k})-\mathbf{z}_{1}^{k}\\ \vdots\\ f_{L}(\mathbf{z}_{L-1}^{k})-\mathbf{z}_{L}^{k}\end{array}\right], \tag{3}\]

that is, the system is _block bidiagonal_. This structure is a direct consequence of the Markovian nature of the sequential operation: since each step relates only two adjacent variables \(\mathbf{z}_{l-1}\) and \(\mathbf{z}_{l}\), only two diagonals appear. The core of DeepPCR lies in applying a specialized parallel algorithm for solving systems with this very structure, as described in Sec. 3.
Parallel Cyclic Reduction for NNs

The solution of a block bidiagonal system is usually obtained via forward substitution: once \(\mathbf{z}_{l}\) is known, it is used to recover \(\mathbf{z}_{l+1}\) and so on, in increasing order in \(l\). This procedures is efficient, but inherently sequential, and as such might represent a bottleneck for large \(L\). Interestingly, there exist alternative algorithms for the solution of such systems, which trade-off more complex instructions and extra memory consumption for a higher degree of parallelization. One such algorithm, and the one our method is based on, is Parallel Cyclic Reduction (PCR) [19]. Originally, PCR was devised to parallelize the solution of tridiagonal systems; in this work, we describe its adaptation for bidiagonal systems such as (3). In a nutshell, PCR works by combining the equations of a system to progressively reduce its dimension, until it becomes easily solvable. Pseudo-code for the adapted algorithm is reported in Alg. 1, and a schematic of how the reduction is performed is outlined in Fig. 2. More details on its functioning are provided next.

We start by noting that systems like (3) can be compactly represented as a set of equations involving only two _adjacent_ variables \(\delta\mathbf{z}_{l-1}\), \(\delta\mathbf{z}_{l}\):

\[\delta\mathbf{z}_{l}-\underbrace{J_{f_{1}}|_{\mathbf{z}_{l-1}}}_{=:A_{l}^{0}}\delta \mathbf{z}_{l-1}-(\underbrace{f_{l}(\mathbf{z}_{l-1})-\mathbf{z}_{l}}_{=:\mathbf{r}_{l}^{0}})=0,\qquad l=1,\dots,L, \tag{4}\]

with \(\delta\mathbf{z}_{0}=f_{0}(\mathbf{x})-\mathbf{z}_{0}^{k}\) known. The \(0\) superscripts in the operators \(A_{l}^{0}\) and vectors \(\mathbf{r}_{l}^{0}\) defined above refer to the current (0-th) PCR step. As a first step for PCR, we substitute the \((l-1)\)-th equation into the \(l\)-th, for each \(l\) in parallel, recovering

\[\delta\mathbf{z}_{l}-\underbrace{A_{l}^{0}A_{l-1}^{0}}_{=:A_{l}^{1}}\delta\mathbf{z}_ {l-2}-\underbrace{\left(\mathbf{r}_{l}^{0}-A_{l}^{0}\mathbf{r}_{l-1}^{0}\right)}_{=: \mathbf{r}_{l}^{1}}=0,\qquad l=2,\dots,L. \tag{5}\]

Notice that the original structure is still preserved, but now the equations relate variables \(l\) to \(l-2\). In other words, the even and the odd variables have become separated, and we have split the original system into two independent subsystems: one involving variables \(\delta\mathbf{z}_{0},\delta\mathbf{z}_{2},\dots\), the other \(\delta\mathbf{z}_{1},\delta\mathbf{z}_{3},\dots\). At the next step, we substitute equations \(l-2\) into \(l\), to recover:

\[\delta\mathbf{z}_{l}-\underbrace{A_{l}^{1}A_{l-2}^{1}}_{=:A_{l}^{2}}\delta\mathbf{z}_ {l-4}-\underbrace{\left(\mathbf{r}_{l}^{1}-A_{l}^{1}\mathbf{r}_{l-2}^{1}\right)}_{=: \mathbf{r}_{l}^{1}}=0,\qquad l=5,\dots,L, \tag{6}\]

so that now only variables at distance 4 are related. Ultimately, at each step of PCR, we are splitting each subsystem into two independent subsystems. If we iterate this procedure for \(\log_{2}L\) steps, we finally obtain \(L\) systems in one variable, which are trivially solvable, thus recovering the solution to the original system.

Figure 2: Left: pseudo-code for PCR algorithm. Right: schematic of row reductions in PCR: green rows are combined pairwise to obtain a system of equations in even unknowns; at the same time, blue rows are combined to obtain a system in odd unknowns only. The result is two independent systems with half the original number of unknowns. The procedure is then repeated for \(\log_{2}L\) steps.
### Limitations of DeepPCR

The main advantage of using DeepPCR for solving (1) lies in the fact that it requires only \(\mathcal{O}(\log_{2}L)\) sequential steps, as opposed to the \(\mathcal{O}(L)\) necessary for traditional forward substitution. However, some conditions must be verified for this procedure to be effective in achieving speedups. We discuss next some recommendations and limitations associated with DeepPCR.

Effective speedup for deep modelsWhile PCR requires fewer sequential steps overall, each step is in principle more computationally intensive than its sequential counterpart, as it requires multiple matrix-matrix multiplications to be conducted concurrently (by comparison, one step of the sequential case requires applying the step function \(f_{l}(\mathbf{z})\)), as per line 6 in Alg. 1. If this cannot be done efficiently, for example because of hardware limitations, then we can expect performance degradation. Moreover, the difference between the linear and logarithmic regimes becomes useful only for large \(L\). Both these facts are investigated in Sec. 4.1.

Controlling Newton iterationsWhenever (1) is nonlinear, the complexity actually becomes \(\mathcal{O}(c_{N}\log_{2}L)\), where \(c_{N}\) identifies the number of Newton iterations necessary for convergence. On the one hand, it is important for \(c_{N}\) to remain (roughly) constant and small, particularly with respect to \(L\), for the logarithmic regime to be preserved and speedups to be attained; on the other hand, there is a positive correlation between \(c_{N}\) and the accuracy of the solution recovered by the Newton solver. Implications of this trade-off are discussed in Sec. 4.4. We also point out that, in general, Newton's method provides no guarantees on _global_ convergence (unlike Jacobi's in [39], which reduces to the sequential solution in the worst-case scenario). Even though in our experiments the method never fails to converge, it is worth keeping in mind that ultimately the solver performance is dependent both on the regularity of the target function (1), and on the initialization choice. In particular, the effect of the latter is investigated in appendix F, but already the simple heuristics employed in our experiments (such as using the average of the train set images as initialization for the output of our DMs) have proven to be effective in providing valid initial guesses for Newton.

Benefits from larger memoryTo apply DeepPCR, it is necessary to store the temporary results from the equation reductions (most noticeably, the operators \(A_{l}\) in line 6 in Alg. 1). The associated memory requirements scale linearly in the number of steps \(L\) and quadratically in the dimension of each step output \(\mathbf{z}\). This results in an increase in memory usage with respect to classical approaches (roughly \(2\times\) as much for forward passes in MLPs, as measured and reported in appendix C.2). We point out that the additional memory requirements of DeepPCR may limit its applications to some distributed training settings where memory is already a bottleneck. Moreover, one can expect additional communication overhead to arise in these settings.

## 4 Results

In this section, we set out to demonstrate the applicability of DeepPCR to a variety of scenarios. We start by investigating the performance characteristics of DeepPCR when applied to the forward and backward passes through a Multi-Layer Perceptron (MLP). Experimenting with this model problem is mostly aimed at identifying regimes where DeepPCR achieves speedup. Specifically, in Sec. 4.1 we show that, when applied to the forward pass, DeepPCR becomes effective in architectures with more than \(2^{7}\) layers. For the backward pass, this regime is reached earlier, in architectures with \(2^{5}\) layers. Next, we explore the effects of applying DeepPCR to speedup the whole training procedure, considering ResNets architectures: in Sec. 4.2 we verify not only that the speedups measured for the single forward and backward passes carry over to this scenario, achieving a \(7\times\) speedup over the sequential implementation, but also that training with DeepPCR results in equivalent models than using sequential passes. In Sec. 4.3, we showcase the flexibility of DeepPCR by using it to speedup another type of sequential operation: the denoising procedure employed by diffusion models in image generation. We consider applications to latent diffusion, and find speedups of up to \(11.2\times\), with negligible error with respect to the sequential counterpart. Lastly, in Sec. 4.4 we focus on the role of the Newton solver in the DeepPCR procedure, establishing that the method remains stable and recovers satisfactory results even by limiting the number of Newton iterations, thus allowing to trade-off additional speedup for an increased approximation error with respect to sequential solutions.
All the experiments in this section were conducted on a V100 GPU with 40GB of RAM; our models are built using the PyTorch framework, without any form of neural network compilation.

### Speeding up forward and backward passes in MLPs: identifying performance regimes

Our first goal is to identify under which regimes DeepPCR can effectively provide a speedup. To this end, we consider a single forward pass through a randomly initialized MLP with a constant number of hidden units (namely, its width \(w\)) at each layer, and profile our algorithm for varying \(w\) and NN depth, \(L\). Notice that these two parameters directly affect the size of (3): \(L\) determines the number of equations, while \(w\) the unknowns in each equation; as such, they can be used as indication of when to expect speedups for more complex problems.

Timing results for these experiments are reported in Fig. 3. The leftmost column refers to the sequential implementation of forward (top) and backward (bottom) pass, and clearly shows the linear complexity in \(L\) of such operations: the curves flatten on a line of inclination 1. Conversely, the graphs in the middle column illustrate DeepPCR's performance, and trace a logarithmic curve for the most part, confirming the theoretical expectations on its \(\mathcal{O}(\log_{2}L)\) complexity. Notice this reduces the wall-clock time for a single forward pass from \(0.55s\) to \(0.015s\), and for a backward pass from \(589ms\) to \(2.45ms\), corresponding to speedups of \(>30\times\) and \(200\times\), respectively, at least for the most favorable architectures - and this despite the fact that there has been more than 20 years of optimization into extracting the best performance from the current GPU hardware when running the sequential forward and backward pass. This result is encouraging as our proposed algorithm can gain from further optimization in each of its steps.

As the MLP grows in width, however, the logarithmic regime is abandoned in favour of a linear regime. This performance degradation is due to the fact that the reductions in line 6 necessary for PCR cannot be performed concurrently anymore. Notice that \(w\) relates directly to the size of the Jacobian blocks in (3), so we can expect similar problems whenever the Jacobian size grows past a given threshold. This issue is caused by hardware limitations, and can be addressed by using dedicated hardware or by optimizing the implementation: evidence of this claim is provided in appendix C.1, where we measure how the threshold for abandoning the logarithmic regime shifts as we use GPUs with different amounts of dedicated memory. Finally, the rightmost graphs in Fig. 3 show the ratio of timings for the sequential versus parallel implementation: any datapoint above 1 indicates effective speedup. The break-even point between the two methods lies around \(L\approx 2^{7}\) for the forward pass.

Figure 3: Time to complete a single forward pass (top) and backward pass (bottom), for MLPs of varying depths \(L\) and widths \(w\), with ReLU activation function. Each datapoint reports the minimum time over 100 runs. The left, center, and right columns refer to the sequential implementation, the DeepPCR implementation, and the ratio between the timings of the two, respectively.
Results for backward pass are qualitatively comparable, but achieve break-even at \(L\approx 2^{5}\): this gain is due to the fact that the backward pass is a linear operation, and as such does not require Newton iterations. For a more in-depth analysis of the role of the Newton solver, we refer to Sec. 4.4.

### Speeding up training of ResNets

The results in Sec. 4.1 identify regimes where one can expect to achieve speedup using DeepPCR, but they only refer to a single forward and backward pass through a freshly initialized model. The results in this section aim to verify that DeepPCR can be used to accelerate forward and backward passes for the whole training procedure, and that the speedup is maintained throughout. To this end, we train a deep ResNet model composed of only fully-connected layers. Each ResNet block consists of 4 layers of width \(2^{4}\) and the ReLU activation function. The models are trained on a classification task on MNIST [8], both using the sequential approach and DeepPCR. We train for 8 epochs using an SGD optimizer with a learning rate of \(10^{-3}\) without a scheduler. We perform training runs with various seeds but report results from only one for readability: the others are comparable, and we show their statistics in appendix D. In Fig. 4 we report the evolution of the wall-clock time measurements for the forward pass throughout the training procedure. We can notice these remain roughly constant, confirming that the speedup achieved by DeepPCR is preserved during training. Notice that using DeepPCR translates into a speedup of \(7\times\) over the sequential implementation: over the whole course of training, this entails a wall-clock time difference of \(3.2h\) versus \(30min\), even without including the gains from the backward pass.

As mentioned in Sec. 3.1, we remind the reader that DeepPCR uses Newton in order to solve (1). Being Newton an approximate solver, one may wonder whether we are accumulating numerical errors with respect to the sequential solution, how does it affect the evolution of the parameters, and what is

Figure 4: Time to complete forward pass during training, for sequential (left) and DeepPCR implementation (center), and ratio between the two (right), for ResNets of varying depths \(L\), with \(w=2^{4}\), skip connection of length \(4\), and ReLU activation function. Each datapoint is an average over 100 optimization steps, and the shaded area spans to \(\pm 1\) standard deviation.

Figure 5: Loss evolution during training with forward and backward passes computed sequentially (left), with DeepPCR (center), and difference between the two (right), for ResNets of varying depths \(L\), with \(w=2^{4}\), skip connection of length \(4\), and ReLU activation function. Each datapoint is an average over 100 optimization steps, and the shaded area spans \(\pm 1\) standard deviation.
the impact on the quality of the final trained model. In our experiments, we measure such impact by comparing the evolution of the loss curves for the models trained sequentially and in parallel with DeepPCR. These are reported in Fig. 5, which shows that, for our experiments, the evolutions are practically equivalent. To further confirm this, we report the accuracy evolution on the test set in appendix D: in both cases, it sits around \(94\%\) at the end of training. The effects of the Newton solver on performance are further discussed in Sec. 4.4.

### Speeding up image generation in Diffusion Models

The experiments in this section showcase the flexibility of DeepPCR in accelerating more general definitions of sequential operations. As an example, we apply DeepPCR to speedup image generation via latent-space diffusion models [37]. Note that we are interested in parallelizing the whole denoising procedure, rather than the single forward pass through the denoiser: we refer to appendix A.4 for the specifics on how this operation falls within the DeepPCR framework. We consider the size of the latent space and the number of denoising steps as the two main parameters which can impact the effectiveness of DeepPCR, and measure how the performance of our method varies according to them. Notice that, in determining the size of system (3), these two parameters cover the same role as \(w\) and \(L\) in Sec. 4.1, respectively, so we identify them using the same notation. Our latent diffusion model considers a simplification of the KL-AutoEncoder introduced by [37] as an encoder, and a custom MLP with residual connections as denoiser: see appendix E for details.

In Fig. 6 (left) we report the average time1 for completing the diffusion procedure, either sequentially or using DeepPCR, for 100 runs on architectures trained on MNIST with various values of \(w\) and \(L\). Notice how even in this case the time for the sequential approach grows linearly with respect to the number of denoising steps, while for DeepPCR the growth is logarithmic for the most part. Increasing \(w\) past \(\sim 2^{6}\), though, results in a speedup reduction for the largest \(L=2^{10}\), matching what is observed in Fig. 3: similarly, this is related to hardware limitations, and we refer again to appendix C.1 for an analysis of the phenomenon. The distributions of the associated speedups are also plotted in Fig. 6 (middle), where we can see that DeepPCR manages to generate images up to \(11\times\) faster, reducing the required time from \(1.3s\) to \(0.12s\) for certain configurations. To ensure the quality of the resulting images, we follow the FID score [16] and measure the Wasserstein-2 distance between the latent distribution of the original test set and the latent distribution of the images recovered, either sequentially or using DeepPCR. The difference of these distances is also reported in Fig. 6, and is consistently close to \(0\), hinting that using either method results in images of similar qualities. Some examples images generated sequentially or using DeepPCR can be seen in Fig. 18, to further confirm that they are hardly distinguishable. We also experimented with diffusion in pixel-space: the corresponding timings can be found in Tab. 2, and their behavior mimics what was observed for latent diffusion.

Footnote 1: We point out that the timings in Fig. 6 and 7 are a proxy, evaluated assuming perfect parallelizability of the Jacobian assembly operation necessary to initialize system (3). We could not measure exact wall-clock time due to incompatibilities between the vmap and autograd functionalities provided in PyTorch. Nonetheless, this proxy is reasonably accurate, as the time required to assemble the Jacobians is negligible with respect to that for the PCR reduction (see appendix E.2, and particularly Fig. 17 for details).

Figure 6: Results from applying DeepPCR to speedup image generation in latent diffusion trained on MNIST, for various latent space dimensions \(w\) and number of denoising steps \(L\). Left: timings using sequential and DeepPCR approaches (average over 100 runs). Middle: violin plots of speedups distribution (ratio of sequential/DeepPCR timings for 100 runs). Right: difference between Wasserstein-2 distances to test distribution of latents recovered sequentially and using DeepPCR.
Finally, in order to provide empirical evidence of the capability of DeepPCR to provide speedup also for other datasets, we experiment with latent diffusion on CIFAR-10 [25] and CelebA [29] as well. The corresponding timings results are reported in Fig. 7. We limit ourselves to \(w>2^{6}\) due to the difficulty of training VAEs for these datasets on smaller latent dimensions. Nonetheless, the timing results are comparable to the ones measured for MNIST in Fig. 6, and even in this case we manage to recover speedups of \(8\times\) and \(9\times\) for CIFAR-10 and CelebA, respectively. We can see that also for these more complex datasets the performance of DeepPCR starts degrading for \(w>2^{7}\), similarly to what is observed in Fig. 6. This observation further confirms that the speedup attained by DeepPCR is influenced by the problem parameters \(w\) and \(L\), but is otherwise dataset-independent.

### Accuracy/Speedup trade-off: analysis on Newton convergence

As outlined in Sec. 2, when system (1) is nonlinear, DeepPCR relies on a Newton solver. This is an iterative solver, which only recovers an _approximate_ solution, correct up to a fixed tolerance. The experiments in the previous sections were conducted with a tolerance of \(10^{-4}\), as we were interested in recovering a solution which would closely match the sequential one. The tolerance of the solver, however, grants us a degree of freedom in trading off accuracy for additional speedup. In this section we investigate in detail the properties of the Newton method when used for the solution of the problems considered in Sec. 4.1 and 4.2.

As a first result, we show that Newton can indeed recover high-quality solutions, within a number of iterations \(c_{N}\) which is small and roughly independent of the configuration considered. To this purpose, we report in Fig. 8 the values of \(c_{N}\) recorded for the experiments in Sec. 4.1 and 4.2. In all configurations considered, they remained bounded below \(c_{N}\leq 6\), and practically independent on the system configuration, particularly of \(L\). In Fig. 8 (first on the left), we see that the performance of the Newton solver is indeed impacted by the type of activation function used in the layers of the MLP: using ReLUs generally requires more iterations for convergence than using a smoother counterpart such as sigmoid. This is in line with the properties of the Newton method which assumes differentiability of the underlying function for fast convergence.

Additionally, for the same set-up, we show (second plot in Fig. 8) the error between the solution recovered via Newton with DeepPCR and the traditional solution, recovered sequentially. This error is expressed in terms of the \(L^{2}\) difference of the NN output (for the experiments in Sec. 4.1) and in terms of the \(L^{\infty}\) difference of the parameters evolution (for the experiments in Sec. 4.2), to better reflect the relevant metrics of the two experiments. The former sits almost always around machine precision, confirming that sequential and DeepPCR solutions are extremely close. For the latter, we see that small numerical errors eventually accumulate throughout the training procedure. Still, the discrepancies are bounded, and this does not affect the final performance of the trained model (as shown also in Fig. 5, and appendix D).

Finally, we conduct an ablation study on the effect of reducing the accuracy of the recovered solution. To this end, we consider again the framework in Sec. 4.2, but this time we fix the number of Newton iterations for solving the forward pass to increasingly small values, and check at which stage training

Figure 7: Results from applying DeepPCR to speedup image generation in latent diffusion, for various latent space dimensions \(w\) and number of denoising steps \(L\). The timings compare sequential (baseline) and DeepPCR approaches, reporting an average over 100 runs, for models trained over the CIFAR-10 (left) and CelebA (right) datasets.
of the ResNets fails. The results reported in appendix F.1 show that, for the problem considered, stopping Newton at \(c_{N}=3\) still results in successful training. This translates into an additional \(2\times\) speedup with respect to the ResNet times reported in Fig. 4, for a total of up to \(14\times\) speedup. For more general problems, we can expect that fine-tuning the Newton solver would play a relevant role in the final speedup attained. Particularly, choosing the correct initial guess for the system and identifying the most apt tolerance level.

## 5 Conclusion, Limitations, and Future Work

We introduced DeepPCR, a method for parallelizing sequential operations which are relevant in NN training and inference. The method relies on the target sequence being Markovian: if this is satisfied, the sequential operation can be interpreted as the solution of a bidiagonal system of equations. The system is then tackled using Parallel Cyclic Reduction, combined with Newton's method. We investigated the effectiveness and flexibility of DeepPCR by applying it to accelerate: i) forward/backward passes in MLPs, ii) training of ResNets, and iii) image generation in diffusion models, attaining speedups of up to \(30\times\), \(7\times\), and \(11\times\) for the three problems, respectively. We identified regimes where the method is effective, and further analyzed trade-offs in terms of speedup, accuracy, and memory consumption.

The main bottleneck for our DeepPCR implementation is represented by the decay in performance associated with the growth in size of the Jacobian blocks in (3). While this can be curbed by using hardware with larger memory and/or better parallelization capabilities, investigating alternative ways to circumvent this issue would greatly benefit the applicability of DeepPCR. Another potential issue is related to the reliance of DeepPCR on a Newton solver for recovering the solution to the target system. While Newton proved to be reasonably robust for the target applications we investigated, in order to achieve best performance one might have to perform _ad-hoc_ adjustments to the solver, depending on the specific sequential operation considered.

Future work will focus on relaxing the limitations outlined above, but also on investigating the applicability of DeepPCR to speedup forward and backward passes through more complex architectures, as well as to speedup different types of sequential operations. In particular, text generation in large language models [4] could be a suitable candidate. Overall, DeepPCR represents a promising method for speeding up training and inference in applications where reducing wall-clock time is critical, and additional computational power is available for parallelization. Furthermore, DeepPCR has the potential to unlock architectures which were not previously experimented upon, due to the long computational time required to perform inference on them.

Figure 8: Newton solver analysis for forward pass through MLP (left), and ResNet training (right).
## Acknowledgements

The authors would like to thank Barry Theobald, David Grangier and Ronan Collobert for their effort and help in proofreading the paper, and Nicholas Apostoloff and Jerremy Holland for supporting this work. The work by Federico Danieli was conducted as part of the AI/ML Residency Program in MLR at Apple.
