DeepPCR: Parallelizing Sequential Operations in Neural Networks
Authors: Federico Danieli
Miguel Sarabia
Xavier Suau
Pau Rodriguez
Luca Zappella
Venue: NeurIPS 2023 poster (2023)
Keywords: Acceleration
layer-parallelization
diffusion
Parallel Cyclic Reduction
TL;DR: We use Parallel Cyclic Reduction to accelerate the computation of Markovian sequences related to NNs inference and training

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_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.

DeepPCR: Parallelizing Sequential Operations in Neural Networks
Federico Danieli
Miguel Sarabia
Xavier Suau
Pau Rodríguez
Luca Zappella
Apple
(f_danieli, miguelsdc, xsuaucuadros, pau.rodriguez, 1zappella/@apple. com
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 computational cost proportional to the number of steps involved, presenting potential bottleneck as the number of steps increases. In this work, we introduce DeepPCR, a novel algorithm which parallelizes lypically sequential operations in order to speed up inference and training of neural networks. DeepPCR based on interpreting a sequence of L steps the solution of specific system of equations, which we recover using the Parallel Cyclic Reduction algorithm. This reduces the complexity of computing the sequential operations from ((L) to Ologg L). thus yielding speedup for large L. To verify the theoretical lower complexity of the algorithm, and to identify regimes for speedup, parallelizing the forward and backward pass in multi-layer perceptrons, and reach speedups of up 30 for the forward, and x for the backward pass. We additionally showcase the flexibility of DeepPCR by parallelizing training of ResNets with as many 1024 layers, and generation in diffusion models, enabling up to 7 x faster training and 11 x faster generation, respectively, when compared to the sequential approach.
test the effectiveness of DeepPCR in
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,
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 slow, particularly when complex architectures are involved [4]. To reduce training times, 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 focus on in this work, consists in par- allelizing 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-paralled paalleaaaalllizaiion. where multiple datapoints are processed simultaneously in batches; and model- model- parallelization, where the model is split among multiple computational units, which perform their evaluations in parallel [1].
Still, certain operations which key for training and inference in NNs have sequential struc- ture. The forward and backward pass of a NN are examples of such operations, where activations
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(or gradients) are computed sequentially, one layer a time. Moreover, some generative mod- els suffer from similar shortcomings: in diffu- sion models (DMs), for example, the output im- age is generated through sequence of denoising steps [36]. Sequential operations such as these require computational effort which grows lin- early with the sequence length L (that is, with the number of layers, denoising steps), which represents a bottleneck when L is large. Given the prevalence of these operations, any effort towards their acceleration can result notice- able 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 se- quentially: for example, extremely deep NNs, or diffusion over tens of thousands of denoising steps.
7
9
7
7
1.224s
Heseline: Sequential
M(OD)
11.2 speedup
DeepPCR: Parallel 60logg Z.) L- 1024 steps
1 0.109s
Figure 1: DeepPCR allows executing sequential operations, such as denoising in latent diffusion, in (logg L.) time, as opposed to the (L)) needed for the traditional approach (T. being the number of steps). In our experiments, DeepPCR achieves a 11.2 x 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, novel method which provides a flexible framework for turning such sequential operations into parallel ones, thus accelerating operations such training, inference, and the denoising procedure in DMs.
The core idea behind DeepPCR lies in interpreting sequential operation of L steps as the solution of 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 Ologz L) steps, rather than the (L) steps required for its sequential counterpart. In our test, this translates into inference speedups of up to 30 for the forward pass and x for the backward pass in certain regimes, and 11.2x 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, 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 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 For the most part,
spearheaded in [13, 32, 24, 31, 41], under the concept of layer-parallelizae these approaches have been limited accelerating the training of deep ResNets [15], since they rely on the interpretation of a ResNet as the discretization of a time-evon evoolving differential equation [6]. whose solution then recovered in time tamelllrlllll fashion [11].
More closely resembling our approach is the work in [39], where the authors start by interpreting sequential operation the solution of 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 PixeICNN+ [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 a Markov chain" [39]: we can expect the convergence of Jacobi to fall to (LL) in that case, thus providing no speedup over the sequential approach. By contrast, our method specifically targets Markov sequences, solving them with complexity (logz L). and is in this sense complementary to theirs. We point out that 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, novel algorithm for parallelizing sequential operations in NN training and inference, reducing the complexity of these processes from (L) to ( L), L being the sequence length.
ii) We analyze DeepPCR speedup of forward and backward passes in MLPs, identify high- performance regimes of the method in terms of simple architecture parameters, and we discuss the trade-off's between memory consumption, accuracy of the final solution, and speedup.
ii) We showcase the flexibility of DeepPCR applying it to accelerate training of deep ResNet [15] on MNIST [8], and generation 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 and 11x x faster, respectively.
2 Turning sequential operations into systems of equations
Our approach is rooted in casting the application of sequence of L steps as the solution of a system of I. equations, which we then proceed to solve all at once, in parallel. In this section, we illustrate general framework to perform this casting and recover the target system. Specific examples for the applications considered 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 ž = ft(2)-1), for / 1, I..... L, starting from Zy = (x)). The various fi could represent, for example, the application of layer to the activations #1-1 (if we considering a forward pass), the application of the /-th denoising step the partially recovered image ø1-1 (if we are considering a diffusion mechanism). Notice we are assuming that the output of cach step z depends only that of the previous step ø1-1 and no past ones: that is, 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 collate this
sequence of operations into a system of equations for the collated variable = = [z] 2[,17 T and obtain:
20 fo(x)
I 10
2
fo(
21 - (1
21
0
F(z)
= 0.
: Z1
1L f((EL-1)
fz() I
0
(1)
Notice that, to better highlight the structure of the operator involved, are abusing matrix notation and considering that the "multiplication" of fi(.) with 21-1 results in its application fi(21-1), although fi is generally nonlinear operator. To tackle the nonlinearity (when present), use Newton's method [34]. In more detail, denoting with superscript the Newton iteration, we start from an initial guess for iteration & O, namely z z and iteratively update the solution *+1 z 8z
by solving the linearized system
J 1z
(2)
until we reach convergence. Jواا* denotes the Jacobian of the global sequential operation F(z) evaluated at the current iteration ² This Jacobian defines the target system we need to solve, and obeys very specific structure: taking the derivative of (1) with respect to z, and expanding (2), we
see that
dz
fo(x) x f1 - 2
Jr
dz
(2)
(3)
J1n 24__
I
82
fz (z 2
that is, the system is block bidiagonal. This structure is direct consequence of the Markovian nature of the sequential operation: since each step relates only two adjacent variables z;-1 and z,, only two diagonals appear. The core of DeepPCR lies 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 usually obtained via forward substitution: once z; is known, it is used to recover ø141 and so on, in increasing order in l. This procedures is efficient, but inherently sequential, and 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 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 nutshell, PCR works by combining the equations of system to progressively reduce its dimension, until it becomes easily solvable. Pseudo-code for the adapted algorithm 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 adjacent variables òz;-1, Ōzy;
8z1 J 821-1 (Ji(z1-1) - 21) = 0, I.,...,,4 L,
=:4
=
with òzg = fo(ar) - = known. The superscripts the operators Ap and vectors ry defined above refer to the current (0-th) PCR step. As first step for PCR, we substitute the (1 1)-th equation into the /-th, for each in parallel, recovering
dz:- APAL 8z1-2 (T)
Apr( O,
(= 2= L.
(5)
=AA
=ir
Notice that the original structure is still preserved, but now the equations relate variables / to / 2. In other words, even and the odd variables have become separated, and we have split the original system into two independent subsystems: one involving variables òzo, ōzz,.... the other 821, dz3..... At the next step, we substitute equations - into I, to recover.
dzz- A 82144 (r) A r| O, 1= 1=.,...,44
(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, steps, we finally obtain I. systems in one variable, which are trivially solvable, thus recovering the solution to the original system.
A
Algorithm 1: PCR for block bidiagonal systems
/* Initialisation:
*/
-A
A
Set sub-diagonal operators AG
A
2 Set right-hand sides rı;
A142
/* PCR steps (sequential)
*/
A1+1
3 for 1 do
/* Reduction steps (parallel)
*/
4
parfor 1 i to T. do Update rı + Ti Am--i
5
System of even unknowns 221
System of odd unknowns 221-1
6
Update A + A:Aiii
7 i 2i;
Return solution r1;
Z4t
Z41+2
241+1
Z4t +3
Figure 2: Left: pseudo-code for PCR algorithm. Right: schematic of row reductions in PCR: green rows are combined pairwise to obtain a system cquations in even unknowns; at the same time, blue combined to obtain 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, steps.

---

3.1 Limitations of DeepPCR
The main advantage of using DeepPCR for solving (1) lies in the fact that it requires only (logg L) sequential steps, as opposed to the (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ı(=)), as per line in Alg. 1. 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 OCey log2 L), where cN identifies the number of Newton iterations necessary for convergence. On the one hand, it is important for oN to remain (roughly) constant and small, particularly with respect L, for the logarithmic regime to be preserved and speedups to be attained; on the other hand, there is a positive correlation between eN and the accuracy of the solution recovered by the Newton solver. Implications of this trade-off 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₁ 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 This results in an increase in memory usage with respect classical approaches (roughly 2x as much for forward passes in MLPs, as measured and reported in appendix C.2). We point that the additional memory requirements of DeepPCR may limit its applications to some distributed training settings where memory 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 27 layers. For the backward pass, this regime is reached earlier, in architectures with 2 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 7x 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.2x. 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.
5

---

All the experiments in this section were conducted on VI00 GPU with 40GB of RAM: our models 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 identify under which regimes DeepPCR can effectively provide a speedup. To this end, we consider a single forward pass through randomly initialized MLP with 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 each equation; as such, they can be used indication of when to expect speedups for more complex problems.
Timing results for these experiments reported in Fig. 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 line of inclination 1. Conversely, the graphs in the middle column illustrate DeepPCR's 's performance, and trace logarithmic curve for the most part, confirming the theoretical expectations on its (log, L) complexity. Notice this reduces the wall-clock time for a single forward pass from 0.55s to 0.015s, and for backward pass from 589ms to 2.45ms, corresponding to speedups of 30x and 200x, 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 abandoned in favour of a linear regime. This performance degradation due the fact that the reductions line 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 given threshold. This issue is caused by hardware limitations, and can be addressed by using dedicated hardware by optimizing the implementation: evidence of this claim is provided appendix C.I. where we measure how the threshold for abandoning the logarithmic regime shifts as we use GPUs with different amounts dedicated memory. Finally, the rightmost graphs Fig. show the ratio of timings for the sequential versus parallel implementation: any datapoint above indicates effective speedup. The break-even point between the two methods lies around 27 for the forward pass.
DESPPCR
10
Pervud pavuuiiggi diningggg
10-1
9-1
30
30
27
10-1
10-2
10P
Bockwand pas linings
10P
10
10
27
20
27
211
233
27
23
21
23
211
29
NN Depth L
NNN Depth
Depth
Figure 3: Time complete 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.

---

BASELINE
RATIO BANILIN/DDEFPCE
12
10
Forwad pase tiinngs
2
Epoch
Eprcl
Epech
Figure 4: Time to complete forward pass during training, for sequential (left) and DeepPCR imple- mentation (center), and ratio between the two (right), for ResNets of varying depths L, with w = 2', skip connection of length 4, and ReLU activation function. Each datapoint is average over 100 optimization steps, and the shaded area spans to +1 standard deviation.
Results for backward pass are qualitatively comparable, but achieve break-even at T. x 2 this gain due the fact that the backward pass is linear operation, and 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 but they only refer to single forward and backward pass through a freshly initialized model. The results in this section aim 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, train deep ResNet model composed of only fully-connected layers. Each ResNet block consists of 4 layers of width 2 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 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 speedup of over the sequential implementation: over the whole course of training, this entails wall-clock time difference of 3.2h versus 30min, even without including the gains from the backward pass.
expect to achieve speedup using DeepPCR,
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 accumulating numerical errors with respect to the sequential solution, how does affect the evolution of the parameters, and what is
BASPLINE
DEEPPRR
DumFannnc BASHUNS-DERPER
Las evounion
00.2
Epoch
Epoch
Epoch
Figure 5: Loss evolution during training with forward and backward passes computed sequentially (left), with DecpPCR (center), and difference between the two (right), for ResNets varying depths L. with 2 skip connection of length 4, and ReLU activation function. Each datapoint is an average over 100 optimization steps, and the shaded area spans +1 standard deviation.
7

---

42
Figure 6: Results from applying DeepPCR to speedup image generation in latent diffusion trained on MNIST, for various latent space dimensions and number of denoising steps L. Left: timings using sequential and DeepPCR approaches (average over 100 runs). Middle: violin plots of speedups distri- bution (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 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 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 and L in Sec. 4.1, respectively, so we identify them using the same notation. Our latent diffusion model considers simplification of the KL- AutoodoEee introduced by [37] as an encoder, and custom MLP with residual connections as denoiser: see appendix for details.
In Fig. (left) we report the average time for completing the diffusion procedure, either sequentially or using DeepPCR, for 100 runs on architectures trained on MNIST with various values of and L. Notice how even in this case the time for the sequential approach grows linearly with respect the number of denoising steps, while for DeepPCR the growth is logarithmic for the most part. Increasing w past ~ 26, though, results in a speedup reduction for the largest L = 20 matching what is observed in Fig. 3: similarly, this related 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. (middle), where we can see that DeepPCR manages generate images up to 11x faster, reducing the required time from 1.3s to 0. 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 seen 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.
'We point that the timings in Fig. and proxy, evaluated assuming perfect parallelizability of the Jacobian assembly operation necessary 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).

---

CIFAR-10
CELERA
1.5
1.5
DDDeyPCR
Baseline
Baeelin
Times
0.5
0.5
0
Figure 7: Results from applying DeepPCR to speedup image generation in latent diffusion, for various latent space dimensions 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-IO [25] and CelebA [29] as well. The corresponding timings results are reported in Fig. 7. We limit ourselves difficulty of training VAEs for these datasets on smaller latent dimensions. Nonetheless, the timing results comparable to the ones measured for MNIST in Fig. 6, and even in this case manage to recover speedups of 8x and 9x for CIFAR-10 10 and CelebA, respectively. We can see that also for these more complex datasets the performance of DeepPCR starts degrading for 27, similarly to what observed in Fig. 6. This observation further confirms that the speedup attained by DeepPCR influenced by the problem parameters and L, but is otherwise
2 due to the
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 approximate solution, correct up to fixed tolerance. The experiments in the previous sections were conducted with a tolerance of 10 -4- as we were interested in recovering 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 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 first result, we show that Newton can indeed recover high-quality solutions, within number of iterations ex which is small and roughly independent of the configuration considered. To this purpose, report in Fig. 8 the values of cN recorded for the experiments in Sec. 4.1 and 4.2. In all configurations considered, they remained bounded below eN S 6, and practically independent on the system configuration, particularly of L. In Fig. 8 (first the leftt, we that the performance of the Newton solver is indeed impacted by the type of activation function used the layers of the MLP: using ReLUs generally requires more iterations for convergence than using smoother counterpart such as sigmoid. This is in line with the properties 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 expressed in terms of the L2 difference of the NN output (for the experiments in Sec. 4. and in terms of the 0e 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, 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 ablation study on the effect of reducing the accuracy of the recovered solution. To this end, we consider again the framework Sec. 4.2, but this time we fix the number of Newton iterations for solving the forward pass increasingly small values, and check at which stage training

---

Pororr pan
locd Tranig (.5.A -O
Dryph
Fpwh
(a) Newton solver analysis for forward pass through an MLP of depth L with layers of width (same config- urations as Fig. 3). Left plot indicates the number of iterations until convergence. Right plot indicates the error between the model output obtained via sequential and DeepPCR approaches.
(b) Newton solver analysis for ResNet training with L layers (same configuration in Fig. 5). Left plot indicates number of itcrations until convergence (mean across 100 optimization steps), and shaded area spans +1 standard deviation. Right plot indicates evolution of error between network parameters obtained training with sequential baseline and DeepPCR.
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 cy 3 still results successful training. This translates into an additional 2x speedup with respect to the ResNet times reported in Fig. 4, for a total of up 14x 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.
Conclusion, Limitations, and Future Work
We introduced DeepPCR, method for parallelizing sequential operations which relevant in NN training and inference. The method relies on the target sequence being Markovian: this is satisfied, the sequential operation
be interpreted the solution of 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 accelerate: i) forward/backwaar passes in MLPs, ii) training of ResNets, and iii) image generation in diffusion models, attaining speedups of up 30x, 7x, and 11x 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 be curbed by using hardware with larger memory and/or better parallelization capabilities, investigating alternative ways 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 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 appli- cability of DeepPCR speedup forward and backward passes through as well as to speedup different types of sequential operations. In particular, text generation in large language models [4] could be 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 them.
complex architectures,
10

---

Acknowledgements
The authors would like 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 the AI/L Residency Program in MLR at Apple.
References
[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.//.
[2] Bini, and Meini, The eyclic reduction algorithm: from Poisson equation 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: 13: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):833-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, cudnn: Efficient primitives for deep learning. arXiv preprint arXiv: 1410.059,, 2014.
[8] Deng, I.. 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 initializa- tion: Why are we missing the mark? In International Conference on Learming Representations, 2021.
[10] Gander, W. Golub, G. H. Cyclic reduction - history and applications. Luk, T. and Plemmons, R. J. (eds.), Workshop Scientific Computing, New York, 1997. Springer Verlag.
[11] Gander, M. 50 years of time parallel time integration. 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-
[12] Germain, M., Gregor, K., Murray, I., and Larochelle, H. Made: Masked autoencoder for distribution estimation. International conference on machine learning, pp. 881-889, PMLR, 2015.
[13] Günther, 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. N0.137/19M1447600 URL https://doi: org/10 10. 1137/19M12476200
[14] He, K., Zhang, X., Ren, S., and Sun, Deep residual learning for image recognition. In Proceedings of the IEEE conference 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.1 10.1109/CVP..20169990

---

[16] Heusel, M.. Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to local nash equilibrium. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. https http:tttt///rrreeednggg neerr cc/ wwwpppppppppp..///e/.. wtterffllessppee/11//11 Paper.pde
[17] Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[18] Ho, Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[19] Hockney, R. W. A fast direct solution of Poisson's equation using Fourier analysis. J. ACM, 12(1):95-113, jan 1965. ISSN 0004-5411,
doi: 10.1145/321250.322259 URL https:
//doi. 1145/321250 321259.
[20] Horton, M., Jin, Y., Farhadi, A., and Rastegari, M. Layer-wise data-free CNN compression. 2022. URL https: http///xxvvvv//o/ .wgz///iirrr/oooooooooo org/ htt//dff22111 /2011. w9.08.005..
[21] Huang, X. and Belongie, S. Arbitrary style transfer in real-time with adaptive instance normal- ization. In Proceedings of the IEEE international conference computer vision, pp. 1501-1510, 2017.
[22] Jouppi, N. P., Kurian, G., Li, S., Ma, P., Nagarajan, R., Nai, L., Patil, N., Subramanian, S., Swing, A., Towles, B., et al. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. arXiv preprint arXiv:2304.0 2023.
[23] Kenyon, C. and Capano, C. Apple silicon performance in scientific computing. In 2022 IEEE High Performance Extreme Computing Conference (HPEC), pp. 1-10. IEEE, 2022.
[24] Kirby, A., Samsi, S., Jones, M., Reuther, A., Kepner, J., and Gadepally, V. Layer-paralle training with GPU concurrency of deep residual neural networks via nonlinear multigrid. In 2020 IEEE High Performance Extreme Computing Conference (HPEC), pp. 1-7, 2020. doi: 10.1 /0CHH/PEEAA677422009999559
[25] Krizhevsky, A. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[26] Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):84-90, 2017.
[27] Lee, N., Ajanthan, T., and Torr, P. H. S. SNIP: single-shot network pruning based on connection sensitivity. In International Conference Learning Representations, 2019.
[28) Li, H., Kadav, A., Durdanovic, I., Samet, H., and Graf, H. P. Pruning filters for efficient ConvNets. In International Conference on Learning Representations, 2017.
[29] Liu, Z., Luo, P.. Wang, X., and Tang, X. Deep learning face attributes in the wild. In ICCV, pp. 3730-3738. IEEE Computer Society, 2015.
[30] Long, J., Shelhamer, E., and Darrell, T. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3431- 3440, 2015.
[31] Massaroli, S., Poli, M., Sonoda, S., Suzuki, T., Park, J., Yamashita, A., and Asama, H. Differen- tiable multiple shooting layers. ArXiv, abs/2106.03885, 2021.
[32] Moon, E. and Cyr, E. C. Parallel training of GRU networks with multi-grid solver for long sequences. In International Conference on Learning Representations, 2022. URL https: //openreview. nett/forum?ko-NLIIIOIII
[33] Naumov, M. Parallel complexity of forward and backward propagation. CoRR, ab/ 2017. URL http: http://arxiv...e org/ abs/1712. 06577.
12

---

[34] Ortega, J. M. and Rheinboldt, W. C. Iterative Solution of Nonlinear Equations in Several Vari- ables. Society for Industrial and Applied Mathematics, 2000. doi: (0.111771.97888887777 URI. (ttpst
/tetttps:/eebbb... siam. wttp////abs/////ooooooooggogaatt 10. 1137/1 97..977/999788887999988.
[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 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 syn- thesis 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. Pixelenn++: Improving the pix- elenn with discretized logistic mixture likelihood and other modifications. arXiv preprint 17X.:0701..5517, 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, Penalty and augmented Lagrangian methods for layer-parallel training of residual networks. ArXiv, abs/2009.01466 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 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: http///eooceeeenseeeeee/e// mlr pwwps/veiivveeee 11/9 19wii///ew20ee html.
13

---

A Specialization of target system for selected applications
In Sec. 2, introduced generic framework for the application of DeepPCR, illustrating how performing sequential operation can be equivalently recast into solving block bidiagonal system such (3). In this section, 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 fi(=) in (I), and recovering their Jacobians Jn, appearing in (3).
A.1 Forward pass in MLPs
The forward pass through a NN can be interpreted a sequential operation whose /-th step (x1-1) corresponds to the application of layer 1, to recover activations z₁. Considering an MLP for simplicity, the corresponding formula for this operation reads
Zo (F (a:) ) = Wox + bo:
21 = P (z = Wior(z + b for 1=
for various weights and biases (Wi, br), and nonlinear activation functions ar(-). From (7), we can recover the Jacobian of each step:
J5PT
W Silz:
(8)
where in turn Stl_ identifies the Jacobian of ar(-) evaluated at z. Since the activation functions are applied wiss this reduces to the diagonal matrix
Silz diag {of ([x1,))
Similar considerations can be extended to the forward pass through more complex layers, such as convolutional layers, normalization layers, attention layers, and on. In those cases, an explicit for- mula such as (8) for the Jacobian calculations might not be readily available, but one can nonetheless recover efficiently using autograd functionalities.
A.2 Backward pass in MLPs
The backward pass can also be re-cast as sequential operation through the layers of a NN, as we illustrate here. The purpose of the backward pass is finding the gradients of target cost function c(y, y) with respect to the hidden variables 2;, which we can then use to recover the gradients with respect to the network parameters. Considering again MLP for simplicity, c(y, g) with respect to weights and biases at each layer can be recovered as
have that gradients of
Vw,c Vz,c (1) and Vb, = 2
(10)
with a @ b bt denoting the outer product operator. The gradients V., C, turn, are given by applying the chain rule and solving
V c= JyPP
C. for 1= 1=
(11)
$1-1
backwards, starting from Vz, C, which can be evaluated directly from the output y. This is the target sequential operation that aim parallelize. Notice that 2; represents the activations
layer recovered during the forward pass, and that J is the transpose of the same operator
appearing in (8). Notice moreover that (11) already represents sequence of linear operations, and as such we can solve by applying PCR directly, without needing Newton iterations. More explicitly, the target system is given by
V V
V
J
I
122-1
C
c
0
(12)
C
J
This property not limited to MI.Ps; indeed, backward passes through any architecture are linear by default.
14

---

A.3 Forward and backward passes in ResNets
ResNets are characterized by skip connections which relate activations from layers specific distance s. As such, they break the Markovian assumption that DeepPCR relies on: in fact, the step function representing the application of a layer has form 2 f.R (z1-1, #1-a) for 1%s 0 (where % denotes the modulo operation), meaning it depends both on the activations at the previous layer and those layers prior. Nonetheless, we can still apply the DeepPCR procedure, we collapse all layers in-between skip connection into a single one. This amounts to considering the following step function:
= pRN (21 1, = S1 (SR (1) J.R 1 INN U. f.. (1-8+(221-5) (1-a) := P. (z1-s);
RN
(13)
where stands for function composition. Notice that, after this modification is done, one can apply DeepPCR to the reduced system composed only of unknowns our experiments, we consider MLPs as intermediate layers, and simple residual addition as skip connection layer. This can be denoted as
= = ... In
Wior(z1-1)++ b Wio)(Zi-1) + bi + 21-a else.
1% + 0
FRN
(14)
The Jacobians of the resulting collapsed forward pass is then given by
J gaxx
JRRN
(15)
21
and Ja shares the same form Jep in (8). A similar reasoning can be applied to the backward pass.
A.4 Denoising procedure in diffusion models
Also image generation in diffusion model is, its heart, the result of a sequential operation: an image is extracted from random noise by having undergo a sequence of denoising steps. As such, this operation too is a target for parallelization by DeepPCR. Following [36], the step function in this case is defined by
21 PM (21-1)
1 Vaz
21
v1 dr 0(E-1.1)) + B.N(O,
1 on
(16)
where z! denotes partially denoised image (which underwent I denoising steps), /N 1) is a Normal sample, while o4, 3r, and ūn are parameters of the noise scheduling. Notice that the particularity of this step function is that it hides a NN application inside it: g(z, () identifies the action of the NN trained to predict the noise in image at step L, Once again, we only need the Jacobians in order to assemble the target system, which implies that we need to compute the Jacobian the NN output with respect to its input. We have in fact
I
1
04 αi
Jap
Jg(( (21-1,0).
(17)
-1-
Luckily, Jyl Jz((z,-)))) .) can be efficiently calculated using autograd functionalities, at a cost comparable of that of applying ((,
Notice this applies equivalently diffusion mechanisms acting in pixel and latent space: in the former, z denotes an actual (noisy) image; in the latter, z denotes (noisy) representation of an image, according given encoder [37].
B Relaxing the Markovian assumption
When introducing our method, in Sec. 2 we remark on the fact that, for DeepPCR to be applicable, we must be tackling Markovian sequence 2 = f1(2)-1). Indeed, this automatically guarantees the block bidiagonal structure of system (3), upon which the PCR routine relies. In this section,
15

---

we comment on the validity of the Markov assumption in applications involving common NN architectures, and provide some workarounds to still guarantee the applicability of DeepPCR in some cases where this assumption does not hold.
As it turns out, many architectures automatically obey the Markovian property: vanilla Feed-forward, Convolutional, and Graph NN all consist in sequence of layers whose application only depends on the result of the previous layer. This is also the case for uni-directional Recurrent NN; while for bi-directional ones, both the forward and the backward legs, considered separately, are still Markovian. In general, autoregressive models whose state depends only on the previous one (such as the Diffusion Model considered in our experiments), all give rise to Markov sequences which can then directly be tackled by DeepPCR. The main breakdown cases we identified which invalidate the Markovian assumption are twofold: (i) when residual connections are present, and (ii) when considering autoregressive models whose state depends on a history of previous states. Even in these cases, though, often one can leverage the available flexibility in defining both the states 2; and the step functions fi composing the sequence, so to revert back to the Markov case, and still recover a sequence which can be tackled by DeepPCR, In the following, we discuss how this can be achieved for some relevant applications.
Presence of residual connections Residual connections break the Markovian assumption because they introduce dependencies between states far away from each other along the target sequence, so that a certain state # ends up becoming a function of multiple previous states x = f((z)-1. =1-e) (see also (I4)). Depending on how these connections are distributed, however, different workarounds are available to recover a Markovian sequence:
If the skip connections length is short with respect to the number of layers (s < L), one can simply collapse together the states inside each skip, and consider only the sequence composed by the states connected by the skips. This is precisely the approach we utilized in our ResNets applications, so we refer to appendix A.3 for more details on how this can be implemented.
Conversely, if the skip length is comparable to that of the whole sequence, one can consider applying DeepPCR to parallelize the application of the layers within the skip. Notice however that this severely caps the potential for parallclization offered by our method.
If multiple skip connections nested within each other, that is if we have z1 ((=1-1,2(-s) and 2; = f) for two (or more) given pairs 1- < 2i <
then the whole sequence can be cut between i 7-3 å and and split into two parts. Since the residuals originating from the "left" part are not involved in the generation of the left sequence, and similarly for the "right" part, the two parts are Markovian when taken indi- vidually, Notice this occurs in UNets, where images at the corresponding resolution level during down- and up-sampling are connected together.
The failure case for DeepPCR is when each layer is connected to every other, as it happens for DenseNets. In this framework, in fact, the target system (3) would be fully lower- triangular, and it is not possible to leverage any sparsity pattern thereof to return to the bidiagonal (Markovian) case.
Autoregressive models whose state depends on a history of previous states If the model relies on a window of s states to produce the next output in the generating sequence (as is the case, for example, with PixeICNN [38]), then we are considering a non-Markovian sequence where z = f1 21- =1 To revert to the Markovian case, we can instead consider the collation of states 1 =
=f=,] 1 and notice that indeed we can write & = f)( Note that this results in a Jacobian with larger blocks (by factor s), but if the interdependence between the states is limited, it will still present a sparse structure, thus reducing the complexity of the necessary PCR reduction operations.
While the ones outlined in this section are all feasible workarounds to apply DeepPCR when the Markovian assumption is broken, we would like to point out that the solutions proposed should only be treated as indications, since only some of them have been tested in our work. Moreover, as we are shifting away from the Markovian case, the Jacobi approach investigated in [39] should also become competitive: a comparison with DeepPCR within this framework, and more in-depth investigation of the solutions outlined here is matter of future research.
16

---

BARLINDDDEVPCR
10
10
Forvadd pas Simings
101
10
10
3."
10
10P
Bidoord pais tinings
10-1
101
10-3
10
NN Depth
NN Depth
NN Depth
Figure 9: Time complete single forward pass (top) and backward pass (bottom). for MLPs of varying depths L and widths w. Same configuration as Fig. 3, but using tanh activation functions. Each datapoint reports the minimum time 100 runs. The left, center, and right columns refer the sequential implementation, the DeepPCR implementation, and the ratio between the timings of the two, respectively.
C Additional results for single pass in MLPs
In this section, we expand on the results from the application DeepPCR to accelerate forward and backward passes through MLPs.
BASEI INF
DESPPCR
RATO RASPUN//DDaWCCR
10
Forward plwaadpasttin pastainnggs
30
Kgg
10-2
Backwadd pas timiage
10-1
200 10
10-9
10-1
21
21
NN Depth
NN
Figure 10: Time to complete a single forward pass (top) and backward pass (bottom), for MLPs of varying depths L and widths w. Same configuration Fig. 3, but using sigmoid activation functions. 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.
17

---

BASILINE
BARILIN/DDuPPRRR
10
21
10
Foraadd pas timing
10
coos
10 -1
10"2
211
10
29
NN Width m
NN Width o7
NN Width a
Figure 11: Time to complete a single forward pass, for ML.s of varying depths I, and widths w. The left column refers to classical sequential implementation, the center column the parallel implemen- tation using DeepPCR, the right column reports the ratio between the two. Same configuration as Fig. but the results are reported as function of width, rather than depth.
In Fig. and 10, we report the results from the same experiments in Sec. 4.1, but using tanh and sig- moid as activation functions, respectively, Notice the overall behaviour is basically indistinguishable from that observed in Fig. 3, proving that DeepPCR manages to speedup the forward pass MLPs regardless of the nonlinearity considered. Notice moreover that the break-even point for achieving speedup in these cases occurs earlier than with ReLU: at L closer to 2 than 2 This is also line with what observed in Sec. 4.4, that that the Newton solver generally requires fewer iterations to reach convergence, when the underlying nonlinearity is smoother.
C.1 DeepPCR falling to linear regime
In Sec. 4.1, we point out how increasing the values the parameters w and past a certain threshold negatively affects the performance of DeepPCR: the middle graph in Fig. 3 shows how, under certain configurations, DeepPCR abandons the logarithmic regime, falling onto linear one. This is first indication that the reduction operations in Alg. stop being conducted in parallel, and start being serialized instead. Here we show evidence that this phenomenon linked memory limitations in the hardware used for our experiments.
In Fig. 11 we report the same timing results as in Fig. 3, but expressed function of the MLPs width w instead. This is to better highlight failure conditions for DeepPCR: since its complexity is only function of timings should not be affected by changes in w, and hence the curves Fig. 11 should be constant. This is however not the case, and particularly from the right plot in Fig. 11 we can infer that is the value of w² L that determines when performance drops. To better highlight this, we include in the graph an isoline of 2 L: notice how the threshold after which DeepPCR performance drops has the same inclination. The quantity w2 L. directly related to the size of the system in (3):
(40G)))
(BOGR)
RATIO V100VA100
1.4
22
1.2
Foroom
211
2.3
NN Depth
NN Depth
Figure 12: Timings for forward pass through MLP of various sizes w, (same configurations as in Fig. 3), for two different GPUs: V100 (left) and A100 (middle), and ratio between the two (right). The ratios cluster around two values: 1.1 (showing the GPUs behave comparably for configurations where DecpPCR is still in logarithmic regime), or 0.65 (when DeepPCR falls to linear regime: the factor of 2 difference stems from the A A100 managing to preserve the logarithmic regime for longer).
18

---

BASELINE
RATIO BARILINE-DDEPPCR A100
10
Forvad pas Simings
10
10
10
10-4
kgg
10
2
10
21
NN Depth 7
NN Depth
NN Depth
Figure 13: Time to complete single forward pass, for MLPs of varying depths L and widths w, with ReLU activation function. Analogously to Fig. 3, each datapoint reports the minimum time over 100 runs, and the left, center, and right columns refer to the sequential implementation, the DeepPCR implementation, and the ratio between the timings of the two, respectively. The sole difference is that the experiments reported here conducted 80GB A100 GPU.
2 determines the size of each Jacobian block appearing in the system, while L defines how many blocks there are in total. This hints at the fact that the issue is due memory limitations. To further verify this, we run the same experiments in Sec. 4.1 on two different Graphical Processing Units (GPUs): an NVIDIA Tesla V100, with 40GB of available memory, and NVIDIA Tesla A000, with 80GB. We observe that indeed the GPU with larger memory manages to delay the decay of DeepPCR linear regime by a factor of 2 the 2 configuration of the MLP considered. Particularly, this illustrated in the right plot of Fig. 12, which reports the ratio of the timings for the experiments in Sec. 4.1, conducted with the two GPUs. Notice that the ratios sit around ~ 1.1 for the contigurations where DeepPCR performance is preserved, and fall to half that value, 0.65, when DeepPCR regime falls to linear, further confirming our hypothesis.
For completeness, in Fig. 13 we also report the comparison result of DeepPCR against the sequential baseline (equivalent the experiment Fig. 3) the 80GB A100 GPU: the overall behaviour is comparable to that on the 40GB V100 GPU.
C.2 Memory consumption in DeepPCR
In Sec. 3.1 we mention the additional memory requirements DeepPCR, and how these are mostly related to the need of storing intermediate operators A resulting from the row reductions in line of Alg. 1. From (4), we can see how these operators have the same size as the Jacobian of the step function (i) of the target sequence considered for parallelization via DeepPCR. When applied
memory DeepPiR
0.7
0.0
23
27
29
13
NN Depth
Figure 14: GPU memory allocation for forward pass through MLP: ratio between baseline and DeepPCR for MLPs of varying depths L and widths Same configurations in Fig. 3, maximum over 100 runs.
19

---

BASHLINE-DErre
10090
0078
Aceincy evelusion
8325
8008
7556
Epoch
Opoch
Accuracy evolution during training with forward and backward passes computed sequen- tially (left) and in parallel with DeepPCR (center), and difference between the two (right), for ResNets of various depths L. Same configuration as in Fig. 5.
Figure 15:
to parallelize the forward pass through a MLP, as in our experiments in Sec. 4.1, these Jacobians have size w², exactly the weight matrices Wi the network. Since in our experiments storing the model itself is the most memory-consuminn consuming operation, we can expect the total memory required by DeepPCR to be roughly twice as much as that required in the sequential approach. This verified in Fig. 14, where we report the ratio of the maximum memory allocated when performing forward pass, using the sequential approach and DeepPCR. The ratio sits consistently around 0.6: other words, DeepPCR requires slightly less than twice the amount of memory, as expected.
For more general application, one can expect the extra memory requirements of DeepPCR grow as w w being the size of the Jacobians of the step function fi(z) considered, and L the total number of steps parallelize over.
D Additional results for training of ResNets
In this section, we complement the results in Sec. 4.2, stemming from the applications of DeepPCR to speedup training of ResNets.
In Fig. 15 we report the evolution of the accuracy throughout the training procedure, for the same NN considered Sec. 4.2. This serves an additional confirmation that training sequentially or using DeepPCR results in comparable performance for the NN: in both cases, the accuracy hits 94%, and their difference is consistently close O.
Additionally, in Fig. 16 we report similar training results stemming from different initializations. Regardless of the seed chosen, the behaviour of both loss and accuracy curves remain consistent, and in particular training sequentially or with DeepPCR still provides comparable results.
20

---

Loss CUIVE PVOOUTOO
ACCORACY ENGLUTION
1005
307
1000
780
Spocc
Figure 16: Loss curve and accuracy curve evolution during training for ResNets of different depths, with forward and backward passes computed sequentially and in parallel with DeepPCR. Same configuration as in Fig. 5, but using different seeds for initialization of the network parameters.
21

---

E
Additional results on generation in Diffusion Models
In this section we provide more information regarding the results reported in Sec. 4.3: we describe in detail the architecture of the diffusion models employed, elaborate how the timings were conducted, and include additional results from pixel-space diffusion, which was not described the main paper.
E.1 Choice of architectures and training configuration
To set-up the latent-space diffusion model used for the experiments in Sec. 4.3, we follow Rombach et al. [37] and first train an autoencoder MNIST by minimizing the loss
LAE = || G(E(z))1 + (L(E()),NV(0 1)) + A D(G(E()))),
(18)
where E' is an encoder, G' is generator/decoder, and D is adversarial discriminator. We set B to 10 and to 0.5 [37]. Differently from Rombach et al. [37], we do not include a perceptual loss for training the autoencoder. Since we are training MNIST, we also simplify the architectures of E, G. and D, which consist of 4 convolutional layers and produce a I-dimensional latent space, rather than maintaining the image topology. We experiment with different latent space sizes: 16, 32, 64, and 128. Note that similarly to what is done with a variational autoencoder, the latents obtained from the encoder are sampled using the reparameterization trick. Details on the architectures are displayed in Tables la and lb.
Once the autoencoder is trained, we freeze its weights and train a diffusion process top of the latent vectors produced by applying the encoder on the data. Here, instead of following Rombach et al. [37], we apply the process described in Appendix A.4. Since the latents produced by our autoencoder are I-dimensional, we use a residual MLP instead of the UNet used by Rombach et al. [37]. See Table I for details the MLP architecture. Finally, we generate on randomly sampled latents as in Section 4.3, and generate the final images by applying the decoder just once the resulting latents. To condition the diffusion model on the current timestep and class, we first embed these two values with a sinusoid positional encoding [42], project them with a 2-layer MLP [37], add them (only if we are conditioning the class), and pass the result through linear layer to predict each of the elementwise affine coefficients of the AdaL ayyeroom [21] layers of the diffusion MLP.
samples by rumning diffusion process
The models are trained with AdamW with a batch size of 4096 and a constant learning rate of 10-4 for 400 epochs. Following [3], we train the discriminator with the hinge loss and double the learning rate the generator. We trained all models on a single A A100 GPU with 80G of vRAM.
We also experiment with diffusion in pixel space. To this purpose, we train UNet to recognize the noise in given image, using MSE as a loss function, following [17]. The training procedure uses AdamW W with a constant learning rate of 10-4 as optimizer, and runs for 500 epochs with batch size of 512. Conditioning with this model is performed the same way as with the Iatent-diffusion model, by adding the sinusoidal embedding of time and class. Additional details on the architecture of the network are reported in Tab. le.
E.2 Remarks on profiling for diffusion experiments
Timing of Jacobian evaluations One of the main sub-tasks in the application of DeepPCR for parallelizing the denoising procedure consists in the assembly of the Jacobians of the corresponding step functions (16), and in particular the Jacobians of the NN g(z, () used as denoiser. This requires assembling one Jacobian per denoising step task which is perfectly parallelizable over l. To perform this parallelization, autograd, which are however a recent addition in Py Torch, and still under development. Indeed, after some testing with this implementation, was order of magnitudes higher than actually required, effectively causing the program to crash due to lack of available memory. After some investigation, suboptimal implementation of vmap for our target application, which ends up spawning unnecessary copies of the NN, thus rapidly occupying all available memory. To cireumvent this issue, and still report a reasonable measure for the performance of DeepPCR, in our experiments we conduct the actual computation of the Jacobians sequentially, and then correct the DeepPCR timing accordingly.
attempted to combine the functionalities provided by vmap and
noticed that the memory usage during this phase
suspect the issue is associated with
22

---

Table Description of latent and pixel diffusion model components. 3 specifies the convolution kernel size. ConvDown is convolution layer with stride 2, thus reducing spatial resolution by a factor of 2. ConvUp is convolution layer followed by nearest-neighbor upsample operation, thus increasing spatial resolution by factor of 2. Adal.ayerNorm corresponds to the 1d version of adaptive instance normalization [21], which re-scales the input based on the diffusion time embedding and an optional class-conditionn embedding, following Rombach et al. [37]. We also use Swish activations according to [37]. The input/output dimensionalities in and out are reported for each layer as in - out. [Z] denotes the latent space dimensionality
(a) Image Encoder (E)
(b) Image decoder/generatoor (G)
Grayscale Image z E IR 32
Latent vector R
x 3 ConvDown 1 + 16, ReLU
Linear |Z] 256
3 ConvDown 16 + 32, ReLU
Linear 256 + 2048
3 x 3 ConyDown 32 + 64, ReLU
Reshape R 2048 + R4x4X125
3 3 ConyDown 61 >> 128, ReLU
3 3 ConvUp 128 - 64, ReLU
Reshape 4xx4x128 + R204
3 3 ConvUp 64 + 32, ReLU
Linear 2048 - 256
3 X 3 ConvUp 32 + 16, ReLU
14, Linear 256 > 2 * [Z|
1 x Conv 16 ReLU
z~ A (11,
Tanh
(c) Image discriminator (D)
(d) Residual MLP
Grayscale Image æ E 32x333
Latent code : E [R/21
x ConvDown + 16, ReI.U
Linear |Z] + 2048
3 X 3 ConyDown 16 - 32, ReLU
AdaLayerNorm, Swish, Linear 2048 + 2048
3 x 3 ConvDown 32 64, ReLU
AdaLayerNorm, Swish. Linear 2048 + 2048
x3
x ConvDown 64 + 128 ReLU
Residual
1 X 1 ConvDown 128 +
LayerNorm, Swish, Linear 2048
|Z|
(c) UNet
Grayscale Image E IR" 3 x 3 Conv
+ 16
3 x 3 Conv c c c, ReLU, BatchNorm
x 3 Conv c c c, ReLU, BatchNorm
x2
3 x 3 ConvDown + 2c
Residual
x 3 Conv c c-+ c, ReLU, BatchNorm
3 x 3 Conv c-+ c, ReLUU, BatchNorm
x2
3 3 ConvUp c c/2 Residual
x Conv +
assuming perfect parallelization. By doing this, we are slightly favoring DeepPCR, as we cannot expect perfect parallelization to hold in practice: there will necessarily be some overhead associated with this operation; however, it is reasonable to assume this overhead to be negligible with respect to the actual assembly of the Jacobians. Moreover, the whole Jacohians assembly operation (which itself can be considered as an overhead associated with the application of DeepPCR) is negligible with respect to the time required to perform the PCR reductions. We report comparison of these timings in Fig. 17, for the latent diffusion experiments described in Sec. 4.3.
Ensuring consistency in sequential and DeepPCR denoising procedure Looking at (16), we can that denoising procedure is stochastic process: at each step, random Gaussian noise is added to
23

---

0.3
Jacobian assembly
DeepPCR
0.2
Times
0.1
0
27
25 20
27
2'
25
20
27
2
210
Figure 17: Comparison of the times required for the assembly of the Jacobians (17), and for the complete DeepPCR routine, for the latent diffusion experiments. Setup in Fig. 6 (left). From this we can infer that overall, the main bottleneck in the application of DeepPCR remains the PCR reduction operations in line of Alg. 1.
the image; this is generally done ensure variability in the output generated. In experiments we are interested in matching the results from applying (16) sequentially and those recovered by parallelizing this procedure using DeepPCR. To ensure this, for each generated image simply pre-sample the noise to include at each step, and feed the same sample to both the sequential and DeepPCR procedures.
E.3
Additional results application of DeepPCR to diffusion
In Tab. 2 we expand on the results on latent diffusion in Fig. 6, by reporting numerical values for the timings and speedups recovered, well as the average number of Newton iterations convergence and 1 norm of the error at convergence. Moreover, we include results from the application of DeepPCR pixel-space diffusion. Overall, also in this case the convergence and the Newton iterations remain bounded, confirming what we observed for Sec. 4. and 4.2 as well. Notice however the increase Newton iterations for the models with the largest sizes, which also has an effect on the speedup recovered: this is investigated more in detail in appendix F.
remains small,
In Fig. 18 also show some examples of the images generated by the latent diffusion process trained on MNIST. We put side-by-side the ones recovered using the baseline sequential procedure and those using DeepPCR, to further showcase their similarity.
2
2
5
5
9
9
S
5
1
2
7
8
8
8
8
0
4
4
7
Figure 18: Example images generated by the diffusion models used for the experiments in Fig. 6. The images generated sequentially and using DeepPCR are put side-by-side, and show no visible difference.

---

Table 2: Times complete diffusion generation in latent-space and pixel-space (minimum over 100 trials). dim(Xo) refers to the input dimensions, L the total number of diffusion steps, SEQUENTIAL the amount of time inference takes in sequential fashion, ADJ. DEEPPCR is the amount of time inference takes with DeepPCR: notice this is an estimate, and has been adjusted for parallelism in the Jacobian computation: see details in appendix E.2. We also report the average Newton iterations required (NEWTON ITS.), and the convergence error expressed as the % norm between the generated image using the sequential model and DeepPCR.
dim(xo)
dim((z)
SEQUENTIAL
ADJ. PCR
SPEED UP
NEWTON TTS.
L 90
FRROR
x2
256
16
0.3063
0.0554
5.53x
7.05
0.00195
32x
512
16
0.5997
0,0721
8.32x
9.06
0.00304
32
1024
16
11.22
it
0.00523
LATENT-SPACE SPACE
32 x2
256
32
0.3092
0.0547
5,65
6.67
0.00252
32 322
512
32
0.6358
0.0757
8.40
8.90 it
0.00298
1024
32
1.2543
0.1143%
10.98
14.06
0.00608
32 32
256
64
0.3093
0,0556
5.56X
6,43
0.00283
8.25x
32
1024
64
1,2509
0.1420
8.81x
17:35 at
0.00271
32x
256
128
0,3151
0,0481
6,55x
0.00282
32
512
128
0.6303
0.0669
9,42x
9.08
0.00364
32 32
1024
128
1,2390
0,2033
6.09x
18.54 it
0.00248
X 1
256
16
0.6629
0.0974
6.83x
6.44 it
0.00069
PIXEL SPACE
1
1024
16
2.6095
0.1884
14.25
12.33 it
0.00247
8
256
64
0.6971
0.1070
6,55
0.00094
8
512
64
1,3112
0.11755
11,18X
7.89
0.00276
X 8
1024
64
2.8571
0.1792
15.95
11.11 it
0.00418
256
3.46x
st
0.00126
16x
512
256
1,3752
0.3940
3,79
15.78 st
0.00235
16x16
1024
256
2,8357
1.1594
2.48x
27.33 a
0.02638
Dependence of Newton solver on configuration details
As we have seen, key step of DeepPCR lies in using Newton's 's method to solve the system of equations (2). In this section we analyze more in detail how changing the Newton solver configuration affects the results in Sec. 4. In particular, we focus on the impact on the solver performance of changing the maximum number of Newton iterations, as well as the choice of initial guess.
F.1 Changing the maximum number of Newton iterations
Throughout the main text of our paper, configure the Newton solver with the parameters shown in Tab. 3. From there, we can see that the chosen tolerance is tight, since our aim is to be as accurate as possible, verifying that DeepPCR achieves solution that differs little from the sequential implementation. One can reasonably expect that decreasing these tolerances will produce a noisier, less-accurate final solution. However, both the stochastic gradient descent algorithm generally used for training NNs, and the denoising procedure used to generate images in diffusion, have built- in noisy component in their design. We can then consider the approximate results stemming from an carly-stopped Newton solver as yet another source of noise.
Table 3: Newton solver parameters and description, along with default values for forward pass computation and diffusion generation.
DIFFUSION
Maximum number of Newton iterations the solver apply.
NUMBER OF ITURATIONS
15
30
ABSOLUTE THRESHOLD
Infinity norm of the Newton solver residuul.
10
10-4
Infinity narm of the Newion sulver residuul
RELATIVE THRESHOLD
normalized the infinity-norm of the first iteration residual.
25

---

In the following subsections, we verify the impact of reducing the tolerance in the Newton solver. To this end, we ablate the number of iterations the Newton solver is allowed to take in order to establish whether it is possible to successfully use DeepPCR with noisier solution for (2): first for ResNets training, and then for diffusion generation.
Effects on ResNets training For these experiments, we set the Newton solver to only stop once the specified number of iterations has been reached eN (1,2,3,5, 5, 8, 13), independently of the absolute or relative errors of the residual.
Our results in Fig. 19 show that it is possible to train ResNet of up to 1024 layers, each layer with a width of 16, and a residual connection every layers, with just 3 Newton iterations for solving each forward pass. This is only half of the 6 iterations reported in Sec. 4.2 to keep low error with respect to the sequential baseline, and as such represent a potential further 2x speedup with respect to the reported one.
Effects on Diffusion generation In Fig. 20 and we show convergence - enror as a function of the number of Newton iterations, for Diffusion models of 256, 512 and 1024 denoising steps, acting both in latent- and pixel-space. For the former, we consider latent dimension w 16, 32, 64, 128, while for the latter we consider images of sizes corresponds to a measurement with the maximum number of iterations limited certain value.
4 4, 8 x 8, 16 x 16. Each marker in the plots
Notice that, for the most part, the - of the error falls below 0.1 very quickly (in fewer than 15 iterations regardless of image size). This value already represents very reasonable convergence, considering that in pixel space the images are scaled within [-1, 1]. Moreover, the convergence behaviour is similar regardless of the parameters of the denoising procedure: the convergence curves share similar inclinations.
One exceptions this is the latent-diffusion model with the most denoising steps in Fig. 20: for the first few Newton iterations there no improvement of the solution, and convergence starts relatively late in the process. This hints at the fact that a better strategy for initialization might be beneficial in this case (see also appendix F.2. Another exception is in pixel-space diffusion: for the model with largest in Fig. 20, the Newton solver stalls after reaching 10 error. While this still results in reasonably good solutions, it might represent a bottleneck for scaling the method to pixel-space diffusion in larger dimensions. Further investigation required in order to identify the causes of this behavior.
F.2 Changing the initial guess for the Newton solver
Choosing suitable initial guess n0 for (2) also have relevant impact the performance of the Newton solver. In conducting our experiments, we investigated also this aspect of the solver configuration, and tried to identify ways initialize the Newton solver which could best leverage the available information. Notice that this additional optimization is only available because DeepPCR resorts to an iterative method to recover the target solution: even if reasonable guesses on the solution are provided, using the sequential approach we cannot, in general, leverage them.
NN DErru L- 206
NN Durm L 512
NN Darmi L. 1024
0.8
Arcuney crsiiton
0.6
0.4
0.2
00
Epoch
Epoch
Epoch
Figure 19: Accuracy of ResNets training using a fixed number of Newton iteration. Lines represent mean accuracy over 5 runs, with shaded spanning two standard deviations. Notice how, for the curves corresponding to ey 2 3, the accuracy evolution is basically indistinguishable.
26

---

Latent dimension 16
Latent dimension or 32
Latent dimension or
Latent dimension o 122
107
10
oror
10
STenN J. 206
10
512
10-2
1021
10
10 15
15
20
20
Newton Iterations
Newton Iterations
New Derations
Newton lierations
Figure 20: Convergence eπror in - norm as function of the Newton iterations, for latent-space diffusion with models using different latent-space dimensions w and number of denoising steps
Our choices for the initial guess z" for the three sets of experiments considered in Sec. 4 are discussed next.
Forward pass through MLPs For the set of experiments in Sec. 4.1, Newton is aimed at recovering the results from the forward pass: namely, the unknowns in z are the activations at each layer. Not much information is available to kickstart the method: as initial guess we simply copy the results from the application of the first layer, but the results (not reported here) are comparable to picking a random initialization. Notice that the backward pass, being a linear operation, does not require Newton iterations, as the discussion in appendix A.2.
Training of ResNets Also for the set of experiments Sec. 4.2 we are aiming at recovering the results from the forward pass, but this time this operation is repeated throughout the training procedure: as such, we can reuse information from the previous optimization step to initialize Newton. For the experiments considered, forward pass at the previous optimization step. The rationale behind this choice is as follows: at each step of the training procedure, the NN parameters are perturbed only slightly (the SGD update is proportional to the learning rate); moreover, the datapoints the batches are chosen uniformly, so on average the inputs at each optimization step should be comparable. From these facts, it follows that also the activations should be, on average, comparable.
2 we use batch-average of the activations resulting from the
This choice seems to make for an effective initial guess: in Fig. 22 we measure its impact on the solver convergence, comparing with zero and random initialization for 20. For the model, we use similar setup as for Sec. 4.2, with a ResNet network 1024 layers deep, each layer being 16 units wide, with residual connections every 4 layers. The results Fig. 22 show that we need more than 5 Newton iterations with the zero and random initial guesses, while only 3 suffice when picking as initial guess the average the activations at the previous optimization step.
Diffusion generation For the set of experiments in Sec. 4.3, the Newton solver is used to recover the noisy images (or their encodings) resulting from each step of the denoising procedure. We can expect the procedure to recover images that similar to the ones seen during train time, since the
pixels
t0 8 pixels
t0 = 16 x 16 pixels
10a
10a
1 error
10a
steps
256
10-1
512 1024
10
10-3
10
10
20
30
40
50
60
10
20
30
40
50
60
0
10
20
30
40
50
60
Newton Iterations
Newton Iterations
Newton lterations
Figure 21: Convergence error in - norm function of the Newton iterations, for pixel-space diffusion with models using different image sizes w and number denoising steps
27

---

PREvIOUS ACTIVATIONS INITIAL GuesS
RANDON INTIAL Guess
Zanos INSIALL Guess
2.8
L200
0.5
11
2
Epoch
Epoch
Figure 22: Training dynamics of 1024 layer ResNet with DeepPCR, different initial guesses, and fixed set of Newton iterations. Left plot shows shows behaviour when initial guess of the Newton solver is the batch-mean of the activations at the previous optimization step. Middle plot shows the training dynamics when the initial guess a random guess from normal distribution. Right plot shows training dynamics when the initial guess is all zeros.
denoiser has been trained for this very purpose. As such, as an initial guess x for these noisy images, pick the average of the train set (either in pixel- or in latent-space).
When considering conditional diffusion, we can further fine-tune the initial guess by taking averages over the images in the train test which belong to the same class. We found that this can generally boost the performance of the Newton solver even further. Results from this comparison are reported in Fig. 23, where we compare the speedup achieved with the conditional and unconditional diffusion models, using their respective initial guesses, for various images sizes. Notice that the conditional models consistently achieve better speedup, thanks to the better initialization for the Newton method.
stps
stops
soeps
Spocdup
speedup
gredip
speedap
condrional
maseniitinall
10-2
5
15
100
75
30 0
10
25
30
70
30
(ps)
sidle
Figure 23: Measured speedup for various pixel diffusion models, as a function of the image size. Notice how conditional models achieve larger speedup due more fine-tuned initialization of the Newton method.
28


Reviews:
[{"review_id": "tWv72jOxkk", "reviewer": "Reviewer_uCYF", "summary": "In this paper, the authors proposed an algorithm (DeepPCR) which parallelizes typically sequential operations used in 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 got recovered using the Parallel Cyclic Reduction algorithm. \nTo verify the the effectiveness of DeepPCR, authors presented the results on multi-layer perceptrons, and reach speedups of up to 30x for forward and 200x for backward pass. In addition, with ResNets with as many as 1024 layers, and generation in diffusion models, DeepPCR enables up to 7x faster training and 11x faster generation, respectively. ", "questions": "1. Which platform are the experiments in Sec. 4.1 performing at, CPU or GPU? In the appendix, deepPCR rerun for two GPUs: V100 and A100, but didn't see the result of baseline on these two GPUs.  \n2. Compared with Jacobi iterations in [36], is there any experimental comparison proving deepPCR have more benefit, such as ResNets? ", "limitations": "As mentioned in the limitation part, the proposed deepPCR has three main limitations:\n1. PCR requires fewer sequential steps overall which needs deployments efficiently, otherwise lead to performance degradation.\n2. The complexity actually becomes O(cN * log2 L), where cN identifies the number of Newton iterations necessary for convergence. There is a trade-off between complexity and accuracy. \n3. Memory increment for the temporary results. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. In this paper, authors proposed to interpret a sequential operation of L steps as the solution of a system of L equations, which can be recovered using the Parallel Cyclic Reduction algorithm. \n2. The paper is organized and written well. \n3. Compared with previous work,  instead of relying on variations of Jacobi iterations,  authors proposed our method specifically targets Markov sequences. \n4. In the experiment result section, authors have provided a comprehensive result to demonstrate the applicability of DeepPCR to a variety of scenarios. ", "weaknesses": "1. The motivation is not very clear. With the constraints of deepPCR, how much real-time saving is not shown in the result. That would be more helpful to understand if the speedup for neural networks sampling on dataset is provided. \n2. Compared with Jacobi iterations", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> W1.1. The motivation is not very clear.\n\nThe main motivation behind the proposed algorithm lies in providing an additional venue for speeding-up training and inference in NNs, to be used in conjunction with already-established parallelisation methods: this is particularly useful in time-sensitive applications, where reducing inference or training time is paramount. Based on the reviewer’s input, we made sure to adapt the abstract to further highlight this, including the sentence “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\", and we remain open to other suggestions.\n\n> W1.2. With the constraints of deepPCR, how much real-time saving is not shown in the result. That would be more helpful to understand if the speedup for neural networks sampling on dataset is provided.\n\nThe reviewer raises a valid point. Please notice that in the graphs showcasing our main results (Figs. 3, 4, 6, 8, 10) we do report the wall-times of both baseline and DeepPCR side-by-side. It is true however that in our discussion we mostly focus on speedup: this choice is dictated by the fact that we are mostly interested in checking that the scaling is as predicted by theory. Nonetheless, following the reviewer’s suggestion, we explicitly included some examples or wall-times in the main text as well:\n\n* In Line 205: “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 $\\gt30\\times$ and $200\\times$, respectively“\n* In Line 235: “Notice that using DeepPCR translates into a speedup of $\\sim7\\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.”\n* In Line 264: “We can see that DeepPCR manages to generate images up to $11\\times$ faster, reducing the required time from $1.3s$ to $0.12s$”\n\nWe hope our reply addresses the reviewer’s concerns. If not, we remain open to provide further clarifications.\n\n> W2, Q2. Compared with Jacobi iterations in [36], is there any experimental comparison proving deepPCR have more benefit, such as ResNets?\n\nWe thank the reviewer for pointing this out. We did not consider an experimental comparison between PCR and Jacobi iterations, but this is because the two methods are designed for fundamentally different target applications; as a consequence, we believe a comparison would not be fair (for Jacobi), nor insightful. More specifically, as the authors in [36] point out, Jacobi is performant only if the target system has a “heavy” lower-triangular structure. If instead the system reduces to block bi-diagonal, Jacobi basically becomes equivalent to forward substitution (that is, it reduces to the sequential case), as is described in [36, Sec4.1, Example 3]. Conversely, DeepPCR specifically relies on the block bi-diagonal assumption (guaranteed by the Markovian property of the target sequence), which is why the two algorithms are best suited for fundamentally different scenarios. We discuss this in lines 83-89 of our manuscript.\n\n> Q1. Which platform are the experiments in Sec. 4.1 performing at, CPU or GPU? In the appendix, deepPCR rerun for two GPUs: V100 and A100, but didn't see the result of baseline on these two GPUs. \n\nThe experiments in Sec4 were run on a V100 GPU with 40GB of memory. Only the experiments described in Fig11 were run on a A100 with 80GB. The purpose of the latter experiment is to investigate the causes for the lower performance of DeepPCR that we observed in Fig3, and ultimately to pinpoint them to memory limitations - which is why we compare GPUs with different available memory. Since this investigation pertained specifically to the properties of DeepPCR, we initially did not deem relevant to include baseline results as well. However, the reviewer is right in that those results are helpful in providing more evidence of DeepPCR working across different devices. The baseline results have now been included in Fig11, and can be seen in Fig1 in the attached pdf. \nAs expected, they are qualitatively similar to the ones recovered on the V100: we believe this further solidifies the theoretical analysis conducted in Sec3.\n\n> L1. PCR requires fewer sequential steps overall which needs deployments efficiently, otherwise lead to performance degradation.\nL2. The complexity actually becomes O(cN * log2 L), where cN identifies the number of Newton iterations necessary for convergence. There is a trade-off between complexity and accuracy. \nL3. Memory increment for the temporary results.\n\nIndeed, these are the main limitations of DeepPCR.\n\n* Regarding L1 and L3: With the experiments in Fig3 and Figs10-11-12, we investigate in detail the effect of memory limitations on DeepPCR, but still identify relevant regimes for speedup.\n* Regarding L2: the trade-off between complexity and accuracy can be seen as both a limitation and an advantage: in our experiments in Sec4.1 we constrain ourselves to maintaining high-accuracy (for a fairer comparison wrt baseline), but one can squeeze additional speedup by sacrificing this, without affecting the final results significantly (this is investigated in Fig.7 and more in detail in App.E)\n\n"}, {"review_id": "wGNzivRDew", "reviewer": "Reviewer_Cpwv", "summary": "This paper presents a method called DeepPCR for converting sequential operations in Deep Neural Networks into parallel ones, thus accelerating DNN training, inference, and denoising procedure in diffusion models. The key idea of DeepPCR is to interpret a sequential operation of L steps that satisfies the Markov property as the solution of a system of L equations. DeepPCR tackles this solution by using the Parallel Cyclic Reduction algorithm (PCR) and Newton’s method in parallel, resulting in O(log L) steps compared to O(L) steps of the sequential operations. Evaluation results show DeepPCR brings up to 30X speedup for the forward pass and 200X speedup for the backward pass for certain cases in MLPs, and 11.2X speedup in image generation via diffusion. DeepPCR also accelerates ResNet training by up to 7X. \n", "questions": "In summary, this is an interesting work that aims to parallelize sequential operations in DNNs from a more theoretical and fundamental perspective. I generally enjoyed my reading and I have some detailed questions as follows besides the items in the weakness section above. It would be very helpful if the authors could provide more justifications:\n\n1. DeepPCR currently works on sequential neural network structures that satisfy the Markov property. Is it a fundamental limitation? It would be helpful to explicitly list some neural architectures that follow (or do not follow) the Markov property to assist in understanding this limitation (if space allows).\n2. DeepPCR focuses on execution performance. Some important evaluation specifications are necessary, e.g., GPU type and DNN training/inference framework/engine.\n3. DeepPCR trains ResNet on the MNIST dataset. Why not use ImageNet that might be more convincing? \n4. What is the relationship between Deep PCR’s design and the effort in operator fusion (and subsequent efficient parallelization)? It would be extremely helpful to include some discussion about this.\n5. PCR requires extra memory (and results in more intensive computation) than its sequential counterpart, thus sensitive to the available hardware. How does it affect parallel efficiency? Does this cause any problems for large-scale training tasks that require distributed settings and intensive communications among nodes?\n", "limitations": "Limitations are carefully explained.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ This is a solid work that improves the parallel execution of DNNs from a more theoretical and fundamental perspective.\n+ This paper is well written. Both the theory and system aspects of DeepPCR are carefully analyzed.\n+ The resulting performance gains are significant.\n", "weaknesses": "- The general idea of interpreting a sequential operation as the solution of a large system of equations is not brand new [36].  \n- It has a relatively strong assumption that the operation sequence needs to satisfy the Markov property, i.e., the output of each step depends only on that of the previous step and no past steps.\n- Some evaluation details seem missing or unclear.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> W1 The general idea of interpreting a sequential operation as the solution of a large system of equations is not brand new [36]\n\nIndeed this intuition is not new: as mentioned in the paper, the first theorisation of this comes from [30] (a work also referenced in [36]). Still, we do not claim that our novelty lies in this intuition, but rather in the operationalisation of this method for Markovian sequences, as well as an in-depth theoretical and experimental analysis of its characteristics. This application is complementary to the one in [36] (as their method instead targets non-Markovian sequences), but is particularly relevant for modern Deep Learning architectures, as we discuss in W2, Q1\n\n> W2,Q1 DeepPCR currently works on sequential neural network structures that satisfy the Markov property. Is it a fundamental limitation? It would be helpful to explicitly list some neural architectures that follow (or do not follow) the Markov property to assist in understanding this limitation\n\nAs correctly stated by the reviewer, our working assumption is that we are dealing with Markovian sequences. We point out that this assumption is typically satisfied: a non-exhaustive list of architectures obeying it is\n* Feed-Forward NN\n* Convolutional NN\n* Graph NN\n* (Unidirectional) Recurrent NN\n* Normalising Flows\n* Diffusion Networks\n* Autoregressive models depending on one previous state\n\nConversely, this assumption breaks down in two main cases\n* Presence of residual connections\n* Autoregressive models relying on a history of previous states\n\nEven if the Markovian assumption does break down, in some cases one can still rely on some workarounds to apply DeepPCR. Indeed we successfully use one such method in our ResNets experiments. We refer to the answer of Q1 of 7Bpu for a list of possible approaches to circumvent some common non-Markovian cases\n\n> W3, Q2 DeepPCR focuses on execution performance. Some important evaluation specifications are necessary, eg, GPU type and DNN training/inference framework/engine\n\nWe thank the reviewer for this remark. We have now included this information in the paper, adding at Line 193: “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.” For the study on the effect of increasing GPU memory in Fig11, we are using a A100 GPU with 80GB, as reported in Line 540.\nWe hope this addresses the reviewer’s concerns. If more details still require clarification, we kindly invite the reviewer to specify them, and we will gladly include them.\n\n> Q3 DeepPCR trains ResNet on the MNIST dataset. Why not use ImageNet that might be more convincing?\n\nWe agree with the reviewer on the relevance of showcasing results also on more complex datasets. For the specific experiments in Sec4.1/4.2 our goal is mainly measuring how DeepPCR performance varies in terms of specific architecture parameters (namely, width and depth of the NN considered). A smaller dataset allowed us to iterate freely across these parameters, and since in this case the computations involved are rather agnostic to the type of data being fed to the NN, we do not expect that changing the dataset will affect our conclusions on the recovered speedup.\nFollowing the reviewer's suggestion, however, we are including additional results on diffusion on CIFAR10 and CelebA. These can be seen in the pdf, and are qualitatively similar to the ones on MNIST, reporting up to 9x and 10x speedups with no loss in FID.\n\n>Q4 What is the relationship between Deep PCR’s design and the effort in operator fusion (and subsequent efficient parallelization)?\n\nWe were not familiar with operation fusion, so we are grateful to the reviewer for bringing this to our attention. From [what we could gather](https://learn.microsoft.com/en-us/windows/ai/directml/dml-fused-activations), it is an optimisation to reduce redundant memory reads, merging certain layer operations (most noticeably, the element-wise nonlinearity which typically follows other layers). This could benefit DeepPCR too, as it will accelerate the two main overheads associated with the Newton solver:\n* Residual computation: The right-hand side of the linearised system (3) is computed by applying each “step function” $f_l$ in parallel. These are the same functions constituting the layers of the NN, so these too can benefit from operation fusion, in a similar way as the sequential forward pass\n* Jacobian assembly: Similarly, the Jacobians in (3) are assembled using autograd in parallel over the individual $f_l$: operator fusion should be able to benefit backward passes as well, following similar principles (since also the derivatives of the nonlinearities are an element-wise operation) - if this is true, then DeepPCR can gain performance also in this situation\n\n> Q5 PCR requires extra memory (and results in more intensive computation) than its sequential counterpart, thus sensitive to the available hardware. How does it affect parallel efficiency? Does this cause any problems for large-scale training tasks that require distributed settings and intensive communications among nodes?\n\nWe thank the reviewer for raising this interesting point. The inter-node communication overhead associated with the DeepPCR operations ultimately depends on the way the model is partitioned:\n* For a data-parallel setting, no additional communication is necessary, but the GPU must be able to accommodate for the extra memory requirements (roughly $2\\times$ the size of the model, as per Line174 and analysis in SecB.2)\n* For a model-parallel setting, there will likely be additional communication during the reduction phase of DeepPCR, depending on how the Jacobians are distributed across the GPUs. Notice that each of the reduction steps (lines 5-6 in Alg1) only involve communication between two nodes ($l$ and $l-i$), although a sync between all nodes is required at the end of each reduction step"}, {"review_id": "DRHHsleWtl", "reviewer": "Reviewer_7Bpu", "summary": "This work introduces DeepPCR, a method for parallelizing sequential computations with an application in neural networks. The method is based on the assumption that the target sequence is Markovian, and if so, DeepPCR can leverage Parallel Cyclic Reduction to convert the sequence into a parallel version that executes in O(log_2 L) instead of O(L). To account for the nonlinearity in neural networks, the authors approximate it using Newton’s method. They show that considerable wall-time speedup can be achieved using DeepPCR in both forward and backward passes on MLP/ResNet/diffusion process. ", "questions": "* One important assumption made by DeepPCR is the Markov property. However, some layers/operators do require additional historical information or outputs of the previous layer. Could you clarify how to handle such cases, for instance, residual connections? Although you may fold the residual blocks(as discussed in the appendix), I wonder what will happen when the sizes of the residual blocks are non-uniform. Will the load of each parallel stage be unbalanced?\n* I am interested in the PCR overhead prior to training. How much time is required to perform PCR?\n", "limitations": "The authors have properly addressed the limitations. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The authors tackle a significant problem using novel approaches. To solve the non-linearity, they propose to use Newton's method and prove that it requires few iterations when the non-linear operations (eg. activation) are smooth.\n* The evaluation section was convincing, and the limitations were properly discussed. \n* This work could have a profound impact on the parallelization techniques \n* The paper was overall well-written, with the main claims clearly articulated and properly supported by references and evaluations.", "weaknesses": "* As one of the tradeoffs here is memory v.s. compute time, it would be great to see an analysis of space complexity with respect to the Jacobian blocks/model width, and a detailed memory profile breakdown at training time (which can be obtained using a profiler).\n* More evaluations on the transformer-based model will be welcome. The architecture of BERT/T5/GPT should satisfy the Markovian property, and I would expect DeepPCR to be able to work on them.\n* To evaluate the quality of stable diffusion models trained using DeepPCR, it would be great to present the FID and the generated images for quantitative and qualitative comparison. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">W1 As one of the tradeoffs here is memory vs compute time, it would be great to see an analysis of space complexity with respect to the Jacobian blocks/model width, and a detailed memory profile breakdown at training time (which can be obtained using a profiler).\n\nThe reviewer raises an excellent point. Due to space limitations, we left the memory complexity analysis in the appendix SecB.2. However, we agree that this should be highlighted in the main text: we hence moved the relevant considerations to Sec3.1. In short, memory usage scales as expected: linearly in the number of steps in the sequence, and quadratically with Jacobian size.\n\n>W2 More evaluations on the transformer-based model will be welcome. The architecture of BERT/T5/GPT should satisfy the Markovian property, and I would expect DeepPCR to be able to work on them\n\nWe thank the reviewer for the suggestion: indeed we agree, as this is the focus of our future work. The main complication lies in that generally the text sequence is generated using argmax to sample tokens over a probability distribution. This is heavily nondifferentiable, so Newton is bound to perform poorly. As an alternative, we are considering instead the sequences of probability distributions over the tokens. Nonetheless, identifying the optimal way to apply the algorithm to LLMs remains a matter of future investigation\n\n>W3 To evaluate the quality of stable diffusion models trained using DeepPCR, it would be great to present the FID and the generated images for quantitative and qualitative comparison\n\nAbsolutely. Notice that as form of quantitive measure, in Fig6 we considered the _difference_ in Wasserstein distances (same distance used in FID) between baseline and DeepPCR-generated images. This difference is less than 0.01, showing that DeepPCR maintains the same quality of the sequential model. Following the reviewer's feedback, we have included examples of actual images in the appendix (and in the pdf), showing that images generated with DeepPCR are qualitatively indistinguishable from those generated sequentially. We have also included the actual Wasserstein distances of both sequentially and DeepPCR-generated images to the original dataset. \n\n>Q1.1 One important assumption made by DeepPCR is the Markov property. However, some layers/operators do require additional historical information or outputs of the previous layer. Could you clarify how to handle such cases, for instance, residual connections? \n\nWe thank the reviewer for the relevant remark. Indeed, the working assumption of our algorithm is that the target sequence $z_l=f_l(z_{l-1})$ is Markovian. Nonetheless, we can leverage the flexibility in defining both the states $z_l$ and the step functions $f_l$ to find workarounds if the Markovian property is not strictly satisfied. These cases are briefly discussed next:\n1. Residual connections: Depending on how the connections are distributed, feasible workarounds are:\n    * a. Short→ collapse blocks: this is the applied in our ResNets experiments\n    * b. Long→ apply DeepPCR inside, then again hierarchically to stitch the states across the various connections (inspired by [13])\n    * c. Nested (UNets)→ split network in half (downsampling/upsampling), apply DeepPCR to each half. A similar approach can be used for bidirectional RNNs, separating forward and backward part\n    * d. Criss-crossed→collapse blocks and consider state as collation of all states inside\n    * e. DenseNets→failure case of DeepPCR: if all layers are connected to all those after them, the Jacobian is fully lower-triangular, and one cannot leverage any sparsity pattern to return to the bidiagonal/Markovian case.\n\n2. Autoregressive models relying on a history of previous states: If the model relies on a window of states to produce the next output in the generating sequence (eg PixelCNN), we can collapse these states together to return to the Markovian case, similarly to 1d). Notice this causes each block in the Jacobian to grow, but if the interdependence between the states is limited, it will present a sparse structure, reducing the complexity of its reduction operations\n\nWe are grateful to both reviewers Cpwv and 7Bpu for triggering this discussion: we believe it makes for a relevant addition to our work, and we have included an expanded version of it in a novel section of the appendix (to be further refined based on the reviewing process)\n\n> Q1.2 Although you may fold the residual blocks, I wonder what will happen when the sizes of the residual blocks are non-uniform. Will the load of each parallel stage be unbalanced?\n\nYes, if length/distribution of the residual connections vary throughout the NN, the cost of applying the corresponding step functions $f_l$ also ends up varying with $l$, causing bottlenecks in the application of each step of DeepPCR. However, one can leverage the flexibility in the definition of the sequence $z_l=f_l(z_{l-1})$ to try and render these costs more uniform. For example, by composing together layers including shorter residual connections so to match the cost of applying a layer presenting a longer one\n\n> Q2 I am interested in the PCR overhead prior to training. How much time is required to perform PCR?\n\nThanks for the relevant remark. The main cost of applying DeepPCR lies in the PCR routine (Alg1). On top of this, the main overhead is associated with assembling system (3), and particularly the Jacobian blocks (but this is a perfectly parallel operation).\n\nThe timing results in Sec4 already report the total time required by DeepPCR (overheads included); following the reviewer's suggestion, though, we further break this down to highlight the cost of the Jacobian assembly for the diffusion experiments in Fig6. These remain roughly constant (at ~3.5ms) regardless of the architecture chosen, and low with respect to the total time taken by DeepPCR (which ranges from 0.05s to 0.25s). The figure is included in the pdf, as well as in the main paper in SecD3\n"}]

Author rebuttal:
We are grateful to the reviewers for the comments on our work, and their appreciation for the soundness of the manuscript and for the impactfulness of our proposed algorithm. We are encouraged that all reviewers found that the paper is well-written, that we “provided a comprehensive result to demonstrate the applicability of DeepPCR to a variety of scenarios” (uCYF), that the “resulting performance gains are significant” (CpWv), and that our “work could have a profound impact on the parallelization techniques” (7Bpu). The remarks provided sparkled interesting discussions among the authors, and will undoubtedly contribute to increasing the overall quality of the paper.

Before answering each individual reviewer, we wanted to summarise the main changes that these reviews have triggered:

* We have expanded the results including applications of DeepPCR to Diffusion Models trained on CIFAR10 and on CelebA. The speedups in these cases are consistent with those shown on MNIST, reaching up to 9x and 10x respectively, depending on the architectures considered. These results are provided in the attached pdf and will be included in the updated submission
* We have provided more details about hardware used and the evaluation performed, which have also been included in the paper
* We have listed common architectures that follow the Markovian assumption underlying our algorithm, as well as cases where this assumption does not hold, and identified workarounds that can be used to circumvent this (some of which were tested in our experiments)
* We have better illustrated the relationship between our algorithm and other approaches to speedup such as Jacobi iterations [36] and operators fusion. We have explained in which way they differ and in which case they are complementary to DeepPCR
* Following a reviewer’s suggestion, we have produced a breakdown of the timings required by DeepPCR for the diffusion experiments, showing how overhead costs (particularly, Jacobian assembly) are negligible with respect to the "core" Alg1. These have been included in the pdf, and the updated paper

We appreciate the reviewer’s insights and questions, which we have aimed to address comprehensively. We kindly ask the reviewers for any further feedback they might have, and if they deem appropriate, we would be grateful for taking our responses into account in their assessment.

Meta-review: All reviewers agree that this is an interesting paper.

One of the concerns was that although the current work is inspired by the prior work by Song et al. [36], there is no direct discussion of pros and cons of Newton iteration compared to Jacobi. I think the authors did explain this well in the rebuttal. While there is no global convergence guarantee, it is nice to see that parallel speedup is possible for the strictly sequential case (which is the worst case for Jacobi).
Decision: Accept (poster)