SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tobit Klug, Dogukan Atik, Reinhard Heckel
School of Computation, Information and Technology
tobit.klug, dogukan.atik, reinhard.heckel {@tum.de}
Supervised training of deep neural networks on pairs of clean image and noisy measurement achieves state-of-the-art performance for many image reconstruction tasks, but such training pairs are difficult to collect. Self-supervised methods enable training based on noisy measurements only, without clean images. In this work, we investigate the cost of self-supervised training in terms of sample complexity for a class of self-supervised methods that enable the computation of unbiased estimates of gradients of the supervised loss, including noise2noise methods. We analytically show that a model trained with such self-supervised training is as good as the same model trained in a supervised fashion, but self-supervised training requires more examples than supervised training. We then study self-supervised denoising and accelerated MRI empirically and characterize the cost of self-supervised training in terms of the number of additional samples required, and find that the performance gap between self-supervised and supervised training vanishes as a function of the training examples, at a problem-dependent rate, as predicted by our theory.
Deep neural networks trained in a supervised fashion to map a noisy measurement to a clean image achieve state-of-the-art performance for image reconstruction tasks including image denoising [48, 12], image super-resolution [8, 26] and accelerated magnetic resonance imaging (MRI) [13, 11].
However, collecting clean training images is sometimes not possible, and is often expensive and time-consuming. For example, to collect clean target images for denoising is difficult since a sensor in a camera only collects noisy images [33]. To collect target images for MRI is difficult as it requires acquiring fully sampled data with long scan times which, introduces difficulties like increased motion artifacts.
This has motivated research on self-supervised methods that enable the training of neural networks from noisy measurements of images only. Self-supervised approaches are generally based on constructing a self-supervised loss. For example, noise2noise [24] constructs a self-supervised loss based on a noisy image and a second noisy observation of the same image. Noisier2noise [30], noise2self [2], and neighbor2neighbor [18] construct losses based on a single noisy image for denoising. One might expect that a model trained with a self-supervised loss computed from noisy images performs worse than the same model trained in a supervised manner. And indeed, for image reconstruction problems a model trained with many of the self-supervised losses (in particular noisier2noise, noise2self, and neighbor2neighbor) performs worse than a model trained with a supervised loss, even when abundant training data is available.
However, with infinite training data, noise2noise [24] self-supervised training promises to achieve the performance of supervised training, since the noise2noise loss enables the computation of unbiased estimates of the gradients of the supervised loss.
In this paper, we study a class of self-supervised methods including noise2noise based on constructing unbiased estimates of the gradients of the supervised loss. We characterize the cost of self-supervised training in terms of the number of training examples required to achieve the same performance as supervised training. Our contributions are:
- Finite sample theory. We start by viewing the noise2noise loss as enabling computation of unbiased estimates of the gradients of the supervised loss. Based on this view, we characterize the sample complexity of noise2noise like self-supervised training. We find that the risk of a method trained in a noise2noise-like fashion as a function of training examples approaches the optimal risk at the same rate as supervised training (i.e., at rate  $1/N$ , where  $N$  is the number of training examples). However, to reach the same performance as with supervised training more training examples are required. The amount of extra training examples required is a function of the self-supervised loss and of the reconstruction problem.
- Empirical sample complexity of self-supervised denoising. We empirically characterize the performance of noise2noise self-supervised training for Gaussian and camera-noise image denoising as a function of the training examples, and find that, as predicted by theory, once the training set size becomes sufficiently large noise2noise training yields a denoiser essentially on par with supervised training. For Gaussian denoising, we also characterize the performance of noisier2noise and neighbor2neighbor like training as a function of the number of training examples, and find, again as expected, models trained with a noisier2noise and neighbor2neighbor self-supervised loss perform worse than models trained in a supervised manner, even when abundant training data is available.
- Empirical sample complexity of self-supervised compressive sensing. Finally, we characterize the performance of noise2noise-like self-supervised training similar to the approaches from Yaman et al. [45] and Millard and Chiew [29] for compressive sensing reconstruction as a function of the training examples. Again the performance gap between models trained with a self-supervised and supervised loss goes to zero as a function of the training examples at a problem dependent rate. Perhaps surprisingly, the performance gap is vanishing already for small training set sizes because the gradients constructed from the self-supervised loss have a similar variance than the gradients constructed from the supervised loss.
Together, our results show that networks trained with a self-supervised loss based on computing unbiased estimates of the gradients of the supervised loss perform as well as model trained in a supervised fashion at the cost of additional training examples required.
Related work. We consider self-supervised learning methods since they enable training neural networks for problems where clean examples are not available or are scarce. We discuss the most related works throughout the paper. Another training technique that is useful when clean data is scarce are data augmentation techniques [7, 10]. Another class of methods that is interesting for the regime, where no clean data is available are zero-shot methods that are not based on any data, such as deep-image-prior based methods [40, 15, 14, 6]. There is also a variety of zero shot methods that rely partly on self-supervised loss functions, such as Yaman et al. [44]'s method for compressive sensing, noise2fast [25], and zero-shot noise2noise [28].
2 Background on self-supervised learning based on estimates of gradients of the supervised loss
Let  $f_{\theta} \colon \mathbb{R}^m \to \mathbb{R}^n$  be a neural network with parameters  $\theta$  for estimating an image  $\mathbf{x} \in \mathbb{R}^n$  based on a measurement  $\mathbf{y} \in \mathbb{R}^m$ . Our goal is to find a neural network  $f_{\theta}$  that has a small risk
$$
R (\boldsymbol {\theta}) = \mathbb {E} _ {(\mathbf {x}, \mathbf {y})} [ \ell (f _ {\boldsymbol {\theta}} (\mathbf {y}), \mathbf {x}) ]. \tag {1}
$$
Here, expectation is over the signal-measurement pairs  $(\mathbf{x},\mathbf{y})$  (for example a clean image  $\mathbf{x}$  and a noisy measurement  $\mathbf{y} = \mathbf{x} + \mathbf{z}$ ), and  $\ell$  is a supervised loss, which we take as the mean-squared error.
Since the data distribution is unknown, we can't compute and minimize the risk. In supervised training, we approximate the risk with an empirical risk computed based on  $N$  pairs of training examples. This requires pairs of ground-truth signal  $\mathbf{x}$  and associated measurements  $\mathbf{y}$ .
We consider self-supervised training with a self-supervised loss that, unlike the supervised loss, does not require ground-truth images. Instead, it depends on another, randomized measurement of the ground-truth image, denoted by  $\mathbf{y}'$ . We are given  $N$  pairs of a randomized measurement and original measurement  $(\mathbf{y}_1', \mathbf{y}_1), \ldots, (\mathbf{y}_N', \mathbf{y}_N)$ , and train a network to reconstruct a clean image from the original measurement by minimizing
$$
\mathcal {L} _ {\mathrm {S S}} (\boldsymbol {\theta}) = \frac {1}{N} \sum_ {i = 1} ^ {N} \ell_ {\mathrm {S S}} \left(f _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {i}\right), \mathbf {y} _ {i} ^ {\prime}\right). \tag {2}
$$
We consider loss functions  $\ell_{\mathrm{SS}}$ , like the noise2noise loss discussed below, with the property that in expectation over the training data, a gradient of the self-supervised loss is also a gradient of the risk  $R(\pmb{\theta})$ , and thus with sufficiently many training examples, we expect a network trained with such a self-supervised loss to achieve the same performance as when trained in a supervised manner.
However, the 'noise' induced by relying on randomized measurements of the original image instead of relying on the original image increases the variance of the gradients computed from the self-supervised loss. Therefore, self-supervised training requires more training examples to yield a network on par with supervised training. Next, we discuss two self-supervised losses, one for denoising and one for compressive sensing in the context of MRI.
2.1 Noise2noise loss for denoising
For denoising, our goal is to train a neural network  $f_{\theta}$  to estimate an image  $\mathbf{x}$  from a noisy observation  $\mathbf{y} = \mathbf{x} + \mathbf{z}$ , where  $\mathbf{z}$  is additive Gaussian, non-Gaussian, or even structured noise. Noise2noise assumes access to pairs of noisy observations  $(\mathbf{y}_1, \mathbf{y}_1'), \ldots, (\mathbf{y}_N, \mathbf{y}_N')$ , where  $\mathbf{y}_i = \mathbf{x}_i + \mathbf{z}_i$  and  $\mathbf{y}_i' = \mathbf{x}_i + \mathbf{e}_i$ . Here,  $\mathbf{e}_i$  is zero-mean noise, independent of but not necessarily of the same distribution as the noise  $\mathbf{z}_i$ . Lehtinen et al. [24] introduced the (noise2noise) self-supervised loss
$$
\ell_ {\mathrm {S S}} \left(f _ {\boldsymbol {\theta}} (\mathbf {y}), \mathbf {y} ^ {\prime}\right) = \left\| f _ {\boldsymbol {\theta}} (\mathbf {y}) - \mathbf {y} ^ {\prime} \right\| _ {2} ^ {2}, \tag {3}
$$
which aims at finding a network that predicts one noisy observation of an image based on another one, therefore the name noise2noise.
In expectation, a minimizer of the self-supervised loss is also a minimizer of the associated risk (1), as formalized by the proposition below. Thus, training with the self-supervised loss (3) with infinitely many training examples is as good as supervised training with infinitely many training examples. The proof is in Appendix A.
Proposition 1. Suppose that a signal  $\mathbf{x}$  and a corresponding measurement  $\mathbf{y}$  are drawn from a joint distribution, and let  $\mathbf{y}' = \mathbf{x} + \mathbf{e}$  be another randomized measurement of the signal. Assume that the noise  $\mathbf{e}$  is uncorrelated with the residual, i.e.,  $\mathbb{E}_{(\mathbf{x},\mathbf{y},\mathbf{e})}\left[(f_{\boldsymbol{\theta}}(\mathbf{y}) - \mathbf{x})^T\mathbf{e}\right] = 0$ , for all  $\boldsymbol{\theta}$ . Then, the minimizer  $\boldsymbol{\theta}$  of the self-supervised risk  $\mathbb{E}_{(\mathbf{y},\mathbf{y}')}\left[\| f_{\boldsymbol{\theta}}(\mathbf{y}) - \mathbf{y}'\|_2^2\right]$  is also a minimizer of the supervised risk  $\mathbb{E}_{(\mathbf{x},\mathbf{y})}\left[\| f_{\boldsymbol{\theta}}(\mathbf{y}) - \mathbf{x}\|_2^2\right]$ .
Noise2noise self-supervised training has been applied to many domains including biomedical imaging [5, 3, 19], channel estimation [51] or acoustic sensing [23].
2.2 Compressive sensing
For accelerated MRI, our goal is to train a network  $f_{\theta}$  to reconstruct an image  $\mathbf{x}$  from an undersampled measurement  $\mathbf{y} = \mathbf{M}\mathbf{F}\mathbf{x}$  in the frequency domain. Here,  $\mathbf{M} \in \mathbb{R}^{n \times n}$  is an undersampling mask (i.e., a diagonal matrix with zeros and ones on its diagonal), and  $\mathbf{F} \in \mathbb{C}^{n \times n}$  is the Fourier transform.
We do not have access to a ground-truth image  $\mathbf{x}$ . Instead, we are given an undersampled measurement  $\mathbf{y} = \mathbf{M}\mathbf{F}\mathbf{x}$  as well as a second randomized measurement of the same image, obtained as  $\mathbf{y}' = \mathbf{M}'\mathbf{F}\mathbf{x}$ , where  $\mathbf{M}' \in \mathbb{R}^{n \times n}$  is a randomized mask which is zero or one on its diagonal, and is one with
non-zero probability, so that  $\mathbf{W} = \mathbb{E}\left[\mathbf{M}'\right]^{-1/2}$  exists. From this measurement, we can compute a self-supervised loss defined as
$$
\ell_ {\mathrm {S S}} \left(f _ {\boldsymbol {\theta}} (\mathbf {y}), \mathbf {y} ^ {\prime}\right) = \left\| \mathbf {W} \left(\mathbf {M} ^ {\prime} \mathbf {F} f _ {\boldsymbol {\theta}} (\mathbf {y}) - \mathbf {y} ^ {\prime}\right) \right\| _ {2} ^ {2}, \tag {4}
$$
where  $\mathbf{W} = \mathbb{E}\left[\mathbf{M}'\right]^{-1/2}$  is a diagonal weighting mask. The loss is evaluated only on the frequencies that are given through the randomized mask  $\mathbf{M}'$  and are weighted by how often a frequency occurs in the randomized mask  $\mathbf{M}'$  (thus, the multiplication with the weighting mask  $\mathbf{W}$ ).
Analogous as for the noise2noise loss for denoising, in expectation, a minimizer of the self-supervised loss is also a minimizer of the associated risk, as formalized by the following proposition. The proof is in Appendix A.
Proposition 2. Suppose that a signal  $\mathbf{x}$  and a corresponding measurement  $\mathbf{y}$  are drawn from some joint distribution (e.g., as  $\mathbf{y} = \mathbf{M}\mathbf{F}\mathbf{x}$ , where  $\mathbf{M}$  is a mask), and let  $\mathbf{y}' = \mathbf{M}'\mathbf{F}\mathbf{x}$  be an independent measurement taken with a randomized mask  $\mathbf{M}'$  with 0's or 1's on its diagonal and non-zero probability of a one on the diagonal. Then a minimizer  $\theta$  of the self-supervised risk  $\mathbb{E}_{(\mathbf{y},\mathbf{y}')}\left[\| \mathbf{W}(\mathbf{M}'\mathbf{F}f_{\theta}(\mathbf{y}) - \mathbf{y}')\| _2^2\right]$  with  $\mathbf{W} = \mathbb{E}\left[\mathbf{M}'\right]^{-1 / 2}$  is also a minimizer of the supervised risk  $\mathbb{E}_{(\mathbf{x},\mathbf{y})}\left[\| f_{\theta}(\mathbf{y}) - \mathbf{x}\| _2^2\right]$ .
A variety of works consider training a neural network with a self-supervised loss similar to (4) [24, 45, 42, 17, 16, 52, 29]. However, none of those works studied the sample complexity. Further, Lehtinen et al. [24] train with a noise2noise loss, but sample new masks  $\mathbf{M}$  and  $\mathbf{M}'$  in every training epoch, which requires access to fully sampled measurements. Similar to our setup, Yaman et al. [45] and Millard and Chiew [29] fix the set of undersampled measurements per image available for creating the masks  $\mathbf{M}$  and  $\mathbf{M}'$ . However, Yaman et al. [45] and Millard and Chiew [29] construct a noisere2noise-like self-supervised loss, i.e., the input mask  $\mathbf{M}$  implements a higher undersampling factor on the training inputs than the undersampling factor used at inference. We assume the mask  $\mathbf{M}$  to have the same distribution at training and inference, which is important for Proposition 2 to hold.
3 Finite sample theory for self-supervised denoising
We start by studying the performance of a denoiser trained with a self-supervised noise2noise loss analytically in the finite sample regime. We measure performance in terms of the risk  $R(\theta) = \mathbb{E}\left[\| f_{\theta}(\mathbf{y}) - \mathbf{x}\| _2^2\right]$  and learn an estimator  $f_{\theta}$  by minimizing the self-supervised objective function (2), where we take the noise2noise loss (3) as the self-supervised loss.
Finite-sample risk bound for linear denoising. We consider a simple linear denoising problem, where the joint distribution of the image and corresponding measurement is as follows. The signal  $\mathbf{x}$  is drawn from a  $d$ -dimensional linear subspace according to  $\mathbf{x} = \mathbf{U}\mathbf{c}$ , where  $\mathbf{U} \in \mathbb{R}^{n \times d}$  is an orthonormal basis for the subspace, and where  $\mathbf{c} \sim \mathcal{N}(0,1 / d\mathbf{I})$ . Thus, for large  $d$ , the vector is drawn approximately uniformly from the intersection of the subspace with the unit sphere. We draw a measurement as  $\mathbf{y} = \mathbf{x} + \mathbf{z}$ , where  $\mathbf{z} \sim \mathcal{N}(0,\sigma_z^2 / n\mathbf{I})$  is Gaussian noise, and then draw a second measurement as  $\mathbf{y}' = \mathbf{x} + \mathbf{e}$ , where  $\mathbf{e} \sim \mathcal{N}(0,\sigma_e^2 / n\mathbf{I})$  is the target noise. With this scaling, the expected SNR of the denoising problem is  $1 / \sigma_z^2$ .
We consider a linear estimator of the form  $f_{\mathbf{W}}(\mathbf{y}) = \mathbf{W}\mathbf{y}$ . Note that the optimal linear estimator (i.e., the estimator that minimizes the risk) is  $\mathbf{W}^{*} = \frac{1}{1 + \sigma_{z}^{2}\frac{d}{n}}\mathbf{U}\mathbf{U}^{T}$ . The optimal estimator projects the measurement onto the subspace and shrinks the projected measurement by a noise-dependent factor.
We provide a bound on the expected risk of the estimator that is learned by minimizing the self-supervised loss function (2) by running the stochastic gradient method for one epoch. Given a dataset  $\mathcal{D} = \{(\mathbf{y}_1,\mathbf{y}_1'),\dots ,(\mathbf{y}_N,\mathbf{y}_N')\}$  consisting of pairs of two noisy measurements  $\mathbf{y}_i = \mathbf{x}_i + \mathbf{z}_i,\mathbf{y}_i' = \mathbf{x}_i + \mathbf{e}_i$  drawn i.i.d. from the distribution specified above, we start the stochastic gradient method at  $\mathbf{W}_0 = 0$  and update
$$
\mathbf {W} _ {k + 1} = \mathbf {W} _ {k} - \alpha_ {k} \nabla_ {\mathbf {W}} \| \mathbf {W} \mathbf {y} _ {k} - \mathbf {y} _ {k} ^ {\prime} \| _ {2} ^ {2},
$$
for  $k = 1,\ldots ,N$ . Here,  $\alpha_{k}$  is the stepsize.
Figure 1: Subspace denoising. Simulated risk of a linear estimator learned with early stopped gradient descent for different levels of target noise  $\sigma_{e}$ , signal dimension  $d = 10$ , ambient dimension  $n = 100$ , and input noise level  $\sigma_{z} = 0.1$ . Left: The larger the target noise level, the more training examples are required for self-supervised training to approach the optimal risk  $R(\mathbf{W}^{*})$ . Right: Plotting the risk minus the optimal risk reveals the convergence rate to be  $1 / N$ . Error bars show the standard deviation over 5 independent runs.
Theorem 1. Consider the estimate  $\mathbf{W}_N$  obtained by running the SGM for  $N$  iterations on the training set  $\mathcal{D}$  with a decaying stepsize  $\alpha_{k} = \frac{1}{c + k}$ , where  $c$  is a constant. Then, the expected generalization error, where expectation is over the random training set  $\mathcal{D}$ , obeys
$$
\mathbb {E} \left[ R \left(\mathbf {W} _ {N}\right)\right) ] \leq R \left(\mathbf {W} ^ {*}\right) + \frac {1 / d + \sigma_ {z} ^ {2} / n}{\left(\sigma_ {z} ^ {2} / n\right) ^ {2}} \frac {1}{N - 2} \left(2 + \left(1 2 \sigma_ {z} ^ {2} \frac {d}{n} + \sigma_ {e} ^ {2} \left(1 + \sigma_ {z} ^ {2}\right)\right) ^ {2}\right). \tag {5}
$$
The proof (detailed in Appendix B), is based on an standard convergence analysis of the stochastic gradient method [32, 43, 35]. In a nutshell, from the noise2noise self-supervised loss we compute stochastic gradients of the risk as  $\nabla \ell_{\mathrm{SS}}(f_{\pmb{\theta}}(\mathbf{y}_i),\mathbf{y}_i')$ , where  $(\mathbf{y}_i,\mathbf{y}_i')$  is a training example. This stochastic gradient is an unbiased estimate of the gradient of the risk,  $\nabla R(\pmb {\theta})$ . The variance of the stochastic gradient determines the rate on the RHS of equation (5).
Theorem 1 establishes that the risk is upper bounded by a problem-dependent noise floor, which is the risk of the optimal estimator,  $R(\mathbf{W}^*)$ , plus a term associated with minimizing a self-supervised loss constructed from finitely many training examples. This term decays as  $c / N$ , which is exactly the rate we see in the simulations below for this setup. Moreover, the term associated with minimizing the self-supervised loss becomes larger in the noise variance  $\sigma_e^2$ , and thus reflects that the self-supervised loss is a worse approximation of the supervised loss, as  $\sigma_e^2$  increases.
Numerical results for linear denoising. Figure 1 shows the risk of the linear denoiser  $\mathbf{W}$  trained on  $N$  examples from the linear subspace model. The linear estimator is learned by applying gradient descent to the self-supervised loss function (2) regularized with early stopping. As predicted by Theorem 1, the performance of the linear estimator  $\mathbf{W}$  converges to the performance of the optimal estimator  $\mathbf{W}^*$  at the rate  $1 / N$  (right plot in Figure 1). Also as predicted by the Theorem, larger levels of target noise  $\sigma_{e}$  require more training examples, reflected by the term that multiplies with  $1 / N$  increasing in  $\sigma_{e}^{2}$ .
4 Empirical results for self-supervised denoising
In this section, we study the performance of a network trained for denoising with the self-supervised objective function (2) and noise2noise loss (3) as a function of the number of training examples, and compare to the performance of the network trained in a supervised fashion. We consider Gaussian denoising and denosing real-world-camera noise.
Throughout this section, we focus on training a U-net [36] with 7.4M network parameters. A U-net is a good choice, since it is widely used and gives near SOTA models for image denoising [4, 12, 50, 49]. However, our qualitative results are independent of the network architecture. In the appendix, we present results for SwinIR [26], an attention-based SOTA architecture, confirming this.
Figure 2: Gaussian image denoising. Left: Noise2noise training approaches the performance of supervised training as the number of training images increases at a rate dependent on the target noise level  $\sigma_{e}$ . Networks trained with noisier2noise and neighbor2neighbor also improve as a function of the training examples, but are far from approaching the performance of networks trained in a supervised fashion. Right: Histogram of the variance of the stochastic gradients. The noisier the gradients, the slower the convergence as a function of the training examples,  $N$ .
4.1 Gaussian denoising
Setup. We consider denoising images  $\mathbf{x}$  from noisy measurements  $\mathbf{y} = \mathbf{x} + \mathbf{z}$  with noise level  $\sigma_z = 25$  (for pixel values in [0, 255]). For the noisy targets  $\mathbf{y}' = \mathbf{x} + \mathbf{e}$ , we study noise levels  $\sigma_e \in \{0, 25, 50\}$ , where  $\sigma_e = 0$  corresponds to supervised training. The target noise level determines how fast as a function of the training set size self-supervised approaches the performance of supervised training. In all our experiments, the measurement noise  $\mathbf{z}_i$  and target noise  $\mathbf{e}_i$  are only sampled once per clean image  $\mathbf{x}_i$  during training, which corresponds to the practical setup in which we are given two noisy measurements per image. We use cropped patches of size  $128 \times 128$  from ImageNet [37] to create training sets with 100 to 300k images. We pick the best run out of different runs with independently sampled network initialization. See Appendix C.1.2 for training details.
Results and discussion. Figure 2 shows the performance in PSNR of U-nets trained in a supervised and self-supervised manner for different training set sizes  $N$ . We find, as expected from the theory in Section 3, that the gap between noise2noise like self-supervised training and supervised training vanishes as a function of the training examples, at a rate that increases as  $\sigma_{e}$  decreases.