[Diffusion] 노이즈 제거 Score Matching (DSM)

안암동컴맹·2026년 1월 6일

Diffusion 101

목록 보기
7/10
post-thumbnail

🤗 소개

앞 글에서는 tractable한 score matching 목적 함수를 통해 score 기반 생성 모델의 이론적 토대를 마련했다. 그러나 해당 목적 함수는 score의 Jacobian과 trace를 포함하므로, 고차원 공간에서는 계산 비용이 크고 구현이 까다롭다는 문제 가 있다. 이번 글에서는 이러한 실용적인 한계를 어떻게 극복할 수 있는지를 살펴본다. 특히 score matching의 구조를 재해석한 Denoising Score Matching(DSM) 을 소개함으로써, 보다 효율적이고 안정적인 학습 방법으로 이어지는 동기를 제시할 예정이다. 이번 글에서는 많은 수학적 증명들이 나타나므로 꽤 긴 글이 될 것으로 예상된다.


🌦️ DSM 파헤치기

1️⃣ DSM이 태동하게 된 배경

이전 글에서 소개한 대안 목적 함수

L~SM(ϕ)=Expdata[Tr(xsϕ(x))+12sϕ(x)22]\tilde{\mathcal{L}}_\text{SM}(\phi)=\mathbb{E}_{\mathbf{x}\sim p_\text{data}}\left[\operatorname{Tr}(\nabla_\mathbf{x}\mathbf{s}_\phi(\mathbf{x}))+\frac{1}{2}\|\mathbf{s}_\phi(\mathbf{x})\|_2^2\right]

tractable 하다는 장점이 있지만, 여전히 Jacobian의 trace Tr(xsϕ(x))\operatorname{Tr}(\nabla_\mathbf{x}\mathbf{s}_\phi(\mathbf{x}))를 계산해야 한다는 문제가 있다. 이 연산은 최악의 경우(worst-case) 계산 복잡도가 O(D2)\mathcal{O}(D^2)에 이르며, 이러한 복잡도는 고차원 데이터로의 확장성을 제한한다.

이 문제를 해결하기 위해 sliced score matching(Song et at., 2020b)무작위 투영에 기반한 확률적 추정 을 사용하여 trace 항을 대체한다. 아래에서는 이 아이디어를 간단히 개괄한다.

🥒 Sliced Score Matching과 Hutchinson 추정기

Sliced score matching은 score matching에서의 trace 항을 무작위로 선택된 "슬라이스" 방향을 따라 계산한 방향 도함수들의 평균으로 대체한다. 예를 들어 uRD\mathbf{u}\in\mathbb{R}^D를 등방적인(isotropic) 확률 벡터라고 하자. 이 벡터는 E[u]=0\mathbb{E}[\mathbf{u}]=0E[uu]=I\mathbb{E}[\mathbf{uu}^\top]=\mathbf{I}를 만족한다.

Hutchinson 항등식에 따르면,

Tr(A)=Eu[uAu],Eu[(usϕ(x))2]=sϕ(x)22\operatorname{Tr}(\mathbf{A})=\mathbb{E}_\mathbf{u}[\mathbf{u}^\top\mathbf{Au}],\quad \mathbb{E}_\mathbf{u}[(\mathbf{u}^\top\mathbf{s}_\phi(\mathbf{x}))^2]=\|\mathbf{s}_\phi(\mathbf{x})\|_2^2

가 성립한다. 이를 이용하면 다음과 같은 정확한 형태의 목적 함수를 얻는다.

L~SM(ϕ)=Ex,u[u(xsϕ(x))u+12(usϕ(x))2]\tilde{\mathcal{L}}_\text{SM}(\phi)=\mathbb{E}_{\mathbf{x},\mathbf{u}}\left[\mathbf{u}^\top(\nabla_\mathbf{x}\mathbf{s}_\phi(\mathbf{x}))\mathbf{u}+\frac{1}{2}(\mathbf{u}^\top\mathbf{s}_\phi(\mathbf{x}))^2\right]

이 목적 함수는 큰 Jacobian이나 헤시안 행렬을 명시적으로 계산하지 않고도, 자동 미분을 이용한 Jacobian-Vector Product(JVP)와 Vector-Jacobian Produuct(VJP) 연산을 통해 효율적으로 계산할 수 있다.

무작위 probe를 KK개 평균내면 분산이 O(1/K)\mathcal{O}(1/K)불편향 추정기(unbiased estimator) 를 얻을 수 있으며, 방향 항 u(xsϕ)u\mathbf{u}^\top(\nabla_\mathbf{x}\mathbf{s}_\phi)\mathbf{u} 역시 명시적인 Jacobian 없이 JVP/VJP 루틴으로 효율적으로 계산된다.

직관적으로 보면, 이는 모델의 움직임을 무작위 방향들에서만 검사하는 것 과 같다. 투영된 score는 데이터 밀도가 더 높은 영역과 정렬되도록 유도되며, 그 결과 데이터 포인트들은 기댓값의 측면에서 정류점 들이 된다.

🧼 Sliced Score Matching에서 Denoising Score Matching으로

Sliced score matching은 Jacobian 계산을 피할 수는 있지만, 여전히 원래의 데이터 분포에 직접적으로 의존한다는 한계를 가진다. 이로 인해 다음과 같은 취약점을 가질 수 있다. 예를 들어 이미지 데이터가 저차원 다양체(manifold) 위에 놓여 있는 경우, xlogpdata(x)\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})정의되지 않거나 매우 불안정 할 수 있다.

또한, 이 방법은 관측된 데이터 지점에서만 벡터장을 제약하므로, 그 주변 영역에 대해서는 약한 제어만 을 제공한다. 추가적으로, 무작위 probe로 인한 분산 문제와 JVP/VJP 연산을 반복적으로 수행해야 하는 비용 역시 존재한다.

이번 글에서 중점적으로 다룰 보다 강력한 대안은 Denoising Score Matching(DSM; Vincent, 2011) 이다. DSM은 원리적으로 타당하면서도 확장성이 뛰어난 해결책을 제공한다.

2️⃣ DSM의 훈련

앞서 소개한 score matching 손실을 다시 살펴보자.

LSM(ϕ)=12Expdata(x)[sϕ(x)xlogpdata(x)22]\mathcal{L}_\text{SM}(\phi)=\frac{1}{2}\mathbb{E}_{\mathbf{x}\sim p_\text{data}(\mathbf{x})}\left[\|\mathbf{s}_\phi(\mathbf{x})-\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})\|_2^2\right]

여기서 문제가 되는 부분은 intractable한 항 xlogpdata(x)\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})이다.

☝🏻 조건화를 통한 해결

xlogpdata(x)\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})의 intractability를 극복하기 위해 Vincent(2011)는 알려진 조건부 분포 pσ(x~x)p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})를 통해 데이터 xpdata\mathbf{x}\sim p_\text{data}노이즈를 주입하는 방법 을 제안했다.

신경망 sϕ(x~;σ)\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)는 다음과 같은 주변 교란 분포(marginal perturbed distribution) 를 근사하도록 학습된다.

pσ(x~)=pσ(x~x)pdata(x)dxp_\sigma(\tilde{\mathbf{x}})=\int p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})d\mathbf{x}

이를 위해 다음 손실을 최소화한다.

LSM(ϕ;σ):=12Ex~pσ[sϕ(x~;σ)x~logpσ(x~)22]\mathcal{L}_\text{SM}(\phi;\sigma):=\frac{1}{2}\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma}\left[\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)-\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\|_2^2\right]

비록 x~logpσ(x~)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})가 일반적으로 intractable 하더라도 Vincent(2011)는 xpdata\mathbf{x}\sim p_\text{data}에 조건화를 하면, 동등하면서 tractable한 목적 함수, 즉 Denoising Score Matching(DSM) 손실을 얻을 수 있음을 보였다.

LDSM(ϕ;σ):=12Expdatax~pσ(x)[sϕ(x~;σ)x~logpσ(x~x)22]\mathcal{L}_\text{DSM}(\phi;\sigma):=\frac{1}{2}\mathbb{E}_{\mathbf{x}\sim p_\text{data}\tilde{\mathbf{x}}\sim p_\sigma(\cdot|\mathbf{x})}\left[\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)-\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\|_2^2\right]

위 식의 최적해 s\mathbf{s}^\ast는 다음을 만족한다.

s(x~;σ)=x~logpσ(x~)\mathbf{s}^\ast(\tilde{\mathbf{x}};\sigma)=\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})

이는 또한 LSM\mathcal{L}_\text{SM}에 대해서도 최적이다.

예를 들어 pσ(x~x)p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})가 분산이 σ2\sigma^2인 Gaussian 노이즈인 경우 pσ(x~x)=N(x~;x,σ2I)p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})=\mathcal{N}(\tilde{\mathbf{x}};\mathbf{x},\sigma^2\mathbf{I}), 기울기 x~logpσ(x~x)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})closed-form 으로 표현될 수 있다. 이로 인해 회귀의 목표가 명시적으로 주어지며 tractable하다. 또한 σ0\sigma\approx0일 때 pσ(x~)pdata(x)p_\sigma(\tilde{\mathbf{x}})\approx p_\text{data}(\mathbf{x})가 되고,

s(x~;σ)=x~logpσ(x~)xlogpdata(x)\mathbf{s}^\ast(\tilde{\mathbf{x}};\sigma)=\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\approx\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})

가 성립한다. 이는 학습된 score가 원래 데이터 분포의 score를 근사하게 됨을 의미하며, 이를 생성 과정에 활용할 수 있음을 보여준다.

이러한 논의를 바탕으로 LSM\mathcal{L}_\text{SM}LDSM\mathcal{L}_\text{DSM} 사이의 기울기 등가성을 다음 정리로 공식화할 수 있다.

정리 3.3.1 – LSM\mathcal{L}_\text{SM}LDSM\mathcal{L}_\text{DSM}의 등가성
임의의 고정된 노이즈 스케일 σ>0\sigma>0에 대해 다음이 성립한다.

LSM(ϕ;σ)=LDSM(ϕ;σ)+C\mathcal{L}_\text{SM}(\phi;\sigma)=\mathcal{L}_\text{DSM}(\phi;\sigma)+C

여기서 CCϕ\phi와 무관한 상수항이다. 또한, 두 손실의 최소해 s(;σ)\mathbf{s}^\ast(\cdot;\sigma)는 다음을 만족한다.

s(;σ)=x~logpσ(x~),for almost every x~\mathbf{s}^\ast(\cdot;\sigma)=\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}),\quad\text{for almost every}~\tilde{\mathbf{x}}

💡 정리 3.3.1에서 등가성에 대한 증명

LSM(ϕ;σ)\mathcal{L}_\text{SM}(\phi;\sigma)LDSM(ϕ;σ)\mathcal{L}_\text{DSM}(\phi;\sigma)를 모두 전개해보면 다음을 얻는다.

LSM(ϕ;σ)=12Ex~pσ(x~)[sϕ(x~;σ)222sϕ(x~;σ)x~logpσ(x~)+x~logpσ(x~)22]LDSM(ϕ;σ)=12Epdata(x)pσ(x~x)[sϕ(x~;σ)222sϕ(x~;σ)x~logpσ(x~x)+x~logpσ(x~x)22]\small \begin{aligned} \mathcal{L}_\text{SM}(\phi;\sigma)&=\frac{1}{2}\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\left[\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2-2\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})+\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\|_2^2\right] \\ \mathcal{L}_\text{DSM}(\phi;\sigma)&=\frac{1}{2}\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\left[\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2-2\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})+\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\|_2^2\right] \\ \end{aligned}

이 두 손실에 대해 뺄셈을 적용하면,

LSM(ϕ;σ)LDSM(ϕ;σ)=12(Ex~pσ(x~)sϕ(x~;σ)22Epdata(x)pσ(x~x)sϕ(x~;σ)22)(Ex~pσ(x~)[sϕ(x~;σ)x~logpσ(x~)]Epdata(x)pσ(x~x)[sϕ(x~;x)x~logpσ(x~x)])+12(Ex~pσ(x~)x~logpσ(x~)22Epdata(x)pσ(x~x)x~logpσ(x~x)22)\begin{aligned} &\quad\mathcal{L}_\text{SM}(\phi;\sigma)-\mathcal{L}_\text{DSM}(\phi;\sigma) \\ &=\frac{1}{2}\bigg(\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2-\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2\bigg)\\ &\small\quad-\bigg(\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\left[\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\right]-\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\left[\mathbf{s}_\phi(\tilde{\mathbf{x}};\mathbf{x})^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\right]\bigg) \\ &\quad+\frac{1}{2}\bigg(\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\|_2^2-\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\|_2^2\bigg) \\ \end{aligned}

이제 이 식을 한번에 한 항씩 살펴보자. 우선 첫번째 항부터 보자면, pσ(x~)=pσ(x~x)pdata(x)dxp_\sigma(\tilde{\mathbf{x}})=\int p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})d\mathbf{x}이므로, 다음과 같이 재서술할 수 있다.

Ex~pσ(x~)sϕ(x~;σ)22=(pσ(x~x)pdata(x)dx)sϕ(x~;σ)22 dx~=pdata(x)pσ(x~x)sϕ(x~;σ)22 dx~dx=Epdata(x)pσ(x~x)sϕ(x~;σ)22\begin{aligned} \mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2&=\int\bigg(\int p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})d\mathbf{x}\bigg)\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2~d\tilde{\mathbf{x}} \\ &=\int p_\text{data}(\mathbf{x})\int p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2~d\tilde{\mathbf{x}}d\mathbf{x} \\ &=\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2 \\ \end{aligned}

즉, 첫 번째 항은 00으로 상쇄할 수 있다. 다음으로 두 번째 항을 살펴보자.

Ex~pσ(x~)[sϕ(x~;σ)x~logpσ(x~)]=pσ(x~)sϕ(x~;σ)(x~pσ(x~)pσ(x~))dx~=sϕ(x~;σ)x~pσ(x~x)pdata(x) dxdx~=sϕ(x~;σ)x~pσ(x~x)pdata(x) dx~dx=Epdata(x)pσ(x~x)[sϕ(x~;σ)x~logpσ(x~x)]\begin{aligned} \mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\left[\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\right]&=\int p_\sigma(\tilde{\mathbf{x}})\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\left(\frac{\nabla_{\tilde{\mathbf{x}}}p_\sigma(\tilde{\mathbf{x}})}{p_\sigma(\tilde{\mathbf{x}})}\right)d\tilde{\mathbf{x}} \\ &=\int\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\int p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})~d\mathbf{x}d\tilde{\mathbf{x}} \\ &=\iint\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})~d\tilde{\mathbf{x}}d\mathbf{x} \\ &=\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\left[\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\right] \\ \end{aligned}

마찬가지로 00으로 상쇄할 수 있다. 마지막으로 세 번째 항은 다음과 같이 정리할 수 있다.

C:=12(Ex~pσ(x~)x~logpσ(x~)22Epdata(x)pσ(x~x)x~logpσ(x~x)22)C:=\frac{1}{2}\bigg(\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma(\tilde{\mathbf{x}})}\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\|_2^2-\mathbb{E}_{p_\text{data}(\mathbf{x})p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})}\|\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})\|_2^2\bigg)

하지만 이는 ϕ\phi에 무관한 상수항이기 때문에 최적화 관점에서 무시할 수 있다. 따라서,

arg minϕLSM(ϕ;σ)=arg minϕLDSM(ϕ;σ)\boxed{ \argmin_\phi\mathcal{L}_\text{SM}(\phi;\sigma)=\argmin_\phi\mathcal{L}_\text{DSM}(\phi;\sigma) }

이므로, 두 손실은 본질적으로 동등함을 증명할 수 있다. _\blacksquare

🪄 정리 3.3.1에서 SM과 DSM의 최소해 유도

최소해 s\mathbf{s}^\ast를 구하기 위해, 먼저 고정된 시간 tt를 도입하고 목적 함수에 나타나는 내부 기댓값을 분석해보자.

J(t,ϕ):=Ex0pdataExtpt(x0)[sϕ(xt,t)xtlogpt(xtx0)22]\mathcal{J}(t,\phi):=\mathbb{E}_{\mathbf{x}_0\sim p_\text{data}}\mathbb{E}_{\mathbf{x}_t\sim p_t(\cdot|\mathbf{x}_0)}\left[\|\mathbf{s}_\phi(\mathbf{x}_t,t)-\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)\|_2^2\right]

이 기댓값을 최소화하기 위해서는 각 xt\mathbf{x}_t에 대해 기대 제곱오차를 최소화하는 sϕ(xt,t)\mathbf{s}_\phi(\mathbf{x}_t,t)를 찾아야 한다. 이를 위해 X0X_0XtX_t의 결합 분포를 사용하여 위 기댓값을 다음과 같이 다시 쓸 수 있다.

J(t,ϕ)=pdata(x0)pt(xtx0)sϕ(xt,t)xtlogpt(xtx0)22 dx0dxt\mathcal{J}(t,\phi)=\iint p_\text{data}(\mathbf{x}_0)p_t(\mathbf{x}_t|\mathbf{x}_0)\|\mathbf{s}_\phi(\mathbf{x}_t, t)-\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)\|_2^2~d\mathbf{x}_0d\mathbf{x}_t

고정된 xt\mathbf{x}_t에 대해 다음 항을 최소화해야 한다.

p(x0Xt=xt)pt(xt)sϕ(xt,t)xtlogpt(xtx0)22 dx0\int p(\mathbf{x}_0|X_t=\mathbf{x}_t)p_t(\mathbf{x}_t)\|\mathbf{s}_\phi(\mathbf{x}_t, t)-\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)\|_2^2~d\mathbf{x}_0

pt(xt)p_t(\mathbf{x}_t)sϕ(xt,t)\mathbf{s}_\phi(\mathbf{x}_t,t)에 대해 상수이므로, 이는 다음을 최소화하는 것과 같다.

p(x0Xt=xt)sϕ(xt,t)xtlogpt(xtx0)22 dx0\int p(\mathbf{x}_0|X_t=\mathbf{x}_t)\|\mathbf{s}_\phi(\mathbf{x}_t, t)-\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)\|_2^2~d\mathbf{x}_0

이는 sϕ(xt,t)\mathbf{s}_\phi(\mathbf{x}_t,t)가 다음 조건부 기댓값과 같을 때 최소가 된다.

s(xt,t)=Ex0p(X0Xt=xt)[xtlogpt(xtX0)]\mathbf{s}^\ast(\mathbf{x}_t,t)=\mathbb{E}_{\mathbf{x}_0\sim p(X_0|X_t=\mathbf{x}_t)}\left[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|X_0)\right]

이제 이것이 xtlogpt(xt)\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)와 같음을 보여야 한다. Bayes' rule과 marginal probability의 정의에 따르면,

pt(xt)=pt(xtx0)pdata(x0) dx0p_t(\mathbf{x}_t)=\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0

여기에 로그를 취한 뒤 xt\mathbf{x}_t에 대해 미분하면,

xtlogpt(xt)=xtpt(xt)pt(xt)=xtpt(xtx0)pdata(x0) dx0pt(xtx0)pdata(x0) dx0\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)=\frac{\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t)}{p_t(\mathbf{x}_t)}=\frac{\nabla_{\mathbf{x}_t}\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0}{\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0}

적절한 정칙성 조건 하에서는 기울기 연산과 적분을 교환할 수 있으므로,

xtlogpt(xt)=xtpt(xtx0)pdata(x0) dx0pt(xtx0)pdata(x0) dx0\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)=\frac{\int\nabla_{\mathbf{x}_t} p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0}{\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0}

이제 s(xt,t)\mathbf{s}^\ast(\mathbf{x}_t,t)를 다시 전개해보자.

s(xt,t)=p(x0xt)xtlogpt(xtx0) dx0\mathbf{s}^\ast(\mathbf{x}_t,t)=\int p(\mathbf{x}_0|\mathbf{x}_t)\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)~d\mathbf{x}_0

여기에도 위에서와 마찬가지로 Bayes'rule과 marginal probability의 정의를 적용하면 다음과 같다.

s(xt,t)=pt(xtx0)pdata(x0)pt(xt)xtlogpt(xtx0) dx0=pt(xtx0)pdata(x0)pt(xt)xtpt(xtx0)pt(xtx0) dx0=1pt(xt)xtpt(xtx0)pdata(x0) dx0=xtpt(xtx0)pdata(x0) dx0pt(xtx0)pdata(x0) dx0\begin{aligned} \mathbf{s}^\ast(\mathbf{x}_t,t)&=\int\frac{p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)}{p_t(\mathbf{x}_t)}\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|\mathbf{x}_0)~d\mathbf{x}_0 \\ &=\int\frac{p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)}{p_t(\mathbf{x}_t)}\cdot\frac{\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)}{p_t(\mathbf{x}_t|\mathbf{x}_0)}~d\mathbf{x}_0 \\ &=\frac{1}{p_t(\mathbf{x}_t)}\int\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0 \\ &=\frac{\int\nabla_{\mathbf{x}_t} p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0}{\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_\text{data}(\mathbf{x}_0)~d\mathbf{x}_0} \\ \end{aligned}

따라서 다음과 같이 SM과 DSM 공통의 최소해 s\mathbf{s}^\ast를 유도할 수 있다.

s(xt,t)=Ex0p(X0Xt=xt)[xtlogpt(xtX0)]=xtlogpt(xt)\boxed{ \mathbf{s}^\ast(\mathbf{x}_t,t)=\mathbb{E}_{\mathbf{x}_0\sim p(X_0|X_t=\mathbf{x}_t)}\left[\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t|X_0)\right]=\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}\quad_\blacksquare

마치 이전에 소개한 DDPM에서처럼, 정리 3.3.1은 조건화로부터 다음과 같은 통찰을 얻을 수 있게 해준다.

통찰 3.3.1 – 조건화 기법
조건화 기법은 DDPM의 변분적 관점에서도 등장한다(정리 2.2.1). 이 경우 데이터 포인트 x\mathbf{x}에 대해 조건화를 수행함으로써, intractable한 손실을 몬테카를로 추정을 통해 tractable한 손실로 변환할 수 있다.

🌟 특수한 경우: 가산적 Gaussian 노이즈

이제 각 데이터 포인트 xpdata\mathbf{x}\sim p_\text{data}에 분산이 σ2\sigma^2인 Gaussian 노이즈 N(0,σ2I)\mathcal{N}(\mathbf{0},\sigma^2\mathbf{I})를 더하는 일반적인 경우를 고려해보자. 즉,

ϵN(0,I),x~=x+σϵ\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I}),\quad\tilde{\mathbf{x}}=\mathbf{x}+\sigma\epsilon

로 설정하면, 노이즈가 섞인 데이터 x~\tilde{\mathbf{x}}는 다음 분포를 따른다.

pσ(x~x)=N(x~;x,σ2I)p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})=\mathcal{N}(\tilde{\mathbf{x}};\mathbf{x},\sigma^2\mathbf{I})

이 설정에서 조건부 score는 해석적으로 다음과 같이 주어진다.

x~logpσ(x~x)=xx~σ2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}|\mathbf{x})=\frac{\mathbf{x}-\tilde{\mathbf{x}}}{\sigma^2}

따라서 DSM 손실은 다음과 같이 단순화된다.

LDSM(ϕ;σ)=12Ex,x~x[sϕ(x~;σ)xx~σ222]=12Ex,ϵ[sϕ(x+σϵ;σ)+ϵσ22]\boxed{ \begin{aligned} \mathcal{L}_\text{DSM}(\phi;\sigma)&=\frac{1}{2}\mathbb{E}_{\mathbf{x},\tilde{\mathbf{x}}|\mathbf{x}}\bigg[\bigg\Vert \mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)-\frac{\mathbf{x}-\tilde{\mathbf{x}}}{\sigma^2}\bigg\Vert_2^2\bigg] \\ &=\frac{1}{2}\mathbb{E}_{\mathbf{x},\epsilon}\bigg[\bigg\Vert \mathbf{s}_\phi(\mathbf{x}+\sigma\epsilon;\sigma)+\frac{\epsilon}{\sigma}\bigg\Vert_2^2\bigg] \\ \end{aligned} }

여기서 ϵN(0,I)\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})이다. 이 목적 함수는 score-based diffusion 모델의 핵심을 이룬다.

노이즈 수준 σ\sigma가 작은 경우, Gaussian으로 평활화된 marginal 분포 pσ=pdataN(0,σ2I)p_\sigma=p_\text{data}\ast\mathcal{N}(\mathbf{0},\sigma^2\mathbf{I})는 원래의 데이터 분포와 거의 동일한 고밀도 영역과 score를 갖게 된다.

x~logpσ(x~)xlogpdata(x)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})\approx\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})

따라서 노이즈가 섞인 score 방향 x~logpσ\nabla_{\tilde{\mathbf{x}}}\log p_\sigma으로 작은 한 step을 이동하면, 노이즈가 포함된 샘플은 원본 분포의 high-likelihood 영역과 거의 동일한 영역으로 이동하게 된다. 이는 앞서 요약한 score matching의 직관과 유사하다.

반대로 σ\sigma가 큰 경우에는, 이러한 평활화가 분포의 지형을 과도하게 단순화 한다. 즉, pσp_\sigma는 국소적인 mode들을 씻어내며, 그 score는 주로 전체 질량이 모여 있는 방향(예를 들어 평균 방향으로의 수축)을 향하게 된다. 그 결과 거친 수준의 denoising이 이루어지며, 과도한 평활화가 발생할 수 있다. 그러나 실제로 DSM에서는 주입되는 노이즈가 작고 완만하다고 가정하는 경우가 일반적이다.

3️⃣ 샘플링

노이즈 수준 σ\sigma에서 학습된 score 모델 sϕ×(x~;σ)\mathbf{s}_{\phi^\times}(\tilde{\mathbf{x}};\sigma)이 주어지면, 실제 score 대신 학습된 모델을 사용하여 Langevin dynamics로 샘플을 생성할 수 있다. 업데이트 규칙은 다음과 같다.

x~n+1=x~n+ηsϕ×(x~n;σ)+2ηϵn,ϵnN(0,I)\tilde{\mathbf{x}}_{n+1}=\tilde{\mathbf{x}}_n+\eta\mathbf{s}_{\phi^\times}(\tilde{\mathbf{x}}_n;\sigma)+\sqrt{2\eta}\epsilon_n,\quad\epsilon_n\sim\mathcal{N}(\mathbf{0},\mathbf{I})

여기서 sϕ×(x~n;σ)x~logpσ(x~n)\mathbf{s}_{\phi^\times}(\tilde{\mathbf{x}}_n;\sigma)\approx\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}_n)이며, n=0,1,2,n=0,1,2,\ldots이고 초기값 x~0\tilde{\mathbf{x}}_0에서 시작한다. σ\sigma가 충분히 작다면, 충분한 반복 횟수 이후 x~n\tilde{\mathbf{x}}_npdatap_\text{data}로부터의 샘플을 근사하게 된다.

😊 노이즈 주입의 장점

위에서 언급한 기본적인 score matching(LSM\mathcal{L}_\text{SM})과 비교할 때, Gaussian 노이즈를 주입하여 pσp_\sigma를 구성하는 것은 두 가지 중요한 장점을 제공한다.

  • 잘 정의된 기울기
    노이즈는 데이터를 저차원 다양체(manifold)로부터 벗어나게 만들어, RD\mathbb{R}^D 전체에 걸쳐 support를 갖는 분포 pσp_\sigma를 형성한다. 그 결과 score 함수 x~logpσ(x~)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})는 모든 지점에서 잘 정의된다.

  • 향상된 커버리지
    노이즈는 mode들 사이의 희소한 영역을 부드럽게 만들어 학습 신호의 품질을 향상시키고, Langevin dynamics가 저밀도 영역을 보다 효과적으로 통과할 수 있도록 돕는다.

4️⃣ DSM이 노이즈 제거인 이유 – Tweedie 공식

해당 이유를 찾기 위해 먼저 Tweeidie 공식(Efron, 2011) 에서 출발해보자. 이 공식은 노이즈가 섞인 관측값만으로도 원리적인 denoising이 가능함을 보여주는 이론적 근거 를 제공한다.

구체적으로, 알려지지 않은 xpdata\mathbf{x}\sim p_\text{data}로부터 생성된 단일 Gaussian 노이즈 관측값 x~N(;αx,σ2I)\tilde{\mathbf{x}}\sim\mathcal{N}(\cdot;\alpha\mathbf{x},\sigma^2\mathbf{I})이 주어졌을 때, x~\tilde{\mathbf{x}}에 조건화된 모든 가능한 깨끗한 신호들의 평균, 즉 denoised 추정치는 노이즈가 섞인 margianl 분포의 score x~logpσ(x~)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})방향으로 크기 σ2\sigma^2만큼 x~\tilde{\mathbf{x}}를 이동시킴으로써 얻어진다.

여기서 노이즈가 섞인 marginal 분포는 다음과 같이 정의된다.

pσ(x~):=N(x~;αx,σ2I)pdata(x) dxp_\sigma(\tilde{\mathbf{x}}):=\int\mathcal{N}(\tilde{\mathbf{x}};\alpha\mathbf{x},\sigma^2\mathbf{I})p_\text{data}(\mathbf{x})~d\mathbf{x}

이를 공식화하면 다음과 같다.

보조정리 3.3.2 – Tweedie 공식
xpdata\mathbf{x}\sim p_\text{data}이고, x\mathbf{x}에 조건화했을 때 x~N(;αx,σ2I)\tilde{\mathbf{x}}\sim\mathcal{N}(\cdot;\alpha\mathbf{x},\sigma^2\mathbf{I})이며 α0\alpha\ne0라고 가정하자. 이때 Tweedie 공식은 다음이 성립함을 보여준다.

αExp(xx~)[xx~]=x~+σ2x~logpσ(x~)\alpha\mathbb{E}_{\mathbf{x}\sim p(\mathbf{x}|\tilde{\mathbf{x}})}[\mathbf{x}|\tilde{\mathbf{x}}]=\tilde{\mathbf{x}}+\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})

여기서 기댓값은 x~\tilde{\mathbf{x}}가 주어졌을 때의 x\mathbf{x}에 대한 사후 분포 p(xx~)p(\mathbf{x}|\tilde{\mathbf{x}})에 대해 취해진다.

✏️ 보조정리 3.3.2의 증명

먼저 시간에 따라 변하는 Gaussian perturbation을 고려하는 Tweedie 공식의 보다 일반적인 형태를 제시하고, 그 뒤에 증명을 이어나갈 예정이다.

xtN(;αtx0,σt2I)\mathbf{x}_t\sim\mathcal{N}(\cdot;\alpha_t\mathbf{x}_0,\sigma^2_t\mathbf{I})인 Gaussian 확률 벡터를 고려해보자. 이때 Tweedie 공식은 다음을 만족한다.

αtEx0p(x0xt)[x0xt]=xt+σt2xtlogpt(xt)\alpha_t\mathbb{E}_{\mathbf{x}_0\sim p(\mathbf{x}_0|\mathbf{x}_t)}[\mathbf{x}_0|\mathbf{x}_t]=\mathbf{x}_t+\sigma_t^2\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)

여기서 기댓값은 관측된 xt\mathbf{x}_t가 주어졌을 때의 x0\mathbf{x}_0에 대한 사후 분포 p(x0xt)p(\mathbf{x}_0|\mathbf{x}_t)에 대해 취해지며, pt(xt)p_t(\mathbf{x}_t)는 아래와 같은 xt\mathbf{x}_t의 marginal 분포이다.

pt(xt)=pt(xtx0)p0(x0) dx0p_t(\mathbf{x}_t)=\int p_t(\mathbf{x}_t|\mathbf{x}_0)p_0(\mathbf{x}_0)~d\mathbf{x}_0

이제 score 함수를 계산해보자.

xtlogpt(xt)=xtpt(xt)pt(xt)=1pt(xt)xtpt(xtx0)p0(x0) dx0\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)=\frac{\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t)}{p_t(\mathbf{x}_t)}=\frac{1}{p_t(\mathbf{x}_t)}\int\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)p_0(\mathbf{x}_0)~d\mathbf{x}_0

따라서 조건부 밀도의 기울기 xtpt(xtx0)\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)를 계산해야 한다.

조건부 Gaussian 밀도의 기울기는 다음과 같다.

xtpt(xtx0)=pt(xtx0)σt2(xtαtx0)\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)=-p_t(\mathbf{x}_t|\mathbf{x}_0)\sigma^{-2}_t(\mathbf{x}_t-\alpha_t\mathbf{x}_0)

이를 앞선 식에 대입하면,

xtpt(xt)=xtpt(xtx0)p0(x0) dx0=σt2(xtαtx0)pt(xtx0)p0(x0) dx0=σt2(xtαtx0)p(x0xt)pt(xt) Bayes’ Rule dx0=σt2pt(xt)(xtαtx0)p(x0xt) dx0=σt2pt(xt)(xtp(x0xt) dx0=xt1αtx0p(x0xt) dx0)=σt2pt(xt)(xtαtEp(x0xt)[x0xt])\begin{aligned} \nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t)&=\int\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t|\mathbf{x}_0)p_0(\mathbf{x}_0)~d\mathbf{x}_0 \\ &=-\sigma^{-2}_t\int(\mathbf{x}_t-\alpha_t\mathbf{x}_0)p_t(\mathbf{x}_t|\mathbf{x}_0)p_0(\mathbf{x}_0)~d\mathbf{x}_0 \\ &=-\sigma^{-2}_t\int(\mathbf{x}_t-\alpha_t\mathbf{x}_0)\underbrace{p(\mathbf{x}_0|\mathbf{x}_t)p_t(\mathbf{x}_t)}_{\because~\text{Bayes' Rule}}~d\mathbf{x}_0 \\ &=-\sigma^{-2}_tp_t(\mathbf{x}_t)\int(\mathbf{x}_t-\alpha_t\mathbf{x}_0)p(\mathbf{x}_0|\mathbf{x}_t)~d\mathbf{x}_0 \\ &=-\sigma^{-2}_tp_t(\mathbf{x}_t)\bigg(\underbrace{\int\mathbf{x}_tp(\mathbf{x}_0|\mathbf{x}_t)~d\mathbf{x}_0}_{=\mathbf{x}_t\cdot1}-\alpha_t\int\mathbf{x}_0 p(\mathbf{x}_0|\mathbf{x}_t)~d\mathbf{x}_0\bigg) \\ &=-\sigma^{-2}_tp_t(\mathbf{x}_t)(\mathbf{x}_t-\alpha_t\mathbb{E}_{p(\mathbf{x}_0|\mathbf{x}_t)}[\mathbf{x}_0|\mathbf{x}_t]) \\ \end{aligned}

양변을 pt(xt)p_t(\mathbf{x}_t)로 나누면 다음을 얻을 수 있다.

xtpt(xt)pt(xt)=xtlogpt(xt)=σt2(xtαtEp(x0xt)[x0xt])\frac{\nabla_{\mathbf{x}_t}p_t(\mathbf{x}_t)}{p_t(\mathbf{x}_t)}=\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)=-\sigma^{-2}_t(\mathbf{x}_t-\alpha_t\mathbb{E}_{p(\mathbf{x}_0|\mathbf{x}_t)}[\mathbf{x}_0|\mathbf{x}_t])

이를 재배열하면,

αtEp(x0xt)[x0xt]=xt+σt2xtlogpt(xt)\boxed{\alpha_t\mathbb{E}_{p(\mathbf{x}_0|\mathbf{x}_t)}[\mathbf{x}_0|\mathbf{x}_t]=\mathbf{x}_t+\sigma^2_t\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)}

여기서 xt\mathbf{x}_t를 노이즈가 주입된 데이터 x~\tilde{\mathbf{x}}로 해석하면 보조정리 3.3.2를 증명할 수 있다. ~_\blacksquare

Tweedie 공식은 DDPM에서와 같이 여러 단계의 노이즈를 도입하는 diffusion 모델에서 중심적인 역할을 한다. 이 공식은 score 함수를 통해 노이즈가 섞인 관측값으로부터 깨끗한 샘플을 추정할 수 있게 해주며, 그로써 score 예측과 denoiser 사이의 근본적인 연결고리 를 확립한다.

E[xx~]Estimated denoiser=1α(x~+σ2x~logpσ(x~))\underbrace{\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]}_\text{Estimated denoiser}=\frac{1}{\alpha}(\tilde{\mathbf{x}}+\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}))

특히, 노이즈가 섞인 로그 likelihood logpσ(x~)\log p_\sigma(\tilde{\mathbf{x}})에 대해 step-size가 정확히 σ2\sigma^2인 단 한번의 gradient ascent step이 곧 denoised 추정치, 즉 깨끗한 조건부 평균 신호에 해당한다. 이로 인해 DSM 학습과 denoising은 매우 밀접하게 연결된다.

만약 DSM을 통해 학습된 sϕ×(x~)\mathbf{s}_{\phi^\times}(\tilde{\mathbf{x}})x~logpσ(x~)\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})를 잘 근사한다면,

1α(x~+σ2sϕ(x~))\frac{1}{\alpha}(\tilde{\mathbf{x}}+\sigma^2\mathbf{s}_{\phi}(\tilde{\mathbf{x}}))

은 곧 denoiser가 된다.

📈 로그 정규화항 λ(x~)\lambda(\tilde{\mathbf{x}})를 갖는 지수족(Exponential Family) 설정

잠재 자연 매개변수(latent natural parameter) ηRD\boldsymbol\eta\in\mathbb{R}^D가 주어졌을 때, x~\tilde{\mathbf{x}}의 조건부 분포가 다음과 같은 자연 지수족(natural exponential family) 에 속한다고 가정하자.

qσ(x~η)=exp(ηx~ψ(η))q0(x~)q_\sigma(\tilde{\mathbf{x}}|\boldsymbol\eta)=\exp(\boldsymbol\eta^\top\tilde{\mathbf{x}}-\psi(\boldsymbol\eta))q_0(\tilde{\mathbf{x}})

여기서 q0(x~)q_0(\tilde{\mathbf{x}})기저 측도(base measurement) 로, η\boldsymbol\eta에 의존하지 않는 항이다. 예를 들어 분산이 σ2I\sigma^2\mathbf{I}인 가산적 Gaussian 노이즈의 경우,

q0(x~)=(2πσ)D/2exp(x~222σ2)q_0(\tilde{\mathbf{x}})=(2\pi\sigma)^{-D/2}\exp\left(-\frac{\|\tilde{\mathbf{x}}\|_2^2}{2\sigma^2}\right)

가 된다. 이제 잠재 자연 매개변수 η\boldsymbol\eta의 사전 분포를 p(η)p(\boldsymbol\eta)로 두자. 이는 재매개변수화된 깨끗한 데이터 분포 로 해석할 수 있으며(Gaussian 위치 모형의 경우 η=x/σ2\boldsymbol\eta=\mathbf{x}/\sigma^2), 관측된 노이즈가 섞인 marginal 분포는 다음과 같이 주어진다.

pσ(x~)=qσ(x~η)p(η) dηp_\sigma(\tilde{\mathbf{x}})=\int q_\sigma(\tilde{\mathbf{x}}|\boldsymbol\eta)p(\boldsymbol\eta)~d\boldsymbol\eta

이제 x~\tilde{\mathbf{x}}에 대한 로그 정규화항을 다음과 같이 정의해보자.

λ(x~):=logpσ(x~)logq0(x~)\lambda(\tilde{\mathbf{x}}):=\log p_\sigma(\tilde{\mathbf{x}})-\log q_0(\tilde{\mathbf{x}})

그러면 x~\tilde{\mathbf{x}}기 주어졌을 때 η\boldsymbol\eta의 사후 분포는

p(ηx~)exp(ηx~ψ(η)λ(x~))p(η)p(\boldsymbol\eta|\tilde{\mathbf{x}})\propto\exp(\boldsymbol\eta^\top\tilde{\mathbf{x}}-\psi(\boldsymbol\eta)-\lambda(\tilde{\mathbf{x}}))p(\boldsymbol\eta)

로 주어지며, 이는 x~\tilde{\mathbf{x}}를 변수로 보았을 때 사후 분포가 자연 매개변수 x~\tilde{\mathbf{x}}, 충분 통계량 η\boldsymbol\eta, 그리고 로그 정규화항 λ(x~)\lambda(\tilde{\mathbf{x}})를 갖는 지수족 형태임을 보여준다.

🥞 사후 누적량(Posterior Cumulants)를 생성하는 λ\lambda의 도함수

λ\lambda의 도함수는 사후 누적량(posterior cumulants) 를 생성하는데, 여기에는 두 가지 간단한 규칙이 적용한다.

  1. 정규화: 모든 x~\tilde{\mathbf{x}}에 대해 다음이 성립한다.

    exp(ηx~ψ(η)λ(x~))p(η) dη=1\int\exp(\boldsymbol\eta^\top\tilde{\mathbf{x}}-\psi(\boldsymbol\eta)-\lambda(\tilde{\mathbf{x}}))p(\boldsymbol\eta)~d\boldsymbol\eta=1

    이 항등식을 x~\tilde{\mathbf{x}}에 대해 미분하면, 지수항으로부터 η\boldsymbol\eta의 거듭제곱이 내려오고 λ(x~)\lambda(\tilde{\mathbf{x}})의 도함수가 함게 등장한다. 그 결과를 00으로 두면, λ\lambda의 도함수들과 η\boldsymbol\eta사후 모멘트(posterior moment) 들 사이의 등식들이 얻어진다.

  2. 지수족의 표준 성질: 로그 정규화 함수는 충분 통계량의 누적 생성 함수(cumulant generating function) 이다. 따라서 다음이 성립한다.

    x~λ(x~)=E[ηx~],x~2λ(x~)=Cov[ηx~],x~(k)λ(x~)=κk(ηx~)(k3)\nabla_{\tilde{\mathbf{x}}}\lambda(\tilde{\mathbf{x}})=\mathbb{E}[\boldsymbol\eta|\tilde{\mathbf{x}}],\quad\nabla_{\tilde{\mathbf{x}}}^2\lambda(\tilde{\mathbf{x}})=\operatorname{Cov}[\boldsymbol\eta|\tilde{\mathbf{x}}],\quad\nabla_{\tilde{\mathbf{x}}}^{(k)}\lambda(\tilde{\mathbf{x}})=\kappa_k(\boldsymbol\eta|\tilde{\mathbf{x}})\quad(k\ge3)

    여기서 κk\kappa_kx~\tilde{\mathbf{x}}가 주어졌을 때의 확률변수 η\boldsymbol\eta에 대한 kk조건부 누적량(conditional cumulants) 이며, 표준적인 모멘트-누적량(moment-cumulant) 관계를 통해 얻어진다. 이는 고차 Tweedie 공식에 해당한다.

이를 η=x/σ2\boldsymbol\eta=\mathbf{x}/\sigma^2인 Gaussian 위치 모형에 특화하면, logpσ(x~)\log p_\sigma(\tilde{\mathbf{x}})의 도함수로 표현되는 익숙한 형태들을 얻는다.

E[xx~]=x~+σ2x~logpσ(x~),Cov[xx~]=σ2I+σ4x~2logpσ(x~)\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]=\tilde{\mathbf{x}}+\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}),\quad\operatorname{Cov}[\mathbf{x}|\tilde{\mathbf{x}}]=\sigma^2\mathbf{I}+\sigma^4\nabla_{\tilde{\mathbf{x}}}^2\log p_\sigma(\tilde{\mathbf{x}})

그리고 더 높은 차수의 누적량들은 logpσ(x~)\log p_\sigma(\tilde{\mathbf{x}})의 더 높은 차수 도함수에 비례하여 스케일된다.

5️⃣ DSM이 노이즈 제거인 이유 – SURE

대략적으로 말해, Stein 불편향 위험 추정기(Stein's Unbiased Risk Estimator; SURE) 는 깨끗한 신호를 알지 못하더라도 denoiser D\mathbf{D}의 평균제곱오차(MSE)를 추정할 수 있게 해주는 기법이다. 다시 말해, SURE는 노이즈가 섞인 데이터만 주어졌을 때도 denoiser를 선택하거나 학습할 수 있는 방법을 제공한다.

명확성을 위해, 가산적 Gaussian 노이즈 설정을 고려해보자.

x~=x+σϵ,ϵN(0,I)\tilde{\mathbf{x}}=\mathbf{x}+\sigma\boldsymbol\epsilon,\quad\boldsymbol\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})

여기서 xRD\mathbf{x}\in\mathbb{R}^D는 알려지지 않은 깨끗한 신호이고, x~\tilde{\mathbf{x}}는 관측된 노이즈 버전이다. Denoiser란 D:RDRD\mathbf{D}:\mathbb{R}^D\to\mathbb{R}^D인 약하게 미분 가능한 임의의 함수로, x\mathbf{x}의 추정치 D(x~)\mathbf{D}(\tilde{\mathbf{x}})를 출력한다.

자연스러운 품질 척도는 조건부 평균제곱오차(MSE) 이다.

R(D;x):=Ex~x[D(x~)x22x]R(\mathbf{D};\mathbf{x}):=\mathbb{E}_{\tilde{\mathbf{x}}|\mathbf{x}}\Big[\|\mathbf{D}(\tilde{\mathbf{x}})-\mathbf{x}\|_2^2\Big|\mathbf{x}\Big]

이 양은 알려지지 않은 실제 정답 x\mathbf{x}에 의존하므로 직접 계산할 수 없다. 그러나 Stein의 항등식을 이용하면, 다음과 같은 관측 가능한 대체량(observable surrogate) 을 얻을 수 있다.

SURE(D;x~)=D(x~)x22+2σ2x~D(x~)Dσ2\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}})=\|\mathbf{D}(\tilde{\mathbf{x}})-\mathbf{x}\|_2^2+2\sigma^2\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})-D\sigma^2

여기서 x~D(x~)\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})D\mathbf{D}의 발산(divergence)을 의미한다. 다시 한번 강조하자면, SURE(D;x~)\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}})는 깨끗한 신호 x\mathbf{x}가 아니라 관측된 노이즈 샘플 x~\tilde{\mathbf{x}}만을 필요로 한다.

📐 Stein 항등식을 이용한 대체량 유도

Stein 항등식은 부분적분 기법을 이용하여, 알려지지 않은 density 하의 기댓값을 관측 가능한 함수와 그 도함수에 대한 기댓값으로 변환하는 방법이다. 이 과정에서 정규화 상수(partition function)가 상쇄되며, 알려지지 않은 density나 정규화 상수를 직접 계산하지 않고도 불편향(unbiased)이고 tractable한 목적 함수와 검증을 가능하게 한다.

우선 가장 단순한 1차원의 경우부터 시작하여, 이후 SURE의 대체 손실을 증명하는 데 필요한 형태로 확장해보겠다.

  • 1차원 표준 정규분포의 경우
    zN(0,1)z\sim\mathcal{N}(0,1)이고 함수 ff가 적절한 감쇠(decay; limzf=0\lim_{|z|\to\infty}f=0) 조건을 만족한다고 하자. 이때 Stein 항등식은 다음과 같이 주어진다.

    E[f(z)]=E[zf(z)]\mathbb{E}[f'(z)]=\mathbb{E}[zf(z)]

    ϕ(z):=12πexp(z2/2)\phi(z):=\frac{1}{\sqrt{2\pi}}\exp(-z^2/2)를 1차원 표준 정규분포의 density라고 하자. 위 식에 대한 증명은 ϕ(z)=zϕ(z)\phi'(z)=-z\phi(z)라는 사실과, 경계 항이 사라진다는 점을 이용한 부분적분으로 이루어진다. 이를 명확히 보기 위해 다음을 계산해보자.

    E[f(Z)]=f(z)ϕ(z) dz\mathbb{E}[f'(Z)]=\int^\infty_{-\infty}f(z)\phi(z)~dz

    부분적분을 적용하여 u=f(z), dv=ϕ(z) dzu=f(z),~dv=\phi'(z)~dz로 두면,

    f(z)ϕ(z) dz=[f(z)ϕ(z)]Boundary term  0f(z)ϕ(z) dz\int f'(z)\phi(z)~dz=\underbrace{\bigg[f(z)\phi(z)\bigg]_{-\infty}^\infty}_{\text{Boundary term}~\to~0}-\int f(z)\phi'(z)~dz

    ϕ(z)=zϕ(z)\phi'(z)=-z\phi(z)이고, z|z|\to\infty일 때 f(z)ϕ(z)0f(z)\phi(z)\to0 이므로 경계 항은 사라진다. 따라서,

    E[f(z)]=f(z)zϕ(z) dz=E[zf(z)]\mathbb{E}[f'(z)]=\int f(z)z\phi(z)~dz=\mathbb{E}[zf(z)]

    이로써 1차원 케이스에서의 Stein 항등식이 증명되었다.

  • 다변량 표준 정규분포의 경우
    zN(0,ID)\mathbf{z}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_D)이고 g:RDRg:\mathbb{R}^D\to\mathbb{R}일 때, Stein 항등식은 다음과 같다.

    E[g(z)]=E[zg(z)]\mathbb{E}[\nabla g(\mathbf{z})]=\mathbb{E}[\mathbf{z}g(\mathbf{z})]

    이에 대한 동치로, u:RDRD\mathbf{u}:\mathbb{R}^D\to\mathbb{R}^D에 대해 다음이 성립한다.

    E[zu(z)]=E[zu(z)]\mathbb{E}[\nabla_\mathbf{z}\cdot\mathbf{u}(\mathbf{z})]=\mathbb{E}[\mathbf{z}^\top\mathbf{u}(\mathbf{z})]

위에서 보인 Stein 항등식을 바탕으로, x~=x+σz\tilde{\mathbf{x}}=\mathbf{x}+\sigma\mathbf{z}이고, zN(0,ID)\mathbf{z}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_D)이며, 적절한 정칙성을 만족하는 임의의 벡터 함수 a\mathbf{a}에 대해 다음이 성립한다.

E[(x~x)a(x~)x]=σE[za(x+σz)x]=σE[za(x+σz)x]=σ2E[x~a(x~)x]\small \mathbb{E}\big[(\tilde{\mathbf{x}}-\mathbf{x})^\top\mathbf{a}(\tilde{\mathbf{x}})\big|\mathbf{x}\big]=\sigma\mathbb{E}\big[\mathbf{z}^\top\mathbf{a}(\mathbf{x}+\sigma\mathbf{z})\big|\mathbf{x}\big]=\sigma\mathbb{E}\big[\nabla_\mathbf{z}\cdot\mathbf{a}(\mathbf{x}+\sigma\mathbf{z})\big|\mathbf{x}\big]=\sigma^2\mathbb{E}\big[\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{a}(\tilde{\mathbf{x}})\big|\mathbf{x}\big]

여기서 z=σx~\nabla_\mathbf{z}=\sigma\nabla_{\tilde{\mathbf{x}}}를 사용하였다. 이제 이를 바탕으로 조건부 MSE로부터 SURE을 유도해보자.

앞서 정의한 조건부 MSE를 x~\tilde{\mathbf{x}}를 기준으로 전개하면,

R(D;x)=E[(D(x~)x~)+(x~x)22x]=E[D(x~)x~22x]+2E[(D(x~)x~)(x~x)x]+E[x~x22x]=E[D(x~)x~22x]+2(E[(x~x)D(x~)x]σ2E[x~D(x~)x]  (1)E[(x~x)x~x]σ2D  (2))+E[x~x22x]σ2D=E[D(x~)x~22+2σ2x~D(x~)Dσ2x]\small \begin{aligned} R(\mathbf{D};\mathbf{x})&=\mathbb{E}\big[\|(\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}})+(\tilde{\mathbf{x}}-\mathbf{x})\|_2^2\big|\mathbf{x}\big] \\ &=\mathbb{E}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}\|_2^2\big|\mathbf{x}\big]+2\mathbb{E}\big[(\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}})^\top(\tilde{\mathbf{x}}-\mathbf{x})\big|\mathbf{x}\big]+\mathbb{E}\big[\|\tilde{\mathbf{x}}-\mathbf{x}\|_2^2\big|\mathbf{x}\big] \\ &=\mathbb{E}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}\|_2^2\big|\mathbf{x}\big]+2\bigg(\underbrace{\mathbb{E}\big[(\tilde{\mathbf{x}}-\mathbf{x})^\top\mathbf{D}(\tilde{\mathbf{x}})\big|\mathbf{x}\big]}_{\sigma^2\mathbb{E}[\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})|\mathbf{x}]~\cdots~(1)}-\underbrace{\mathbb{E}\big[(\tilde{\mathbf{x}}-\mathbf{x})^\top\tilde{\mathbf{x}}\big|\mathbf{x}\big]}_{\sigma^2D~\cdots~(2)}\bigg)+\underbrace{\mathbb{E}\big[\|\tilde{\mathbf{x}}-\mathbf{x}\|_2^2\big|\mathbf{x}\big]}_{\sigma^2D} \\ &=\mathbb{E}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}\|_2^2+2\sigma^2\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})-D\sigma^2\big|\mathbf{x}\big] \\ \end{aligned}

여기서 (1)(1), (2)(2)에서 Stein 항등식을 사용하였다. 따라서 관측 가능한 대체량

SURE(D;x~):=D(x~)x~22+2σ2x~D(x~)Dσ2\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}}):=\|\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}\|_2^2+2\sigma^2\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})-D\sigma^2

은 다음을 만족한다.

E[SURE(D;x~)x]=R(D;x)\mathbb{E}\big[\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}})\big|\mathbf{x}\big]=R(\mathbf{D};\mathbf{x})

즉, SURE를 (기댓값에서 혹은 경험적으로) 최소화하는 것은 노이즈 관측값만을 사용하여 실제 조건부 MSE를 최소화하는 것과 동등하다. 이로써 Stein 항등식으로부터 SURE 대체량을 유도하였다. ~_\blacksquare

직관적으로 보면 SURE는 서로를 보완하는 두 가지 항으로 이루어져 있다.

  1. D(x~)x~22\|\mathbf{D}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}\|_2^2는 denoiser의 출력이 노이즈가 섞인 입력 x~\tilde{\mathbf{x}}로부터 얼마나 떨어져 있는지를 측정한다. 그러나 x~\tilde{\mathbf{x}} 자체가 이미 손상된 관측값이기 때문에 이 항만으로는 실제 오차를 과소평가 하게 된다.

  2. 발산항 x~D(x~)\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{D}(\tilde{\mathbf{x}})은 이러한 과소평가를 보정 하는 역할을 한다. 이 항은 denoiser가 입력에 가해진 작은 섭동(perturbation)에 얼마나 민감하게 반응하는지를 포착하며, 그 결과 노이즈로 인해 추가된 분산(variance)를 효과적으로 반영하게 된다.

🔗 Tweedie 공식 및 Bayes 최적성과의 연결

예를 들어 노이즈가 섞인 marginal 분포를 pσ(x~)=(pdataN(0,σ2I))(x~)p_\sigma(\tilde{\mathbf{x}})=(p_\text{data}\ast\mathcal{N}(\mathbf{0},\sigma^2\mathbf{I}))(\tilde{\mathbf{x}})라 설정해보자. 앞서 언급한 것 처럼, SURE는 x\mathbf{x}에 조건화된 상태에서, 노이즈에 대한 평균제곱오차의 불편향 추정량이다.

Ex~x[SURE(D;x~)]=Ex~x[D(x~)x22]\mathbb{E}_{\tilde{\mathbf{x}}|\mathbf{x}}\big[\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}})\big]=\mathbb{E}_{\tilde{\mathbf{x}}|\mathbf{x}}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\mathbf{x}\|_2^2\big]

따라서 기댓값에서의 SURE를 최소화하는 것은 Bayes 위험(Bayes' risk)

E(x,x~)[D(x~)x22]=Ex~[Exx~[D(x~)x22]]\mathbb{E}_{(\mathbf{x},\tilde{\mathbf{x}})}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\mathbf{x}\|_2^2\big]=\mathbb{E}_{\tilde{\mathbf{x}}}\Big[\mathbb{E}_{\mathbf{x}|\tilde{\mathbf{x}}}\big[\|\mathbf{D}(\tilde{\mathbf{x}})-\mathbf{x}\|_2^2\big]\Big]

을 최소화하는 것과 동등하다. 여기서 마지막 등식은 전체 기댓값의 법칙(law of total expectation) 을 사용하였다.

이러한 분해는 안쪽 기댓값 Exx~\mathbb{E}_{\mathbf{x}|\tilde{\mathbf{x}}}에 대해 pointwise 최적화 minzRDExx~[zx22]\min_{\mathbf{z}\in\mathbb{R}^D}\mathbb{E}_{\mathbf{x}|\tilde{\mathbf{x}}}[\|\mathbf{z}-\mathbf{x}\|_2^2]를 가능하게 해준다. 거의 모든 x~\tilde{\mathbf{x}}에 대해,

D(x~)=arg minzExx~[zx22]=E[xx~]\mathbf{D}^\ast(\tilde{\mathbf{x}})=\argmin_\mathbf{z}\mathbb{E}_{\mathbf{x}|\tilde{\mathbf{x}}}\big[\|\mathbf{z}-\mathbf{x}\|_2^2\big]=\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]

따라서 SURE-최적 denoiser는 앞 절에서의 Bayes 추정기와 일치하며, Tweedie 공식에 의해 다음과 같이 주어진다.

D(x~)=E[xx~]=x~+σ2x~logpσ(x~)\mathbf{D}^\ast(\tilde{\mathbf{x}})=\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]=\tilde{\mathbf{x}}+\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})

🫱🏻‍🫲🏾SURE과 Score Matching의 관계

바로 위 식에 주어진 항등식은 denoiser D\mathbf{D}를 score field를 통해 매개변수화하도록 동기를 부여한다. 즉,

D(x~)=x~+σ2sϕ(x~;σ)\mathbf{D}(\tilde{\mathbf{x}})=\tilde{\mathbf{x}}+\sigma^2\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)

여기서 sϕ(x~;σ)\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)는 노이즈가 섞인 score x~logpσ()\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\cdot)를 근사하는 함수로 작용한다.

이제 D(x~)=x~+σ2sϕ(x~;σ)\mathbf{D}(\tilde{\mathbf{x}})=\tilde{\mathbf{x}}+\sigma^2\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)SURE(D;x~)\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}}) 식에 대입하면 다음을 얻을 수 있다..

12σ4SURE(D;x~)=Tr(x~sϕ(x~;σ))+12sϕ(x~;σ)22+Cσ\frac{1}{2\sigma^4}\operatorname{SURE}(\mathbf{D};\tilde{\mathbf{x}})=\operatorname{Tr}(\nabla_{\tilde{\mathbf{x}}}\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma))+\frac{1}{2}\|\mathbf{s}_\phi(\tilde{\mathbf{x}};\sigma)\|_2^2+C_\sigma

여기서 CσC_\sigma는 상수이다. 따라서 x~pσ\tilde{\mathbf{x}}\sim p_\sigma에 대해 기댓값을 취하면, SURE를 최소화하는 것은 분포 pσp_\sigma하에서의 기댓값으로 정의된 Hyvärinen의 대안적 score matching 목적 함수,

L~SM(ϕ):=Expdata[Tr(xsϕ(x))+12sϕ(x)22]\tilde{\mathcal{L}}_\text{SM}(\phi):=\mathbb{E}_{\mathbf{x}\sim p_\text{data}}\Big[\operatorname{Tr}(\nabla_\mathbf{x}\mathbf{s}_\phi(\mathbf{x}))+\frac{1}{2}\|\mathbf{s}_\phi(\mathbf{x})\|_2^2\Big]

를 최소화하는 것과 동등하다. 그 결과 두 목적 함수는 동일한 최소해를 가지며, 이는 위에서 주어진 denoiser E[xx~]\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]이다.

6️⃣ 일반화된 Score Matching

고전적인 score matching, denoising score matching, 그리고 그 고차 변형들은 모두 공통적으로 어떤 density pp에 대해

Lp(x)p(x),for some density p\frac{\mathcal{L}p(\mathbf{x})}{p(\mathbf{x})},\quad\text{for some density}~p

의 형태를 목표로 한다. 여기서 L\mathcal{L}은 density에 적용하는 선형 연산자(linear operator) 이다. 고전적인 경우에는 L=x\mathcal{L}=\nabla_\mathbf{x}이며, 이때

xp(x)p(x)=xlogp(x)\frac{\nabla_\mathbf{x} p(\mathbf{x})}{p(\mathbf{x})}=\nabla_\mathbf{x}\log p(\mathbf{x})

를 얻는다.

이와 같은 Lpp\frac{\mathcal{L}p}{p} 구조는 부분적분을 통해 정규화 상수를 제거할 수 있게 해주며, 그 결과 분포 pp로부터의 샘플과 학습된 벡터장 sϕ\mathbf{s}_\phi에만 의존하는 tractable한 목적 함수를 얻을 수 있다. 이러한 관점이 바로 일반화된 score matching 프레임워크를 도입하게 되는 동기를 제공한다.

🎣 일반화된 Fisher 발산

데이터 분포 pp, 임의의 모델 분포를 qq라고 하자. x\mathbf{x}의 스칼라 함수에 작용하는 선형 연산자 L\mathcal{L}에 대해 일반화된 Fisher 발산을 다음과 같이 정의한다.

DL(pq):=p(x)Lp(x)p(x)Lq(x)q(x)22 dx~D_\mathcal{L}(p\Vert q):=\int p(\mathbf{x})\bigg\Vert\frac{\mathcal{L}p(\mathbf{x})}{p(\mathbf{x})}-\frac{\mathcal{L}q(\mathbf{x})}{q(\mathbf{x})}\bigg\Vert_2^2~d\tilde{\mathbf{x}}

만약 L\mathcal{L}완전(complete) 하다면, 즉

Lp1p1=Lp2p2 a.e.impliesp1=p2 a.e.\frac{\mathcal{L}p_1}{p_1}=\frac{\mathcal{L}p_2}{p_2}~\text{a.e.}\quad\text{implies}\quad p_1=p_2~\text{a.e.}

가 성립할 경우, DL(pq)=0D_\mathcal{L}(p\Vert q)=0q=pq=p임을 나타낸다. 특히 L=x~\mathcal{L}=\nabla_{\tilde{\mathbf{x}}}인 경우에는 고전적인 Fisher 발산을 다시 얻게 된다.

🎛️ Score 매개변수화

실제로는 정규화된 밀도 qq직접 모델링하지 않는다. 대신, 일반화된 score Lp(x)p(x)\frac{\mathcal{L}p(\mathbf{x})}{p(\mathbf{x})}를 근사하기 위해 벡터장 sϕ(x)\mathbf{s}_\phi(\mathbf{x})를 직접 매개변수화한다. 우선 다음을 한번 살펴보자.

DL(psϕ):=Exp[sϕ(x)Lp(x)p(x)22]\mathcal{D}_\mathcal{L}(p\Vert\mathbf{s}_\phi):=\mathbb{E}_{\mathbf{x}\sim p}\left[\bigg\Vert\mathbf{s}_\phi(\mathbf{x})-\frac{\mathcal{L}p(\mathbf{x})}{p(\mathbf{x})}\bigg\Vert_2^2\right]

비록 Lp(x)p(x)\frac{\mathcal{L}p(\mathbf{x})}{p(\mathbf{x})}는 알려져 있지 않지만, 부분적분을 사용하면 이 손실은 score에만 의존하도록 만들 수 있다. L\mathcal{L}수반 연산자(adjoint)L\mathcal{L}^\dag라 하고, 다음과 같이 정의해보자.

(Lf)g=f(Lg)\int(\mathcal{L}f)^\top g=\int f(\mathcal{L}^\dag g)

이는 경계항이 사라질 때 적분 안에서 L\mathcal{L}형식적으로 반대편으로 옮기는 역할 을 한다. 제곱을 전개하고 이 항등식을 적용하면, 다음과 같은 tractable한 목적 함수를 얻는다.

LGSM(ϕ)=Exp[12sϕ(x)22(Lsϕ)(x)]+C\boxed{ \mathcal{L}_\text{GSM}(\phi)=\mathbb{E}_{\mathbf{x}\sim p}\left[\frac{1}{2}\|\mathbf{s}_\phi(\mathbf{x})\|_2^2-(\mathcal{L}^\dagger\mathbf{s}_\phi)(\mathbf{x})\right]+C }

여기서 상수항 CCϕ\phi에 의존하지 않는다. 실제로는 분포 pp를 기댓값을 통해서만 사용하므로, 일반화된 score matching 손실은 고전적인 score matching과 마찬가지로 학습 데이터로부터 경험적 추정량을 갖는다.

L=\mathcal{L}=\nabla인 경우, 그 수반 연산자는 L=\mathcal{L}^\dag=-\nabla가 된다. 이로부터 Hyvärinen의 score matching 목적 함수 Ep[12sϕ22+sϕ]\mathbb{E}_p\big[\frac{1}{2}\|\mathbf{s}_\phi\|_2^2+\nabla\cdot\mathbf{s}_\phi\big]를 그대로 다시 얻게 된다.

💁🏻‍♂️ 선형 연산자와 수반 연산자 – Linear Operator & Adjoint Operator
유한차원 내적공간 VV와 선형 연산자 T:VVT:V\to V를 생각하자. 이때 다음이 성립한다.

TS(V),  !TS(V)  s.t.  x,yV,  Tx,y=x,Ty\forall\, T \in \mathcal{S}(V),\; \exists!\, T^\dagger \in \mathcal{S}(V) \;\text{s.t.}\; \forall x,y \in V,\; \langle T x, y\rangle = \langle x, T^\dagger y\rangle

여기서 S(V)\mathcal{S}(V)VV위의 존재 가능한 모든 선형 연산자들의 집합이고, T:VVT^\dag:V\to V는 선형 연산자 TT에 대한 수반 연산자(adjoint operator)이다.

1차원 함수공간에서의 수반 연산자
함수 ff, gg에 대해 두 함수의 내적은 다음과 같이 정의된다.

f,g:=abf(x)g(x) dx\langle f,g\rangle:=\int_a^b f(x)g(x)~dx

이때 선형 연산자 T=ddxT=\frac{d}{dx}에 대응하는, Tf,g=f,Tg\langle Tf,g\rangle=\langle f,T^\dag g\rangle를 만족하는 수반 연산자 TT^\dag를 찾기 위해서는 아래와 같이 부분적분을 적용하면 된다.

Tf,g=abdfdxg dx=[fg]abRfdgdx dx\langle Tf,g\rangle=\int_a^b \frac{df}{dx}\cdot g~dx=\big[f\cdot g\big]_a^b-\int_\mathbb{R}f\cdot\frac{dg}{dx}~dx

여기서 x{a,b}:f(x)=g(x)=0\forall x\in\{a,b\}:f(x)=g(x)=0를 가정하면 경계항은 사라지고 dfg=fdg\int df\cdot g=-\int f\cdot dg이므로, 수반 연산자는 다음과 같이 음의 미분연산자가 된다.

T=(ddx)=(ddx)T^\dag=\left(\frac{d}{dx}\right)^\dag=-\left(\frac{d}{dx}\right)

다차원 함수공간에서의 수반 연산자
다차원 함수에 대해서는 내적과 연산자가 함께 바뀐다. 스칼라 함수 f:RDRf:\mathbb{R}^D\to\mathbb{R}와 벡터장 v:RDRD\mathbf{v}:\mathbb{R}^D\to\mathbb{R}^D에 대한 내적은 fv dx\int f^\top\mathbf{v}~d\mathbf{x}로 정의되고, 선형 연산자 T=T=\nabla를 적용한 내적은 다음과 같다.

f,v=ΩRD(f(x))v(x) dx=i=1DΩfxivi dx\langle\nabla f,\mathbf{v}\rangle=\int_{\Omega\subset\mathbb{R}^D}(\nabla f(\mathbf{x}))^\top\mathbf{v}(\mathbf{x})~d\mathbf{x}=\sum_{i=1}^D\int_\Omega\frac{\partial f}{\partial x_i}v_i~d\mathbf{x}

이제 각 ii에 대해 1차원 부분적분을 적용해보면,

Ωfxivi dx=Ωf vi ni dSΩfvixi dx\int_\Omega\frac{\partial f}{\partial x_i}v_i~d\mathbf{x}=\int_{\partial\Omega}f~v_i~n_i~dS-\int_\Omega f\frac{\partial v_i}{\partial x_i}~d\mathbf{x} \\

이 성립한다. 여기서 n=(n1,n2,,nD)ΩRD\mathbf{n}=(n_1,n_2,\ldots,n_D)\in\partial\Omega\to\mathbb{R}^D는 경계 Ω\partial\Omega의 바깥 방향 법선 벡터(normal vector)이다. 이제 이 결과를 모두 합치면 다음과 같이 정리된다.

Ω(f(x))v(x) dx=Ωf(vn) dSΩf(v) dx\int_\Omega(\nabla f(\mathbf{x}))^\top\mathbf{v}(\mathbf{x})~d\mathbf{x}=\int_{\partial\Omega}f(\mathbf{v}\cdot\mathbf{n})~dS-\int_\Omega f(\nabla\cdot\mathbf{v})~d\mathbf{x}

여기에 경계 조건 xΩ:f(x)=0  v(x)n(x)=0\forall\mathbf{x}\in\partial\Omega:f(\mathbf{x})=0~\vee~\mathbf{v}(\mathbf{x})\cdot\mathbf{n}(\mathbf{x})=0을 가정하면 경계항 Ωf(vn) dS\int_{\partial\Omega}f(\mathbf{v}\cdot\mathbf{n})~dS은 사라지게 된다.

즉, 결론적으로

Ω(f)v dx=Ωf(v) dx\int_\Omega(\nabla f)^\top\mathbf{v}~d\mathbf{x}=-\int_\Omega f(\nabla\cdot\mathbf{v})~d\mathbf{x}

이 되어 f,v=f,v\langle \nabla f,v\rangle=\langle f,-\nabla\cdot\mathbf{v}\rangle이 성립하고, 따라서 T=T=\nabla의 수반 연산자는 다음과 같이

T=()=T^\dag=(\nabla)^\dag=-\nabla

음의 그래디언트(발산)이 된다.

📏 선형 연산자들의 예시

  • 고전적 Score Matching
    고전 SM은 score logp(x)=p(x)p(x)\nabla\log p(\mathbf{x})=\frac{\nabla p(\mathbf{x})}{p(\mathbf{x})}를 벡터장 sϕ(x)\mathbf{s}_\phi(\mathbf{x})로 맞추는 문제로 볼 수 있다. 이때 선형 연산자는 L=\mathcal{L}=\nabla로 설정한다. 이를 일반화된 SM에 대입하면,

    LGSM(ϕ)=Exp[12sϕ(x)22(sϕ(x))]+C=L~SM(ϕ)\mathcal{L}_\text{GSM}(\phi)=\mathbb{E}_{\mathbf{x}\sim p}\bigg[\frac{1}{2}\|\mathbf{s}_\phi(\mathbf{x})\|_2^2-(-\nabla\cdot\mathbf{s}_\phi(\mathbf{x}))\bigg]+C=\tilde{\mathcal{L}}_\text{SM}(\phi)

    즉, 수반 연산자 L=\mathcal{L}^\dag=-\nabla를 대입하기만 하면 바로 Hyvärinen의 score matching의 목적 함수를 얻을 수 있다.

  • 노이즈 제거 Score Matching (DSM)
    가산적 Gaussian 노이즈 하에서의 DSM의 선형 연산자는 L=x~+σ2\mathcal{L}=\tilde{\mathbf{x}}+\sigma^2\nabla로 설정한다. 이에 대한 수반 연산자는 L=x~σ2\mathcal{L}^\dag=\tilde{\mathbf{x}}-\sigma^2\nabla로 정의된다. 이를 일반화된 SM에 대입하면,

    LGSM(ϕ)=Ex~pσ[12sϕ(x~)22x~sϕ(x~)+σ2x~sϕ(x~)]+C\mathcal{L}_\text{GSM}(\phi)=\mathbb{E}_{\tilde{\mathbf{x}}\sim p_\sigma}\bigg[\frac{1}{2}\|\mathbf{s}_\phi(\tilde{\mathbf{x}})\|_2^2-\tilde{\mathbf{x}}\mathbf{s}_\phi(\tilde{\mathbf{x}})+\sigma^2\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{s}_\phi(\tilde{\mathbf{x}})\bigg]+C

    이제 Euler-Lagrange 최적화 를 이용해 이 목적 함수의 최소해를 구해보자. 그러기 위해선 우선 최소화하려는 함수형(functional) 을 정의해야 한다.

    J[sϕ]:=RDpσ(x~)(12sϕ(x~)22x~sϕ(x~)+σ2x~sϕ(x~)) dx~\mathcal{J}[\mathbf{s}_\phi]:=\int_{\mathbb{R}^D}p_\sigma(\tilde{\mathbf{x}})\left(\frac{1}{2}\|\mathbf{s}_\phi(\tilde{\mathbf{x}})\|_2^2-\tilde{\mathbf{x}}^\top\mathbf{s}_\phi(\tilde{\mathbf{x}})+\sigma^2\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{s}_\phi(\tilde{\mathbf{x}})\right)~d\tilde{\mathbf{x}}

    이제 이 함수형에 대해 변분 미분 을 취해보면,

    δJδsϕ=RDpσ(x~) (s(x~)x~σ2x~logpσ(x~))h(x~) dx~\frac{\delta\mathcal{J}}{\delta\mathbf{s}_\phi}=\int_{\mathbb{R}^D} p_\sigma(\tilde{\mathbf{x}})~(\mathbf{s}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}-\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}))^\top\mathbf{h}(\tilde{\mathbf{x}})~d\tilde{\mathbf{x}}

    여기서 h\mathbf{h}sϕ\mathbf{s}_\phi에 대한 변분 방향을 나타내며, 특히 세 번째 항(발산항) 은 다음과 같은 세부 과정을 통해 변분 미분이 이루어졌다.
    우선 함수형의 세 번째 항은 J3[sϕ]=σ2RDpσ(x~)x~sϕ(x~) dx~\mathcal{J}_3[\mathbf{s}_\phi]=\sigma^2\int_{\mathbb{R}^D}p_\sigma(\tilde{\mathbf{x}})\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{s}_\phi(\tilde{\mathbf{x}})~d\tilde{\mathbf{x}}이다. 이제 이 항을 εh\varepsilon\mathbf{h}만큼 이동시키면 다음과 같다.

    J3[sϕ+εh]=σ2RDpσ(x~) (x~sϕ+εx~h) dx~\mathcal{J}_3[\mathbf{s}_\phi+\varepsilon\mathbf{h}]=\sigma^2\int_{\mathbb{R}^D}p_\sigma(\tilde{\mathbf{x}})~(\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{s}_\phi+\varepsilon\nabla_{\tilde{\mathbf{x}}}\cdot\mathbf{h})~d\tilde{\mathbf{x}}

    따라서 변분법의 정의에 따라 이를 ε\varepsilon에 대해 미분하면,

    ddεε=0J3[sϕ+εh]=σ2RDpσ(x~)x~h(x~) dx~\frac{d}{d\varepsilon}\bigg\vert_{\varepsilon=0}\mathcal{J}_3[\mathbf{s}_\phi+\varepsilon\mathbf{h}]=\sigma^2\int_{\mathbb{R}^D}p_\sigma(\tilde{\mathbf{x}})\nabla_{\tilde{\mathbf{x}}}\mathbf{h}(\tilde{\mathbf{x}})~d\tilde{\mathbf{x}}

    이제 위 식에 부분적분(발산 정리)를 적용한다.

    pσh=hpσ,ddεε=0J3[sϕ+εh]=σ2RDh(x~)x~pσ(x~) dx~=σ2RDpσ(x~)h(x~)x~logpσ(x~) dx~\small \begin{aligned} \int p_\sigma\nabla\cdot\mathbf{h}=-\int\mathbf{h}^\top\nabla p_\sigma,\quad\frac{d}{d\varepsilon}\bigg\vert_{\varepsilon=0}\mathcal{J}_3[\mathbf{s}_\phi+\varepsilon\mathbf{h}]&=-\sigma^2\int_{\mathbb{R}^D}\mathbf{h}(\tilde{\mathbf{x}})^\top\nabla_{\tilde{\mathbf{x}}}p_\sigma(\tilde{\mathbf{x}})~d\tilde{\mathbf{x}} \\ &=-\sigma^2\int_{\mathbb{R}^D}p_\sigma(\tilde{\mathbf{x}})\mathbf{h}(\tilde{\mathbf{x}})^\top\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})~d\tilde{\mathbf{x}} \\ \end{aligned}

    이로써 δJ3/δsϕ\delta\mathcal{J}_3/\delta\mathbf{s}_\phi 항을 계산하였다. δJ/δsϕ\delta\mathcal{J}/\delta\mathbf{s}_\phi이 모든 시험 함수 h\mathbf{h}에 대해 00이 되려면 다음이 성립해야 한다(Euler-Lagrange 조건 ).

    pσ(x~) (s(x~)x~σ2x~logpσ(x~))=0a.e.p_\sigma(\tilde{\mathbf{x}})~(\mathbf{s}(\tilde{\mathbf{x}})-\tilde{\mathbf{x}}-\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}}))=0\quad\text{a.e.}

    하지만, 거의 대부분 x~\tilde{\mathbf{x}}에서 pσ(x~)>0p_\sigma(\tilde{\mathbf{x}})>0이므로,

    sϕ(x~)=x~+σ2x~logpσ(x~)\mathbf{s}_\phi(\tilde{\mathbf{x}})=\tilde{\mathbf{x}}+\sigma^2\nabla_{\tilde{\mathbf{x}}}\log p_\sigma(\tilde{\mathbf{x}})

    즉, Tweedie 공식을 유도하게 된 것이다. 여기서 앞서 소개한 denoiser E[xx~]\mathbb{E}[\mathbf{x}|\tilde{\mathbf{x}}]를 얻을 수 있다.

일반화된 score matching은 연속 변수에 국한되지 않고, 언어 모델링을 포함한 이산 변수로까지 확장된다. 또한 이는 denoising 스타일의 목적 함수를 산출하는 score-inspired된 학습을 정당화한다. 이러한 연산자 관점은 다양한 목적 함수들을 하나의 틀로 통합하고, 데이터로부터의 경험적 추정을 가능하게 하며, 선형 연산자 L\mathcal{L}의 적절한 선택을 통해 손실 함수를 설계하는 일반적인 원리를 제공한다.


✅ 정리

이번 글에서는 score matching의 계산적 한계를 출발점으로 삼아, DSM(Denoising Score Matching) 이 어떻게 보다 안정적이고 확장 가능한 학습 프레임워크로 이어지는지를 체계적으로 정리했다. Sliced score matching을 통해 Jacobian 계산을 회피하는 아이디어를 살펴본 뒤, 노이즈 주입과 조건화를 통해 intractable한 score 학습을 tractable한 회귀 문제로 변환하는 DSM의 핵심 원리를 도입했다. 특히 Gaussian 노이즈 설정에서는 DSM의 최적해가 노이즈가 섞인 marginal 분포의 score와 일치함을 보였고, 이는 Tweedie 공식을 통해 denoising, score 예측, Bayes 최적 추정이 하나로 연결됨을 드러낸다. 더 나아가 SURE와의 등가성, 그리고 일반화된 score matching(GSM) 관점은 score matching, DSM, denoising 목적함수들을 하나의 연산자 기반 프레임워크로 통합하며, 적절한 선형 연산자 L\mathcal{L} 선택을 통해 새로운 손실 함수를 설계할 수 있는 원리적 기반을 제공한다. 결과적으로 DSM은 단순한 테크닉을 넘어, 확산 모델과 score 기반 생성 모델 전반을 관통하는 이론적 중심축임을 확인할 수 있다.


📄 출처

[1] Lai, Chieh-Hsin, et al. The Principles of Diffusion Models. arXiv, 24 Oct. 2025, arXiv:2510.21890.

profile
Korea Univ. Computer Science & Engineering

0개의 댓글