Inject Noise to Remove Noise: A Deep Dive into Score-Based Generative Modeling Techniques

A look at the recent Score-Based Generative Modeling through Stochastic Differential Equations paper by Yang Song et al.
Sayantan Das
Created on April 8|Last edited on December 10
Comment
﻿
﻿Try it in Colab﻿
ContentsWhy Score Based?
Score Estimation and Score Matching
Sample Generation and Langevin Dynamics
Gaussian Perturbations of data
Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021)
Experiments
Conclusions
Appendix
1. Why Score Based?
What is a Score?Before we get into score-based modeling, it's important to first understand what a "score" is in this context. Given a probability density function p(x)p(x)p(x)﻿ the 'score' is defined as 
						∇xlog⁡p(x),\nabla_\mathbf{x} \log p(\mathbf{x}),∇x​logp(x),﻿﻿
or the gradient of the log-likelihood of the object x with respect to the input dimensions x. In this case, the object x is input dimensions. 
Instead of working with a probability density function (pdf) ﻿, we work with the gradients and the gradients here are with respect to the input dimensions and notably not w.r.t the model parameters θ.\theta.θ.﻿﻿
In this report, we will assume that pdfs are continuous random variables.
The score is a vector field of the gradient at any point x. This gradient of log⁡p(x)\log p(x) logp(x)﻿ tells us the directions in which to move if we want to increase the likelihood as much as possible.
Score-based generative models are trained to estimate ∇xlog⁡p(x).\nabla_\mathbf{x} \log p(\mathbf{x}).  ∇x​logp(x).﻿ 
Why use Scores?Unlike Likelihood-based models like Normalizing Flows or Autoregressive models, these models are easier to parameterize and do not need to be normalized.
Moreover, ∇xlog⁡p(x)\nabla_\mathbf{x} \log p(\mathbf{x})∇x​logp(x)﻿ behaves like an unconstrained function so modeling it is easier.
Success of Score-Based techniques with Deep Energy-Based ModelsIn EBMs, we treat pθ(x)p_\theta(x)pθ​(x)﻿ as a function 
pθ(x)=e−fθ(x)Zθp_\theta(x) = \frac{e^{-f_\theta(x)}}{Z_\theta}pθ​(x)=Zθ​e−fθ​(x)​﻿
where fθ(x)f_\theta(x)  fθ​(x)﻿ is an energy function parameterized by θ\thetaθ﻿ and ZθZ_\thetaZθ​﻿ is the normalization constant aka partition function, which is introduced here for the same reasons.
	Learning Parameters via Maximum Likelihood EstimationEpdata[−log⁡pθ(x)]=Epdata[log⁡fθ(x)−log⁡Zθ]\mathbb{E}_{p_{data}}[-\log p_\theta (x)] = \mathbb{E}_{p_{data}}[\log f_\theta (x) - \log Z_\theta] Epdata​​[−logpθ​(x)]=Epdata​​[logfθ​(x)−logZθ​]﻿
where log⁡Zθ\log Z_\thetalogZθ​﻿ is intractable.
        The Score does not depend on the partition function ZθZ_\thetaZθ​﻿﻿∇xlog⁡p(x)=−∇xlog⁡f(x)−∇xlog⁡Zθ\nabla_\mathbf{x} \log p(\mathbf{x}) = - \nabla_\mathbf{x} \log f(\mathbf{x}) - \sout{ \nabla_\mathbf{x} \log Z_\theta}∇x​logp(x)=−∇x​logf(x)−∇x​logZθ​​﻿
Idea:  Learn θ\thetaθ﻿ by fitting ∇xlog⁡pθ(x)\nabla_\mathbf{x} \log p_\theta(\mathbf{x})∇x​logpθ​(x)﻿ to ∇xlog⁡pdata(x)\nabla_\mathbf{x} \log p_{data}(\mathbf{x})∇x​logpdata​(x)﻿.
This is also called score estimation.
CelebA sampling. Extracted from https://github.com/ermongroup/ncsn﻿
CIFAR10 sampling. Extracted from https://github.com/ermongroup/ncsn﻿
2. Score Estimation and Score Matching
                     Left: ∇xlog⁡pdata(x)\nabla_x \log p_{data} (x)∇x​logpdata​(x)﻿ , Right: Sθ(x)S_\theta(x)Sθ​(x)﻿. The data density is encoded using an orange colormap: darker colour implies higher density. Red rectangles highlight regions where model scores are close to the data scores. Source: https://arxiv.org/abs/1907.05600 
The caveat to all of this is : we don't have pdatap_{data}pdata​﻿. In this article,
we assume that  
{x1,x2,...xN}∼iidpdata(x)=p(x).\{{x1,x2,...x_N}\}  
\stackrel{iid}{\sim} p_{data}(x) = p(x).{x1,x2,...xN​}∼iidpdata​(x)=p(x).﻿
Task: Estimate the score ∇xlog⁡p(x)\nabla_x \log p(x)∇x​logp(x)﻿﻿
Score Model: A vector-valued function Sθ(x):RD→RDS_\theta(x) : \R^D \rightarrow \R^DSθ​(x):RD→RD﻿.
Objective: How to compare two vector fields of scores?
           		
		Score MatchingAverage Euclidean distance over the whole space.
﻿12Ep(x)[∥∇xlog⁡p(x)−Sθ(x)∥22]\frac{1}{2} \mathbb{E}_{p(x)}[\lVert\nabla_x \log p(x) - S_\theta(x)\rVert_2^2]21​Ep(x)​[∥∇x​logp(x)−Sθ​(x)∥22​]﻿ which is also referred to as Fisher Divergence.
This expression is simplified using Integration By Parts method which gives us
﻿Ep(x)[12∥Sθ(x)∥22+tr(∇xSθ(x))]\mathbb{E}_{p(x)}[\frac{1}{2}\lVert S_\theta(x)\rVert_2^2 + \mathrm{tr}(\nabla_x S_\theta(x))]Ep(x)​[21​∥Sθ​(x)∥22​+tr(∇x​Sθ​(x))]﻿﻿
This was derived by Aapo Hyvarinen in 2005. On further simplification by lifting the Expectation terms, we obtain
﻿≈12∑i=1N[12∥Sθ(xi)∥22+tr(∇xSθ(xi))]\approx \frac{1}{2} \sum_{i=1}^{N}[\frac{1}{2}\lVert S_\theta(x_i)\rVert_2^2 + \mathrm{tr}(\nabla_x S_\theta(x_i))]≈21​∑i=1N​[21​∥Sθ​(xi​)∥22​+tr(∇x​Sθ​(xi​))]﻿,
﻿{x1,x2,...,xN}∼iidp(x)\{x_1,x_2,...,x_N\} \stackrel{iid}{\sim} p(x){x1​,x2​,...,xN​}∼iidp(x)﻿.
Some takeaways:
Trace of the jacobian of the score is used in order to make local maxima.
The score objective says "Choose θ\thetaθ﻿ such that every data point in the training set becomes a local maxima of our estimated density".
3. Sample Generation and Langevin DynamicsVanilla Score Matching has its drawbacks:
"Not Scalable"
  	Let's take an example of a simple score model that computes ∥Sθ(x)∥22\lVert S_\theta (x)\rVert^2_2∥Sθ​(x)∥22​﻿ on forward pass.
								Score  Matching is not Scalable
The second of our score matching objective requires O(D) backpropagations which is expensive in practice. 
A remedy that the community has come up with: Sliced Score Matching
Project the vectors onto random directions and attach an Expectation over these directions over the original objective.
This is cheaper in practice than the vanilla formulation. This sliced fisher divergence is as follows:
12EpvEpdata[(vT∇xlog⁡pθ(x)−vT∇xlog⁡pdata(x))2],\frac{1}{2}\mathbb{E}_{p_v}\mathbb{E}_{p_{data}}[(v^T\nabla_x \log p_\theta(x)-v^T\nabla_x \log p_{data}(x))^2  ],21​Epv​​Epdata​​[(vT∇x​logpθ​(x)−vT∇x​logpdata​(x))2],﻿
where v is a random direction based on some distribution pvp_vpv​﻿. It has been found that Multivariate Rademacher Distribution is a good candidate for this distribution other than Multivariate Standard Normal.
Langevin DynamicsIn computational statistics and recently in generative modeling, Langevin sampling has had great success. Langevin Monte Carlo is a Markov Chain Monte Carlo (MCMC) method for obtaining random samples from probability distributions for which direct sampling is difficult. The goal is to "follow the gradient but add a bit of noise" so as to not get stuck at the local optima regions and thus we are able to explore the distribution and sample from it.
	ProcedureSample from p(x)p(x)p(x)﻿ using only the score ∇xlog⁡p(x).\nabla_x \log p(x).∇x​logp(x).﻿﻿
Initialize x0∼∼π(x).\stackrel{\sim}{x_0} \sim \pi(x).x0​∼​∼π(x).﻿﻿
Repeat for t←1,2,...,T:t \leftarrow 1,2,...,T:t←1,2,...,T:﻿﻿
﻿zt∼N(0,1)z_t \sim \mathcal{N}(0,1)zt​∼N(0,1)﻿﻿
﻿xt∼←x∼t−1+ϵ2∇xlog⁡p(x∼t−1)+ϵzt\stackrel{\sim}{x_t}\leftarrow\stackrel{\sim}{x}_{t-1} + \frac{\epsilon}{2} \nabla_x \log p(\stackrel{\sim}{x}_{t-1}) + \sqrt[]{\epsilon}z_txt​∼​←x∼t−1​+2ϵ​∇x​logp(x∼t−1​)+ϵ​zt​﻿﻿
						gradient                gaussian noise
Score Based Workflow﻿
				Created by Sayantan Das
4. Gaussian Perturbations of dataScore Matching has various pitfalls:
Pitfall #1: Manifold Hypothesis
This hypothesis states that high dimensional data often tends to lie in a low dimensional manifold -- which makes ∇xlog⁡pdata(x)\nabla_x \log p_{data}(x)∇x​logpdata​(x)﻿ undefined in some regions.
Pitfall #2: Inaccurate Score Estimation in Low Data-Density regions
When most of the samples will be around the modes of the distribution, the vector field can be misled around low density regions where not enough samples can guide the correct direction to the score vector.
Pitfall #3: Slow Mixing of Langevin Dynamics between data modes
When the distribution has disconnected supports -- or is a mixture of two disjoint components with a weighting coefficient π\piπ﻿, scores cannot recover this coefficient and is invariant towards mode weights. This is a failure mode of Langevin Dynamics as the density is not supported over the whole space.
Solution - Multiple levels of gaussian perturbations of the dataThis is also known as "Annealed Langevin Dynamics"
 Sample using σ1,σ2,...,σL\sigma_1,\sigma_2,...,\sigma_Lσ1​,σ2​,...,σL​﻿ sequentially with Langevin Dynamics(LD) where,
		σ1>σ2>...>σL\sigma_1>\sigma_2>...>\sigma_Lσ1​>σ2​>...>σL​﻿﻿
Run LD1LD_1LD1​﻿ with σ1\sigma_1σ1​﻿﻿
Run LD2LD_2LD2​﻿ with particles from the previous LD run i.e LD1LD_1LD1​﻿and σ2\sigma_2σ2​﻿﻿
and so on.
This was presented in the paper - Noise Conditional Score Networks (NeurIPS 2019).
5. Score-Based Generative Modeling through Stochastic Differential Equations﻿Link to Paper →﻿﻿﻿
﻿Link to Code →﻿﻿﻿
﻿Link to Colab Notebook →﻿
Introduction     As we saw in the previous section,  gaussian perturbations corrupt the data into random noise. In order to generate samples with score-based models, we need to consider a diffusion process. Scores arise when we reverse this diffusion process leading to sample generation. Let x(t)x(t)x(t)﻿ be a diffusion process, indexed by a continuous random variable of time t:0,...,Tt:0,...,Tt:0,...,T﻿. A diffusion process is governed by a stochastic differential equation (SDE), in the following form
dx=f(x,t)dt+g(t)dw,d\textbf{x} = \textbf{f}(x,t)dt + g(t)d\textbf{w},dx=f(x,t)dt+g(t)dw,﻿
where f(.,t)\textbf{f}(.,t)f(.,t)﻿ is the drift coefficient﻿ of the SDE , g(t) is the diffusion coefficient, w represents the standard brownian motion. It is this w that makes SDEs the stochastic generalization of an Ordinary Differential Equation (ODE) -- since the particles not only follow the deterministic drift guided by f(x,t), but are also affected by the random noise coming from g(t)dw. 
Reversing the SDEFor score based generative modeling, the diffusion process needs to be bounded i.e. x(0)∼p0x(0) \sim p_0x(0)∼p0​﻿ and x(T)∼pTx(T) \sim p_Tx(T)∼pT​﻿ where pt(x)p_t(x)pt​(x)﻿ is used to denote the distribution for x(t)x(t)x(t)﻿. Here  p0p_0p0​﻿ is the data distribution where we have a dataset of i.i.d. samples, and pTp_TpT​﻿ is the prior distribution that has a tractable form and easy to sample from. The noise perturbation by the diffusion process is large enough to ensure  pTp_TpT​﻿ does not depend on  p0p_0p0​﻿ .
By starting out with a sample from pTp_TpT​﻿ and reversing this diffusion SDE, we can obtain samples from our data distribution p0.p_0.p0​.﻿ The reverse-time SDE is given as:
dx=[f(x,t)−g2(t)∇xlog⁡pt(x)]dt+g(t)dwˉd\textbf{x} = [\textbf{f}(\textbf{x},t) - g^2(t)\nabla_x \log p_t(x)]dt + g(t)d\bar{\textbf{w}}dx=[f(x,t)−g2(t)∇x​logpt​(x)]dt+g(t)dwˉ﻿
where wˉ\bar{w}wˉ﻿ is the brownian motion in the reverse time direction and dt is the infinitesimal negative time step.
									Extracted from the paper.
Score EstimationA time-dependent score function sθ(x,t)s_\theta(\textbf{x},t)sθ​(x,t)﻿ is required to approximate ∇xlog⁡pt(x)\nabla_x \log p_t(x)∇x​logpt​(x)﻿ in order to numerically solve the reverse-time SDE. This score model is trained using denoising objective:
min⁡θEt∼U(0,T)[λ(t)Ex(0)∼p0(x)Ex(t)∼p0t(x(t)∣x(0))[∥sθ(x(t),t)−∇x(t)log⁡p0t(x(t)∣x(0))∥22]]\min_\theta \mathbb{E}_{t\sim \mathcal{U}(0, T)} [\lambda(t) \mathbb{E}_{\mathbf{x}(0) \sim p_0(\mathbf{x})}\mathbf{E}_{\mathbf{x}(t) \sim p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0))}[ \|s_\theta(\mathbf{x}(t), t) - \nabla_{\mathbf{x}(t)}\log p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0))\|_2^2]]minθ​Et∼U(0,T)​[λ(t)Ex(0)∼p0​(x)​Ex(t)∼p0t​(x(t)∣x(0))​[∥sθ​(x(t),t)−∇x(t)​logp0t​(x(t)∣x(0))∥22​]]﻿
where U(0,T)\mathcal{U}(0,T)U(0,T)﻿ is a uniform distribution over [0,T][0, T][0,T]﻿, p0t(x(t)∣x(0))p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0))p0t​(x(t)∣x(0))﻿ denotes the transition probability from x(0)\mathbf{x}(0)x(0)﻿ to x(t)\mathbf{x}(t)x(t)﻿, and λ(t)∈R>0\lambda(t) \in \mathbb{R}_{>0}λ(t)∈R>0​﻿ denotes a positive weighting function.
In the objective, the expectation over x(0)\mathbf{x}(0)x(0)﻿ can be estimated with empirical means over data samples from p0p_0p0​﻿. The expectation over x(t)\mathbf{x}(t)x(t)﻿ can be estimated by sampling from p0t(x(t)∣x(0))p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0))p0t​(x(t)∣x(0))﻿, which is efficient when the drift coefficient f(x,t)\mathbf{f}(\mathbf{x}, t)f(x,t)﻿ is affine. The weight function λ(t)\lambda(t)λ(t)﻿ is typically chosen to be inverse proportional to E[∥∇xlog⁡p0t(x(t)∣x(0))∥22]\mathbb{E}[\|\nabla_{\mathbf{x}}\log p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0)) \|_2^2]E[∥∇x​logp0t​(x(t)∣x(0))∥22​]﻿.
Sampling Procedures in the paperThe authors discuss various sampling methods in their paper which are as follows:
Sampling using Numerical Solvers -- such as the Euler-Maruyama approach where dt is approximated as Δt\Delta tΔt﻿.
Sampling using Predictor-Corrector Methods : These methods use one step of Numerical Solvers like the one discussed above to obtain xt−Δt\textbf{x}_{t-\Delta t}xt−Δt​﻿ from xt\textbf{x}_{t}xt​﻿ which is called the 'predictor' step. This is followed by applying several steps of Langevin MCMC to refine xt\textbf{x}_{t}xt​﻿ such that it becomes a more accurate sample from pt−Δt(x)p_{t-\Delta t}(x)pt−Δt​(x)﻿. This is the 'corrector' step as the MCMC helps reduce the error of the numerical SDE solver.
	Reason this is popular is that Score-based MCMC approaches can produce samples from underlying distributions when ∇xlog⁡p(x)\nabla_x \log p(x)∇x​logp(x)﻿ is known -- which in our case is approximated well by the time-dependent score-model.
3.  Sampling with Numerical ODE Solvers: For any SDE of the form
﻿dx=f(x,t)dt+g(t)dw,d \mathbf{x} = \mathbf{f}(\mathbf{x}, t) d t + g(t) d \mathbf{w},dx=f(x,t)dt+g(t)dw,﻿﻿
       there exists an associated ordinary differential equation (ODE)
﻿dx=[f(x,t)−12g(t)2∇xlog⁡pt(x)]dtd \mathbf{x} = \bigg[\mathbf{f}(\mathbf{x}, t) - \frac{1}{2}g(t)^2 \nabla_\mathbf{x} \log p_t(\mathbf{x})\bigg] dtdx=[f(x,t)−21​g(t)2∇x​logpt​(x)]dt﻿﻿
such that their trajectories have the same mariginal probability density  𝑝𝑡(𝐱) . Therefore, by solving this ODE in the reverse time direction, we can sample from the same distribution as solving the reverse-time SDE. We call this ODE the probability flow ODE.
This approach is also very popular since when the formulation turns into an ODE there is a possibility to estimate the true likelihood of data or in our case p0(x)p_0(\textbf{x})p0​(x)﻿, using the change of variable formula. More on this in the ICLR 21 paper.
6. ExperimentsWe reproduce experiments for Controllable Generation from the score_sde repository.
More details about them can be found in Appendix 1.2 in the ICLR paper.
﻿
Run set1
﻿
7. ConclusionThis report presents and summarizes the latest developments in score-based generative models -- with a goal to enable better understanding of existing approaches, new sampling algorithms, exact likelihood computation and conditional generation abilities to the family of score based generative models.﻿﻿
8. Appendix﻿