SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {W _ {l}} c = \nabla_ {\boldsymbol {z} _ {l}} c \otimes \sigma_ {l} (\boldsymbol {z} _ {l - 1}), \quad \text {a n d} \quad \nabla_ {\boldsymbol {b} _ {l}} c = \nabla_ {\boldsymbol {z} _ {l}} c, \tag {10}
$$
with  $\pmb{a} \otimes \pmb{b} = \pmb{a} \cdot \pmb{b}^T$  denoting the outer product operator. The gradients  $\nabla_{\pmb{z}_l} c$ , in turn, are given by applying the chain rule and solving
$$
\nabla_ {\boldsymbol {z} _ {l - 1}} c = J _ {f _ {l} ^ {F P}} \left| _ {\hat {\boldsymbol {z}} _ {l - 1}} ^ {T} \nabla_ {\boldsymbol {z} _ {l}} c, \quad \text {f o r} \quad l = 1, \dots , L \right. \tag {11}
$$
backwards, starting from  $\nabla_{\boldsymbol{z}_L}c$ , which can be evaluated directly from the output  $\boldsymbol{y}$ . This is the target sequential operation that we aim to parallelize. Notice that  $\hat{\boldsymbol{z}}_l$  represents the activations at layer  $l$  recovered during the forward pass, and that  $J_{f_l^{FP}}\Big|_{\hat{\boldsymbol{z}}_{l - 1}}^T$  is the transpose of the same operator appearing in (8). Notice moreover that (11) already represents a sequence of linear operations, and as such we can solve it by applying PCR directly, without needing Newton iterations. More explicitly, the target system is given by
$$
\left[ \begin{array}{c c c c} I & & & \\ - J _ {f _ {L} ^ {F P}} \Big | _ {\hat {\mathbf {z}} _ {L - 1}} ^ {T} & I & & \\ & \ddots & \ddots & \\ & & - J _ {f _ {1} ^ {F P}} \Big | _ {\hat {\mathbf {z}} _ {0}} ^ {T} & I \end{array} \right] \left[ \begin{array}{c} \nabla_ {\boldsymbol {z} _ {L}} c \\ \nabla_ {\boldsymbol {z} _ {L - 1}} c \\ \vdots \\ \nabla_ {\boldsymbol {z} _ {0}} c \end{array} \right] = \left[ \begin{array}{c} \nabla_ {\boldsymbol {z} _ {L}} c \\ \boldsymbol {0} \\ \vdots \\ \boldsymbol {0} \end{array} \right]. \tag {12}
$$
This property is not limited to MLPs: indeed, backward passes through any architecture are linear by default.
A.3 Forward and backward passes in ResNets
ResNets are characterized by skip connections which relate activations from layers at a 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  $z_{l} = f_{l}^{RN}(z_{l - 1},z_{l - s})$  for  $l\% s = 0$  (where  $\%$  denotes the modulo operation), meaning it depends both on the activations at the previous layer and those  $s$  layers prior. Nonetheless, we can still apply the DeepPCR procedure, if we collapse all layers in-between a skip connection into a single one. This amounts to considering the following step function:
$$
\begin{array}{l} \boldsymbol {z} _ {l} = f _ {l} ^ {R N} \left(\boldsymbol {z} _ {l - 1}, \boldsymbol {z} _ {l - s}\right) = f _ {l} ^ {R N} \left(f _ {l - 1} ^ {R N} \left(\boldsymbol {z} _ {l - 2}\right), \boldsymbol {z} _ {l - s}\right) \tag {13} \\ = \dots = f _ {l} ^ {R N} \left(f _ {l - 1} ^ {R N} \circ \dots \circ f _ {l - s + 1} ^ {R N} \left(\boldsymbol {z} _ {l - s}\right), \boldsymbol {z} _ {l - s}\right) := \hat {f} _ {l} ^ {R N} \left(\boldsymbol {z} _ {l - s}\right), \\ \end{array}
$$
where  $\circ$  stands for function composition. Notice that, after this modification is done, one can apply DeepPCR to the reduced system composed only of unknowns  $z = [z_0^T, z_s^T, z_{2s}^T, \ldots]$ . In our experiments, we consider MLPs as intermediate layers, and simple residual addition as a skip connection layer. This can be denoted as
$$
f _ {l} ^ {R N} = \left\{ \begin{array}{l l} W _ {l} \sigma_ {l} \left(\boldsymbol {z} _ {l - 1}\right) + \boldsymbol {b} _ {l} & \text {if} \quad l \% s \neq 0 \\ W _ {l} \sigma_ {l} \left(\boldsymbol {z} _ {l - 1}\right) + \boldsymbol {b} _ {l} + \boldsymbol {z} _ {l - s} & \text {else}. \end{array} \right. \tag{14}
$$
The Jacobians of the resulting collapsed forward pass is then given by
$$
\left. J _ {\hat {f} _ {l} ^ {R N}} \right| _ {\boldsymbol {z} _ {l}} = I + \prod_ {i = 0} ^ {s - 1} \left. J _ {f _ {l - i} ^ {R N}} \right| _ {\boldsymbol {z} _ {l - i}}, \tag {15}
$$
and  $J_{f_{l - i}^{RN}}$  shares the same form as  $J_{f_{l - i}^{FP}}$  in (8). A similar reasoning can be applied to the backward pass.
A.4 Denoising procedure in diffusion models
Also image generation in a diffusion model is, at its heart, the result of a sequential operation: an image is extracted from random noise by having it 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
$$
\boldsymbol {z} _ {l} = f _ {l} ^ {D M} \left(\boldsymbol {z} _ {l - 1}\right) = \frac {1}{\sqrt {\alpha_ {l}}} \left(\boldsymbol {z} _ {l} - \frac {1 - \alpha_ {l}}{\sqrt {1 - \bar {\alpha} _ {l}}} g \left(\boldsymbol {z} _ {l - 1}, l\right)\right) + \sqrt {\beta_ {l}} \mathcal {N} (0, 1), \tag {16}
$$
where  $\mathbf{z}_l$  denotes a partially denoised image (which underwent  $l$  denoising steps),  $\mathcal{N}(0,1)$  is a Normal sample, while  $\alpha_l, \beta_l$ , and  $\bar{\alpha}_l$  are parameters of the noise scheduling. Notice that the particularity of this step function is that it hides a NN application inside of it:  $g(\mathbf{z},l)$  identifies the action of the NN trained to predict the noise in image  $\mathbf{z}$  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 of the NN output with respect to its input. We have in fact
$$
\left. J _ {f _ {l} ^ {D M}} \right| _ {\boldsymbol {z} _ {l}} = \frac {I}{\sqrt {\alpha_ {l}}} - \frac {1 - \alpha_ {l}}{\sqrt {1 - \bar {\alpha} _ {l}}} \left. J _ {g} \right| _ {\left(\boldsymbol {z} _ {l - 1}, l\right)} \left(\boldsymbol {z} _ {l - 1}, l\right). \tag {17}
$$
Luckily,  $J_{g}|_{(\pmb{z}_{l-1}, l)}$  can be efficiently calculated using autograd functionalities, at a cost comparable of that of applying  $g(\pmb{z}, l)$ .
Notice this applies equivalently to diffusion mechanisms acting in pixel and latent space: in the former,  $z_{l}$  denotes an actual (noisy) image; in the latter,  $z_{l}$  denotes a (noisy) representation of an image, according to a 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 a Markovian sequence  $z_{l} = f_{l}(z_{l - 1})$ . Indeed, this automatically guarantees the block bidiagonal structure of system (3), upon which the PCR routine relies. In this section,
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 a 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  $z_{l}$  and the step functions  $f_{l}$  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  $z_{l}$  ends up becoming a function of multiple previous states  $z_{l} = f_{l}(z_{l - 1},z_{l - s})$  (see also (14)). 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 \ll 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 parallelization offered by our method.
- If multiple skip connections are nested within each other, that is if we have  $z_{l} = f_{l}(z_{l - 1},z_{l - s})$  and  $z_{\hat{l}} = f_{\hat{l}}(z_{\hat{l} -1},z_{\hat{l} -\hat{s}})$  for two (or more) given pairs  $l - s < \hat{l} -\hat{s} < \hat{l} < l$ , then the whole sequence can be cut between  $\hat{l} -\hat{s}$  and  $\hat{l}$  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 individually. 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 PixelCNN [38]), then we are considering a non-Markovian sequence where  $\boldsymbol{z}_l = f_l(\boldsymbol{z}_{l-1}, \boldsymbol{z}_{l-2}, \ldots, \boldsymbol{z}_{l-s})$ . To revert to the Markovian case, we can instead consider the collation of states  $\hat{\boldsymbol{z}}_{l-1} = \left[\boldsymbol{z}_{l-1}^T, \boldsymbol{z}_{l-2}^T, \ldots, \boldsymbol{z}_{l-s}^T\right]^T$  and notice that indeed we can write  $\hat{\boldsymbol{z}}_l = \hat{f}_l(\hat{\boldsymbol{z}}_{l-1})$ . Note that this results in a Jacobian with larger blocks (by a 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 a more in-depth investigation of the solutions outlined here is matter of future research.
Figure 9: Time to complete a 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 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.
C Additional results for single pass in MLPs
In this section, we expand on the results from the application of DeepPCR to accelerate forward and backward passes through MLPs.
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 as 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.
Figure 11: Time to complete a single forward pass, for MLPs of varying depths  $L$  and widths  $w$ . The left column refers to classical sequential implementation, the center column to the parallel implementation using DeepPCR, the right column reports the ratio between the two. Same configuration as Fig. 3, but the results are reported as a function of width, rather than depth.
In Fig. 9 and 10, we report the results from the same experiments in Sec. 4.1, but using tanh and sigmoid 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 in 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^{6}$  than  $2^{7}$ . This is also in line with what observed in Sec. 4.4, that is 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 of the parameters  $w$  and  $L$  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 a linear one. This is a first indication that the reduction operations in Alg. 1 stop being conducted in parallel, and start being serialized instead. Here we show evidence that this phenomenon is linked to memory limitations in the hardware used for our experiments.
In Fig. 11 we report the same timing results as in Fig. 3, but expressed as a function of the MLPs width  $w$  instead. This is to better highlight failure conditions for DeepPCR: since its complexity is only a function of  $L$ , timings should not be affected by changes in  $w$ , and hence the curves in Fig. 11 should be constant. This is however not the case, and particularly from the right plot in Fig. 11 we can infer that it is the value of  $w^2 L$  that determines when performance drops. To better highlight this, we include in the graph an isoline of  $w^2 L$ : notice how the threshold after which DeepPCR performance drops has the same inclination. The quantity  $w^2 L$  is directly related to the size of the system in (3):
Figure 12: Timings for forward pass through MLP of various sizes  $w$ ,  $L$  (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 DeepPCR is still in logarithmic regime), or 0.65 (when DeepPCR falls to linear regime: the factor of 2 difference stems from the A100 managing to preserve the logarithmic regime for longer).
Figure 13: Time to complete a 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 are conducted on a 80GB A100 GPU.
$w^{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 to 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 an NVIDIA Tesla A100, with 80GB. We observe that indeed the GPU with larger memory manages to delay the decay of DeepPCR to linear regime by a factor of 2 in the  $w^{2}L$  configuration of the MLP considered. Particularly, this is 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  $\sim 1.1$  for the configurations where DeepPCR performance is preserved, and fall to half that value,  $\sim 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 to the experiment in Fig. 3) on 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 of DeepPCR, and how these are mostly related to the need of storing intermediate operators  $A_{l}$  resulting from the row reductions in line 6 of Alg. 1. From (4), we can see how these operators have the same size as the Jacobian of the step function  $f_{l}(z)$  of the target sequence considered for parallelization via DeepPCR. When applied
Maximum memory allocation - ratio Baseline/DeepPCR
Figure 14: GPU memory allocation for forward pass through MLP: ratio between baseline and DeepPCR for MLPs of varying depths  $L$  and widths  $w$ . Same configurations as in Fig. 3, maximum over 100 runs.
Figure 15: Accuracy evolution during training with forward and backward passes computed sequentially (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.
to parallelize the forward pass through a MLP, as in our experiments in Sec. 4.1, these Jacobians have size  $w^2$ , exactly as the weight matrices  $W_{l}$  in the network. Since in our experiments storing the model itself is the most memory-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 is verified in Fig. 14, where we report the ratio of the maximum memory allocated when performing a forward pass, using the sequential approach and DeepPCR. The ratio sits consistently around 0.6: in 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 to grow as  $w^{2}L$ ,  $w$  being the size of the Jacobians of the step function  $f_{l}(z)$  considered, and  $L$  the total number of steps to 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 in Sec. 4.2. This serves as 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 to 0.
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.
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 5 different seeds for initialization of the network parameters.
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 on how the timings were conducted, and include additional results from pixel-space diffusion, which was not described in 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 on MNIST by minimizing the loss
$$
\mathcal {L} _ {A E} = \left\| x - G (E (x)) \right\| _ {2} ^ {2} + \beta \cdot K L (E (x), \mathcal {N} (0, 1)) + \lambda \cdot D (G (E (x))), \tag {18}
$$
where  $E$  is an encoder,  $G$  is a generator/decoder, and  $D$  is an adversarial discriminator. We set  $\beta$  to  $10^{-6}$  and  $\lambda$  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 on MNIST, we also simplify the architectures of  $E$ ,  $G$ , and  $D$ , which consist of 4 convolutional layers and produce a 1-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 1a and 1b.
Once the autoencoder is trained, we freeze its weights and train a diffusion process on 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 1-dimensional, we use a residual MLP instead of the UNet used by Rombach et al. [37]. See Table 1d for details on the MLP architecture. Finally, we generate new samples by running a diffusion process on randomly sampled latents as in Section 4.3, and generate the final images by applying the decoder just once on 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 on the class), and pass the result through a linear layer to predict each of the elementwise affine coefficients of the AdaLayerNorm [21] layers of the diffusion MLP.
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 of the generator. We trained all models on a single A100 GPU with 80G of vRAM.
We also experiment with diffusion in pixel space. To this purpose, we train a UNet to recognize the noise in a given image, using MSE as a loss function, following [17]. The training procedure uses AdamW with a constant learning rate of  $10^{-4}$  as optimizer, and runs for 500 epochs with a batch size of 512. Conditioning with this model is performed in the same way as with the latent-diffusion model, by adding the sinusoidal embedding of time and class. Additional details on the architecture of the network are reported in Tab. 1e.
E.2 Remarks on profiling for diffusion experiments