SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Saurabh Saxena, Charles Herrmann, Junhwa Hur, Abhishek Kar, Mohammad Norouzi, Deqing Sun, David J. Fleet*
{srbs,irwinherrmann,junhwahur,abhiskar,deqingsun,davidfleet}@google.com
Google DeepMind and Google Research
Denoising diffusion probabilistic models have transformed image generation with their impressive fidelity and diversity. We show that they also excel in estimating optical flow and monocular depth, surprisingly, without task-specific architectures and loss functions that are predominant for these tasks. Compared to the point estimates of conventional regression-based methods, diffusion models also enable Monte Carlo inference, e.g., capturing uncertainty and ambiguity in flow and depth. With self-supervised pre-training, the combined use of synthetic and real data for supervised training, and technical innovations (infilling and step-unrolled denoising diffusion training) to handle noisy-incomplete training data, and a simple form of coarse-to-fine refinement, one can train state-of-the-art diffusion models for depth and optical flow estimation. Extensive experiments focus on quantitative performance against benchmarks, ablations, and the model's ability to capture uncertainty and multimodality, and impute missing values. Our model, DDVM (Denoising Diffusion Vision Model), obtains a state-of-the-art relative depth error of 0.074 on the indoor NYU benchmark and an Fl-all outlier rate of  $3.26\%$  on the KITTI optical flow benchmark, about  $25\%$  better than the best published method. For an overview see diffusion-vision.github.io
Diffusion models have emerged as powerful generative models for high fidelity image synthesis, capturing rich knowledge about the visual world [21, 48, 55, 62]. However, at first glance, it is unclear whether these models can be as effective on many classical computer vision tasks. For example, consider two dense vision estimation tasks, namely, optical flow, which estimates frame-to-frame correspondences, and monocular depth perception, which makes depth predictions based on a single image. Both tasks are usually treated as regression problems and addressed with specialized architectures and task-specific loss functions, e.g., cost volumes, feature warps, or suitable losses for depth. Without these specialized components or the regression framework, general generative techniques may be ill-equipped and vulnerable to both generalization and performance issues.
In this paper, we show that these concerns, while valid, can be addressed and that, surprisingly, a generic, conventional diffusion model for image to image translation works impressively well on both tasks, often outperforming the state of the art. In addition, diffusion models provide valuable benefits over networks trained with regression; in particular, diffusion allows for approximate inference with multi-modal distributions, capturing uncertainty and ambiguity (e.g. see Figure 1).
One key barrier to training useful diffusion models for monocular depth and optical flow inference concerns the amount and quality of available training data. Given the limited availability of labelled
*DF is also affiliated with the University of Toronto and the Vector Institute.
Figure 1: Examples of multi-modal prediction on depth (NYU) and optical flow (Sintel and KITTI). Each row shows an input image (or overlayed image pair for optical flow), a variance heat map from 8 samples, and 3 individual samples. Our model captures multi-modality in uncertain/ambiguous cases, such as reflective (e.g. mirror on NYU), transparent (e.g. vehicle window on KITTI), and translucent (e.g. fog on Sintel) regions. High variance also occurs at object boundaries, which are often challenging cases for optical flow, and also partially originate from noisy ground truth measurements for depth. See Fig. 8, 9, 10 and 11 for more examples.
training data, we propose a training pipeline comprising multi-task self-supervised pre-training followed by supervised pre-training using a combination of real and synthetic data. Multi-task self-supervised pre-training leverages the strong performance of diffusion models on tasks like colorization and inpainting [e.g., 54]. We also find that supervised (pre-)training with a combination of real and large-scale synthetic data improves performance significantly.
A further issue concerns the fact that many existing real datasets for depth and optical flow have noisy and incomplete ground truth annotations. This presents a challenge for the conventional training framework and iterative sampling in diffusion models, leading to a problematic distribution shift between training and inference. To mitigate these issues we propose the use of an  $L_{1}$  loss for robustness, infilling missing depth values during training, and the introduction of step-unrolled denoising diffusion. These elements of the model are shown through ablations to be important for both depth and flow estimation.
1. We formulate optical flow and monocular depth estimation as image to image translation with generative diffusion models, without specialized loss functions and model architectures.
2. We identify and propose solutions to several important issues w.r.t data. For both tasks, to mitigate distribution shift between training and inference with noisy, incomplete data, we propose infilling, step-unrolling, and an  $L_{1}$  loss during training. For flow, to improve generalization, we introduce a new dataset mixture for pre-training, yielding a RAFT [74] baseline that outperforms all published methods in zero-shot performance on the Sintel and KITTI training benchmarks.
3. Our diffusion models is competitive with or surpasses SOTA for both tasks. For monocular depth estimation we achieve a SOTA relative error of 0.074 on the NYU dataset and perform competitively on KITTI. For flow, diffusion surpasses the stronger RAFT baseline by a large margin in pre-training and our fine-tuned model achieves an F1-all outlier rate of  $3.26\%$  on the public KITTI test benchmark,  $\sim 25\%$  lower than the best published method [70].
4. Our diffusion model is also shown to capture flow and depth uncertainty, and the iterative denoising process enables zero-shot, coarse-to-fine refinement, and imputation.
Optical flow and depth estimation have been extensively studied. Here we briefly review only the most relevant work, and refer the interested readers to the references cited therein.
Optical flow. The predominant approach to optical flow is regression-based, with a focus on specialized network architectures to exploit domain knowledge, e.g., cost volume construction [14, 22, 23, 38, 68, 81, 83, 85], coarse-to-fine estimation [68, 77, 82], occlusion handling [24, 27, 67], or iterative refinement [25, 26, 74], as evidenced by public benchmark datasets [4, 44]. Some
Figure 2: Training architecture. Given ground truth flow/depth, we first infill missing values using interpolation. Then, we add noise to the label map and train a neural network to model the conditional distribution of the noise given the RGB image(s), noisy label, and time step. One can optionally unroll the denoising step(s) during training (with stop gradient) to bridge the distribution gap between training and inference for  $y_{t}$ .
recent work has also advocated for generic architectures: Perceiver IO [28] introduces a generic transformer-based model that works for any modality, including optical flow and language modeling. Regression-based methods, however, only give a single prediction of the optical flow and do not readily capture uncertainty or ambiguity in the flow. Our work introduces a surprisingly simple, generic architecture for optical flow using a denoising diffusion model.
We find that this generic generative model is surprisingly effective for optical flow, recovering fine details on motion boundaries, while capturing multi-modality of the motion distribution.
Monocular depth. Monocular depth estimation has been a long-standing problem in computer vision [58, 59] with recent progress focusing on specialized loss functions and architectures [1, 5, 15, 31] such as the use of multi-scale networks [12, 13], adaptive binning [3, 35] and weighted scale-shift invariant losses [13]. Large-scale in-domain pre-training has also been effective for depth estimation [49, 50, 52], which we find to be the case here as well. We build on this rich literature, but with a simple, generic architecture, leveraging recent advances in generative models.
Diffusion models. Diffusion models are latent-variable generative models trained to transform a sample of a Gaussian noise into a sample from a data distribution [21, 62]. They comprise a forward process that gradually annihilates data by adding noise, as 'time'  $t$  increases from 0 to 1, and a learned generative process that reverses the forward process, starting from a sample of random noise at  $t = 1$  and incrementally adding structure (attenuating noise) as  $t$  decreases to 0. A conditional diffusion model conditions the steps of the reverse process (e.g., on labels, text, or an image).
Central to the model is a denoising network  $f_{\theta}$  that is trained to take a noisy sample  $y_{t}$  at some time-step  $t$ , along with a conditioning signal  $x$ , and predict a less noisy sample. Using Gaussian noise in the forward process, one can express the training objective over the sequence of transitions (as  $t$  slowly decreases) as a sum of non-linear regression objectives, with the L2 loss (here with the  $\epsilon$ -parameterization):
$$
\mathbb {E} _ {(\boldsymbol {x}, \boldsymbol {y})} \mathbb {E} _ {(t, \epsilon)} \left\| f _ {\theta} (\boldsymbol {x}, \underbrace {\sqrt {\gamma_ {t}} \boldsymbol {y} + \sqrt {1 - \gamma_ {t}}} _ {\boldsymbol {y} _ {t}} \epsilon , t) - \epsilon \right\| _ {2} ^ {2} \tag {1}
$$
where  $\epsilon \sim \mathcal{N}(0, I)$ ,  $t \sim \mathcal{U}(0, 1)$ , and where  $\gamma_t > 0$  is computed with a pre-determined noise schedule. For inference (i.e., sampling), one draws a random noise sample  $\pmb{y}_1$ , and then iteratively uses  $f_\theta$  to estimate the noise, from which one can compute the next latent sample  $\pmb{y}_s$ , for  $s < t$ .
Self-supervised pre-training. Prior work has shown that self-supervised tasks such as colorization [33, 86] and masked prediction [80] serve as effective pre-training for downstream vision tasks. Our work also confirms the benefit of self-supervised pre-training [54] for diffusion-based image-to-image translation, by establishing a new SOTA on optical flow and monocular depth estimation while also representing multi-modality and supporting zero-shot coarse-to-fine refinement and imputation.
Figure 3: Effects of adding synthetic datasets in pretraining. Diffusion models trained only with AutoFlow (AF) tend to provide very coarse flow estimates and can hallucinate shapes. The addition of FlyingThings (FT), Kubric (KU), and TartanAir (TA) remove the AF-induced bias toward polygonal-shaped regions, and significantly improve flow quality on fine detail, e.g. trees, thin structures, and motion boundaries.
3 Model Framework
In contrast to the conventional monocular depth and optical flow methods, with rich usage of specialized domain knowledge on their architecture designs, we introduce simple, generic architectures and loss functions. We replace the inductive bias in state-of-the-art architectures and losses with a powerful generative model along with a combination of self-supervised pre-training and supervised training on both real and synthetic data.
The denoising diffusion model (Figure 2) takes a noisy version of the target map (i.e., a depth or flow) as input, along with the conditioning signal (one RGB image for depth and two RGB images for flow). The denoiser effectively provides a noise-free estimate of the target map (i.e., ignoring the specific loss parameterization used). The training loss penalizes residual error in the denoised map, which is quite distinct from typical image reconstruction losses used in optical flow estimation.
3.1 Synthetic pre-training data and generalization
Given that we train these models with a generic denoising objective, without task-specific inductive biases in the form of specialized architectures, the choice of training data becomes critical. Below we discuss the datasets used and their contributions in detail. Because training data with annotated ground truth is limited for many dense vision tasks, here we make extensive use of synthetic data in the hope that the geometric properties acquired from synthetic data during training will transfer to different domains, including natural images.
AutoFlow [69] has recently emerged as a powerful synthetic dataset for training flow models. We were surprised to find that training on AutoFlow alone is insufficient, as the diffusion model appears to devote a significant fraction of its representation capacity to represent the shapes of AutoFlow regions, rather than solving for correspondence. As a result, models trained on AutoFlow alone exhibit a strong bias to generate flow fields with polygonal shaped regions, much like those in AutoFlow, often ignoring the shapes of boundaries in the two-frame RGB inputs (e.g. see Figure 3).
To mitigate bias induced by AutoFlow in training, we further mix in three synthetic datasets during training, namely, FlyingThings3D [40], Kubric [19] and TartanAir [76]. Given a model pre-trained on AutoFlow, for compute efficiency, we use a greedy mixing strategy where we fix the relative ratio of the previous mixture and tune the proportion of the newly added dataset. We leave further exploration of an optimal mixing strategy to future work. Zero-shot testing of the model on Sintel and KITTI (see Table 1 and Fig. 3) shows substantial performance gains with each additional synthetic dataset.
We find that pre-training is similarly important for depth estimation (see Table 7). We learn separate indoor and outdoor models. For the indoor model we pre-train on a mix of ScanNet [9] and SceneNet RGB-D [41]. The outdoor model is pre-trained on the Waymo Open Dataset [71].
3.2 Real data: Challenges with noisy, incomplete ground truth
Ground truth annotations for real-world depth or flow data are often sparse and noisy, due to highly reflective surfaces, light absorbing surfaces [65], dynamic objects [43], etc. While regression-based
Algorithm 1 Denoising diffusion train step with infilling and step unrolling
1:  $x\gets$  conditioning images,  $y\gets$  flow or depth map, mask  $\leftarrow$  binary mask of known values   
2:  $t\sim U(0,1),\epsilon \sim N(0,1)$    
3:  $y =$  fill_holes_with_interpolation(y)   
4:  $y_{t} = \sqrt{\gamma_{t}} *y + \sqrt{1 - \gamma_{t}} *\epsilon$    
5: if unroll_step then   
6:  $\epsilon_{pred} =$  stop_gradient  $(f_{\theta}(x,y_t,t))$    
7:  $y_{pred} = (y_t - \sqrt{1 - \gamma_t} *\epsilon_{pred}) / \sqrt{\gamma_t}$    
8:  $y_{t} = \sqrt{\gamma_{t}} *y_{pred} + \sqrt{1 - \gamma_{t}} *\epsilon$    
9:  $\epsilon = (y_t - \sqrt{\gamma_t} *y) / \sqrt{1 - \gamma_t}$    
10: end if   
11:  $\epsilon_{pred} = f_{\theta}(x,y_t,t)$    
12: loss  $=$  reduce_mean(|ε-εpred|[mask])
methods can simply compute the loss on pixels with valid ground truth, corruption of the training data is more challenging for diffusion models. Diffusion models perform inference through iterative refinement of the target map  $\mathbf{y}$  conditioned on RGB image data  $\mathbf{x}$ . It starts with a sample of Gaussian noise  $\mathbf{y}_1$ , and terminates with a sample from the predictive distribution  $p(\mathbf{y}_0 \mid \mathbf{x})$ . A refinement step from time  $t$  to  $s$ , with  $s < t$ , proceeds by sampling from the parameterized distribution  $p_{\theta}(\mathbf{y}_s \mid \mathbf{y}_t, \mathbf{x})$ ; i.e., each step operates on the output from the previous step. During training, however, the denoising steps are decoupled (see Eqn. 1), where the denoising network operates on a noisy version of the ground truth depth map instead of the output of the previous iteration (reminiscent of teaching forcing in RNN training [79]). Thus there is a distribution shift between marginals over the noisy target maps during training and inference, because the ground truth maps have missing annotations and heavily-tailed sensor noise while the noisy maps obtained from the previous time step at inference time should not. This distribution shift has a very negative impact on model performance. Nevertheless, with the following modifications during training we find that the problems can be mitigated effectively.
Infilling. One way to reduce the distribution shift is to impute the missing ground truth. We explored several ways to do this, including simple interpolation schemes, and inference using our model (trained with nearest neighbor interpolation). We find that nearest neighbor interpolation is sufficient to impute missing values in the ground truth maps in the depth and flow field training data.
Despite the imputation of missing ground truth depth and flow values, note that the training loss is only computed and backpropagated from pixels with known (not inflated) ground truth depth. We refer to this as the masked denoising loss (see Figure 2).
Step-unrolled denoising diffusion training. A second way to mitigate distribution shift in the  $y_{t}$  margins in training and inference, is to construct  $y_{t}$  from model outputs rather than ground truth maps. One can do this by slightly modifying the training procedure (see Algorithm 1) to run one forward pass of the model and build  $y_{t}$  by adding noise to the model's output rather than the training map. We do not propagate gradients for this forward pass. This process, called step-unrolled denoising diffusion, slows training only marginally ( $\sim 15\%$  on a TPU v4). This step-unrolling is akin to the predictor-corrector sampler of [63] which uses an extra Langevin step to improve the target marginal distribution of  $y_{t}$ . Interestingly, the problem of training / inference distribution shift resembles that of exposure bias [51] in autoregressive models, for which the mismatch is caused by teacher forcing during training [79]. Several solutions have been proposed for this problem in the literature [2, 32, 84]. Step-unrolled denoising diffusion also closely resembles the approach in [57] for training denoising autoencoders on text.
We only perform step-unrolled denoising diffusion during model fine-tuning. Early in training the denoising predictions are inaccurate, so the latent marginals over the noisy target maps will be closer to the desired true marginals than those produced by adding noise to denoiser network outputs. One might consider the use of a curriculum for gradually introducing step-unrolled denoising diffusion in the later stages of supervised pre-training, but this introduces additional hyper-parameters, so we simply invoke step-unrolled denoising diffusion during fine-tuning, and leave an exploration of curricula to future work.
$L_{1}$  denoiser loss. While the  $L_{2}$  loss in Eqn. 1 is ideal for Gaussian noise and noise-free ground truth maps, in practice, real ground truth depth and flow fields are noisy and heavy tailed; e.g., for distant objects, near object boundaries, and near pixels with missing annotations. We hypothesize
Figure 4: Visual results comparing RAFT with our method after pretraining. Note that our method does much better on fine details and ambiguous regions.
that the robustness afforded by the  $L_{1}$  loss may therefore be useful in training the neural denoising network. (See Tables 11 and 12 in the supplementary material for an ablation of the loss function for monocular depth estimation.)
3.3 Coarse-to-fine refinement
Training high-resolution diffusion models is often slow and memory intensive but increasing the image resolution of the model has been shown to improve performance on vision tasks [18]. One simple solution, yielding high-resolution output without increasing the training cost, is to perform inference in a coarse-to-fine manner, first estimating flow over the entire field of view at low resolution, and then refining the estimates in a patch-wise manner. For refinement, we first up-sample the low-resolution map to the target resolution using bicubic interpolation. Patches are cropped from the up-scaled map, denoted  $z$ , along with the corresponding RGB inputs. Then we run diffusion model inference starting at time  $t'$  with a noisy map  $y_{t'} \sim \mathcal{N}(y_{t'}, \sqrt{\gamma_{t'}} z, (1 - \gamma_{t'}) I)$ . For simplicity,  $t'$  is a fixed hyper-parameter, set based on a validation set. This process is carried out for multiple overlapping patches. Following Perceiver IO [28], the patch estimates are merged using weighted masks with lower weight near the patch boundaries since predictions at boundaries are more prone to errors. (See Section H.5 for more details.)
As our denoiser backbone, we adopt the Efficient UNet architecture [55], pretrained with Palette [54] style self-supervised pretraining, and slightly modified to have the appropriate input and output channels for each task. Since diffusion models expect inputs and generate outputs in the range  $[-1., 1.]$ , we normalize depths using max depth of 10 meters and 80 meters respectively for the indoor and outdoor models. We normalize the flow using the height and width of the ground truth. Refer to Section H for more details on the architecture, augmentations and other hyper-parameters.
Optical flow. We pre-train on the mixture described in Section 3.1 at a resolution of  $320 \times 448$  and report zero-shot results on the widely used Sintel [4] and KITTI [44] datasets. We further fine-tune this model on the standard mixture consisting of AutoFlow [69], FlyingThings [40], VIPER [53], HD1K [30], Sintel and KITTI at a resolution of  $320 \times 768$  and report results on the test set from the public benchmark. We use a standard average end-point error (AEPE) metric that calculates L2
Figure 5: Visual results comparing RAFT with our method after finetuning. Ours does much better on fine details and ambiguous regions.
Table 1: Zero-shot optical flow estimation results on Sintel and KITTI. We provide a new RAFT baseline using our proposed pre-training mixture and substantially improve the accuracy over the original. Our diffusion model outperforms even this much stronger baseline and achieves state-of-the-art zero-shot results on Sintel.final and KITTI.
Model	Dataset	Sintel.clean Sintel.final	KITTI
AEPE	AEPE	Fl-all
FlowFormer	Chairs→Things	1.01	2.40	4.09	14.72%
RAFT	Chairs→Things	1.68	2.80	5.92	-
Perceiver IO	AutoFlow	1.81	2.42	4.98	-
RAFT	AutoFlow	1.74	2.41	4.18	13.41%
RAFT (ours)	AF→AF+FT+KU+TA	1.27	2.28	2.71	9.16%
DDVM (ours)	AF→AF+FT+KU+TA	1.24	2.00	2.19	7.58%
Table 2: Optical flow finetuning evaluation on public benchmark datasets (AEPE↓ for Sintel and Fl-all↓ for KITTI). Bold indicates the best and underline the  $2^{\text{nd}}$ -best.  $\S$  uses extra datasets (AutoFlow and VIPER) on top of defaults (FlyingThings, HD1K, KITTI, and Sintel). *uses warm start on Sintel.
Method	Sintel.clean	Sintel(final	KITTI
SKFlow [72]*	1.30	2.26	4.84%
CRAFT [66]*	1.44	2.42	4.79%
FlowFormer [23]	1.14	2.18	4.68%
RAFT-OCTC [29]*	1.51	2.57	4.33%
RAFT-it [70]$	1.55	2.90	4.31%
DDVM (ours)$	1.75	2.48	3.26%
distance between ground truth and prediction. On KITTI, we additionally use the outlier rate, Fl-all, which reports the outlier ratio in  $\%$  among all pixels with valid ground truth, where an estimate is considered as an outlier if its error exceeds 3 pixels and  $5\%$  w.r.t. the ground truth.
Depth. We separately pre-train indoor and outdoor models on the respective pre-training datasets described in Section 3.1. The indoor depth model is then finetuned and evaluated on the NYU depth v2 dataset [61] and the outdoor model on the KITTI depth dataset [17]. We follow the standard evaluation protocol used in prior work [35]. For both NYU depth v2 and KITTI, we report the absolute relative error (REL), root mean squared error (RMS) and accuracy metrics  $(\delta_{1} < 1.25)$ .
4.1 Evaluation on benchmark datasets