Original: R EMOVING S TRUCTURED N OISE WITH D IFFUSION M ODELS
Revised: R EMOVING S TRUCTURED N OISE WITH D IFFUSION M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Solving ill-posed inverse problems requires careful formulation of prior beliefs over the signals of interest and an accurate description of their manifestation into noisy measurements.
Revised: Solving ill-posed inverse problems requires careful formulation of prior beliefs over the signals of interest and an accurate description of their manifestation into noisy measurements.

Original: Handcrafted signal priors based on e.g. sparsity are increasingly replaced by data-driven deep generative models, and several groups have recently shown that state-of-the-art score-based diffusion models yield particularly strong performance and flexibility.
Revised: Handcrafted signal priors based on e.g. sparsity are increasingly replaced by data-driven deep generative models, and several groups have recently shown that state-of-the-art score-based diffusion models yield particularly strong performance and flexibility.

Original: In this paper, we show that the powerful paradigm of posterior sampling with diffusion models can be extended to include rich, structured, noise models.
Revised: In this paper, we show that the powerful paradigm of posterior sampling with diffusion models can be extended to include rich, structured, noise models.

Original: To that end, we propose a joint conditional reverse diffusion process with learned scores for the noise and signal-generating distribution.
Revised: To that end, we propose a joint conditional reverse diffusion process with learned scores for the noise and signal-generating distribution.

Original: We demonstrate strong performance gains across various inverse problems with structured noise, outperforming competitive baselines that use normalizing flows and adversarial networks.
Revised: We demonstrate strong performance gains across various inverse problems with structured noise, outperforming competitive baselines that use normalizing flows and adversarial networks.

Original: This opens up new opportunities and relevant practical applications of diffusion modeling for inverse problems in the context of non-Gaussian measurements.
Revised: This opens up new opportunities and relevant practical applications of diffusion modeling for inverse problems in the context of non-Gaussian measurements.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Many signal and image processing problems, such as denoising, compressed sensing, or phase retrieval, can be formulated as inverse problems that aim to recover unknown signals from (noisy) observations.
Revised: Many signal and image processing problems, such as denoising, compressed sensing, or phase retrieval, can be formulated as inverse problems that aim to recover unknown signals from (noisy) observations.

Original: These ill-posed problems are, by definition, subject to many solutions under the given measurement model.
Revised: These ill-posed problems are, by definition, subject to many solutions under the given measurement model.

Original: Therefore, prior knowledge is required for a meaningful and physically plausible recovery of the original signal.
Revised: Therefore, prior knowledge is required for a meaningful and physically plausible recovery of the original signal.

Original: Bayesian inference and maximum a posteriori (MAP) solutions incorporate both signal priors and observation likelihood models.
Revised: Bayesian inference and maximum a posteriori (MAP) solutions incorporate both signal priors and observation likelihood models.

Original: Choosing an appropriate statistical prior is not trivial and is often dependent on both the application as well as the recovery task.
Revised: Choosing an appropriate statistical prior is not trivial and is often dependent on both the application as well as the recovery task.

Original: Before deep learning, sparsity in some transform domain has been the go-to prior in compressed sensing (CS) methods (Eldar & Kutyniok, 2012), such as iterative thresholding (Beck & Teboulle, 2009) or wavelet decomposition (Mallat, 1999).
Revised: Before deep learning, sparsity in some transform domain has been the go-to prior in compressed sensing (CS) methods (Eldar & Kutyniok, 2012), such as iterative thresholding (Beck & Teboulle, 2009) or wavelet decomposition (Mallat, 1999).

Original: At present, deep generative modeling has established itself as a strong mechanism for learning such priors for inverse problem-solving.
Revised: At present, deep generative modeling has established itself as a strong mechanism for learning such priors for inverse problem-solving.

Original: Both generative adversarial networks (GANs) (Bora et al., 2017) and normalizing flows (NFs) (Asim et al., 2020; Wei et al., 2022) have been applied as natural signal priors for inverse problems in image recovery.
Revised: Both generative adversarial networks (GANs) (Bora et al., 2017) and normalizing flows (NFs) (Asim et al., 2020; Wei et al., 2022) have been applied as natural signal priors for inverse problems in image recovery.

Original: These data-driven methods are more powerful compared to classical methods, as they can accurately learn the natural signal manifold and do not rely on assumptions such as signal sparsity or hand-crafted basis functions.
Revised: These data-driven methods are more powerful compared to classical methods, as they can accurately learn the natural signal manifold and do not rely on assumptions such as signal sparsity or hand-crafted basis functions.

Original: Recently, diffusion models have shown impressive results for both conditional and unconditional image generation and can be easily fitted to a target data distribution using score matching (Vincent, 2011; Song et al., 2020).
Revised: Recently, diffusion models have shown impressive results for both conditional and unconditional image generation and can be easily fitted to a target data distribution using score matching (Vincent, 2011; Song et al., 2020).

Original: These deep generative models learn the score of the data manifold and produce samples by reverting a diffusion process, guiding noise samples towards the target distribution.
Revised: These deep generative models learn the score of the data manifold and produce samples by reverting a diffusion process, guiding noise samples towards the target distribution.

Original: Diffusion models have achieved state-of-the-art performance in many downstream tasks and applications, ranging from state-of-the-art text-to-image models such as
Revised: Diffusion models have achieved state-of-the-art performance in many downstream tasks and applications, ranging from state-of-the-art text-to-image models such as

Original: DALL-E 2 (Ramesh et al., 2022) to medical imaging (Song et al., 2021b;
Revised: DALL-E 2 (Ramesh et al., 2022) to medical imaging (Song et al., 2021b;

Original: Jalal et al., 2021a; Chung & Ye, 2022).
Revised: Jalal et al., 2021a; Chung & Ye, 2022).

Original: Furthermore, understanding of diffusion models is rapidly improving and progress in the field is extremely fast-paced (Chung et al., 2022a; Bansal et al., 2022; Daras et al., 2022a;
Revised: Furthermore, understanding of diffusion models is rapidly improving and progress in the field is extremely fast-paced (Chung et al., 2022a; Bansal et al., 2022; Daras et al., 2022a;

Original: Karras et al., 2022; Luo, 2022).
Revised: Karras et al., 2022; Luo, 2022).

Original: The iterative nature of the sampling procedure used by diffusion models renders inference slow compared to GANs and VAEs.
Revised: The iterative nature of the sampling procedure used by diffusion models renders inference slow compared to GANs and VAEs.

Original: However, many recent efforts have shown ways to significantly improve the sampling speed by accelerating the diffusion process.
Revised: However, many recent efforts have shown ways to significantly improve the sampling speed by accelerating the diffusion process.

Original: Inspired by momentum methods in sampling, Daras et al.
Revised: Inspired by momentum methods in sampling, Daras et al.

Original: (2022b) introduces a momentum sampler for diffusion models, which leads to increased sample quality with fewer function evaluations.
Revised: (2022b) introduces a momentum sampler for diffusion models, which leads to increased sample quality with fewer function evaluations.

Original: Chung et al.
Revised: Chung et al.

Original: (2022b) offers a new sampling strategy, namely Come-Closer-Diffuse-Faster (CCDF), which leverages the conditional quality of inverse problems.
Revised: (2022b) offers a new sampling strategy, namely Come-Closer-Diffuse-Faster (CCDF), which leverages the conditional quality of inverse problems.

Original: The reverse diffusion can be initialized from the observation instead of a sample from the base distribution, which leads to faster convergence for conditional sampling.
Revised: The reverse diffusion can be initialized from the observation instead of a sample from the base distribution, which leads to faster convergence for conditional sampling.

Original: Salimans & Ho (2021) proposes a progressive distillation method that augments the training of the diffusion models with a student-teacher model setup.
Revised: Salimans & Ho (2021) proposes a progressive distillation method that augments the training of the diffusion models with a student-teacher model setup.

Original: In doing this, they were able to drastically reduce the number of sampling steps.
Revised: In doing this, they were able to drastically reduce the number of sampling steps.

Original: Lastly, many methods aim to execute the diffusion process in a reduced space to accelerate the diffusion process.
Revised: Lastly, many methods aim to execute the diffusion process in a reduced space to accelerate the diffusion process.

Original: While Jing et al.
Revised: While Jing et al.

Original: (2022) restricts diffusion through projections onto subspaces, Vahdat et al.
Revised: (2022) restricts diffusion through projections onto subspaces, Vahdat et al.

Original: and Rombach et al.
Revised: and Rombach et al.

Original: (2022) run the diffusion in the latent space.
Revised: (2022) run the diffusion in the latent space.

Original: Despite this promise, current score-based diffusion methods for inverse problems are limited to measurement models with unstructured noise.
Revised: Despite this promise, current score-based diffusion methods for inverse problems are limited to measurement models with unstructured noise.

Original: In many image processing tasks, corruptions are however highly structured and spatially correlated.
Revised: In many image processing tasks, corruptions are however highly structured and spatially correlated.

Original: Relevant examples include interference, speckle, or haze.
Revised: Relevant examples include interference, speckle, or haze.

Original: Nevertheless, current conditional diffusion models naively assume that the noise follows some basic tractable distribution (e.g. Gaussian or Poisson).
Revised: Nevertheless, current conditional diffusion models naively assume that the noise follows some basic tractable distribution (e.g. Gaussian or Poisson).

Original: Beyond the realm of diffusion models, Whang et al.
Revised: Beyond the realm of diffusion models, Whang et al.

Original: (2021) extended normalizing flow (NF)-based inference to structured noise applications.
Revised: (2021) extended normalizing flow (NF)-based inference to structured noise applications.

Original: However, compared to diffusion models, NFs require specialized network architectures, which are computationally and memory expensive.
Revised: However, compared to diffusion models, NFs require specialized network architectures, which are computationally and memory expensive.

Original: Given the promising outlook of diffusion models, we propose to learn score models for both the noise and the desired signal and perform joint inference of both quantities, coupled via the observation model.
Revised: Given the promising outlook of diffusion models, we propose to learn score models for both the noise and the desired signal and perform joint inference of both quantities, coupled via the observation model.

Original: The resulting sampling scheme enables solving a wide variety of inverse problems with structured noise.
Revised: The resulting sampling scheme enables solving a wide variety of inverse problems with structured noise.

Original: The main contributions of this work are as follows:
Revised: The main contributions of this work are as follows:

Original: • We propose a novel joint conditional posterior sampling method to efficiently remove structured noise using diffusion models.
Revised: • We propose a novel joint conditional posterior sampling method to efficiently remove structured noise using diffusion models.

Original: Our formulation is compatible with many existing iterative sampling methods for score-based generative models.
Revised: Our formulation is compatible with many existing iterative sampling methods for score-based generative models.

Original: • We show strong performance gains across various challenging inverse problems involving structured noise compared to competitive state-of-the-art methods based on NFs and GANs.
Revised: • We show strong performance gains across various challenging inverse problems involving structured noise compared to competitive state-of-the-art methods based on NFs and GANs.

Original: • We demonstrate improved robustness on out
Revised: • We demonstrate improved robustness on out

Original: -of-distribution signals compared to baselines.
Revised: -of-distribution signals compared to baselines.

Original:  P ROBLEM S TATEMENT
Revised:  P ROBLEM S TATEMENT

Original: Many image reconstruction tasks can be formulated as an inverse problem with the basic form
Revised: Many image reconstruction tasks can be formulated as an inverse problem with the basic form

Original: [Equation]
Revised: [Equation]

Original: where y ∈ R m is the noisy observation, x ∈ R d the desired signal or image, and n ∈ R m the additive noise.
Revised: where y ∈ R m is the noisy observation, x ∈ R d the desired signal or image, and n ∈ R m the additive noise.

Original: The linear forward operator A ∈ R m × d captures the deterministic transformation of x .
Revised: The linear forward operator A ∈ R m × d captures the deterministic transformation of x .

Original: Maximum a posteriori (MAP) inference is typically used to find an optimal solution ˆ x MAP that maximizes posterior density p X | Y ( x | y ) :
Revised: Maximum a posteriori (MAP) inference is typically used to find an optimal solution ˆ x MAP that maximizes posterior density p X | Y ( x | y ) :

Original: [Equation]
Revised: [Equation]

Original: where p Y | X ( y | x ) is the likelihood according to the measurement model and log p X ( x ) the signal prior.
Revised: where p Y | X ( y | x ) is the likelihood according to the measurement model and log p X ( x ) the signal prior.

Original: Assumptions on the stochastic corruption process n are of key importance too, in particular for applications for which this process is highly structured.
Revised: Assumptions on the stochastic corruption process n are of key importance too, in particular for applications for which this process is highly structured.

Original: However, most methods assume i.i.d.
Revised: However, most methods assume i.i.d.

Original: Gaussian distributed noise, such that the forward model becomes p Y | X ( y | x ) ∼ N ( Ax , σ 2 N I ) .
Revised: Gaussian distributed noise, such that the forward model becomes p Y | X ( y | x ) ∼ N ( Ax , σ 2 N I ) .

Original: This naturally leads to the following simplified problem:
Revised: This naturally leads to the following simplified problem:

Original: [Equation]
Revised: [Equation]

Original: However, as mentioned, this naive assumption can be very restrictive as many noise processes are much more structured and complex.
Revised: However, as mentioned, this naive assumption can be very restrictive as many noise processes are much more structured and complex.

Original: A myriad of problems can be addressed under the formulation of equation 1, given the freedom of choice for the noise source n .
Revised: A myriad of problems can be addressed under the formulation of equation 1, given the freedom of choice for the noise source n .

Original: Therefore, in this work, our aim is to solve a more broad class of inverse problems defined by any arbitrary noise distribution n ∼ p N ( n ) ̸ = N and signal prior x ∼ p X ( x ) , resulting in the following, more general, MAP estimator proposed by Whang et al.
Revised: Therefore, in this work, our aim is to solve a more broad class of inverse problems defined by any arbitrary noise distribution n ∼ p N ( n ) ̸ = N and signal prior x ∼ p X ( x ) , resulting in the following, more general, MAP estimator proposed by Whang et al.

Original: [Equation]
Revised: [Equation]

Original: In this paper, we propose to solve this class of problems using flexible diffusion models.
Revised: In this paper, we propose to solve this class of problems using flexible diffusion models.

Original: Furthermore, diffusion models naturally enable posterior sampling, allowing us to take advantage of the benefits thereof (Jalal et al., 2021b; Kawar et al., 2021; Daras et al., 2022a).
Revised: Furthermore, diffusion models naturally enable posterior sampling, allowing us to take advantage of the benefits thereof (Jalal et al., 2021b; Kawar et al., 2021; Daras et al., 2022a).

Original: 2. R ELATED W ORK 2.1. N ORMALIZING F LOWS
Revised: 2. R ELATED W ORK 2.1. N ORMALIZING F LOWS

Original: Whang et al.
Revised: Whang et al.

Original: (2021) propose to use normalizing flows (NFs) to model both the data and the noise distributions.
Revised: (2021) propose to use normalizing flows (NFs) to model both the data and the noise distributions.

Original: Normalizing flows are a special class of likelihood-based generative models that make use of an invertible mapping G : R d → R d to transform samples from a base distribution p Z ( z ) into a more complex multimodal distribution x = G ( z ) ∼ p X ( x ) .
Revised: Normalizing flows are a special class of likelihood-based generative models that make use of an invertible mapping G : R d → R d to transform samples from a base distribution p Z ( z ) into a more complex multimodal distribution x = G ( z ) ∼ p X ( x ) .

Original: The invertible nature of the mapping
Revised: The invertible nature of the mapping

Original: G allows for exact density evaluation through the change of variables formula:
Revised: G allows for exact density evaluation through the change of variables formula:

Original: [Equation]
Revised: [Equation]

Original: where J is the Jacobian that accounts for the change in volume between densities.
Revised: where J is the Jacobian that accounts for the change in volume between densities.

Original: Since exact likelihood computation is possible through the flow direction G − 1 , the parameters of the generator network can be optimized to maximize likelihood of the training data.
Revised: Since exact likelihood computation is possible through the flow direction G − 1 , the parameters of the generator network can be optimized to maximize likelihood of the training data.

Original: Subsequently, the inverse task is solved using the MAP estimation in equation 5:
Revised: Subsequently, the inverse task is solved using the MAP estimation in equation 5:

Original: [Equation]
Revised: [Equation]

Original: where G N and G X are generative flow models for the noise and data respectively.
Revised: where G N and G X are generative flow models for the noise and data respectively.

Original: Analog to that, the solution can be solved in the latent space rather than the image space as follows:
Revised: Analog to that, the solution can be solved in the latent space rather than the image space as follows:

Original: [Equation]
Revised: [Equation]

Original: Note that in equation 8 a smoothing parameter λ is added to weigh the prior and likelihood terms, as was also done in Whang et al.
Revised: Note that in equation 8 a smoothing parameter λ is added to weigh the prior and likelihood terms, as was also done in Whang et al.

Original: 2.1. G ENERATIVE A DVERSARIAL N ETWORKS
Revised: 2.1. G ENERATIVE A DVERSARIAL N ETWORKS

Original: Generative adversarial networks (GANs) are implicit generative models that can learn the data manifold in an adversarial manner (Goodfellow et al., 2020).
Revised: Generative adversarial networks (GANs) are implicit generative models that can learn the data manifold in an adversarial manner (Goodfellow et al., 2020).

Original: The generative model is trained with an auxiliary discriminator network that evaluates the generator’s performance in a minimax game.
Revised: The generative model is trained with an auxiliary discriminator network that evaluates the generator’s performance in a minimax game.

Original: The generator G ( z ) : R l → R d maps latent vectors z ∈ R l ∼ N ( 0 , I ) to the data distribution of interest.
Revised: The generator G ( z ) : R l → R d maps latent vectors z ∈ R l ∼ N ( 0 , I ) to the data distribution of interest.

Original: The structure of the generative model can also be used in inverse problem solving (Bora et al., 2017).
Revised: The structure of the generative model can also be used in inverse problem solving (Bora et al., 2017).

Original: The objective can be derived from equation 3 and is given by:
Revised: The objective can be derived from equation 3 and is given by:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where λ weights the importance of the prior with the measurement error.
Revised: where λ weights the importance of the prior with the measurement error.

Original: The ℓ 2 regularization term on the latent variable is proportional to negative log-likelihood under the prior defined by G X , where the subscript denotes the density that the generator is approximating.
Revised: The ℓ 2 regularization term on the latent variable is proportional to negative log-likelihood under the prior defined by G X , where the subscript denotes the density that the generator is approximating.

Original: While this method does not explicitly model the noise, it remains an interesting comparison, as the generator cannot reproduce the noise found in the measurement and can only recover signals that are in the range of the generator.
Revised: While this method does not explicitly model the noise, it remains an interesting comparison, as the generator cannot reproduce the noise found in the measurement and can only recover signals that are in the range of the generator.

Original: Therefore, due to the limited support of the learned distribution, GANs can inherently remove structured noise.
Revised: Therefore, due to the limited support of the learned distribution, GANs can inherently remove structured noise.

Original: However, the representation error (i.e. observation lies far from the range of the generator (Bora et al., 2017)) imposed by the structured noise comes at the cost of recovery quality.
Revised: However, the representation error (i.e. observation lies far from the range of the generator (Bora et al., 2017)) imposed by the structured noise comes at the cost of recovery quality.

Original: 2. B ACKGROUND ON S CORE - BASED D IFFUSION M ODELS
Revised: 2. B ACKGROUND ON S CORE - BASED D IFFUSION M ODELS

Original: One class of deep generative models is known as diffusion models.
Revised: One class of deep generative models is known as diffusion models.

Original: These generative models have been introduced independently as score-based models (Song & Ermon, 2019; 2020) and denoising diffusion probabilistic modeling (DDPM)
Revised: These generative models have been introduced independently as score-based models (Song & Ermon, 2019; 2020) and denoising diffusion probabilistic modeling (DDPM)

Original: (Ho et al., 2020).
Revised: (Ho et al., 2020).

Original: In this work, we will consider the formulation introduced in Song et al.
Revised: In this work, we will consider the formulation introduced in Song et al.

Original: (2020), which unifies both perspectives on diffusion models by expressing diffusion as a continuous process through stochastic differential equations (SDE).
Revised: (2020), which unifies both perspectives on diffusion models by expressing diffusion as a continuous process through stochastic differential equations (SDE).

Original: Diffusion models produce samples by reversing a corruption process.
Revised: Diffusion models produce samples by reversing a corruption process.

Original: In essence these models are networks trained to denoise its input.
Revised: In essence these models are networks trained to denoise its input.

Original: Through iteration of this process, samples can be drawn from a learned data distribution, starting from random noise.
Revised: Through iteration of this process, samples can be drawn from a learned data distribution, starting from random noise.

Original: The diffusion process of the data (cid:8) x t ∈ R d (cid:9) t ∈ [0 , 1] is characterized by a continuous sequence of
Revised: The diffusion process of the data (cid:8) x t ∈ R d (cid:9) t ∈ [0 , 1] is characterized by a continuous sequence of

Original: Gaussian perturbations of increasing magnitude indexed by time t ∈ [0 , 1] .
Revised: Gaussian perturbations of increasing magnitude indexed by time t ∈ [0 , 1] .

Original: Starting from the data distribution at t = 0 , clean images are defined by x 0 ∼ p ( x 0 ) ≡ p ( x ) .
Revised: Starting from the data distribution at t = 0 , clean images are defined by x 0 ∼ p ( x 0 ) ≡ p ( x ) .

Original: Forward diffusion can be described using an SDE as follows:
Revised: Forward diffusion can be described using an SDE as follows:

Original: [Equation]
Revised: [Equation]

Original: where w ∈ R d is a standard Wiener process, f ( t ) : [0 , 1] → R and g ( t ) : [0 , 1] → R are the drift and diffusion coefficients, respectively.
Revised: where w ∈ R d is a standard Wiener process, f ( t ) : [0 , 1] → R and g ( t ) : [0 , 1] → R are the drift and diffusion coefficients, respectively.

Original: Moreover, these coefficients are chosen so that the resulting distribution p 1 ( x ) at the end of the perturbation process approximates a predefined base distribution p 1 ( x ) ≈ π ( x ) .
Revised: Moreover, these coefficients are chosen so that the resulting distribution p 1 ( x ) at the end of the perturbation process approximates a predefined base distribution p 1 ( x ) ≈ π ( x ) .

Original: Furthermore, the transition kernel of the diffusion process is defined as q ( x t | x ) ∼ N ( x t | α ( t ) x , β 2 ( t ) I ) , where α ( t ) and β ( t ) can be analytically derived from the SDE.
Revised: Furthermore, the transition kernel of the diffusion process is defined as q ( x t | x ) ∼ N ( x t | α ( t ) x , β 2 ( t ) I ) , where α ( t ) and β ( t ) can be analytically derived from the SDE.

Original: Naturally, we are interested in reversing the diffusion process, so that we can sample from x 0 ∼ p 0 ( x 0 ) .
Revised: Naturally, we are interested in reversing the diffusion process, so that we can sample from x 0 ∼ p 0 ( x 0 ) .

Original: The reverse diffusion process is also a diffusion process given by the reverse-time SDE (Anderson, 1982; Song et al., 2020):
Revised: The reverse diffusion process is also a diffusion process given by the reverse-time SDE (Anderson, 1982; Song et al., 2020):

Original: [Equation]
Revised: [Equation]

Original: where ¯ w t is the standard Wiener process in the reverse direction.
Revised: where ¯ w t is the standard Wiener process in the reverse direction.

Original: The gradient of the log-likelihood of the data with respect to itself, a.k.a.
Revised: The gradient of the log-likelihood of the data with respect to itself, a.k.a.

Original: the score function , arises from the reverse-time SDE.
Revised: the score function , arises from the reverse-time SDE.

Original: The score function is a gradient field pointing back to the data manifold and can intuitively be used to guide a random sample from the base distribution π ( x ) to the desired data distribution.
Revised: The score function is a gradient field pointing back to the data manifold and can intuitively be used to guide a random sample from the base distribution π ( x ) to the desired data distribution.

Original: Given a dataset X = (cid:8) x (1) , x (2) , . . .
Revised: Given a dataset X = (cid:8) x (1) , x (2) , . . .

Original: , x ( |X| ) (cid:9) ∼ p ( x ) , scores can be estimated by training a neural network s θ ( x t , t ) parameterized by weights θ , with score-matching techniques such as the denoising score matching (DSM) objective (Vincent, 2011):
Revised: , x ( |X| ) (cid:9) ∼ p ( x ) , scores can be estimated by training a neural network s θ ( x t , t ) parameterized by weights θ , with score-matching techniques such as the denoising score matching (DSM) objective (Vincent, 2011):

Original: [Equation]
Revised: [Equation]

Original: Given a sufficiently large dataset X and model capacity, DSM ensures that the score network converges to s θ ( x t , t ) ≃ ∇ x t log p ( x t ) .
Revised: Given a sufficiently large dataset X and model capacity, DSM ensures that the score network converges to s θ ( x t , t ) ≃ ∇ x t log p ( x t ) .

Original: After training the time-dependent score model s θ , it can be used to calculate the reverse-time diffusion process and solve the trajectory using numerical samplers such as the Euler-Maruyama algorithm.
Revised: After training the time-dependent score model s θ , it can be used to calculate the reverse-time diffusion process and solve the trajectory using numerical samplers such as the Euler-Maruyama algorithm.

Original: Alternatively, more sophisticated samplers, such as
Revised: Alternatively, more sophisticated samplers, such as

Original: ALD (Song & Ermon, 2019), probability flow ODE (Song et al., 2020), and Predictor-Corrector sampler (Song et al., 2020), can be used to further improve sample quality.
Revised: ALD (Song & Ermon, 2019), probability flow ODE (Song et al., 2020), and Predictor-Corrector sampler (Song et al., 2020), can be used to further improve sample quality.

Original: These iterative sampling algorithms discretize the continuous time SDE into a sequence of time steps { 0 = t 0 , t 1 , . .
Revised: These iterative sampling algorithms discretize the continuous time SDE into a sequence of time steps { 0 = t 0 , t 1 , . .

Original: , t T = 1 } , where a noisy sample
Revised: , t T = 1 } , where a noisy sample

Original: ˆ x t i is denoised to produce a sample for the next time step ˆ x t i − 1 .
Revised: ˆ x t i is denoised to produce a sample for the next time step ˆ x t i − 1 .

Original: The resulting samples { ˆ x t i } Ti =0 constitute an approximation of the actual diffusion process { x t } t ∈ [0 , 1] .
Revised: The resulting samples { ˆ x t i } Ti =0 constitute an approximation of the actual diffusion process { x t } t ∈ [0 , 1] .

Original:  M ETHOD 3. C ONDITIONAL P OSTERIOR S AMPLING UNDER S TRUCTURED N OISE
Revised:  M ETHOD 3. C ONDITIONAL P OSTERIOR S AMPLING UNDER S TRUCTURED N OISE

Original: We are interested in posterior sampling under structured noise.
Revised: We are interested in posterior sampling under structured noise.

Original: We recast this as a joint optimization problem with respect to the signal x and noise n given by:
Revised: We recast this as a joint optimization problem with respect to the signal x and noise n given by:

Original: [Equation]
Revised: [Equation]

Original: Solving inverse problems using diffusion models requires conditioning of the diffusion process on the observation y , such that we can sample from the posterior p X | Y ( x , n | y ) .
Revised: Solving inverse problems using diffusion models requires conditioning of the diffusion process on the observation y , such that we can sample from the posterior p X | Y ( x , n | y ) .

Original: Therefore, we construct a joint conditional diffusion process { x t , n t | y } t ∈ [0 , 1] , in turn producing a joint conditional reversetime SDE:
Revised: Therefore, we construct a joint conditional diffusion process { x t , n t | y } t ∈ [0 , 1] , in turn producing a joint conditional reversetime SDE:

Original: [Equation]
Revised: [Equation]

Original: We would like to factorize the posterior using our learned unconditional score model and tractable measurement model, given the joint formulation.
Revised: We would like to factorize the posterior using our learned unconditional score model and tractable measurement model, given the joint formulation.

Original: Consequently, we construct two separate diffusion processes, defined by separate score models but entangled through the measurement model p Y | X,N ( y | x , n ) . 
Revised: Consequently, we construct two separate diffusion processes, defined by separate score models but entangled through the measurement model p Y | X,N ( y | x , n ) . In addition to the original score model s θ ( x , t ) , we introduce a second score model s ϕ [Equation]

Original: In addition to the original score model s θ ( x , t ) , we introduce a second score model s ϕ ( n  , t ) ≃ ∇ n  log p N ( n  ) , parameterized by weights ϕ , to model the expressive noise component n . 
Revised: n t , t ) [Equation] n t log p N ( n t ) , parameterized by weights ϕ , to model the expressive noise component n . These two score networks can be trained independently on datasets for x and n , respectively, using the objective in equation 12.

Original: The gradients of the posterior with respect to x and n are now given by:
Revised: The gradients of the posterior with respect to x and n are now given by:

Original: [Equation]
Revised: [Equation]

Original: where ˆ y t is a sample from p ( y t | y ) , and { y t } t ∈ [0 , 1] is an additional stochastic process that essentially corrupts the observation along the SDE trajectory together with x t .
Revised: where ˆ y t is a sample from p ( y t | y ) [Equation] and { y t } t ∈ [0 , 1] is an additional stochastic process that essentially corrupts the observation along the SDE trajectory together with x t .

Original: As p ( y t | y ) is tractable, we can easily compute ˆ y t = α ( t ) y + β ( t ) Az , using the reparameterization trick with z ∈ R d ∼ N ( 0 , I ) , see Song et al.
Revised: As p ( y t | y ) is tractable, we can easily compute ˆ y t = α ( t ) y + β ( t ) Az , using the reparameterization trick with z ∈ R d ∼ N ( 0 , I ) , see Song et al.

Original: (2021b).
Revised: (2021b).

Original: Subsequently, the approximation in equation 15 and equation 16 can be substituted for the conditional score in equation 14, resulting in two entangled diffusion processes:
Revised: Subsequently, the approximation in equation 15 and equation 16 can be substituted for the conditional score in equation 14, resulting in two entangled diffusion processes:

Original: [Equation]
Revised: [Equation]

Original: which allows us to perform posterior sampling for both the signal, such that x ≡ x 0 ∼ p X | Y ( x 0 | y ) , as well as the structured noise, such that n ≡ n 0 ∼ p N | Y ( n 0 | y ) .
Revised: which allows us to perform posterior sampling for both the signal, such that x ≡ x 0 ∼ p X | Y ( x 0 | y ) , as well as the structured noise, such that n ≡ n 0 ∼ p N | Y ( n 0 | y ) .

Original: To solve the approximated joint conditional reverse-time SDE, we resort to the aforementioned iterative scheme in Section 2.2, however, now incorporating the observation via a data-consistency step.
Revised: To solve the approximated joint conditional reverse-time SDE, we resort to the aforementioned iterative scheme in Section 2.2, however, now incorporating the observation via a data-consistency step.

Original: This is done by taking gradient steps that minimize the ℓ 2 norm between the true observation and its model prediction given current estimates of x and n .
Revised: This is done by taking gradient steps that minimize the ℓ 2 norm between the true observation and its model prediction given current estimates of x and n .

Original: Ultimately, this results in solutions that are consistent with the observation y and have high likelihood under both prior models.
Revised: Ultimately, this results in solutions that are consistent with the observation y and have high likelihood under both prior models.

Original: The data-consistency update steps for both x and n are derived as follows:
Revised: The data-consistency update steps for both x and n are derived as follows:

Original: [Equation]
Revised: [Equation]

Original: where the time difference between two steps ∆ t = 1 /T and λ and µ are weighting coefficients for the signal and noise gradient steps, respectively.
Revised: where the time difference between two steps ∆ t = 1 /T and λ and µ are weighting coefficients for the signal and noise gradient steps, respectively.

Original: An example of the complete sampling algorithm is shown in Algorithm 1, which adapts the Euler-Maruyama sampler (Song et al., 2020) to jointly find the optimal data sample and the optimal noise sample while taking into account the measurement model in line 7 and 8 using the outcome of equation 18 and equation 19, respectively.
Revised: An example of the complete sampling algorithm is shown in Algorithm 1, which adapts the Euler-Maruyama sampler (Song et al., 2020) to jointly find the optimal data sample and the optimal noise sample while taking into account the measurement model in line 7 and 8 using the outcome of equation 18 and equation 19, respectively.

Original: Although we show the Euler-Maruyama method, our addition is applicable to a large family of iterative sampling methods for score-based generative models.
Revised: Although we show the Euler-Maruyama method, our addition is applicable to a large family of iterative sampling methods for score-based generative models.

Original: Require: T, s θ , s ϕ , λ, µ, y 

Revised: Require: T, s θ , s ϕ , λ, µ, y

Original: ˆ x 1 ∼ π ( x ) , ˆ n 1 ∼ π ( n ) , ∆ t ← 1 T 3 for i
Revised: ˆ x 1 ∼ π ( x ) , ˆ n 1 ∼ π ( n ) , ∆ t ← 1 T 3 for i

Original: = T − 1 to 0 do t ← i +1 T
Revised: = T − 1 to 0 do t ← i +1 T

Original: ˆ y t ∼ p 0 t ( y t | y ) //
Revised: ˆ y t ∼ p 0 t ( y t | y ) //

Original: Data consistency steps [Equation] [Equation]
Revised: Data consistency steps ˆ x t − ∆ t ← ˆ x t − λ A T (

Original: [Equation]
Revised: A ˆ x t − ˆ y t + ˆ n t ) ˆ n t − ∆ t ← ˆ n t − µ ( A ˆ x t − ˆ y t + ˆ n t ) 9 ...

Original: ˆ x t − ∆ t ← ˆ x t − f
Revised: ˆ x t − ∆ t ← ˆ x t − f

Original: ( t )
Revised: ( t )

Original: ˆ x t ∆ t ˆ x t − ∆ t ← ˆ x t − ∆ t + g ( t ) 2 s ∗ θ (ˆ x t , t )∆ t z ∼ N ( 0 , I )ˆ x t − ∆ t ← ˆ x t −
Revised: ˆ x t ∆ t ˆ x t − ∆ t ← ˆ x t − ∆ t + g ( t ) 2 s ∗ θ (ˆ x t , t )∆ t z ∼ N ( 0 , I )

Original: [Table]

Original:  t + g ( t ) √∆ t z  ˆ n [Table] t ← ˆ n
Revised: ˆ x t − ∆ t ← ˆ x t − ∆ t + g ( t ) √

Original: [Table] − f ( t )ˆ n t ∆ t ˆ n t − ∆ t ← ˆ n t − ∆ t + g ( t ) 2 s ∗ ϕ (ˆ n t , t )∆ t z ∼ N ( 0 , I ) ˆ n t − ∆ t ← ˆ n t − ∆ t + g ( t ) √ ∆ t z return: ˆ x19 end
Revised: ∆ t z  ˆ n t − ∆ t ← ˆ n t − f ( t )ˆ n t ∆ t ˆ n t − ∆ t ← ˆ n t − ∆ t + g ( t ) 2 s ∗ ϕ (ˆ n t , t ) ∆ t z ∼ N ( 0 , I ) ˆ n t − ∆ t ← ˆ n t − ∆ t + g ( t ) √ ∆ t z return: ˆ x19 end

Original: 3. T RAINING AND I NFERENCE S ETUP
Revised: 3. T RAINING AND I NFERENCE S ETUP

Original: For training the score models, we use the NCSNv2 architecture as introduced in Song & Ermon (2020) in combination with the Adam optimizer and a learning rate of 5e − 4 until convergence.
Revised: For training the score models, we use the NCSNv2 architecture as introduced in Song & Ermon (2020) in combination with the Adam optimizer and a learning rate of 5e − 4 until convergence.

Original: For simplicity, no exponential moving average (EMA) filter on the network weights is applied.
Revised: For simplicity, no exponential moving average (EMA) filter on the network weights is applied.

Revised: Given two separate datasets, one for the data and one for the structured noise, two separate score models can be trained independently.

Revised: This allows for easy adaptation of our method, since many existing trained score models can be reused.

Revised: Only during inference, the two priors are combined through the proposed sampling procedure as described in Algorithm 1, using the adapted Euler-Maruyama sampler.

Original: We use the variance preserving (VP) SDE ( β 0 = 0 .
Revised: We use the variance preserving (VP) SDE ( β 0 = 0 .

Original: 1 , β 1 = 7 . 0 )
Revised: 1 , β 1 = 7 . 0 )

Original: (Song et al., 2020) to define the diffusion trajectory.
Revised: (Song et al., 2020) to define the diffusion trajectory.

Original: Furthermore, we use the exact sampling procedure as described in Algorithm 1, using the adapted Euler-Maruyama sampler.

Original: During each experiment, we run the sampler for T =iterations.
Revised: During each experiment, we run the sampler for T = 600 iterations.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: All models are trained on the CelebA dataset (Liu et al., 2015) and the MNIST dataset withand 27000 training samples, respectively.
Revised: All models are trained on the CelebA dataset (Liu et al., 2015) and the MNIST dataset withand 27000 training samples, respectively.

Original: We downsize the images to 64 × 64 pixels.
Revised: We downsize the images to 64 × 64 pixels.

Original: Due to computational constraints, we test on a randomly selected subset of 100 images.
Revised: Due to computational constraints, we test on a randomly selected subset of 100 images.

Original: We use both the peak signal-to noise ratio (PSNR) and structural similarity index (SSIM) to evaluate our results.
Revised: We use both the peak signal-to noise ratio (PSNR) and structural similarity index (SSIM) to evaluate our results.

Original: 4. B ASELINE M ETHODS
Revised: 4. B ASELINE M ETHODS

Original: The closest to our work is the flow-based noise model proposed by Whang et al.
Revised: The closest to our work is the flow-based noise model proposed by Whang et al.

Original: (2021), discussed in
Revised: (2021), discussed in

Original: Section 2.1.1, which will serve as our main baseline.
Revised: Section 2.1.1, which will serve as our main baseline.

Original: To boost the performance of this baseline and to make it more competitive we moreover replace the originally-used RealNVP (Dinh et al., 2016) with the Glow architecture (Kingma & Dhariwal, 2018).
Revised: To boost the performance of this baseline and to make it more competitive we moreover replace the originally-used RealNVP (Dinh et al., 2016) with the Glow architecture (Kingma & Dhariwal, 2018).

Original: Glow is a widely used flow model highly inspired by RealNVP, with the addition of 1 × 1 convolutions before each coupling layer.
Revised: Glow is a widely used flow model highly inspired by RealNVP, with the addition of 1 × 1 convolutions before each coupling layer.

Original: We use the exact implementation found in Asim et al.
Revised: We use the exact implementation found in Asim et al.

Original: (2020), with a flow depth of K = 18 , and number of levels
Revised: (2020), with a flow depth of K = 18 , and number of levels

Original: L = 4 .
Revised: L = 4 , which has been optimized for the same CelebA dataset used in this work and thus should provide a fair comparison with the proposed method.

Original: Secondly, GANs as discussed in Section 2.1.2 are used as a comparison.
Revised: Secondly, GANs as discussed in Section 2.1.2 are used as a comparison.

Original: We train a DCGAN (Radford et al., 2015), with a generator latent input dimension of l = 100 .
Revised: We train a DCGAN (Radford et al., 2015), with a generator latent input dimension of l = 100 .

Original: The generator architecture consists of 4 strided 2D transposed convolutional layers, having 4 × 4 kernels yielding feature maps of 512, 256, 128 and 64.
Revised: The generator architecture consists of 4 strided 2D transposed convolutional layers, having 4 × 4 kernels yielding feature maps of 512, 256, 128 and 64.

Original: Each convolutional layer is followed by a batch normalization layer and
Revised: Each convolutional layer is followed by a batch normalization layer and

Original: ReLU activation.
Revised: ReLU activation.

Original: Lastly, depending on the reconstruction task, classical non-data-driven methods are used as a comparison.
Revised: Lastly, depending on the reconstruction task, classical non-data-driven methods are used as a comparison.

Original: For denoising experiments, we use the block-matching and 3D filtering algorithm (BM3D) (Dabov et al., 2006), and in compressed sensing experiments, LASSO with wavelet basis (Tibshirani, 1996).
Revised: For denoising experiments, we use the block-matching and 3D filtering algorithm (BM3D) (Dabov et al., 2006), and in compressed sensing experiments, LASSO with wavelet basis (Tibshirani, 1996).

Original: B M 3 D
Revised: B M 3 D

Original: G A N
Revised: G A N

Original: F L O W
Revised: F L O W

Original: D I FF U S
Revised: D I FF U S

Original: I
Revised: I

Original: O N
Revised: O N

Original: PS N R 11.
Revised: PS N R 11.

Original: SS I M (a) CelebA
Revised: SS I M (a) CelebA

Original: B M 3 D
Revised: B M 3 D

Original: G A N
Revised: G A N

Original: F L O W
Revised: F L O W

Original: D I FF U S
Revised: D I FF U S

Original: I
Revised: I

Original: O N
Revised: O N

Original: PS N R 10.
Revised: PS N R 10.

Original: SS I M
Revised: SS I M

Original: [Figure]
Revised: [Figure]

Original: Except for the flow-based method of Whang et al.
Revised: Except for the flow-based method of Whang et al.

Original: (2021), none of these methods explicitly model the noise distribution.
Revised: (2021), none of these methods explicitly model the noise distribution.

Original: Still, they are a valuable baseline, as they demonstrate the effectiveness of incorporating a learned structured noise prior rather than relying on simple noise priors.
Revised: Still, they are a valuable baseline, as they demonstrate the effectiveness of incorporating a learned structured noise prior rather than relying on simple noise priors.

Original: Automatic hyperparameter tuning for optimal inference was performed for all baseline methods on a small validation set of only 5 images.
Revised: Automatic hyperparameter tuning for optimal inference was performed for all baseline methods on a small validation set of only 5 images.

Original: For both GAN and flow-based methods, we anneal the step size during inference based on stagnation of the objective.
Revised: For both GAN and flow-based methods, we anneal the step size during inference based on stagnation of the objective.

Original: 4. R ESULTS 4.2. R EMOVING MNIST DIGITS
Revised: 4. R ESULTS 4.2. R EMOVING MNIST DIGITS

Original: For comparison with Whang et al. (2021), we recreate an experiment introduced in their work, where MNIST digits are added to CelebA faces.
Revised: For comparison with Whang et al. (2021), we recreate an experiment introduced in their work, where MNIST digits are added to CelebA faces.

Original: Moreover, the experiment is easily reproducible as both CelebA and MNIST datasets are publicly available.
Revised: Moreover, the experiment is easily reproducible as both CelebA and MNIST datasets are publicly available.

Original: The corruption process is defined by y = 0 .
Revised: The corruption process is defined by y = 0 .

Original: 5 · x CelebA + 0 . 5 · n MNIST .
Revised: 5 · x CelebA + 0 . 5 · n MNIST .

Revised: In this experiment, the score network s ϕ is trained on the MNIST dataset.

Original: Fig.
Revised: Fig.

Original: 1a shows a quantitative comparison of our method with all baselines.
Revised: 1a shows a quantitative comparison of our method with all baselines.

Original: Furthermore, a random selection of test samples is shown in Fig.
Revised: Furthermore, a random selection of test samples is shown in Fig.

Original: 2 for qualitative analysis.
Revised: 2 for qualitative analysis.

Original: Both our method and the flow-based method are able to recover the data, and remove most of the structured noise.
Revised: Both our method and the flow-based method are able to recover the data, and remove most of the structured noise.

Original: However, more details are preserved using the diffusion method.
Revised: However, more details are preserved using the diffusion method.

Original: In contrast, the flow-based method cannot completely remove the digits in some cases and is unable to reconstruct some subtle features present in the original images.
Revised: In contrast, the flow-based method cannot completely remove the digits in some cases and is unable to reconstruct some subtle features present in the original images.

Original: Furthermore, we observe that for the flow-based method, initialization from the measurement is necessary to reproduce the results in Whang et al.
Revised: Furthermore, we observe that for the flow-based method, initialization from the measurement is necessary to reproduce the results in Whang et al.

Original: (2021) since random initialization does not converge.
Revised: (2021) since random initialization does not converge.

Original: The GAN method is also able to remove the digits, but cannot accurately reconstruct the faces as it is unable to project the observation onto the range of the generator.
Revised: The GAN method is also able to remove the digits, but cannot accurately reconstruct the faces as it is unable to project the observation onto the range of the generator.

Original: Similarly, the BM3D denoiser fails to recover the underlying signal, confirming the importance of prior knowledge of the noise in this experiment.
Revised: Similarly, the BM3D denoiser fails to recover the underlying signal, confirming the importance of prior knowledge of the noise in this experiment.

Original: The metrics in Fig.
Revised: The metrics in Fig.

Original: 1a support these observations.
Revised: 1a support these observations.

Original: See Table 1 for the extended results.
Revised: See Table 1 for the extended results.

Original: Additionally, we expose the methods in a similar experiment to out-of-distribution (OoD) data.
Revised: Additionally, we expose the methods in a similar experiment to out-of-distribution (OoD) data.

Original: The images from this dataset not found in the CelebA dataset, which is the data used for training the models.
Revised: The images from this dataset not found in the CelebA dataset, which is the data used for training the models.

Original: In fact, the out-of-distribution data is generated using the stable-diffusion text-to-image model Rombach et al.
Revised: In fact, the out-of-distribution data is generated using the stable-diffusion text-to-image model Rombach et al.

Original: We use the exact same hyperparameters as during the experiment on the CelebA dataset.
Revised: We use the exact same hyperparameters as during the experiment on the CelebA dataset.

Original: Quantitative and qualitative results are shown in Fig.
Revised: Quantitative and qualitative results are shown in Fig.

Original: 1b and Fig.
Revised: 1b and Fig.

Original: 3, respectively.
Revised: 3, respectively.

Original: Similarly to the findings of Whang et al.
Revised: Similarly to the findings of Whang et al.

Original: ; Asim et al.
Revised: ; Asim et al.

Original: (2020), the flow-based method is robust to OoD data, due to their inherent invertibility.
Revised: (2020), the flow-based method is robust to OoD data, due to their inherent invertibility.

Original: We empirically show that the diffusion method is also resistant to OoD data in inverse tasks with complex noise structures and even outperforms the flowbased methos.
Revised: We empirically show that the diffusion method is also resistant to OoD data in inverse tasks with complex noise structures and even outperforms the flowbased methos.

Original: Unsurprisingly, the GAN method performs even more poorly when subjected to OoD data.
Revised: Unsurprisingly, the GAN method performs even more poorly when subjected to OoD data.

Original: More experiments, covering different inverse problem settings can be found in Appendix A.
Revised: More experiments, covering different inverse problem settings can be found in Appendix A.

Original: [Figure]
Revised: [Figure]

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: npu
Revised: npu

Original: [Figure]
Revised: [Figure]

Original: I
Revised: I

Original: [Figure]
Revised: [Figure]

Original: U
Revised: U

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: npu
Revised: npu

Original: [Figure]
Revised: [Figure]

Original: 4.2. P ERFORMANCE
Revised: 4.2. P ERFORMANCE

Original: To highlight the difference in inference time between our method and the baselines, benchmarks are performed on a single 12GBytes NVIDIA GeForce RTX 3080 Ti, see Table 3 in Appendix B.2.
Revised: To highlight the difference in inference time between our method and the baselines, benchmarks are performed on a single 12GBytes NVIDIA GeForce RTX 3080 Ti, see Table 3 in Appendix B.2.

Original: Although this is not an extensive benchmark, a quick comparison of inference times reveals a 50 × difference in speed between ours and the flow-based method.
Revised: Although this is not an extensive benchmark, a quick comparison of inference times reveals a 50 × difference in speed between ours and the flow-based method.

Original: All the deep generative models need approximately an equal amount of iterations ( T ≈ 600) to converge.
Revised: All the deep generative models need approximately an equal amount of iterations ( T ≈ 600) to converge.

Original: However, for the same modeling capacity, the flow model requires a substantial higher amount of trainable parameters compared to the diffusion method.
Revised: However, for the same modeling capacity, the flow model requires a substantial higher amount of trainable parameters compared to the diffusion method.

Original: This is mainly due to the restrictive requirements imposed on the architecture to ensure tractable likelihood computation.
Revised: This is mainly due to the restrictive requirements imposed on the architecture to ensure tractable likelihood computation.

Original: It should be noted that no improvements to speed up the diffusion process, such as CCDF (Chung et al., 2022b) are applied for the diffusion method, giving room for even more improvement in future work.
Revised: It should be noted that no improvements to speed up the diffusion process, such as CCDF (Chung et al., 2022b) are applied for the diffusion method, giving room for even more improvement in future work.

Original:  D ISCUSSION AND C ONCLUSIONS
Revised:  D ISCUSSION AND C ONCLUSIONS

Original: In this work, we present a framework for removing structured noise using diffusion models.
Revised: In this work, we present a framework for removing structured noise using diffusion models.

Original: Our work provides an efficient addition to existing score-based conditional sampling methods incorporating knowledge of the noise distribution.
Revised: Our work provides an efficient addition to existing score-based conditional sampling methods incorporating knowledge of the noise distribution.

Original: We demonstrate our method on natural and out-ofdistribution data and achieve increased performance over the state-of-the-art.
Revised: We demonstrate our method on natural and out-ofdistribution data and achieve increased performance over the state-of-the-art and established conventional methods for complex inverse tasks.

Revised: Additionally, the diffusion based method is substantially easier to train using the score matching objective compared to other deep generative methods and furthermore allows for posterior sampling.

Original: While our method is considerably faster and better in removing structured noise compared to the flow-based method (Whang et al., 2021), it is not ready (yet) for real-time inference and still slow compared to GANs (Bora et al., 2017) and classical methods.
Revised: While our method is considerably faster and better in removing structured noise compared to the flow-based method (Whang et al., 2021), it is not ready (yet) for real-time inference and still slow compared to GANs (Bora et al., 2017) and classical methods.

Original: Luckily, research into accelerating the diffusion process are well on their way.
Revised: Luckily, research into accelerating the diffusion process are well on their way.

Original: In addition, although a simple sampling algorithm was adapted in this work, many more sampling algorithms for score-based diffusion models exist, each of which introduces a new set of hyperparameters.
Revised: In addition, although a simple sampling algorithm was adapted in this work, many more sampling algorithms for score-based diffusion models exist, each of which introduces a new set of hyperparameters.

Original: For example, the predictor-corrector (PC) sampler has been shown to improve sample quality (Song et al., 2020).
Revised: For example, the predictor-corrector (PC) sampler has been shown to improve sample quality (Song et al., 2020).

Original: Future work should explore this wide increase in design space to understand limitations and possibilities of more sophisticated sampling schemes in combination with the proposed joint diffusion method.
Revised: Future work should explore this wide increase in design space to understand limitations and possibilities of more sophisticated sampling schemes in combination with the proposed joint diffusion method.

Original: Furthermore, the range of problems to which we can apply the proposed method, can be expanded into non-linear likelihood models and extend beyond the additive noise models.
Revised: Furthermore, the range of problems to which we can apply the proposed method, can be expanded into non-linear likelihood models and extend beyond the additive noise models.

Original: Lastly, the connection between diffusion models and continuous normalizing flows through the neural ODE formulation (Song et al., 2021a) is not investigated, but greatly of interest given the comparison with the flow-based method in this work
Revised: Lastly, the connection between diffusion models and continuous normalizing flows through the neural ODE formulation (Song et al., 2021a) is not investigated, but greatly of interest given the comparison with the flow-based method in this work.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Revised: All code used to train and evaluate the models as presented in this paper can be found at https: //anonymous.4open.science/r/iclr2023-joint-diffusion .

Original: The codebase in https://github.com/yang-song/score_sde_pytorch of Song et al.
Revised: Essentially, the codebase in https://github.com/yang-song/score_sde_pytorch of Song et al.

Original: (2020) can be used to train the score-based diffusion networks.
Revised: (2020) is used to train the score-based diffusion networks, for both data and structured noise, independently.

Original: During inference, the lines in Algorithm 1 should be adapted to adjust the sampler.
Revised: To implement the proposed inference scheme, the lines in Algorithm 1 should be adapted to create a sampler that includes both trained diffusion models.

Original: Details regarding the training and inference settings used to reproduce the results in this work can be found in Section 3.2.
Revised: Details regarding the training and inference settings used to reproduce the results in this work can be found in Section 3.2.

Original: A A DDITIONAL E XPERIMENTS
Revised: A A DDITIONAL E XPERIMENTS

Original: The following section explores additional inverse problems with compressed sensing and structured noise.
Revised: The following section explores additional inverse problems with compressed sensing and structured noise.

Original: The goal is to show the performance of the proposed method in a variety of settings.
Revised: The goal is to show the performance of the proposed method in a variety of settings.

Original: A. S TRUCTURED N OISE WITH C OMPRESSED S ENSING
Revised: A. S TRUCTURED N OISE WITH C OMPRESSED S ENSING

Original: The corruption process is defined by y
Revised: The corruption process is defined by y

Original: =
Revised: =

Original: Ax + n sine with a random Gaussian measurement matrix
Revised: Ax + n sine with a random Gaussian measurement matrix

Original: A ∈ R m × d and a noise with sinusoidal variance σ k ∝ exp (cid:0) sin (cid:0) 2 πk 16 (cid:1)(cid:1) for each pixel k .
Revised: A ∈ R m × d and a noise with sinusoidal variance σ k ∝ exp (cid:0) sin (cid:0) 2 πk 16 (cid:1)(cid:1) for each pixel k .

Original: The subsampling factor is defined by the size of the measurement matrix d/m .
Revised: The subsampling factor is defined by the size of the measurement matrix d/m .

Revised: In this experiment, the score network s ϕ is trained on a dataset generated with sinusoidal noise samples n sine .

Original: In Fig.
Revised: In Fig.

Original: 5 the results of the compressed sensing experiment and the comparison with the baselines are shown for an average standard deviation of σ N = 0 . 2 and subsampling of factor d/m = 2 .
Revised: 5 the results of the compressed sensing experiment and the comparison with the baselines are shown for an average standard deviation of σ N = 0 . 2 and subsampling of factor d/m = 2 .

Original: Given the same hyperparameter settings, we repeat the experiment on the out-of-distribution (OoD) dataset, shown in Fig.
Revised: Given the same hyperparameter settings, we repeat the experiment on the out-of-distribution (OoD) dataset, shown in Fig.

Original: Similar to the results found in Section 4.2.1, the diffusion method is more robust to the shift in distribution and is able to deliver high quality recovery under the structured noise setting.
Revised: Similar to the results found in Section 4.2.1, the diffusion method is more robust to the shift in distribution and is able to deliver high quality recovery under the structured noise setting.

Original: In contrast, the flow-based method under-performs when subjected to the OoD data.
Revised: In contrast, the flow-based method under-performs when subjected to the OoD data.

Original: Quantitative results on both CelebA and OoD are found in Fig.
Revised: Quantitative results on both CelebA and OoD are found in Fig.

Original: 4 as well as Table 1 and Table 2, respectively.
Revised: 4 as well as Table 1 and Table 2, respectively.

Original: A. R EMOVING S INUSOIDAL N OISE
Revised: A. R EMOVING S INUSOIDAL N OISE

Original: The corruption process is defined by y
Revised: The corruption process is defined by y

Original: = x + n sine where the noise variance σ k ∝ exp (cid:0) sin (cid:0) 2 πk 16 (cid:1)(cid:1) follows a sinusoidal pattern along each row of the image
Revised: = x + n sine where the noise variance σ k ∝ exp (cid:0) sin (cid:0) 2 πk 16 (cid:1)(cid:1) follows a sinusoidal pattern along each row of the image

Original: k .
Revised: k .

Revised: In this experiment, the score network s ϕ is trained on a dataset generated with 1D sinusoidal noise samples n sine .

Original: See Fig.
Revised: See Fig.

Original: 8 for a comparison of our method to the flow-based method for varying noise variances.
Revised: 8 for a comparison of our method to the flow-based method for varying noise variances.

Original: Both methods perform quite well, with the diffusion method having a slight edge.
Revised: Both methods perform quite well, with the diffusion method having a slight edge.

Original: A visual comparison in Fig.
Revised: A visual comparison in Fig.

Original: 7, however, reveals that the diffusion method preserves more detail in general.
Revised: 7, however, reveals that the diffusion method preserves more detail in general.

Original: G r ound T r u t h
Revised: G r ound T r u t h

Original: D I FF
Revised: D I FF

Original: U S I
Revised: U S I

Original: ON
Revised: ON

Original: F L O
Revised: F L O

Original: W
Revised: W

Original: GAN
Revised: GAN

Original: L A SS
Revised: L A SS

Original: O
Revised: O

Original: G r ound T r u t h
Revised: G r ound T r u t h

Original: D I FF
Revised: D I FF

Original: U S I
Revised: U S I

Original: ON
Revised: ON

Original: F L O
Revised: F L O

Original: W
Revised: W

Original: GA N
Revised: GA N

Original: L A SS
Revised: L A SS

Original: O
Revised: O

Original: [Figure]
Revised: [Figure]

Original: ound T r u t h
Revised: ound T r u t h

Original: N o i s y
Revised: N o i s y

Original: I npu
Revised: I npu

Original: t
Revised: t

Original: D I FF
Revised: D I FF

Original: U S I
Revised: U S I

Original: ON
Revised: ON

Original: F L O
Revised: F L O

Original: W
Revised: W

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: B E XTENDED R ESULTS B. M ETRICS
Revised: B E XTENDED R ESULTS B. M ETRICS

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: C OMPUTATIONAL PERFORMANCE

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: C OMPUTATIONAL PERFORMANCE
Revised: C P SEUDO -C ODE

Original: [Table]
Revised: In this section, we provide pseudo-code for the proposed joint conditional diffusion sampler with the Euler-Maruyama sampling algorithm as basis. Furthermore, we use the SDE formulation for the diffusion process which is denoted as an sde object with drift , diffusion and marginal_prob methods. The latter computes the mean and standard deviation of the diffusion transition kernel at a certain time t . Lastly, there are two trained score networks (NCSNv2) score_data and score_noise for the data and structured noise respectively.

Original: [Table]
Revised: def [Equation] for t in linspace(1, 0, num_steps): # corrupt observation along the diffusion process [Equation]

Revised: # data consistency step for

Revised: [Equation]

Revised: data

Revised: [Equation]

Revised: hat

Revised: [Equation]

Revised: n

Revised: [Equation]

Revised: # data consistency step for n (noise) n

Revised: [Equation]

Revised: n

Revised: [Equation]

Revised: mu_coeff

Revised: [Equation]

Revised: n

Revised: [Equation]

Revised: y_hat + A @ x) # reverse diffusion step for x (data)

Revised: z

Revised: [Equation]

Revised: random.normal(x.shape) x_hat

Original: [Table]
Revised: = x [Equation] sde.drift(t) * x * dt x_hat

Revised: = x_hat + sde.diffusion(t)**2 * score_data(x, t)

Revised: * dt x_hat

Revised: = x_hat + sde.diffusion(t)

Revised: * sqrt(dt) *

Revised: z x = x_hat

Revised: # reverse diffusion step for n (noise)

Revised: z = random.normal(n.shape) n_hat = n - sde.drift(t)

Revised: * n * dt n_hat = n_hat + sde.diffusion(t)**2 * score_noise(n, t) * dt n_hat = n_hat + sde.diffusion(t)

Revised: * sqrt(dt) *

Revised: z n