🌉

[23CVPR] BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models

Diffusion Series

Diffusion Theory Summary

Image-to-image Translation

EGSDE: Unpaired I2I Translation via Energy-Guided SDEs.

SDDM: Score-Decomposed Diffusion Models on Manifolds for Unpaired I2I Translation

Schrodinger Bridge Series

[0] Schrödinger Bridge Overview

[21NeurIPS] DSB: Schrödinger Bridge

[23CVPR] BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models

BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models

Bo Li, Kaitao Xue, Bin Liu (Nanchang Hangkong University, CN)
[paper][code]

Intro & Overview

Prelinamaries: DDPMs

DDPM에 대한 간략한 요약을 하며 시작하겠습니다.

DDPM의 Forward Process와 Reverse process 그리고 Training Objective는 다음과 같이 정의됩니다.

DDPM: Forward Process

•

Forward Process는 Data에 노이즈를 점진적으로 더해가는 과정입니다.

•

Data x0∼qdata(x0)\boldsymbol{x}_0 \sim q_{d a t a}\left(\boldsymbol{x}_0\right)x0​∼qdata​(x0​) 에서 시작합니다.

•

최종 Latent variable xT\boldsymbol x_TxT​이며  platent (xT)=N(0,I)p_{\text {latent }}\left(\boldsymbol{x}_T\right)=\mathcal{N}(\mathbf{0}, \boldsymbol{I})platent ​(xT​)=N(0,I)인 isotropic Gaussian distribution입니다.

•

x0\boldsymbol x_0x0​에서 xT\boldsymbol x_TxT​로 변환하는 과정은 Markov chain으로 정의됩니다.

\begin{equation} q\left(\boldsymbol{x}_1, \ldots, \boldsymbol{x}_T \mid \boldsymbol{x}_0\right)=\prod_{t=1}^T q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}\right) \end{equation}

•

βt\beta_tβt​는 아주 작은 상수.

DDPM: Reverse Process.

•

Reverse Process는 Latent variable xT\boldsymbol x_TxT​로부터 데이터 x0\boldsymbol x_0x0​를 추론하는 과정입니다.

•

이는 다음과 같은 Markov chain으로 정의됩니다.

p_\theta\left(\boldsymbol{x}_0, \ldots, \boldsymbol{x}_{T-1} \mid \boldsymbol{x}_T\right)=\prod_{t=1}^T p_\theta\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)

DDPM: Training Objective

•

DDPM은 Evidence Lower Bound (ELBO)를 이용하여 최적화됩니다.

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{\epsilon}}\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)\right\|_2^2\end{equation}

•

ϵ\boldsymbol{\epsilon}ϵ: xt\boldsymbol x_txt​ 의 가우시안 노이즈. ∇xtln⁡q(xt∣x0)\nabla_{\boldsymbol{x}_t} \ln q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0\right)∇xt​​lnq(xt​∣x0​).

•

ϵθ\boldsymbol{\epsilon}_\thetaϵθ​: denoising model.

Conditional Diffusion.

•

대부분의 conditional diffusion model들은 condition을  Objective에 직접적으로 주입합니다.

\begin{equation} \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{\epsilon}}\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, \boldsymbol{y}, t\right)\right\|_2^2 \end{equation}

•

하지만, p(xt∣y)p\left(\boldsymbol{x}_t \mid y\right)p(xt​∣y)가 이 Objective에 명확히 정의되지 않으므로 생성 결과가 condition에 부합하는지 단언할 수 없습니다.

Prelinamaries: Brownian Bridge

•

Brownian bridge 모델은 continous-time stochastic 모델입니다.

•

이때 condtion은 start point와 ending point에 의해 정해집니다. 

•

즉, 노이즈에서 데이터로 변했던 기존 모델과 달리 데이터에서 데이터로 변하게 됩니다.

•

Brownian Bridge Process는 다음과 같이 정의됩니다.

\begin{equation}p\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0, \boldsymbol{x}_T\right)=\mathcal{N}\left(\left(1-\frac{t}{T}\right) \boldsymbol{x}_0+\frac{t}{T} \boldsymbol{x}_T, \frac{t(T-t)}{T} \boldsymbol{I}\right)\end{equation}

•

식은 간단합니다. T에 대한 t의 비율로 x0\boldsymbol x_0x0​와 xT\boldsymbol x_TxT​를 Linear하게 섞어줍니다.

•

노이즈는 양 극단에서 최소가 되고, 중간에서 최대가 됩니다.

Methodology

Brownian Bridge Diffusin Model 위 Brownian Bridge Process를 적용하면 다음과 같습니다.

\begin{equation}\begin{gathered}q_{B B}\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0, \boldsymbol{y}\right)=\mathcal{N}\left(\boldsymbol{x}_t ;\left(1-m_t\right) \boldsymbol{x}_0+m_t \boldsymbol{y}, \delta_t \boldsymbol{I}\right) \\\boldsymbol{x}_0=\boldsymbol{x}, \quad m_t=\frac{t}{T}\end{gathered}\end{equation}

•

이때, 식 (4)에 따라 variance는 δt=t(T−t)T\delta_t=\frac{t(T-t)}{T}δt​=Tt(T−t)​인데 t=12Tt=\frac{1}{2}Tt=21​T 일때  δ12t=T4\delta_{\frac{1}{2}t}=\frac{T}{4}δ21​t​=4T​ 가 됩니다.

•

이는 T가 매우 커짐에 따라 더욱 큰 값이 됩니다.

•

따라서 variance를 보존 (Preserving) 할 수 있도록 다음과 같이 디자인합니다.

\begin{equation}\delta_t=2 s\left(m_t-m_t^2\right)\end{equation}

•

이때 s는 scale factor 로써 sampling diversity를 컨트롤합니다. (s=1s=1s=1 by default)

Forward Process:

q_{B B}\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}, \boldsymbol{y}\right)

•

위 식 (5) 를 통해 xt,xt−1\boldsymbol{x}_t,\boldsymbol{x}_{t-1}xt​,xt−1​을 구하여 대입해주면 다음과 같습니다.

\begin{equation} \begin{aligned} q_{B B}\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}, \boldsymbol{y}\right)=&\mathcal{N}\left(\boldsymbol{x}_t ; \frac{1-m_t}{1-m_{t-1}} \boldsymbol{x}_{t-1}\right. \\ & \left.\quad+\left(m_t-\frac{1-m_t}{1-m_{t-1}} m_{t-1}\right) \boldsymbol{y}, \delta_{t \mid t-1} \boldsymbol{I}\right) \end{aligned} \end{equation}

\begin{equation} \delta_{t \mid t-1}=\delta_t-\delta_{t-1} \frac{\left(1-m_t\right)^2}{\left(1-m_{t-1}\right)^2} \end{equation}

Reverse Process

x_T=y

에서 출발해

x_0

에 도달하는 reverse process는 다음과 같습니다.

\begin{equation}p_\theta\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{y}\right)=\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\boldsymbol{x}_t, t\right), \tilde{\delta}_t \boldsymbol{I}\right)\end{equation}

•

이때, μθ(xt,t)\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t)μθ​(xt​,t): predicted mean value of noise
δ~t\tilde{\delta}_tδ~t​: variance of noise at each step.

Objective ELBO를 사용합니다.

\begin{equation}\begin{aligned}E L B O & =-\mathbb{E}_q\left(D_{K L}\left(q_{B B}\left(\boldsymbol{x}_T \mid \boldsymbol{x}_0, \boldsymbol{y}\right) \| p\left(\boldsymbol{x}_T \mid \boldsymbol{y}\right)\right)\right. \\& +\sum_{t=2}^T D_{K L}\left(q_{B B}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0, \boldsymbol{y}\right) \| p_\theta\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{y}\right)\right) \\& \left.-\log p_\theta\left(\boldsymbol{x}_0 \mid \boldsymbol{x}_1, \boldsymbol{y}\right)\right)\end{aligned}\end{equation}

•

첫 항은 constant가 되어 KL Divergence는 0이 됩니다.

•

두번째 항은 식 (5)와 (7) 그리고 Bayes’ theorem, Markov chain property를 이용해 다음과 같이 나타냅니다.

\begin{equation}\begin{aligned}q_{B B}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0, \boldsymbol{y}\right) & =\frac{q_{B B}\left(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}, \boldsymbol{y}\right) q_{B B}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_0, \boldsymbol{y}\right)}{q_{B B}\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0, \boldsymbol{y}\right)} \\& =\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{x}_0, \boldsymbol{y}\right), \tilde{\delta}_t \boldsymbol{I}\right)\end{aligned}\end{equation}

\begin{equation}\begin{aligned}\tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{x}_0, \boldsymbol{y}\right) & =\frac{\delta_{t-1}}{\delta_t} \frac{1-m_t}{1-m_{t-1}} \boldsymbol{x}_t \\& +\left(1-m_{t-1}\right) \frac{\delta_{t \mid t-1}}{\delta_t} \boldsymbol{x}_0 \\& +\left(m_{t-1}-m_t \frac{1-m_t}{1-m_{t-1}} \frac{\delta_{t-1}}{\delta_t}\right) \boldsymbol{y}\end{aligned}\end{equation}

\begin{equation} \tilde{\delta}_t=\frac{\delta_{t \mid t-1} \cdot \delta_{t-1}}{\delta_t} \end{equation}

•

이때, x0\boldsymbol{x}_0x0​가 unknown이므로 reparametrization trick을 사용합니다.

\begin{equation}\tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t, \boldsymbol{y}\right)=c_{x t} \boldsymbol{x}_t+c_{y t} \boldsymbol{y}+c_{\epsilon t}\left(m_t\left(\boldsymbol{y}-\boldsymbol{x}_0\right)+\sqrt{\delta_t} \boldsymbol{\epsilon}\right)\end{equation}

\begin{aligned}c_{x t} & =\frac{\delta_{t-1}}{\delta_t} \frac{1-m_t}{1-m_{t-1}}+\frac{\delta_{t \mid t-1}}{\delta_t}\left(1-m_{t-1}\right) \\c_{y t} & =m_{t-1}-m_t \frac{1-m_t}{1-m_{t-1}} \frac{\delta_{t-1}}{\delta_t} \\c_{\epsilon t} & =\left(1-m_{t-1}\right) \frac{\delta_{t \mid t-1}}{\delta_t}\end{aligned}

•

여기서 μ~θ\tilde{\mu}_\thetaμ~​θ​가 아닌 노이즈 ϵθ\epsilon_\thetaϵθ​를 예측하도록 하면 다음과 같은 Linear combindation이 됩니다.

\begin{equation}\boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, \boldsymbol{y}, t\right)=c_{x t} \boldsymbol{x}_t+c_{y t} \boldsymbol{y}+c_{\epsilon t} \boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right) \end{equation}

•

따라서, ELBO (식 10)은 다음과 같습니다.

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{y}, \boldsymbol{\epsilon}}\left[c_{\epsilon t}\left\|m_t\left(\boldsymbol{y}-\boldsymbol{x}_0\right)+\sqrt{\delta_t} \boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right)\right\|^2\right]\end{equation}

수식으로 보나 코드로 보나 사실 별거 없습니다.

m_t(y-x_0)

샘플에 노이즈를 더한 후 (target) L2 로스를 계산하면 끝입니다.

Sampling

\begin{equation}\begin{array}{r}q_{B B}\left(\boldsymbol{x}_{\tau_{s-1}} \mid \boldsymbol{x}_{\tau s}, \boldsymbol{x}_0, \boldsymbol{y}\right)=\mathcal{N}\left(\left(1-m_{\tau_{s-1}}\right) \boldsymbol{x}_0+m_{\tau_{s-1}} \boldsymbol{y}+\right. \\\left.\sqrt{\delta_{\tau_{s-1}}-\sigma_{\tau_s}^2} \frac{1}{\sqrt{\delta_{\tau_s}}}\left(\boldsymbol{x}_{\tau_s}-\left(1-m_{\tau_s}\right) \boldsymbol{x}_0-m_{\tau_s} \boldsymbol{y}\right), \sigma_{\tau_s}^2 \boldsymbol{I}\right)\end{array}\end{equation}

사실 큰 의미는 없습니다. 샘플러 step = 200 에 대해 DDPM 샘플링을 하겠다는 뜻입니다.

Experimental Results

Conclusion

•

굉장히 간단한 논문이지만, 기존의 Bridge Model과 다르게 간단한 수식만으로 DDPM의 VP 프레임워크로 명료화했다는 점에서 다른 Application으로 발판을 마련했다고 볼 수 있습니다.

•

최근의 많은 Architectural 한 논문들은 거진 DDPM을 기반으로 하고 있는 것 처럼요.

•

Bridge Diffusion Model은 Condition의 개입 없이도 Image-to-image task를 수행할 수 있다는 점에서 모델 복잡도, 훈련 시간, 차원 복잡도에서 큰 이득이 있습니다. (논문에는 언급되지 않았지만 이러한 이점이 Bridge Model의 가장 큰 장점이라고 생각합니다.)

•

이를 베이스로 한 저의 논문이 곧 publish 될 예정이니 기대해주시면 감사드리겠습니다.