[비전 생성 모델] 디퓨전 모델에서의 Guidance
Guidance란?
디퓨전 모델은 완전한 노이즈 상태에서 출발하여 점진적으로 노이즈를 제거함으로써, 거친 (coarse) 구조에서 정교한 (fine) 세부 묘사까지 순차적으로 복원해 내는 방식으로 데이터를 생성한다.
가이던스 (Guidance)는 이러한 생성 모델이 텍스트 프롬프트나 이미지와 같은 사용자의 조건을 반영하여 원하는 샘플을 얻어내는 제어 가능한 생성 (Controllable generation)을 위한 기법이다.
가이던스는 샘플링 과정에서 생성 궤적을 특정 조건 공간으로 유도하는 역할을 하며, 구체적으로 reverse 과정에 guidance 항을 추가하여 생성 궤적을 원하는 방향으로 조종한다.
우리가 원하는 conditional score는 아래와 같이 분해할 수 있다.
\[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t\mid\mathbf{c}) =\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}_{\text{unconditional direction}} + \underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t)}_{\text{guidance direction}} \tag{1}\]즉, conditional sampling은 기존의 unconditional score에 데이터를 조건 $\mathbf{c}$ 방향으로 밀어주는 guidance 항을 더하는 방식으로 구현할 수 있다.
이때 생성 과정에서 가이던스를 얼마나 반영할 것인지를 조절하는 guidance scale $\omega$를 사용하여, 식 (1)을 아래와 같이 표현할 수 있다.
\[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t\mid\mathbf{c},\omega) =\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}_{\text{unconditional direction}} + \omega\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t)}_{\text{guidance direction}} \tag{2}\]하지만 식 (1)을 계산하는 데에는 현실적인 한계가 존재한다.
우리는 사전학습된 이미지 분류기나 CLIP 모델을 통해 원본 이미지에 대한 조건부 확률인 $p(\mathbf{c}\mid\mathbf{x}_0)$를 구할 수 있고, 디퓨전모델을 통해 노이즈 이미지로부터 원본 이미지를 추정하는 $p(\mathbf{x}_0\mid\mathbf{x}_t)$를 알 수 있다.
그러나 노이즈가 낀 상태의 이미지 $\mathbf{x}_t$가 조건 $\mathbf{c}$와 얼마나 잘 맞는지를 나타내는 $p(\mathbf{c}\mid\mathbf{x}_t)$는 직접적으로 계산할 수 없다.
이를 수식으로 전개하면 다음과 같다.
\[\begin{aligned} p(\mathbf{c}\mid\mathbf{x}_t) &=\int p(\mathbf{c},\mathbf{x}_0\mid\mathbf{x}_t)d\mathbf{x}_0 &&{\text{(Marginalization)}}\\ &=\int p(\mathbf{c}\mid\mathbf{x}_0,\mathbf{x}_t)p(\mathbf{x}_0\mid\mathbf{x}_t)d\mathbf{x}_0 &&{\text{(Chain Rule)}}\\ &=\int p(\mathbf{c}\mid\mathbf{x}_0)p(\mathbf{x}_0\mid\mathbf{x}_t)d\mathbf{x}_0 &&{\text{(Conditional Independence)}} \end{aligned}\]위 수식에서 Conditional Independence가 성립하는 이유는, 조건 $\mathbf{c}$는 원본 이미지 $\mathbf{x}_0$의 내용에 의해 결정되는 것이지 추가된 노이즈와는 무관하기 때문이다.
그렇다면 왜 $p(\mathbf{c}\mid\mathbf{x}_t)$를 직접 계산할 수 없을까?
바로 마지막 식에 등장하는 적분 $\int\cdots p(\mathbf{x}_0\mid\mathbf{x}_t) d\mathbf{x}_0$ 때문이다. 이는 현재의 노이즈 이미지 $\mathbf{x}_t$로부터 나올 수 있는 모든 가능한 원본 이미지 $\mathbf{x}_0$의 분포에 대해 적분해야 함을 의미한다.
이미지와 같은 고차원 공간에서 이러한 적분을 수행하는 것은 사실상 계산이 불가능하다. (이를 Intractable하다고 함)
예를 들어, 해상도가 $32\times32$인 RGB 이미지 $\mathbf{x}$는 3,072차원 공간에 존재하게 된다. 이에 대한 적분은 다음과 같이 표현된다.
\[\int p(\mathbf{c}\mid\mathbf{x}_0)p(\mathbf{x}_0\mid\mathbf{x}_t)d\mathbf{x}_0 =\underbrace{\int\int\cdots\int}_{\text{3,072번}} p(\mathbf{c}\mid\mathbf{x}_0)p(\mathbf{x}_0\mid\mathbf{x}_t) dx_0^{(1)}dx_0^{(2)}\cdots dx_0^{(3072)}\]따라서 실제 가이던스 기법들에서는 이 불가능한 적분 계산을 피하기 위해, $\mathbf{x}_t$로부터 예측한 하나의 대표 원본 이미지(예: $\hat{\mathbf{x}}_0$)를 활용하여 근사하는 방식을 취하게 된다.
Classifier Guidance (CG)
노이즈가 섞인 데이터 $\mathbf{x}_t$로부터 조건 $\mathbf{c}$를 예측하는 time-conditional 분류기 $p_\psi(\mathbf{c}\mid \mathbf{x}_t,t)$를 별도로 학습시키는 방식이다.
\[\nabla_{\mathbf{x}_t}\log p_\psi(\mathbf{c}\mid \mathbf{x}_t,t) \approx \nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t)\]
이후 샘플링 단계에서 이 분류기의 입력에 대한 그라디언트를 가이던스로 활용하여 unconditional score에 더해준다.Classifier-Free Guidance (CFG)
별도의 분류기를 학습시키지 않고, 단일 디퓨전 모델 내에서 conditional score와 unconditional score를 동시에 학습하는 방식이다.
\[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t\mid\mathbf{c}) \approx \mathbf{s}_\phi(\mathbf{x}_t,t,\mathbf{c})\]
학습 과정 중 일정 확률로 조건 $\mathbf{c}$ 자리에 null token을 입력하여 모델이 unconditional 생성 능력을 함께 익히도록 유도한다.Training-Free Guidance (TFG)
디퓨전 모델이나 분류기의 추가 학습 없이도 생성 과정을 제어할 수 있는 방식이다.
\[\tilde{p}_t(\mathbf{c}\mid\mathbf{x}_t)\propto \exp(-\tau \ell(\mathbf{x}_t,\mathbf{c}; t))\]
앞서 $p_t(\mathbf{c}\mid\mathbf{x}_t)$가 계산 불가능하다고 했기 때문에, 이를 대신할 대리 분포 (Surrogate) $\tilde{p}_t(\mathbf{c}\mid\mathbf{x}_t)$를 새로 정의한다.여기서 $\ell$은 노이즈 이미지 $\mathbf{x}_t$가 조건 $\mathbf{c}$와 부합할수록 값이 작아지는 손실 함수이다.
위 식에 $\nabla_{\mathbf{x}_t} \log$를 취하면 아래와 같이 정의되는데, 결과적으로 선택한 손실 함수 $\ell$의 그라디언트만을 계산함으로써 가이던스 방향을 구할 수 있다.
\[\nabla_{\mathbf{x}_t} \log \tilde{p}_t(\mathbf{c}\mid\mathbf{x}_t) = -\tau \nabla_{\mathbf{x}_t} \ell(\mathbf{x}_t,\mathbf{c}; t)\]
Classifier Guidance (CG)
수식에서 $\phi$는 unconditional 디퓨전 모델의 가중치, $\psi$는 분류기의 가중치를 의미하고, 첨자 $^*$는 학습이 완료됨을 의미한다.
사전학습된 unconditional 디퓨전 모델은 현재의 노이즈 상태 $\mathbf{x}_t$에서 노이즈를 어느 방향으로 제거해야 원본 데이터 매니폴드로 갈 수 있는지를 알려주는 score를 예측한다.
즉, 수학적으로 $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)$를 근사한다고 볼 수 있다.
\[\mathbf{s}_{\phi^*}(\mathbf{x}_t,t) \approx \nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t) \tag{3}\]CG에서는 식 (1)의 guidance 항을 근사하기 위해, time-dependent classifier $p_{\psi^*}(\mathbf{c}\mid\mathbf{x}_t,t)$를 별도로 학습시킨다.
이 분류기는 스텝 $t$와 그에 해당하는 노이즈 이미지 $\mathbf{x}_t$를 입력 받아, $\mathbf{c}$를 예측하도록 cross entropy loss를 최소화하며 학습된다.
$\mathbf{c}$는 클래스 라벨, 텍스트 캡션, 참조 이미지 등이 될 수 있다.
학습이 완료된 분류기의 입력 이미지에 대한 그래디언트는, 이미지가 조건에 부합하도록 유도하는 surrogate score 역할을 할 수 있다.
\[\nabla_{\mathbf{x}_t}p_{\psi^*}(\mathbf{c}\mid\mathbf{x}_t,t) \approx \nabla_{\mathbf{x}_t}p_t(\mathbf{c}\mid\mathbf{x}_t) \tag{4}\]
Sampling
결과적으로, 실제 샘플링 과정에서는 디퓨전 모델을 다시 학습할 필요 없이 식 (3)과 (4)를 결합하기만 하면 된다.
즉, 기존 무조건부 디퓨전 모델이 예측한 방향에 별도로 학습시킨 분류기의 그래디언트 방향을 더해주는 방식으로 생성 과정을 원하는 방향으로 조종한다.
Advantages & Limitations
Advantages
- 이미 잘 학습된 unconditional 디퓨전 모델을 직접 수정할 필요 없이, 가벼운 분류기 하나만 추가로 붙이면 원하는 대로 궤적을 조종할 수 있기 때문에 유연하다.
Limitations
- 분류기를 $t=0$부터 $t=T$까지 모든 노이즈 레벨에서 작동하도록 훈련해야 한다.
특히 노이즈가 극심한 구간 ($t$가 큰 경우)에서는 이미지를 판별하기 어렵기 때문에, 신경망 모델이 의미 있는 그라디언트를 추출하기는 어렵다. - 생성 모델과 분류기가 서로 다른 시점에 분리되어 학습되었기 때문에, 두 모델이 이해하는 데이터의 구조나 manifold가 어긋나 생성 품질이 저하될 수 있다.
Classifier-Free Guidance (CFG)
CG에서는 조건 $\mathbf{c}$로 유도하기 위해 분류기의 그라디언트 $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t)$를 따로 구해서 더해주었다.
반면, CFG에서는 이 항목을 conditional score와 unconditional score의 차이로 표현하였다.
이 수식을 원래의 가이던스 공식 식 (2)에 대입하면 아래와 같다.
\[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t,\omega) = \underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}_{\text{unconditional score}} +\omega(\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t\mid\mathbf{c})}_{\text{conditional score}} -\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}_{\text{unconditional score}}) \tag{7}\]즉, unconditional score를 기준점으로 삼고 conditional score와 unconditional score의 방향 차이를 guidance scale $\omega$만큼 증폭시키는 외삽을 통해 원하는 가이던스를 생성할 수 있다는 뜻이다.
내삽 (Interpolation, $0 \le \omega \le 1$): 두 점 사이에서 값을 찾는 것
외삽 (Extrapolation, $\omega > 1$): 두 점을 연결한 선을 따라가되, 두 점의 범위를 벗어나 바깥쪽으로 연장해서 값을 찾는 것
아래의 그림에서 점 a와 b를 기준으로 볼 때, 내두 점 사이에 있는 c는 내삽을 통해 구한 점이고, 범위를 벗어난 d와 e는 외삽을 통해 구한 점이다.
식 (7)을 다시 정리하면 아래의 수식이 유도된다.
\[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t,\omega) =\omega\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t\mid\mathbf{c})}_{\text{conditional score}} +(1-\omega)\underbrace{\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}_{\text{unconditional score}} \tag{8}\]즉, 식 (7)을 프롬프트 조건에 부합하는 특징은 강하게 증폭시키고, 조건 없이도 생성될 수 있는 일반적인 특징들은 생성 과정에서 억제하는 과정으로도 해석할 수 있다.
- $\omega=0$: 모델이 조건을 완전히 무시하고 아무 이미지나 생성한다.
- $\omega=1$: 일반적인 conditional 생성 모델과 동일하게 동작한다.
- $\omega>1$: conditional score에 더 가중치를 두어 생성 결과가 더 조건에 부합하도록 하지만 생성 결과의 다양성은 떨어지게 된다.
출처: The Principles of Diffusion Models
Training
식 (8)을 풀기 위해서는 conditional score와 unconditional score 두 가지가 모두 필요하지만, 두 개의 모델을 따로 학습시키는 것은 비용이 많이 든다.
이를 해결하기 위해, CFG는 단일 모델 안에서 두 가지를 동시에 학습시키는 방식을 사용한다. 구체적으로 다음과 같다.
Unconditional training: 일정 확률 $p_{\text{uncond}}$로 조건 $\mathbf{c}$를 지워버리고, 빈 조건을 의미하는 null token $\emptyset$을 입력으로 준다.
\[\mathbf{s}_\phi(\mathbf{x}_t, t,\emptyset)\]Conditional training: 조건 $\mathbf{c}$를 입력으로 준다.
\[\mathbf{s}_\phi(\mathbf{x}_t, t, \mathbf{c})\]
일반적으로 $p_{\text{uncond}}=0.1$으로 설정한다.
이렇게 하면 모델은 90%의 상황에서는 조건에 부합하도록 이미지를 생성하는 법을 학습하고, 10%의 상황에서는 조건 없이도 자연스러운 이미지를 생성하는 법을 학습하게 된다.
이때, 학습 단계에서는 가중치 $\omega$를 사용하지 않는다는 것을 주의해야 한다.
Sampling
학습이 완료된 모델을 이용해 실제 이미지를 생성할 때는, 식 (7)에 모델의 예측값을 그대로 대입한다.
\[\mathbf{s}_{\phi^*}^{\text{CFG}}(\mathbf{x}_t,t,\mathbf{c};\omega) =\omega\cdot\mathbf{s}_{\phi^*}(\mathbf{x}_t, t, \mathbf{c})+(1-\omega)\cdot\mathbf{s}_{\phi^*}(\mathbf{x}_t, t,\emptyset) \tag{9}\]즉, 조건을 주고 예측한 결과와 빈 조건을 주고 예측한 결과를 $\omega$를 통해 가중합한다.
Training-Free Guidance (TFG)
학습 없는 가이던스를 위해서는 사전학습된 unconditional 디퓨전 모델과 현재의 샘플 $\mathbf{x}$가 조건 $\mathbf{c}$에 얼마나 잘 부합하는지를 판별하는 손실 함수가 필요하다.
DDIM의 업데이트 수식을 살펴보면, 다음 스텝 $\mathbf{x}_{t-1}$을 만들 때 두 가지 공간의 정보를 활용한다는 것을 알 수 있다.
\[\mathbf{x}_{t-1}=\alpha_{t-1}\underbrace{\hat{\mathbf{x}_0}(\mathbf{x}_t)}_{\text{in data space}}-\sigma_{t-1}\sigma_t\underbrace{\hat{\mathbf{s}}(\mathbf{x}_t)}_{\text{in noise space}} \tag{10}\]- $\hat{\mathbf{x}}_0(\mathbf{x}_t)$: 모델이 노이즈 이미지 $\mathbf{x}_t$를 보고 예측해 낸 데이터 공간에서의 clean 이미지
- $\hat{\mathbf{s}}(\mathbf{x}_t)$: 노이즈를 어느 방향으로 제거해야 하는지를 알려주는 노이즈 공간에서의 방향
대부분의 TFG 기법들은 식 (10)을 변형하여, 데이터 공간 또는 노이즈 공간에서 조건을 만족시키는 방향으로 유도되도록 correction 항을 추가하는 방식을 사용한다.
\[\mathbf{x}_{t-1}=\alpha_{t-1}\underbrace{\left(\hat{\mathbf{x}}(\mathbf{x}_t)+\omega_t^{\text{data}}\mathcal{G}_0\right)}_{\text{in data space}}-\sigma_{t-1}\sigma_t\underbrace{\left(\hat{\mathbf{s}}(\mathbf{x}_t)+\omega_t^{\text{latent}}\mathcal{G}_t\right)}_{\text{in noise space}} \tag{11}\]위 수식에서 $\mathcal{G}_0$와 $\mathcal{G}_t$가 correction 항이고, $\omega_t^{\text{data}}$와 $\omega_t^{\text{latent}}$가 guidance scale이다.
- Guidance in Data Space: 목적지 $\hat{\mathbf{x}}_0$의 좌표를 살짝 변화시켜, 그 목적지를 향하는 $\mathbf{x}_{t-1}$의 궤적도 간접적으로 수정되도록 유도함
- Guidance in Noise Space: U-Net이 예측하는 노이즈 $\epsilon_\theta$ (또는 스코어) 자체를 수정하여, 궤적을 직접 수정함
Guidance in Data Space
이 방식은 디퓨전 모델이 매 스텝 $t$에서 예측한 clean 이미지 $\hat{\mathbf{x}}_0$를 직접 수정하는 방식이다.
구체적으로 모델이 예측한 깨끗한 이미지 $\hat{\mathbf{x}}_0$를 우리가 정한 손실 함수 $\ell$에 넣은 뒤, 이 손실을 줄이기 위해 $\hat{\mathbf{x}}_0$ 공간에 대해 미분하여 그라디언트를 구한다.
즉, correction 항은 아래와 같이 정의된다.
최종적으로 그라디언트 방향을 따라 $\hat{\mathbf{x}}_0$를 픽셀 단위에서 살짝 이동시킨다.
\[\hat{\mathbf{x}}_0'=\hat{\mathbf{x}}_0 + \omega_t^{\text{data}}\mathcal{G}_0 \tag{13}\]이렇게 업데이트된 $\hat{\mathbf{x}}_0’$를 식 (10)에 대입해서 새로운 $\mathbf{x}_{t-1}$을 계산한다.
Guidance in Noise Space
이 방식은 노이즈를 제거해 나가는 방향 자체를 틀어버리는 방식이다.
우리가 가이던스를 주려면 $p_t(\mathbf{c}\mid\mathbf{x}_t)$가 필요하지만, 앞서 말했듯이 이 조건부 확률은 계산이 불가능하다. 대신 우리는 손실 함수 $\ell$을 이용해 근사할 수 있다.
확률 $p_t(\mathbf{c}|\mathbf{x}_t)$는 조건에 맞을수록 값이 커져야 하고, 손실 $\ell$은 조건에 맞을수록 값이 작아져야 한다.
이런 동작을 위해, 손실 함수에 $\exp(-x)$를 씌운다.
이렇게 하면, 손실 함수를 surrogate likelihood로 둔갑시킬 수 있다.
우리는 가이던스를 위해 $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{c}\mid\mathbf{x}_t)$가 필요하기 때문에, correction 항은 아래와 같이 정의된다.
\[\mathcal{G}_t:=-\omega\nabla_{\mathbf{x}_t}\ell(\hat{\mathbf{x}}_0(\mathbf{x}_t),\mathbf{c}) \tag{15}\]위 수식에서 왜 손실 함수의 입력값이 해당 스텝의 노이즈 이미지 $\mathbf{x}_t$가 아니라, 예측한 clean 이미지 $\hat{\mathbf{x}}_0(\mathbf{x}_t)$인지 궁금할 것이다.
이는 손실 함수 계산할 때 사용하는 CLIP 등의 사전학습된 모델이 노이즈 이미지에 대해 학습되지 않았기 때문이다. 따라서 노이즈 이미지를 입력하면, 모델은 무의미한 값을 출력해 올바른 가이던스 방향을 제공할 수 없다.
최종적으로 그라디언트 방향을 따라 $\tilde{\mathbf{s}}({\mathbf{x}}_t)$를 노이즈 공간에서 살짝 이동시킨다.
\[\tilde{\mathbf{s}}({\mathbf{x}}_t) + \omega_t^{\text{latent}}\mathcal{G}_t \tag{16}\]Computational Bottleneck in Noise Space
노이즈 공간에서 가이던스를 적용할 때 발생하는 치명적인 문제는 막대한 연산량이다.
식 (15)를 보면, 노이즈 공간에서 가이던스 $\mathcal{G}_t$를 구하기 위해서는 손실 함수를 노이즈 이미지 $\mathbf{x}_t$로 미분해야 한다.
하지만 손실 함수 $\ell$은 $\mathbf{x}_t$를 직접 평가할 수 없어서 그 대신 $\hat{\mathbf{x}}_0$를 입력으로 주는데, 이를 다시 원래 미분 변수인 $\mathbf{x}_t$에 대해 미분하는 과정에서 chain Rule이 발생한다.
$\nabla_{\hat{\mathbf{x}}_0} \ell(\hat{\mathbf{x}}_0)\in\mathbb{R}^D$
모델이 예측한 원본 이미지 $\hat{\mathbf{x}}_0$가 변할 때, 손실 함수 $\ell$이 얼마나 변하는지를 나타냄
벡터 입력을 받아 스칼라 출력을 내는 함수에 대한 미분이므로, 결과값은 입력과 동일한 크기의 그라디언트 벡터가 됨$\nabla_{\mathbf{x}_t}\hat{\mathbf{x}}_0(\mathbf{x}_t)\in\mathbb{R}^{D\times D}$
노이즈가 낀 이미지 $\mathbf{x}_t$가 변할 때, 예측된 원본 이미지 $\hat{\mathbf{x}}_0$가 어떻게 변할지를 나타냄
벡터 입력을 받아 벡터 출력을 내는 함수에 대한 미분이므로, 결과값은 $D\times D$ 크기의 자코비안 행렬이 됨
이미지 차원 $D$에 대해 $D\times D$ 크기의 행렬을 구해야하기 때문에, 메모리에 할당하는 것이 사실상 불가능하다.
식 (17) 유도
스칼라 함수를 벡터로 미분할 때, 합성함수의 미분은 아래와 같이 표현된다.
\[\frac{\partial\ell}{\partial\mathbf{x}_t} =\left(\frac{\partial\mathbf{y}}{\partial\mathbf{x}_t}\right)^\top \cdot\frac{\partial\ell}{\partial\mathbf{y}} ~~\to~~ \nabla_{\mathbf{x}_t}\ell=J^\top\cdot\nabla_{\hat{\mathbf{x}}_0}\ell\]여기서 $\mathbf{y}=\hat{\mathbf{x}}_0(\mathbf{x}_t)$로 치환하면 식 (17) 형태를 유도할 수 있다.
따라서 노이즈 공간에서의 가이던스는 매 스텝마다 U-Net 전체를 통과하는 역전파를 수행해야되기 때문에, VRAM과 연산 시간이 많이 소모된다.
그런데 DDIM 샘플링 과정에서 $\mathbf{x}_t$로부터 $\mathbf{x}_0$를 예측하는 수식은 아래와 같은 단순 일차식 형태인데, 왜 무거울까?
\[\hat{\mathbf{x}_0}(\mathbf{x}_t)=\frac{\mathbf{x}_t-\sigma_t\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)}{\alpha_t}\]위 수식에 전체에 미분 연산자 $\nabla_{\mathbf{x}_t}$를 씌우면 아래와 같이 정리된다.
\[\begin{aligned} \nabla_{\mathbf{x}_t}\hat{\mathbf{x}_0}(\mathbf{x}_t) &=\nabla_{\mathbf{x}_t}\left(\frac{\mathbf{x}_t-\sigma_t\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)}{\alpha_t}\right)\\ &=\frac{1}{\alpha_t}\left(\nabla_{\mathbf{x}_t}\mathbf{x}_t-\sigma_t\nabla_{\mathbf{x}_t}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)\right)\\ &=\frac{1}{\alpha_t}\left(I-\sigma_t\nabla_{\mathbf{x}_t}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)\right) \end{aligned}\]정리된 수식에서 $\nabla_{\mathbf{x}_t}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)$는 입력 이미지가 $\mathbf{x}_t$가 변할 때, U-Net이 출력하는 예측 노이즈 맵이 얼마나 변하는지를 나타낸다.
그리고 $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)$는 벡터 입력을 받아 벡터 출력을 내는 함수이므로, 이에 대한 미분은 $D\times D$ 크기의 자코비안 행렬이 된다.
결국 $\mathbf{x}_0$를 구하는 연산 자체는 단순한 일차식이지만, 이를 $\mathbf{x}_t$로 미분하는 순간 거대한 U-Net을 통과하는 미분 연산이 포함되게 된다.