SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ayush Tewari $^{1*}$  Tianwei Yin $^{1*}$  George Cazenavette $^{1}$  Semon Rezchikov $^{4}$  Joshua B. Tenenbaum $^{1,2,3}$  Frédo Durand $^{1}$  William T. Freeman $^{1}$  Vincent Sitzmann $^{1}$
$^{1}$ MIT CSAIL  $^{2}$ MIT BCS  $^{3}$ MIT CBMM  $^{4}$ Princeton IAS
Denoising diffusion models have emerged as a powerful class of generative models capable of capturing the distributions of complex, real-world signals. However, current approaches can only model distributions for which training samples are directly accessible, which is not the case in many real-world tasks. In inverse graphics, for instance, we seek to sample from a distribution over 3D scenes consistent with an image but do not have access to ground-truth 3D scenes, only 2D images. We present a new class of conditional denoising diffusion probabilistic models that learn to sample from distributions of signals that are never observed directly, but instead are only measured through a known differentiable forward model that generates partial observations of the unknown signal. To accomplish this, we directly integrate the forward model into the denoising process. At test time, our approach enables us to sample from the distribution over underlying signals consistent with some partial observation. We demonstrate the efficacy of our approach on three challenging computer vision tasks. For instance, in inverse graphics, we demonstrate that our model in combination with a 3D-structured conditioning method enables us to directly sample from the distribution of 3D scenes consistent with a single 2D input image.
Consider the problem of reconstructing a 3D scene from a single picture. Since much of the 3D scene is unobserved, there are an infinite number of 3D scenes that could have produced the image, due to the 3D-to-2D projection, occlusion, and limited field-of-view that leaves a large part of the 3D scene unobserved. Given the ill-posedness of this problem, it is desirable for a reconstruction algorithm to be able to sample from the distribution over all plausible 3D scenes that are consistent with the 2D image, generating unseen parts in plausible manners. Previous data-completion methods, such as in-painting in 2D images, are trained on large sets of ground-truth output images along with their incomplete (input) counterparts. Such techniques do not easily extend to 3D scene completion, since curating a large dataset of ground-truth 3D scene representations is very challenging.
This 3D scene completion problem, known as inverse graphics, is just one instance of a broad class of problems often referred to as Stochastic Inverse Problems, which arise across scientific disciplines whenever we capture partial observations of the world through a sensor. In this paper, we introduce a diffusion-based framework that can tackle this problem class, enabling us to sample from a distribution of signals that are consistent with a set of partial observations that are generated from the signal by a non-invertible, generally nonlinear, forward model. For instance, in inverse graphics, we learn to sample 3D scenes given an image, yet never observe paired observations of images and 3D scenes at training time, nor observe 3D scenes directly.
* Equal Contribution. Project page: diffusion-with-forward-models.github.io
While progress in deep learning for generative modeling has been impressive, this problem remains unsolved. In particular, variational autoencoders and conditional neural processes are natural approaches but have empirically fallen short of modeling the multi-modal distributions required in, for instance, inverse graphics. They have so far been limited to simple datasets. Emerging diffusion models [1], in contrast, enable sampling from highly complex conditional distributions but require samples from the output distribution that is to be modeled for training, e.g. full 3D models. Some recent work in inverse graphics has resorted to a two-stage approach, where one first reconstructs a large dataset of 3D scenes to then train an image-conditional diffusion model to sample from the conditional distribution over these scenes [2, 3]. To avoid a two-stage approach, another recent line of work trains a conditional diffusion model to sample from the distribution over novel views of a scene, only requiring image observations at training time [4, 5]. However, such methods do not model the distribution over 3D scenes directly and therefore cannot sample from the distribution over 3D scenes consistent with an image observation. Thus, a multi-view consistent 3D scene can only be obtained in a costly post-processing stage [6]. A notable exception is the recently proposed RenderDiffusion [7], demonstrating that it is possible to train an unconditional diffusion model over 3D scenes from observing only monocular images. While one can perform conditional sampling even with unconditional models, they are fundamentally limited to simple distributions, in this case, single objects in canonical orientations.
Our core contribution is a novel approach for integrating any differentiable forward model that describes how partial observations are obtained from signals, such as 2D image observations and 3D scenes, with conditional denoising diffusion models. By sampling an observation from our model, we jointly sample the signal that gave rise to that observation. Our approach has a number of advantages that make it highly attractive for solving complex Stochastic Inverse Problems. First, our model is trained end-to-end and does away with two-stage approaches that first require reconstruction of a large dataset of signals. Second, our model directly yields diverse samples of the signal of interest. For instance, in the inverse graphics setting, our model directly yields highly diverse samples of 3D scenes consistent with an observation that can then be rendered from novel views with guaranteed multi-view consistency. Finally, our model naturally leverages domain knowledge in the form of known forward models, such as differentiable rendering, with all guarantees that such forward models provide. We validate our approach on three challenging computer vision tasks: inverse graphics (the focus of this paper), as well as single-image motion prediction and GAN inversion.
1. We propose a new method that integrates differentiable forward models with conditional diffusion models, replacing prior two-step approaches with a conditional generative model trained end-to-end.
2. We apply our framework to build the first conditional diffusion model that learns to sample from the distribution of 3D scenes trained only on 2D images. In contrast to prior work, we directly learn image-conditional 3D radiance field generation, instead of sampling from the distribution of novel views conditioned on a context view. Our treatment of inverse graphics exceeds a mere application of the proposed framework, contributing a novel, 3D-structured denoising step that leverages differentiable rendering both for conditioning and for the differentiable forward model.
3. We formally prove that under natural assumptions, as the number of observations of each signal in the training set goes to infinity, the proposed model maximizes not only the likelihood of observations, but also the likelihood of the unobserved signals.
4. We demonstrate the efficacy of our model for two more downstream tasks with structured forward models: single-image motion prediction, where the forward model is a warping operation, and GAN inversion, where the forward model is a pretrained StyleGAN [8] generator.
Consider observations  $(\mathbf{O}_j^i,\phi_j^i)$  that are generated from underlying signals  $\mathbf{S}_j$  according to a known forward model forward(), i.e.,  $\mathbf{O}_j^i = \mathrm{forward}(\mathbf{S}_j,\phi_j^i)$ , where  $\phi_j^i$  are parameters of the forward model corresponding to observation  $\mathbf{O}_j^i$ . Each observation can be partial. Specifically, given a single observation, there is an infinite number of signals that could have generated this observation. However, we assume that given a hypothetical set of all possible observations, the signal is fully
Figure 1: Overview of our proposed method. (a) We assume a dataset of tuples of observations  $(\mathbf{O},\phi)^i$  generated from unobserved signals  $\mathbf{S}$  via a differentiable forward model. (b) We propose to integrate the forward model directly into the denoising step of a diffusion model: given a pair of observations of the same signal, we designate context  $\mathbf{O}^{\mathrm{ctx}}$  and target  $\mathbf{O}^{\mathrm{trgt}}$ . We add noise to  $\mathbf{O}^{\mathrm{trgt}}$ , then feed  $(\mathbf{O}^{\mathrm{ctx}},\phi^{\mathrm{ctx}},\mathbf{O}_t^{\mathrm{trgt}},\phi^{\mathrm{trgt}})$  to a neural network denoise to estimate the signal  $\mathbf{S}_{t - 1}$ . We then apply the forward model to obtain an estimate of the clean target observation,  $\hat{\mathbf{O}}_{t - 1}^{\mathrm{trgt}}$ . (c) The graphical model of the diffusion process.
determined. In the case of inverse graphics,  $\mathbf{O}_j^i$  are image observations of 3D scenes  $\mathbf{S}_j$  and  $\phi_j^i$  are the camera parameters, where we index scenes with  $j$  and observations of the  $j$ -th scene via  $i$ . forward() is the rendering function. Note that if we were to capture every possible image of a 3D scene, the 3D scene is uniquely determined, but given a single image, there are an infinite number of 3D scenes that could have generated that image, both due to the fact that rendering is a projection from 3D and 2D, and due to the fact that a single image only constrains the visible part of the 3D scene. We will drop the subscript  $j$  in the following, and leave it implied that we always consider many observations generated from many signals. Fig. 1 provides an illustration of the data.
We are now interested in training a model that, at test time, allows us to sample from the distribution of signals that are consistent with a previously unseen observation  $\mathbf{O}$ . Formally, we aim to model the conditional distribution  $p(\mathbf{S}|\mathbf{O},\phi)$ . We make the following assumptions:
- We have access to a differentiable implementation of forward().
- We have access to a large dataset of observations and corresponding parameters of the forward model,  $\{(\mathbf{O}^i,\phi^i)\}_{i}^{N}$ .
- In our training set, we have access to several observations per signal.
Crucially, we do not assume that we have direct access to the underlying signal that gave rise to a particular observation, i.e., we do not assume access to tuples of  $(\mathbf{O},\phi ,\mathbf{S})$ . Further, we also do not assume that we have access to any prior distribution over the signal of interest, i.e., we never observe a dataset of signals of the form  $\{\mathbf{S}^j\}_{j}$ , and thus cannot train a generative model to sample from an unconditional distribution over signals.
Recent advances in deep-learning-based generative modeling have seen the emergence of denoising diffusion models as powerful generative models that can be trained to generate highly diverse samples from complex, multi-modal distributions. We are thus motivated to leverage denoising diffusion probabilistic models to model  $p(\mathbf{S}|\mathbf{O},\phi)$ . However, existing approaches cannot be trained if we do not have access to signals S. In the following, we give background on denoising diffusion models and discuss the limitation.
2.1 Background: Denoising Diffusion Probabilistic Models and their Limitation
Denoising diffusion probabilistic models are a class of generative models that learn to sample from a distribution by learning to iteratively denoise samples. Consider the problem of modeling the distribution  $p_{\theta}(\mathbf{x})$  over samples  $\mathbf{x}$ . A forward Markovian process  $q(\mathbf{x}_{0:T})$  adds noise to the data as
$$
q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) = \mathcal {N} \left(\mathbf {x} _ {t}; \sqrt {1 - \beta_ {t}} \mathbf {x} _ {t - 1}, \beta_ {t} \mathbf {I}\right). \tag {1}
$$
Here,  $\beta_{t}, t \in 1 \dots T$  are the hyperparameters that control the variance schedule. A denoising diffusion model learns the reverse process, where samples from a distribution  $p(x_{T}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$  are transformed incrementally into the data manifold as  $p_{\theta}(\mathbf{x}_{0:T}) = p(x_{T}) \prod_{t=1}^{T} p_{\theta}(\mathbf{x}_{t-1} \mid \mathbf{x}_{t})$ , where
$$
p _ {\theta} \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}\right) = \mathcal {N} \left(\mathbf {x} _ {t - 1}; \mu \left(\mathbf {x} _ {t}, t\right), \Sigma \left(\mathbf {x} _ {t}, t\right)\right). \tag {2}
$$
A neural network denoise  $\theta (\cdot)$  with learnable parameters  $\theta$  learns to reverse the diffusion process. It is also possible to model conditional distributions  $p_{\theta}(\mathbf{x}_{0:T}\mid \mathbf{c})$ , where the output is computed as
denoise $_{\theta}(\mathbf{x}_t, t, \mathbf{c})$ . The forward process does not change in this case; in practice, we merely add the conditional signal as input to the denoising model.
Unfortunately, we cannot train existing denoising diffusion models to sample from  $p(\mathbf{S} \mid \mathbf{O}, \phi)$ , or, in fact, even from an unconditional distribution  $p(\mathbf{S})$ . This would require computation of the Markovian forward process in Eq. 1. However, recall that we do not have access to any signals  $\{\mathbf{S}^j\}_j$  - we thus can not add any noise to any signals to then train a denoising neural network. In other words, since no  $\mathbf{S}$  is directly observed, we cannot compute  $q(\mathbf{S}_t \mid \mathbf{S}_{t-1})$ .
2.2 Integrating Denoising Diffusion with Differentiable Forward Models
We now introduce a class of denoising diffusion models that we train to directly model the distribution  $p(\mathbf{S} \mid \mathbf{O}^{\mathrm{ctx}}; \phi^{\mathrm{ctx}})$  over signals  $\mathbf{S}$  given an observation  $(\mathbf{O}^{\mathrm{ctx}}, \phi^{\mathrm{ctx}})$ . Our key contribution is to directly integrate the differentiable forward model forward() into the iterative conditional denoising process. This enables us to add noise to and denoise the observations, while nevertheless sampling the underlying signal that generates that observation.
Our model is trained on pairs of "context" and "target" observations of the same signal, denoted as  $\mathbf{O}^{\mathrm{ctx}}$  and  $\mathbf{O}^{\mathrm{trgt}}$ . As in conventional diffusion models, for the forward process, we have  $q(\mathbf{O}_t^{\mathrm{trgt}}|\mathbf{O}_{t - 1}^{\mathrm{trgt}}) = \mathcal{N}(\mathbf{O}_t^{\mathrm{trgt}};\sqrt{1 - \beta_t}\mathbf{O}_{t - 1}^{\mathrm{trgt}},\beta_t\mathbf{I})$ . In the reverse process, we similarly denoise  $\mathbf{O}^{\mathrm{trgt}}$  conditional on  $\mathbf{O}^{\mathrm{ctx}}$ :
$$
p _ {\theta} \left(\mathbf {O} _ {0: T} ^ {\operatorname {t r g t}} \mid \mathbf {O} ^ {\operatorname {c t x t}}; \phi^ {\operatorname {c t x t}}, \phi^ {\operatorname {t r g t}}\right) = p \left(\mathbf {O} _ {T} ^ {\operatorname {t r g t}}\right) \prod_ {t = 0} ^ {T} p _ {\theta} \left(\mathbf {O} _ {t - 1} ^ {\operatorname {t r g t}} \mid \mathbf {O} _ {t} ^ {\operatorname {t r g t}}, \mathbf {O} ^ {\operatorname {c t x t}}; \phi^ {\operatorname {c t x t}}, \phi^ {\operatorname {t r g t}}\right), \tag {3}
$$
However, unlike conventional diffusion models, we implement  $p_{\theta}(\mathbf{O}_{t - 1}^{\mathrm{trgt}} \mid \mathbf{O}_t^{\mathrm{trgt}}, \mathbf{O}^{\mathrm{ctxt}}, \phi^{\mathrm{ctxt}}, \phi^{\mathrm{trgt}})$  by first predicting an estimate of the underlying signal  $\mathbf{S}_{t - 1}$  and then mapping it to an estimate of the denoised observations via the differentiable forward:
$$
\mathbf {S} _ {t - 1} = \operatorname {d e n o i s e} _ {\theta} \left(\mathbf {O} ^ {\mathrm {c t x t}}, \mathbf {O} _ {t} ^ {\mathrm {t r g t}}; t, \phi^ {\mathrm {c t x t}}, \phi^ {\mathrm {t r g t}}\right), \tag {4}
$$
$$
\hat {\mathbf {O}} _ {t - 1} ^ {\mathrm {t r g t}} = \operatorname {f o r w a r d} \left(\mathbf {S} _ {t - 1}, \phi^ {\mathrm {t r g t}}\right) \tag {5}
$$
$$
\mathbf {O} _ {t - 1} ^ {\operatorname {t r g t}} \sim \mathcal {N} \left(\mathbf {O} _ {t - 1} ^ {\operatorname {t r g t}}; C _ {t - 1} \hat {\mathbf {O}} _ {t - 1} ^ {\operatorname {t r g t}}, \hat {\beta} _ {t - 1} \mathbf {I}\right) \tag {6}
$$
Here,  $\hat{\mathbf{O}}_{t - 1}^{\mathrm{trgt}}$  is an estimate of the clean observation, and the constants  $C_{t - 1}$  and  $\hat{\beta}_{t - 1}$  are chosen to match the total noise added by the forward process at time  $t - 1$ . See Fig. 1 for an overview. At test time, a signal is sampled by iterating Eq. 4, 5, and 6 starting with  $p(\mathbf{O}_{t = T}^{\mathrm{trgt}})\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . Importantly, our models define a generative model over the underlying signal via Eq. 4:
$$
p _ {\theta , \phi^ {\mathrm {t r g t}}} \left(\mathbf {S} _ {0: T} \mid \mathbf {O} ^ {\mathrm {c t x t}}; \phi^ {\mathrm {c t x t}}\right) = \prod_ {t = 1} ^ {T} p _ {\theta} \left(\mathbf {S} _ {t - 1} \mid \mathbf {O} _ {t} ^ {\mathrm {t r g t}}, \mathbf {O} ^ {\mathrm {c t x t}}; \phi^ {\mathrm {c t x t}}, \phi^ {\mathrm {t r g t}}\right). \tag {7}
$$
We will suppress the subscript in the notation, and refer to this distribution as  $p(\mathbf{S}_{0:T} \mid \mathbf{O}^{\mathrm{ctx}}; \phi^{\mathrm{ctx}})$  for brevity from now.
Loss Function. We train to minimize the following two loss terms:
$$
\mathcal {L} _ {\theta} ^ {\operatorname {t r g t}} = \mathbb {E} _ {\mathbf {O} ^ {\mathrm {c t x t}}, \mathbf {O} ^ {\mathrm {t r g t}}, \phi^ {\mathrm {c t x t}}, \phi^ {\mathrm {t r g t}}, t} \left[ \| \mathbf {O} ^ {\mathrm {t r g t}} - \underbrace {\text {f o r w a r d} \left(\text {d e n o i s e} _ {\theta} \left(\mathbf {O} ^ {\mathrm {c t x t}}, \mathbf {O} _ {t} ^ {\mathrm {t r g t}} ; t , \phi^ {\mathrm {c t x t}} , \phi^ {\mathrm {t r g t}}\right) , \phi^ {\mathrm {t r g t}}\right)} \| ^ {2} \right], \tag {8}
$$
$$
\mathcal {L} _ {\theta} ^ {\text {n o v e l}} = \mathbb {E} _ {\mathbf {O} ^ {\mathrm {c t x}}, \mathbf {O} ^ {\text {n o v e l}}, \phi^ {\mathrm {c t x}}, \phi^ {\mathrm {t r g t}}, \phi^ {\text {n o v e l}}, t} \left[ \| \mathbf {O} ^ {\text {n o v e l}} - \underbrace {\text {f o r w a r d} \left(\operatorname {d e n o i s e} _ {\theta} \left(\mathbf {O} ^ {\mathrm {c t x}}, \mathbf {O} _ {t} ^ {\mathrm {t r g t}} ; t , \phi^ {\mathrm {c t x t}}, \phi^ {\mathrm {t r g t}}\right) , \phi^ {\text {n o v e l}}\right)} \| ^ {2} \right]. \tag {9}
$$
Here, we compute the estimate of the observation from the target, as well as a separate, novel forward model parameter  $\phi^{\mathrm{novel}}$ . In the supplemental document, we show that these losses approximate a total observation loss, maximizing the likelihood of all possible observations of the signal S.
Characterizing the Conditional Distribution Over Signals. Due to the complexity of the reverse process, it may not be clear that the learned distribution over signals will agree with the true distribution, even in the limit of infinite data. However, this model will indeed asymptotically learn the true conditional distribution over signals, as we formally prove in the supplement:
Figure 2: Overview of 3D Generative Modeling. We build a 3D-structured denoise operator on top of pixelNeRF [9] that learns to sample from the distribution of 3D scenes from image observations only. Given a context image  $\mathbf{O}^{\mathrm{ctx}}$  with camera pose  $\phi^{\mathrm{ctx}}$ , we pick a target pose  $\phi^{\mathrm{tgt}}$ . We render out a deterministic estimate of the depth, RGB, and features of the target view  $\mathbf{O}_{\mathrm{det}}^{\mathrm{tgt}}$  using pixel-aligned features  $\mathbf{f}^{\mathrm{ctx}}$  extracted from the context view with encoder  $\mathsf{enc}_{t=0}$  (left, only RGB shown here). To generate a 3D scene, we concatenate the deterministic estimate with noise  $\mathbf{O}_{t}^{\mathrm{tgt}}$ , and extract features  $\mathbf{f}_{t}^{\mathrm{tgt}}$  for the target view with  $\mathsf{enc}_{t}$ .  $\mathbf{f}_{t}^{\mathrm{tgt}}$  and  $\mathbf{f}^{\mathrm{ctx}}$  now jointly parameterize the radiance field of the generated scene  $\mathbf{S}_{t-1}$ , and we may render an estimate of the clean target view  $\hat{\mathbf{O}}_{t-1}^{\mathrm{tgt}}$ . The model is trained end-to-end via a re-rendering loss.
Proposition 1. Suppose that any signal  $\mathbf{S}$  can be reconstructed from the set of all all possible observations of  $\mathbf{S}$ . Under this assumption, if in the limit as the number of known observations per signal goes to infinity, there are parameters  $\theta$  such that  $\mathcal{L}_{\theta}^{\mathrm{trgt}} + \mathcal{L}^{\mathrm{novel}}$  is minimized, then the conditional probability distribution over signals discovered by our model  $p(\mathbf{S} \mid \mathbf{O}^{\mathrm{ctxt}}; \phi^{\mathrm{ctxt}})$  agrees with the true distribution  $p^{\mathrm{true}}(\mathbf{S} \mid \mathbf{O}^{\mathrm{ctxt}}; \phi^{\mathrm{ctxt}})$ .
The proof follows by showing that our losses implicitly minimize a diffusion model loss over total observations, which are collections of all possible observations of our signal. As such, when the observations suffice to completely reconstruct the signal, the correctness of the estimated distribution over total observations forces the estimated distribution over signals to be correct, as well.
3 Prior Work on Latent Variable Models for Inverse Problems
Variational Autoencoders [10, 11], normalizing flows [12], conditional [13] and attentive neural processes [14] are latent-variable models that can be combined with forward models to learn to sample from the distribution of unobserved signals from observations [15, 16]. However, they empirically fall short of accurately modeling complex signal distributions - in inverse graphics, for instance, such models have so far been limited to synthetic 3D scenes. Generative Adversarial Networks can be trained with differentiable forward models in-the-loop, and have yielded impressive results in unconditional generative modeling of unobserved signals [17-19]. Similarly, in concurrent work, diffusion models have been leveraged for unconditional generative modeling through differentiable forward models [2, 7, 20]. However, unconditional models are limited to tight distributions, and no conditional generative modeling of similar quality has been demonstrated. Diffusion models trained directly on signals have been effectively applied to diverse inverse problems such as superresolution [21-25], inpainting [21, 23-26], and medical imaging [27]. These works utilize the learned prior of the data distribution to recover the latent signal through a "plug and play" approach [28-30], integrating the diffusion model with a forward measurement process according to Bayes' rule. These approaches are versatile and can easily adapt to new inverse problems without retraining. However, unlike our models, they rely on direct supervision over the signals in the form of large datasets.
We now apply our framework to three stochastic inverse problems. We focus on applications in computer vision, where we tackle the problems of inverse graphics, single-image motion prediction, and GAN inversion. For each application, we give a detailed description of the forward model, the dataset and baselines, as well as a brief description of prior work.
Figure 3: Sample Diversity. We illustrate different 3D scenes sampled from the same context image for RealEstate10k and Co3D datasets. Unlike deterministic methods like pixelNeRF [9], our method generates diverse and distinct 3D scenes that all align with the context image. Co3D results are generated using autoregressive sampling, where a 360 degree trajectory can be generated by iteratively sampling target images. Note the photorealism and diversity of the generated structures for the indoor scene, such as doors and cabinets. Also note the high-fidelity geometry of the occluded parts of the hydrant and the diverse background appearance.
4.1 Inverse Graphics