变分扩散模型.

(1)前向时间扩散过程

定义一个高斯扩散过程,它从输入数据\(\mathbf{x}_0\)开始,并定义了一系列噪声越来越大的\(\mathbf{x}_t\),其中$t$从$t=0$(噪声最小)到$t=1$(噪声最大)。对于任何$t∈[0, 1]$,以\(\mathbf{x}_0\)为条件的隐变量\(\mathbf{x}_t\)的分布由下式给出:

\[q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}(\alpha_t \mathbf{x}_0, \sigma_t^2 \mathbf{I})\]

定义信噪比(signal-to-noise ratio, SNR):

\[\text{SNR}(t) = \frac{\alpha_t^2}{\sigma_t^2}\]

信噪比是时间$t$的单调递减函数。在实验中设置$\alpha_t = \sqrt{1-\sigma_t^2}$。

(2)噪声策略

噪声$\sigma_t$参数化为:

\[\sigma_t^2 = \text{sigmoid}(\gamma_{\eta}(t))\]

其中$\gamma(\cdot)$是一个单调神经网络。进一步有:

\[\alpha_t^2 = 1-\sigma_t^2 = \text{sigmoid}(-\gamma_{\eta}(t)) \\ \text{SNR}(t) = \frac{\alpha_t^2}{\sigma_t^2} = \exp(-\gamma_{\eta}(t))\]

(3)反向时间生成模型

反向扩散过程一个层次化生成模型,该模型对一系列隐变量\(\mathbf{x}_t\)进行采样,时间从$t=1$反向运行到$t=0$。

\[p(\mathbf{x}) = \int_{\mathbf{x}} p(\mathbf{x}_1)p(\mathbf{x} \mid \mathbf{x}_0) \prod_{i=1}^T p(\mathbf{x}_{(i-1)/T}\mid \mathbf{x}_{i/T})\]

参数化一个噪声预测模型\(\hat{\mathbf{\epsilon}}_{\theta}(\mathbf{x}_t;t)\),使得:

\[\hat{\mathbf{x}}_{\theta}(\mathbf{x}_t;t) = \frac{\mathbf{x}_t-\sigma_t\hat{\mathbf{\epsilon}}_{\theta}(\mathbf{x}_t;t)}{\alpha_t}\]

为了捕捉数据的精细尺度细节,作者建议在噪声预测模型的输入中添加一组模糊特征。设\(\mathbf{x}_0\)为原始数据,缩放到范围$[-1,1]$,设\(\mathbf{x}_t\)为隐变量,为其附加通道\(\sin(2^n\pi \mathbf{x})\)和\(\cos(2^n\pi \mathbf{x})\),其中$n$取整数范围\(\{n_{\min},...,n_{\max}\}\)。这些特征能够放大输入数据\(\mathbf{x}_t\)中的小变化的高频周期函数,在去噪模型的输入中包含这些特征会大大提高似然性。

损失函数通过分布似然的变分下界构造:

\[-\log p(\mathbf{x}) \leq \underbrace{D_{KL}(q(\mathbf{x}_1\mid \mathbf{x})||p(\mathbf{x}_1))}_{\text{prior loss}} + \underbrace{\mathbb{E}_{q(\mathbf{x}_0\mid \mathbf{x})}[-\log p(\mathbf{x}\mid \mathbf{x}_0)]}_{\text{reconstruction loss}} + \underbrace{\mathcal{L}_T(\mathbf{x})}_{\text{diffusion loss}}\]

⚪ 离散时间模型

对于离散时间模型,损失\(\mathcal{L}_T(\mathbf{x})\)构造为:

\[\begin{aligned} \mathcal{L}_T(\mathbf{x}) &= \frac{T}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I}),i\sim U[1,T]} [(\text{SNR}((i-1)/T)-\text{SNR}(i/T)) \left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_t;t)\right\|_2^2] \\ &= \frac{T}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I}),i\sim U[1,T]} [(\exp(\gamma_{\eta}(i/T)-\gamma_{\eta}((i-1)/T))-1 )\left\| \mathbf{\epsilon}-\hat{\mathbf{\epsilon}}_{\theta}(\mathbf{x}_t;t)\right\|_2^2] \end{aligned}\]

⚪ 连续时间模型

采取更多的时间步长会产生更好的似然结果,不妨采取$T→ ∞$, 有效地将时间$t$视为连续的而不是离散的。

\[\begin{aligned} \mathcal{L}_\infty(\mathbf{x}) &= -\frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I})} \int_0^1 \text{SNR}'(t)\left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_t;t)\right\|_2^2dt \\ &= -\frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I}),t\sim \mathcal{U}(0,1)} [\text{SNR}'(t)\left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_t;t)\right\|_2^2]\\ &= \frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I}),t\sim \mathcal{U}(0,1)} [\gamma_{\eta}'(t)\left\| \mathbf{\epsilon}-\hat{\mathbf{\epsilon}}_{\theta}(\mathbf{x}_t;t)\right\|_2^2] \end{aligned}\]

根据单调性假设,信噪比函数\(\text{SNR}(t)\)是可逆的,记\(v=\text{SNR}(t)\),则有:

\[\begin{aligned} \mathcal{L}_\infty(\mathbf{x}) &= \frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I})} \int_{\text{SNR}(1)}^{\text{SNR}(0)} \left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_v;v)\right\|_2^2dv \\ &= \frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I})} \int_{\text{SNR}_{\min}}^{\text{SNR}_{\max}} \left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_v;v)\right\|_2^2dv \\ \end{aligned}\]

进一步在损失中引入加权系数$w$:

\[\begin{aligned} \mathcal{L}_\infty(\mathbf{x},w) &= \frac{1}{2} \mathbb{E}_{\mathbf{\epsilon}\sim \mathcal{N}(0, \mathbf{I})} \int_{\text{SNR}_{\min}}^{\text{SNR}_{\max}} w(v)\left\| \mathbf{x}-\hat{\mathbf{x}}_{\theta}(\mathbf{x}_v;v)\right\|_2^2dv \\ \end{aligned}\]