Noise Conditional Score Network

NCSN
$$p_{\sigma }(\tilde{\mathrm{x}}|\mathrm{x}):=\mathcal{N}(\tilde{\mathrm{x}};\mathrm{x},{\sigma }^2\mathbf{I})$$
$$p_{\sigma }(\tilde{\mathrm{x}}):=\int p_{data}(\mathrm{x})p_{\sigma }(\tilde{\mathrm{x}}|\mathrm{x})d\mathrm{x}$$
$p_{data}(x)$表示目标数据分布。 ${\sigma }_{\min }={\sigma }_1<{\sigma }_2<\cdot \cdot \cdot <{\sigma }_N={\sigma }_{\max }$
${\sigma }_{\min }$足够小，以至于$p_{{\sigma }_{\min }}(\mathrm{x})\approx p_{data}(\mathrm{x}))$, ${\sigma }_{\max }$足够大，以至于$p_{{\sigma }_{\min }}(\mathrm{x}))\approx \mathcal{N}(\mathbf{x};0,{\sigma }_{\max }^2\mathbf{I})$

$${\theta }^{\ast }=\underset{\theta }{\mathrm{argmin}}\sum _{i=1}^N{\sigma }_i^2{\mathbb{E}}_{\mathrm{x}\sim p_{data}(\mathrm{x})}{\mathbb{E}}_{\tilde{\mathrm{x}}\sim p_{{\sigma }_i}(\tilde{\mathrm{x}}|\mathrm{x})}[||s_{\theta }(\tilde{\mathrm{x}},{\sigma }_i)-{▽}_{\tilde{\mathrm{x}}}\log p_{\tilde{\sigma }}(\tilde{\mathrm{x}}|\mathrm{x})|{|}_2^2]$$

模型训练完毕后，执行M步Langevin MCMC采样：
$$x_i^m=x_i^{m-1}+{ϵ}_i{s}_{{\theta }^{\ast }}(x_i^{m-1},{\sigma }_i)+\sqrt{2{ϵ}_i}{z}_i^m,m=1,2,\cdot \cdot \cdot ,M$$
${ϵ}_i>0$为步长，$z_i^m$是标准正态分布。上述采样过程重复$i=N,N-1,\cdot \cdot \cdot ,1$，也就是说对于每个noise level，执行N步，直至样本收敛至当前noise level的最佳位置。
$$x_N^0\sim \mathcal{N}(\mathrm{x}|0,{\sigma }_{max}^2\mathbf{I}),x_i^0=x_{i+1}^M\mathrm{when}i<N$$

样例代码如下：

import torch
def langevin_sampling(score_network, noise_levels, num_steps, step_size, batch_size, device):
	# 初始化样本，从标准正态分布中采样
	x = torch.randn(batch_size, 3, 32, 32).to(device)
	# 多级噪声采样，从高噪声到低噪声逐步采样，最终得到逼近目标分布的样本
	for sigma in noise_levels:
		print(f"Sampling at noise level: {sigma}")
		# Langevin动力学迭代
		for _ in range(num_steps):
			# 计算分布函数梯度（由分数网络预测）
			with torch.no_grad():
				grad = score_network(x, sigma) # 输入当前样本和噪声水平
			
			# 梯度上升步
			x = x + step_size * grad
			# 添加随机噪声步
			noise = torch.randn_like(x) * (2 * step_size) ** 0.5
			x = x + noise
		
		return x
				
			
class DummyScoreNet(nn.Module):
	def forward(self, x, sigma):
		return -x / (sigma ** 2)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
score_network = DummyScoreNet().to(device)

noise_levels = [50.0, 25.0, 10.0, 5.0, 1.0]
num_steps = 50
step_size = 0.1
batch_size = 64

samples = langevin_sampling(score_network,
							noise_levels,
							num_steps,
							step_size,
							batch_size,
							device)