SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Federico Danieli Miguel Sarabia Xavier Suau Pau Rodriguez Luca Zappella Apple
{f_danieli, miguelsdc, xsuaucuadros, pau.rodriguez, lzappella}@apple.com
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.
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.
Figure 1: DeepPCR allows executing sequential operations, such as denoising in latent diffusion, in  $\mathcal{O}(\log_2L)$  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  $11.2\times$  speedup for image generation with latent diffusion with respect to the sequential baseline, with comparable quality in the recovered result.
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_2L)$  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 Work The 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.
Main Contributions The main contributions of this work can be summarized as follows:
i) 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.
ii) 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.
iii) 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  $\boldsymbol{z}_l = f_l(\boldsymbol{z}_{l-1})$ , for  $l = 1, \dots, L$ , starting from  $\boldsymbol{z}_0 = f_0(\boldsymbol{x})$ . The various  $f_l$  could represent, for example, the application of layer  $l$  to the activations  $\boldsymbol{z}_{l-1}$  (if we are considering a forward pass), or the application of the  $l$ -th denoising step to the partially recovered image  $\boldsymbol{z}_{l-1}$  (if we are considering a diffusion mechanism). Notice we are assuming that the output of each step  $\boldsymbol{z}_l$  depends only on that of the previous step  $\boldsymbol{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  $\boldsymbol{z} = [z_0^T, \dots, z_L^T]^T$ , and obtain:
$$
\mathcal {F} (\boldsymbol {z}) = \left[ \begin{array}{c} \boldsymbol {z} _ {0} - f _ {0} (\boldsymbol {x}) \\ \boldsymbol {z} _ {1} - f _ {1} (\boldsymbol {z} _ {0}) \\ \vdots \\ \boldsymbol {z} _ {L} - f _ {L} (\boldsymbol {z} _ {L - 1}) \end{array} \right] = \left[ \begin{array}{c c c c} I & & & \\ - f _ {1} (\cdot) & I & & \\ & \ddots & \ddots & \\ & & - f _ {L} (\cdot) & I \end{array} \right] \left[ \begin{array}{c} \boldsymbol {z} _ {0} \\ \boldsymbol {z} _ {1} \\ \vdots \\ \boldsymbol {z} _ {L} \end{array} \right] - \left[ \begin{array}{c} f _ {0} (\boldsymbol {x}) \\ \boldsymbol {0} \\ \vdots \\ \boldsymbol {0} \end{array} \right] = \boldsymbol {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  $z = z^{0}$ , and iteratively update the solution  $z^{k+1} = z^{k} + \delta z^{k}$  by solving the linearized system
$$
\left. J _ {\mathcal {F}} \right| _ {\boldsymbol {z} ^ {k}} \delta \boldsymbol {z} ^ {k} = - \mathcal {F} (\boldsymbol {z} ^ {k}), \tag {2}
$$
until we reach convergence.  $J_{\mathcal{F}}|_{\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}{c c c c} I & & & \\ - J _ {f _ {1}} | _ {\boldsymbol {z} _ {0} ^ {k}} & I & & \\ & \ddots & \ddots & \\ & & - J _ {f _ {L}} | _ {\boldsymbol {z} _ {L - 1} ^ {k}} & I \end{array} \right] \left[ \begin{array}{c} \delta \boldsymbol {z} _ {0} ^ {k} \\ \delta \boldsymbol {z} _ {1} ^ {k} \\ \vdots \\ \delta \boldsymbol {z} _ {L} ^ {k} \end{array} \right] = \left[ \begin{array}{c} f _ {0} (\boldsymbol {x}) - \boldsymbol {z} _ {0} ^ {k} \\ f _ {1} (\boldsymbol {z} _ {0} ^ {k}) - \boldsymbol {z} _ {1} ^ {k} \\ \vdots \\ f _ {L} (\boldsymbol {z} _ {L - 1} ^ {k}) - \boldsymbol {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.
3 Parallel Cyclic Reduction for NNs
The solution of a block bidiagonal system is usually obtained via forward substitution: once  $z_{l}$  is known, it is used to recover  $z_{l+1}$  and so on, in increasing order in  $l$ . This procedure 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 z_{l-1}, \delta z_l$ :
$$
\delta \boldsymbol {z} _ {l} - \underbrace {J _ {f _ {l}} \mid_ {\boldsymbol {z} _ {l - 1}}} _ {=: A _ {l} ^ {0}} \delta \boldsymbol {z} _ {l - 1} - \underbrace {(f _ {l} (\boldsymbol {z} _ {l - 1}) - \boldsymbol {z} _ {l})} _ {=: \boldsymbol {r} _ {l} ^ {0}} = 0, \quad l = 1, \dots , L, \tag {4}
$$
with  $\delta z_0 = f_0(x) - z_0^k$  known. The 0 superscripts in the operators  $A_l^0$  and vectors  $\boldsymbol{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 \boldsymbol {z} _ {l} - \underbrace {A _ {l} ^ {0} A _ {l - 1} ^ {0}} _ {=: A _ {l} ^ {1}} \delta \boldsymbol {z} _ {l - 2} - \underbrace {\left(\boldsymbol {r} _ {l} ^ {0} - A _ {l} ^ {0} \boldsymbol {r} _ {l - 1} ^ {0}\right)} _ {=: \boldsymbol {r} _ {l} ^ {1}} = 0, \quad 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 z_0, \delta z_2, \ldots$ , the other  $\delta z_1, \delta z_3, \ldots$ . At the next step, we substitute equations  $l - 2$  into  $l$ , to recover:
$$
\delta \boldsymbol {z} _ {l} - \underbrace {A _ {l} ^ {1} A _ {l - 2} ^ {1}} _ {=: A _ {l} ^ {2}} \delta \boldsymbol {z} _ {l - 4} - \underbrace {\left(\boldsymbol {r} _ {l} ^ {1} - A _ {l} ^ {1} \boldsymbol {r} _ {l - 2} ^ {1}\right)} _ {=: \boldsymbol {r} _ {l} ^ {2}} = 0, \quad 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_2L$  steps, we finally obtain  $L$  systems in one variable, which are trivially solvable, thus recovering the solution to the original system.
Algorithm 1: PCR for block bidiagonal systems
/* Initialization: */  
1 Set sub-diagonal operators  $A_{l}$ ;  
2 Set right-hand sides  $r_{l}$ ;  
/* PCR steps (sequential) */  
3 for  $i = 1$  to  $L$  do  
/* Reduction steps (parallel) */  
4 parfor  $l = i + 1$  to  $L$  do  
5 Update  $r_{l} \gets r_{l} - A_{l}r_{l - i}$ ;  
6 Update  $A_{l} \gets -A_{l}A_{l - i}$ ;  
7  $i \gets 2i$ ;  
8 Return solution  $r_{l}$ ;
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_2L$  steps.
3.1 Limitations of DeepPCR
The main advantage of using DeepPCR for solving (1) lies in the fact that it requires only  $\mathcal{O}(\log_2L)$  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 models While 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}(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 iterations Whenever (1) is nonlinear, the complexity actually becomes  $\mathcal{O}(c_N\log_2L)$ , 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 memory To 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  $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.
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.
4.1 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_2L)$  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.