SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
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.
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.
4.2 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  $30\mathrm{min}$ , 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 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.
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.
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.
4.3 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 time $^{1}$  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.
<sup>1</sup>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 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.
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.
4.4 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 8: Newton solver analysis for forward pass through MLP (left), and ResNet training (right).
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.
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 Jeremy Holland for supporting this work. The work by Federico Danieli was conducted as part of the AI/ML Residency Program in MLR at Apple.
[1] Ben-Nun, T. and Hoefler, T. Demystifying parallel and distributed deep learning: An in-depth concurrency analysis. ACM Comput. Surv., 52(4), aug 2019. ISSN 0360-0300. doi: 10.1145/3320060. URL https://doi.org/10.1145/3320060.
[2] Bini, D. and Meini, B. The cyclic reduction algorithm: from Poisson equation to stochastic processes and beyond. Numerical Algorithms, 51:23-60, 2008.
[3] Brock, A., Donahue, J., and Simonyan, K. Large scale gan training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
[4] Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[5] Chen, L.-C., Papandreou, G., Kokkinos, I., Murphy, K., and Yuille, A. L. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. IEEE transactions on pattern analysis and machine intelligence, 40(4):834-848, 2017.
[6] Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Neural Information Processing Systems, 2018.
[7] Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Catanzaro, B., and Shelhamer, E. codnn: Efficient primitives for deep learning. arXiv preprint arXiv:1410.0759, 2014.
[8] Deng, L. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[9] Frankle, J., Dziugaite, G. K., Roy, D. M., and Carbin, M. Pruning neural networks at initialization: Why are we missing the mark? In International Conference on Learning Representations, 2021.
[10] Gander, W. Golub, G. H. Cyclic reduction – history and applications. In Luk, F. T. and Plemmons, R. J. (eds.), Workshop on Scientific Computing, New York, 1997. Springer Verlag.
[11] Gander, M. J. 50 years of time parallel time integration. In Carraro, T., Geiger, M., Körkel, S., and Rannacher, R. (eds.), Multiple Shooting and Time Domain Decomposition Methods, pp. 69-113, Cham, 2015. Springer International Publishing. ISBN 978-3-319-23321-5.
[12] Germain, M., Gregor, K., Murray, I., and Larochelle, H. Made: Masked autoencoder for distribution estimation. In International conference on machine learning, pp. 881-889. PMLR, 2015.
[13] Gunther, S., Ruthotto, L., Schroder, J. B., Cyr, E. C., and Gauger, N. R. Layer-parallel training of deep residual neural networks. SIAM Journal on Mathematics of Data Science, 2(1):1-23, 2020. doi: 10.1137/19M1247620. URL https://doi.org/10.1137/19M1247620.
[14] He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
[15] He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016. doi: 10.1109/CVPR.2016.90.
[34] Ortega, J. M. and Rheinboldt, W. C. Iterative Solution of Nonlinear Equations in Several Variables. Society for Industrial and Applied Mathematics, 2000. doi: 10.1137/1.9780898719468. URL https://epubs.siam.org/doi/abs/10.1137/1.9780898719468.
[35] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
[36] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Computer Vision and Pattern Recognition, 2022.
[37] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
[38] Salimans, T., Karpathy, A., Chen, X., and Kingma, D. P. PixelCNN++: Improving the pixelCNN with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:1701.05517, 2017.
[39] Song, Y., Meng, C., Liao, R., and Ermon, S. Accelerating feedforward computation via parallel nonlinear equation solving. In International Conference on Machine Learning (ICML), 2021.
[40] Suau, X., Zappella, L., Palakkode, V., and Apostoloff, N. Principal filter analysis for guided network compression. In Winter Conference on Applications of Computer Vision, 2020.
[41] Sun, Q., Dong, H., Chen, Z., Dian, W., Sun, J., Sun, Y., Li, Z., and Dong, B. Penalty and augmented Lagrangian methods for layer-parallel training of residual networks. ArXiv, abs/2009.01462, 2020.
[42] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[43] Wang, C., Zhang, G., and Grosse, R. B. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations, 2019.
[44] Wiggers, A. and Hoogeboom, E. Predictive sampling with forecasting autoregressive models. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10260-10269. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/wiggers20a.html.
A Specialization of target system for selected applications
In Sec. 2, we introduced a generic framework for the application of DeepPCR, illustrating how performing a sequential operation can be equivalently recast into solving a block bidiagonal system such as (3). In this section, we further specialize this target system, explicitly recovering its form for the applications considered in our work. Notice that this amounts to prescribing a formula for the step operations  $f_{l}(z)$  in (1), and recovering their Jacobians  $J_{f_l}|_z$  appearing in (3).
A.1 Forward pass in MLPs
The forward pass through a NN can be interpreted as a sequential operation whose  $l$ -th step  $f_{l}^{FP}(z_{l-1})$  corresponds to the application of layer  $l$ , to recover activations  $z_{l}$ . Considering an MLP for simplicity, the corresponding formula for this operation reads
$$
\left\{ \begin{array}{l} \boldsymbol {z} _ {0} = f _ {0} ^ {F P} (\boldsymbol {x}) = W _ {0} \boldsymbol {x} + \boldsymbol {b} _ {0}, \\ \boldsymbol {z} _ {l} = f _ {l} ^ {F P} \left(\boldsymbol {z} _ {l - 1}\right) = W _ {l} \sigma_ {l} \left(\boldsymbol {z} _ {l - 1}\right) + \boldsymbol {b} _ {l} \quad \text {f o r} l = 1, \dots , L, \end{array} \right. \tag {7}
$$
for various weights and biases  $(W_{l},\pmb{b}_{l})$ , and nonlinear activation functions  $\sigma_l(\cdot)$ . From (7), we can recover the Jacobian of each step:
$$
\left. J _ {f _ {l} ^ {F P}} \right| _ {\boldsymbol {z}} = W _ {l} \left. S _ {l} \right| _ {\boldsymbol {z}}, \tag {8}
$$
where in turn  $S_{l}|_{z}$  identifies the Jacobian of  $\sigma_l(\cdot)$  evaluated at  $z$ . Since the activation functions are applied element-wise, this reduces to the diagonal matrix
$$
\left. S _ {l} \right| _ {z} = \operatorname {d i a g} _ {i} \left\{\sigma_ {l} ^ {\prime} ([ z ] _ {i}) \right\}. \tag {9}
$$
Similar considerations can be extended to the forward pass through more complex layers, such as convolutional layers, normalization layers, attention layers, and so on. In those cases, an explicit formula such as (8) for the Jacobian calculations might not be readily available, but one can nonetheless recover it efficiently using autograd functionalities.
A.2 Backward pass in MLPs
The backward pass can also be re-cast as a sequential operation through the layers of a NN, as we illustrate here. The purpose of the backward pass is finding the gradients of a target cost function  $c(\pmb{y}, \bar{\pmb{y}})$  with respect to the hidden variables  $z_l$ , which we can then use to recover the gradients with respect to the network parameters. Considering again a MLP for simplicity, we have that gradients of  $c(\pmb{y}, \bar{\pmb{y}})$  with respect to weights and biases at each layer can be recovered as