Accepted to ICLR 24 :)

Author: HilaManor

README.md

@@ -1,11 +1,12 @@
-[![Python 3.8.0](https://img.shields.io/badge/python-3.8.10+-blue?logo=python&logoColor=white)](https://www.python.org/downloads/release/python-3810/)
+[![Python 3.8.10](https://img.shields.io/badge/python-3.8.10+-blue?logo=python&logoColor=white)](https://www.python.org/downloads/release/python-3810/)
 [![NumPy](https://img.shields.io/badge/numpy-1.24.3+-green?logo=numpy&logoColor=white)](https://pypi.org/project/numpy/1.24.3/)
 [![Matplotlib](https://img.shields.io/badge/matplotlib-3.7.1+-green?logo=plotly&logoColor=white)](https://pypi.org/project/matplotlib/3.7.1)
 [![Notebook](https://img.shields.io/badge/notebook-6.5.4+-green?logo=jupyter&logoColor=white)](https://pypi.org/project/notebook/6.5.4)
 [![torch](https://img.shields.io/badge/torch-2.0.0+-green?logo=pytorch&logoColor=white)](https://pytorch.org/)
 [![torchvision](https://img.shields.io/badge/torchvision-0.15.1+-green?logo=pytorch&logoColor=white)](https://pytorch.org/)
 
-# On the Posterior Distribution in Denoising: Application to Uncertainty Quantification Official Implementation
+# On the Posterior Distribution in Denoising: Application to Uncertainty Quantification Official Implementation [ICLR 2024]
+
 
 ###  [Project page](https://HilaManor.github.io/GaussianDenoisingPosterior) | [Arxiv](https://arxiv.org/abs/2309.13598) | [Supplementary materials](https://hilamanor.github.io/GaussianDenoisingPosterior/resources/supp.pdf)
 
@@ -92,10 +93,12 @@ https://github.com/HilaManor/GaussianDenoisingPosterior/assets/53814901/7d209056
 If you use this code for your research, please cite our paper:
 
 ```
-@article{manor2023posterior,
+@inproceedings{
+    manor2024posterior,
     title={On the Posterior Distribution in Denoising: Application to Uncertainty Quantification},
-    author={Hila Manor and Tomer Michaeli}, journal={arXiv preprint arXiv:2305.10124},
-    journal={arXiv preprint arXiv:2309.13598},
-    year={2023},
+    author={Hila Manor and Tomer Michaeli},
+    booktitle={The Twelfth International Conference on Learning Representations},
+    year={2024},
+    url={https://openreview.net/forum?id=adSGeugiuj}
 }
 ```

docs/index.html

@@ -53,7 +53,7 @@
 		<!-- top - author names, affiliation and links -->
 		<center>
 			<h1>On the Posterior Distribution in Denoising:<br> Application to Uncertainty Quantification</h1>
-			<h4>Under Review</h4>
+			<h4>ICLR 2024</h4>
 			<div class="container-xl">
 				<div class="row justify-content-center">
 					<div class="col-md-auto">
@@ -72,9 +72,6 @@ <h4>Under Review</h4>
 					<div class="col-md-auto">
 						<span style="font-size:24px"><a href='resources/paper.pdf' target="_blank" rel="noopener noreferrer">[<i class="fas fa-file-pdf"></i> Paper]</a></span>
 					</div>
-					<div class="col-md-auto">
-						<span style="font-size:24px"><a href='resources/supp.pdf' target="_blank" rel="noopener noreferrer">[<i class="fas fa-file-pdf"></i> Supplementary]</a></span>
-					</div>
 					<div class="col-md-auto">
 						<span style="font-size:24px"><a href="https://github.com/HilaManor/GaussianDenoisingPosterior/" target="_blank" rel="noopener noreferrer">[<i class="fab fa-github"></i> Code]</a></span>
 					</div>
@@ -441,20 +438,23 @@ <h1 style="margin-top: 10px;">Abstract</h1>
 					border-radius: 20px;
 					padding: 0;;">
 
-                @article{manor2023posterior,
+                @inproceedings{
+                    manor2024posterior,
                     title={On the Posterior Distribution in Denoising: Application to Uncertainty Quantification},
-                    author={Hila Manor and Tomer Michaeli}, journal={arXiv preprint arXiv:2305.10124},
-                    journal={arXiv preprint arXiv:2309.13598},
-                    year={2023},
+                    author={Hila Manor and Tomer Michaeli},
+                    booktitle={The Twelfth International Conference on Learning Representations},
+                    year={2024},
+                    url={https://openreview.net/forum?id=adSGeugiuj}
                 }
+
 			</pre>
 		</tbody>
 		</table>
 		<!-- end bibtex citation -->
 
 		<div class="container-xl">
 			<p class="text-left">
-				More results and further discussion about our method can be found in the <a href="resources/supp.pdf" target="_blank" rel="noopener noreferrer">supplementary material</a>.<br>
+				More results and further discussion about our method can be found in the supplementary material (included in the <a href="resources/paper.pdf" target="_blank" rel="noopener noreferrer">paper</a>).<br>
 				Our official code implementation can be found in the official <a href="https://github.com/HilaManor/GaussianDenoisingPosterior" target="_blank" rel="noopener noreferrer">github repository</a>.
 			</p>
 		</div>
@@ -463,21 +463,13 @@ <h1 style="margin-top: 10px;">Abstract</h1>
 			<table class="table" width="1" style="max-width: 400px;">
 				<tbody>
 					<tr>
-						<td width="50%" style="text-align: center;">
-							<a target="_blank" rel="noopener noreferrer" href="resources/supp.pdf" style="color:black !important;">
-								<img alt="" id="supp" class="round" width="100%" style="max-width:150px" border="1px solid" src="./resources/supp.jpg"/>
-							</a>
-						</td>
 						<td width="50%" style="text-align: center;">
 							<a href="https://github.com/HilaManor/GaussianDenoisingPosterior" target="_blank" rel="noopener noreferrer">
 								<img alt="" id="github_logo" class="round" width="100%" style="max-width:150px" src="./resources/github_logo.png"/>
 							</a>
 						</td>
 					</tr>
 					<tr>
-						<td style="font-size:14pt; text-align: center;">
-							<a href='resources/supp.pdf' target="_blank" rel="noopener noreferrer">[Supplementary]</a>
-						</td>
 						<td style="font-size:14pt; text-align: center;">
 							<a target="_blank" rel="noopener noreferrer" href='https://github.com/HilaManor/GaussianDenoisingPosterior'>[Code]</a>
 						</td>
@@ -488,15 +480,36 @@ <h1 style="margin-top: 10px;">Abstract</h1>
 
 		<hr>
 		<br>
+        <!-- NPPC -->
+        <div class="container-xl">
+			<center><h3>Related Work</h3></center>
+			<p style="font-size: 11pt;">
+                The field of Uncertainty quantification has seen increased interest over the last few years.
+                Apart from related works discussed in the paper, we would like to highlight a concurrent related work by
+                <a target="_blank" rel="noopener noreferrer" href="https://eliasnehme.github.io/">Elias Nehme</a>,
+                <a target="_blank" rel="noopener noreferrer" href="https://www.linkedin.com/in/yairomer/">Omer Yair</a>,
+                and <a href="https://tomer.net.technion.ac.il/" target="_blank" rel="noopener noreferrer">Tomer Michaeli</a>:
+            </p>
+            <center>
+                <p style="font-size:14pt; margin-top: 5px;"><strong>Uncertainty Quantification via Neural Posterior Principal Components</strong><br>
+                    Elias Nehme, Omer Yair, Tomer Michaeli. 
+                    <a href="https://arxiv.org/abs/2309.15533" target="_blank" rel="noopener noreferrer">[<i class="ai ai-arxiv"></i> ArXiv]</a>
+                    <a href="https://eliasnehme.github.io/NPPC/" target="_blank" rel="noopener noreferrer">[Project Webpage]</a>
+                    <br>
+                </p>
+            </center>
+        </div>
+
+		<hr><br>
 
 		<!-- acknowledgement -->
 		<div class="container-xl">
 			<center><h3>Acknowledgements</h3></center>
 			<p style="font-size: 11pt;">
 				This webpage was originally made by <a href="https://www.linkedin.com/in/matan-kleiner/" target="_blank" rel="noopener noreferrer">Matan Kleiner</a> with the 
 				help of <a target="_blank" rel="noopener noreferrer" href="https://www.linkedin.com/in/hilamanor/">Hila Manor</a> 
-				for <a target="_blank" rel="noopener noreferrer" href="https://matankleiner.github.io/sinddm/">SinDDM</a> and can be used as a template.<br>
-				It is inspired by the template that was originally made by <a target="_blank" rel="noopener noreferrer" href="http://web.mit.edu/phillipi/">Phillip Isola</a> and
+                for <a target="_blank" rel="noopener noreferrer" href="https://matankleiner.github.io/sinddm/">SinDDM</a> and can be used as a <a target="_blank" rel="noopener noreferrer" href="https://github.com/matankleiner/academic-webpage-template/tree/main">template</a>.<br>
+                It is inspired by the template that was originally made by <a target="_blank" rel="noopener noreferrer" href="http://web.mit.edu/phillipi/">Phillip Isola</a> and
 				<a target="_blank" rel="noopener noreferrer" href="http://richzhang.github.io/">Richard Zhang</a> for a <a target="_blank" rel="noopener noreferrer" href="http://richzhang.github.io/colorization/">colorful</a> ECCV project;
 				the code for the original template can be found <a target="_blank" rel="noopener noreferrer" href="https://github.com/richzhang/webpage-template">here</a>.<br>
 				A lot of features are taken from <a target="_blank" rel="noopener noreferrer" href="https://getbootstrap.com/">bootstrap</a>. All icons are taken from <a target="_blank" rel="noopener noreferrer" href="https://fontawesome.com/">font awesome</a>.

docs/resources/paper.pdf

@@ -150,7 +150,8 @@
                     sigma = args.noise_std
                 else:
                     sigma = (args.noise_std / 255.)
-                nim = im + sigma * torch.randn_like(im)
+                noisemap = torch.randn_like(im)
+                nim = im + sigma * noisemap
                 # if 'DDPM_FFHQ' not in args.denoiser:
                 if 'MNIST' in args.denoiser:
                     nim = nim.clip(0, 1)
@@ -186,20 +187,23 @@
         os.makedirs(outdir, exist_ok=True)
 
         # pylint: disable=unbalanced-tuple-unpacking
-        eigvecs, eigvals, mmse, sigma, subspace_corr = get_eigvecs(model,
+        eigvecs, eigvals, masked_mmse, sigma, subspace_corr, mmse = get_eigvecs(model,
                                                                    nim,  # .unsqueeze(0),
                                                                    mask,  # .unsqueeze(0),
                                                                    args.n_ev,
                                                                    sigma,
                                                                    device,
                                                                    c=args.const, iters=args.iters,
                                                                    double_precision=args.double_precision)
+        save_im = model.save_im
+        toim = model.toim
+        save_im(toim(mmse.cpu()), os.path.join(outdir, 'mmse.png'))
         save_eigvecs(eigvecs, eigvals, outdir)
 
         max_entropy_params = None
         if args.marginal_dist:
             moments = calc_moments(model, nim, mask, sigma, device,
-                                   mmse, eigvecs, eigvals,
+                                   masked_mmse, eigvecs, eigvals,
                                    var_c=args.var_c, skew_c=args.skew_c, kurt_c=args.kurt_c,
                                    use_poly=args.use_poly, poly_deg=args.poly_deg,
                                    poly_bound=args.poly_bound, poly_pts=args.poly_pts,
@@ -209,7 +213,7 @@
 
         patch = im[mask.to(torch.bool).cpu()].reshape(
             mask.shape[0], int(np.floor(selected_coords['y2'])) - int(np.ceil(selected_coords['y1'])), -1)
-        rpatch = mmse[mask.to(torch.bool)].reshape(
+        rpatch = masked_mmse[mask.to(torch.bool)].reshape(
             mask.shape[0], int(np.floor(selected_coords['y2'])) - int(np.ceil(selected_coords['y1'])), -1)
         npatch = nim[mask.to(torch.bool).cpu()].reshape(
             mask.shape[0], int(np.floor(selected_coords['y2'])) - int(np.ceil(selected_coords['y1'])), -1)
@@ -224,4 +228,5 @@
                      name, outdir, args.denoiser, path_name,
                      amount=args.amount,
                      max_entropy_params=max_entropy_params,
-                     subspace_corr=subspace_corr)
+                     subspace_corr=subspace_corr,
+                     noisemap=noisemap)

docs/resources/supp.jpg

@@ -63,6 +63,10 @@ def _validate_fit_converged(moments: Moments, ev_idx: int, coeffs: List[float],
 
     est_vmu1 = delta * sum(xs * pdf)
     est_vmu2 = delta * sum((xs - est_vmu1)**2 * pdf)
+    est_vmu3 = delta * sum((xs - est_vmu1)**3 * pdf)
+    est_vmu4 = delta * sum((xs - est_vmu1)**4 * pdf)
 
     return (abs(est_vmu1 - moments.vmu1[ev_idx]) / abs(moments.vmu1[ev_idx])) < 0.5 and \
-                    (abs(est_vmu2 - moments.vmu2[ev_idx]) / abs(moments.vmu2[ev_idx])) < 0.5
+                    (abs(est_vmu2 - moments.vmu2[ev_idx]) / abs(moments.vmu2[ev_idx])) < 0.5 and \
+                        (abs(est_vmu3 - moments.vmu3[ev_idx]) / abs(moments.vmu3[ev_idx])) < 0.5 and \
+                            (abs(est_vmu4 - moments.vmu4[ev_idx]) / abs(moments.vmu4[ev_idx])) < 0.5

docs/resources/supp.pdf

@@ -23,8 +23,8 @@ def get_eigvecs(model: torch.nn.Module,
                 iters: int = 5,
                 double_precision: bool = False,
                 compare_backprop: bool = False) -> Union[
-                    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, Optional[List]],
-                    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, Optional[List],
+                    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, Optional[List]], torch.Tensor,
+                    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, Optional[List], torch.Tensor,
                           List[Tuple[float, float]], torch.Tensor]]:
     r"""Calculate the eigvecs of the covariance matrix using power iterations
     $\mu_1(y + \epsilon) - \mu_1(y) ~ \frac{d\mu_1(y)}{dy}\cdot\epsilon$$
@@ -47,6 +47,7 @@ def get_eigvecs(model: torch.nn.Module,
     nim = nim.to(device)
     with torch.no_grad():
         im_mmse = _forward_directional(model, nim.unsqueeze(0), device, 0, 0, double_precision).squeeze(0)
+
     print(f'MSE (x_hat, y): {torch.nn.functional.mse_loss(im_mmse, nim)}')
     print(f'sqrt (MSE (x_hat, y)): {torch.nn.functional.mse_loss(im_mmse, nim).sqrt()}')
     print(f'assumed noise sigma: {sigma}')
@@ -59,7 +60,7 @@ def get_eigvecs(model: torch.nn.Module,
         sigma = torch.nn.functional.mse_loss(im_mmse, nim).sqrt().item()
     print(f'Using assumed noise sigma: {sigma}')
     b_masked_mmse = (im_mmse * mask).repeat(n_ev, *[1]*len(nim.shape)).to(device)
-    del im_mmse
+    # del im_mmse
     torch.cuda.empty_cache()
     nim = nim.repeat(n_ev, *[1]*len(nim.shape)).to(device)
     eigvecs = torch.randn(*nim.shape, device=device, dtype=nim.dtype) * mask * c
@@ -112,9 +113,9 @@ def get_eigvecs(model: torch.nn.Module,
     eigvecs = eigvecs[indices, ...]  # TODO ?????? swap places??
 
     if compare_backprop:
-        return eigvecs, eigvals, b_masked_mmse[0], sigma, corr, cosine_similarity, timing, eigvecs_backprop
+        return eigvecs, eigvals, b_masked_mmse[0], sigma, corr, im_mmse, cosine_similarity, timing, eigvecs_backprop
 
-    return eigvecs, eigvals, b_masked_mmse[0], sigma, corr
+    return eigvecs, eigvals, b_masked_mmse[0], sigma, corr, im_mmse
 
 
 def calc_moments(model: torch.nn.Module,

main.py

@@ -50,7 +50,8 @@ def plot_eigvecs(model: ModelWrapper,
                  amount: int = 1,
                  subspace_corr: Optional[List] = None,
                  max_axis: float = 3,
-                 delta: float = 0.01):
+                 delta: float = 0.01,
+                 noisemap: Optional[torch.Tensor] = None):
 
     n_ev = eigvecs.shape[0]
 
@@ -151,6 +152,8 @@ def plot_eigvecs(model: ModelWrapper,
     save_im(toim(im), os.path.join(outdir, 'im.png'))
     save_im(toim(nim), os.path.join(outdir, 'nim.png'))
     save_im(toim(fullnim), os.path.join(outdir, 'fullnim.png'))
+    if noisemap is not None:
+        torch.save(noisemap.cpu(), os.path.join(outdir, 'noisemap.pth'))
 
     plt.close(fig)