SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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, l)$  used as denoiser. This requires assembling one Jacobian per denoising step  $l$ : a task which is perfectly parallelizable over  $l$ . To perform this parallelization, we attempted to combine the functionalities provided by vmap and autograd, which are however a recent addition in PyTorch, and still under development. Indeed, after some testing with this implementation, we noticed that the memory usage during this phase was order of magnitudes higher than actually required, effectively causing the program to crash due to lack of available memory. After some investigation, we suspect the issue is associated with a suboptimal implementation of vmap for our target application, which ends up spawning unnecessary copies of the NN, thus rapidly occupying all available memory. To circumvent 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,
Table 1: Description of latent and pixel diffusion model components.  $3 \times 3$  specifies the convolution kernel size. ConvDown is a convolution layer with stride 2, thus reducing spatial resolution by a factor of 2. ConvUp is a convolution layer followed by a nearest-neighbor upsample operation, thus increasing spatial resolution by a factor of 2. AdaLayerNorm 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-conditional 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  $\rightarrow$  out.  $|Z|$  denotes the latent space dimensionality
(a) Image Encoder (E)
Grayscale Image x ∈ R32×32
3 × 3 ConvDown 1 → 16, ReLU
3 × 3 ConvDown 16 → 32, ReLU
3 × 3 ConvDown 32 → 64, ReLU
3 × 3 ConvDown 64 → 128, ReLU
Reshape R4×4×128 → R2048
Linear 2048 → 256
μ, σ = Linear 256 → 2 * |Z|
z ~ N(μ, σ)
(b) Image decoder/generator (G)
Latent vector z ∈ R|Z|
Linear |Z| → 256
Linear 256 → 2048
Reshape R2048 → R4×4×128
3 × 3 ConvUp 128 → 64, ReLU
3 × 3 ConvUp 64 → 32, ReLU
3 × 3 ConvUp 32 → 16, ReLU
1 × 1 Conv 16 → 1, ReLU
Tanh
(c) Image discriminator (D)
Grayscale Image x ∈ R32×32
3 × 3 ConvDown 1 → 16, ReLU
3 × 3 ConvDown 16 → 32, ReLU
3 × 3 ConvDown 32 → 64, ReLU
3 × 3 ConvDown 64 → 128 ReLU
1 × 1 ConvDown 128 → 1
(d) Residual MLP
Latent code z ∈ R|Z|
Linear |Z| → 2048
AdaLayerNorm, Swish, Linear 2048 → 2048
AdaLayerNorm, Swish, Linear 2048 → 2048
Residual	×3
LayerNorm, Swish, Linear 2048 → |Z|
(e) UNet
Grayscale Image x ∈ Rw
3 × 3 Conv 1 → 16	×2
3 × 3 Conv c → c, ReLU, BatchNorm
3 × 3 Conv c → c, ReLU, BatchNorm
3 × 3 ConvDown c → 2c
Residual
3 × 3 Conv c → c, ReLU, BatchNorm	×2
3 × 3 Conv c → c, ReLU, BatchNorm
3 × 3 ConvUp c → c/2
Residual
3 × 3 Conv 16 → 1
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 Jacobians 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 a 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 see that denoising procedure is a stochastic process: at each step, random Gaussian noise is added to
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 as 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 6 of Alg. 1.
the image; this is generally done to ensure variability in the output generated. In our 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 we 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 on 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, as well as the average number of Newton iterations to convergence and  $L^{\infty}$  norm of the error at convergence. Moreover, we include results from the application of DeepPCR to pixel-space diffusion. Overall, also in this case the convergence error remains small, and the Newton iterations remain bounded, confirming what we observed for Sec. 4.1 and 4.2 as well. Notice however the increase in 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.
In Fig. 18 we 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.
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 to complete diffusion generation in latent-space and pixel-space (minimum over 100 trials).  $\dim (\mathbf{x}_0)$  refers to the input dimensions,  $L$  is the total number of diffusion steps, SEQUENTIAL is the amount of time inference takes in a 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  $L^{\infty}$  norm between the generated image using the sequential model and DeepPCR.
dim(x0)	L	dim(z1)	SEQUENTIAL	ADJ. DEEP-PCR	SPEED-UP	NEWTON ITS.	L∞ERROR
LATENT-SPACE	32×32	256	16	0.3063 s	0.0554 s	5.53×	7.05 it	0.00195
32×32	512	16	0.5997 s	0.0721 s	8.32×	9.06 it	0.00304
32×32	1024	16	1.2241 s	0.1091 s	11.22×	13.40 it	0.00523
32×32	256	32	0.3092 s	0.0547 s	5.65×	6.67 it	0.00252
32×32	512	32	0.6358 s	0.0757 s	8.40×	8.90 it	0.00298
32×32	1024	32	1.2543 s	0.1143 s	10.98×	14.06 it	0.00608
32×32	256	64	0.3093 s	0.0556 s	5.56×	6.43 it	0.00283
32×32	512	64	0.6175 s	0.0748 s	8.25×	9.24 it	0.00324
32×32	1024	64	1.2509 s	0.1420 s	8.81×	17.35 it	0.00271
32×32	256	128	0.3151 s	0.0481 s	6.55×	6.09 it	0.00282
32×32	512	128	0.6303 s	0.0669 s	9.42×	9.08 it	0.00364
32×32	1024	128	1.2390 s	0.2033 s	6.09×	18.54 it	0.00248
PIXEL-SPACE	4×4	256	16	0.6629 s	0.0974 s	6.85×	6.44 it	0.00069
4×4	512	16	1.3610 s	0.1323 s	10.37×	8.22 it	0.00143
4×4	1024	16	2.6695 s	0.1884 s	14.25×	12.33 it	0.00247
8×8	256	64	0.6971 s	0.1070 s	6.55×	6.89 it	0.00094
8×8	512	64	1.3112 s	0.1175 s	11.18×	7.89 it	0.00276
8×8	1024	64	2.8571 s	0.1792 s	15.95×	11.11 it	0.00418
16×16	256	256	0.6942 s	0.2018 s	3.46×	11.33 it	0.00126
16×16	512	256	1.3752 s	0.3940 s	3.79×	15.78 it	0.00235
16×16	1024	256	2.8357 s	1.1594 s	2.48×	27.33 it	0.02638
F Dependence of Newton solver on configuration details
As we have seen, a key step of DeepPCR lies in using Newton'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, we 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 a 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 a built-in noisy component in their design. We can then consider the approximate results stemming from an early-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.
PARAMETER	DESCRIPTION	FORWARD PASS	DIFFUSION
NUMBER OF ITERATIONS CN	Maximum number of Newton iterations the solver will apply.	15	30
ABSOLUTE THRESHOLD	Infinity-norm of the Newton solver residual.	10-4	10-4
RELATIVE THRESHOLD	Infinity-norm of the Newton solver residual normalized by the infinity-norm of the first iteration residual.	10-4	10-4
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 a 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  $c_{N} \in \{1,2,3,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 a ResNet of up to 1024 layers, each layer with a width of 16, and a residual connection every 4 layers, with just 3 Newton iterations for solving each forward pass. This is only half of the 6 iterations we reported in Sec. 4.2 to keep low error with respect to the sequential baseline, and as such represent a potential further  $2 \times$  speedup with respect to the reported one.
Effects on Diffusion generation In Fig. 20 and 21 we show convergence  $L^{\infty}$  error 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 of  $w = 16, 32, 64, 128$ , while for the latter we consider images of sizes  $w = 4 \times 4, 8 \times 8, 16 \times 16$ . Each marker in the plots corresponds to a measurement with the maximum number of iterations limited to a certain value.
Notice that, for the most part, the  $L^{\infty}$ -norm of the error falls below 0.1 very quickly (in fewer than 15 iterations regardless of image size). This value already represents a 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 to this is the latent-diffusion model with the most denoising steps in Fig. 20: for the first few Newton iterations there is 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  $L$  in Fig. 20, the Newton solver stalls after reaching  $10^{-1}$  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 is required in order to identify the causes of this behavior.
F.2 Changing the initial guess for the Newton solver
Choosing a suitable initial guess  $z^0$  for (2) can also have a relevant impact on the performance of the Newton solver. In conducting our experiments, we investigated also this aspect of the solver configuration, and tried to identify ways to 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.
Figure 19: Accuracy of ResNets training using a fixed number of Newton iteration. Lines represent mean accuracy over 5 runs, with shaded area spanning two standard deviations. Notice how, for the curves corresponding to  $c_{N} \geq 3$ , the accuracy evolution is basically indistinguishable.
Figure 20: Convergence error in  $L^{\infty}$  norm as a function of the Newton iterations, for latent-space diffusion with models using different latent-space dimensions  $w$  and number of denoising steps  $L$ .
Our choices for the initial guess  $z^0$  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 per the discussion in appendix A.2.
Training of ResNets Also for the set of experiments in 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, as  $z^0$  we use a batch-average of the activations resulting from the 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 in 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.
This choice seems to make for an effective initial guess: in Fig. 22 we measure its impact on the solver convergence, comparing it with zero and random initialization for  $z^0$ . For the model, we use a 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 in 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 of 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 are similar to the ones seen during train time, since the
Figure 21: Convergence error in  $L^{\infty}$  norm as a function of the Newton iterations, for pixel-space diffusion with models using different image sizes  $w$  and number of denoising steps  $L$ .
Figure 22: Training dynamics of a 1024 layer ResNet with DeepPCR, different initial guesses, and a fixed set of Newton iterations. Left plot 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 is a random guess from a 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  $z^0$  for these noisy images, we 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.
Figure 23: Measured speedup for various pixel diffusion models, as a function of the image size. Notice how conditional models achieve larger speedup due to a more fine-tuned initialization of the Newton method.