학습 가능한 확산항과 밀도 기반 손실로 생성 SDE의 이탈 샘플을 줄인다

무엇인가

이 논문은 데이터 분포를 생성하는 신경 SDE에서 드리프트뿐 아니라 확산 계수까지 학습하고, 생성된 최종 밀도 자체를 벌점으로 삼는 손실로 훈련하는 문제를 다룬다. 기존 score 기반 확산 모델은 보통 시간 의존 score를 회귀하고 확산 계수는 미리 정하며, flow matching은 ODE 속도를 학습하므로 표준 회귀 목적이 최종 밀도의 KL 오차를 직접 최소화하지 않는다. 저자는 이탈 샘플과 분포 오차를 줄이려면 최종 밀도에 직접 의존하는 손실과 학습 가능한 확산이 필요하다고 본다.

어떻게 동작하나

모델은 이토 SDE dx=Fγ(x)dt+σγ(x)dB_t, x0~N(m0,s0^2 I)이다. Fγ와 σγ는 신경망으로 두고 m0와 공유 표준편차 s0도 함께 학습한다. 드리프트는 Fmax 상한을 갖도록 유계 변환과 반경 1~2의 매끄러운 radial cutoff를 사용해 안쪽은 신경망, 바깥은 안쪽으로 향하는 일정 크기 항으로 잇는다. 확산은 sigmoid로 σ_min과 σ_max 사이에 제한되며 바깥에서는 σ_min으로 수렴한다. DK 업데이트가 상태에 대한 2차 도함수를 필요로 해 SiLU 같은 매끄러운 활성화를 쓰고, 실험에서는 시간 독립적인 장만 학습한다.

무엇과 다른가

핵심 훈련법은 divergence-kernel(DK) 공식으로 최종 로그밀도의 파라미터 도함수를 경로별 응답의 조건부 기대값으로 표현하는 것이다. 각 샘플 경로에서 score 정보를 담는 covector ν_t와 응답 T_t를 전진 적분하고, δ log h_T(x)=E[T_T|x_T=x]를 이용해 forward KL, backward KL, 둘의 평균(제프리스 발산의 절반)의 gradient를 만든다. terminal coefficient는 forward KL이 p/h_T, backward KL이 log(h_T/p), 평균이 1/2(p/h_T+log(p/h_T))이며 이전 업데이트 평균을 baseline으로 뺀다. 상태·covector·응답은 같은 Brownian increment로 오일러 이산화해 누적하고, 응답 증가분은 국소 자동미분으로 계산해 전체 궤적을 저장하지 않는다. 최종 밀도 h_T는 Gaussian KDE로 추정해 terminal weight에 쓴다.

어떻게 쓰나

2차원 실험은 확률 0.1~0.4, 평균 반경 3, 폭과 방향이 다른 네 가우시안 혼합이다. DK는 드리프트 3×24 은닉층과 확산 2×18 은닉층, 초기 가우시안을 합쳐 1,740개 파라미터를 학습했고, EDM score 기반 베이스라인은 1,827개로 조금 더 많다. 1,024개 경로, T=1, Δt=0.002(500스텝), Adam lr 1e-3, 5,000 업데이트다. DK는 EDM보다 이탈 샘플이 적고 남은 이탈 샘플도 평균적으로 정답 데이터에 더 가까웠다. DK 내에서는 backward KL이 평균 외부 거리를 가장 작게, averaged KL이 모드 확률 오차를 가장 작게 만들었고, 세 DK 모두 EDM보다 validation/test averaged KL이 낮았으며 averaged KL DK가 test averaged KL 최저였다. 목표 분포에서 직접 뽑은 두 개의 10,000 샘플 참조 세트는 추정 averaged KL 0.00415, precision 98.4%였고 averaged-KL DK의 precision이 이에 거의 맞먹었다. 5,000 업데이트 벽시계 시간은 averaged-KL DK 924초, EDM 159초였고 네트워크만의 median forward 시간은 둘 다 27.9마이크로초/배치였다.

전제와 한계

8차원 실험에서 DK는 2,148개, EDM은 2,270개 파라미터를 썼다. DK가 EDM보다 이탈 샘플을 훨씬 적게 내고 평균 외부 거리도 작았다. 세 DK 모델 모두 EDM보다 validation averaged KL과 test backward KL이 낮았다. 반면 test forward KL과 test averaged KL, 모드 확률 일치에서는 EDM이 더 좋았다. 목표 분포의 두 10,000 샘플 참조는 averaged KL 2.09, precision 90.1%였고 DK는 precision에서는 참조를 넘었지만 averaged KL은 참조보다 높았다. 학습 시간은 averaged-KL DK 1.42×10^3초, EDM 167초였고 forward는 31.2 대 28.8마이크로초/배치였다. 학습된 장에서는 드리프트가 모드 쪽으로 샘플을 끌고, 확산은 모드 사이 빈 공간에서 커지고 모드 중심 근처에서는 하한에 가까워져 과잉 확률을 분산시키면서 모드 주변 퍼짐을 제한하는 패턴이 관찰됐다.

개발자에게 이 논문의 의미는 생성 SDE의 확산 계수를 고정하지 않고 최종 밀도 기반 KL로 함께 튜닝할 수 있다는 점이다. 특히 목표 분포의 모드 사이 빈 영역에서 생기는 이탈 샘플을 줄이는 것이 목적이면 DK 손실과 확산장 학습을 검토할 수 있다. 다만 현재 결과는 2차원·8차원 합성 가우시안 혼합에 한정되고, KDE 기반 밀도 추정과 여러 경계값(Fmax, σ_min, σ_max, α_target), KL 선택에 따라 결과가 달라질 수 있다. 고차원 이미지·텍스트 생성에 그대로 적용할 근거는 아직 없고, 학습 벽시계 시간이 EDM보다 2차원에서 약 5.8배, 8차원에서 약 8.5배 길었다는 비용도 확인해야 한다.

저자가 밝힌 한계와 전제는 분명하다. 실험은 시간 독립적인 신경장과 상태 의존 스칼라 확산 계수만 사용했고, 일반 확산 행렬은 향후 과제로 남겼다. DK 공식은 요구되는 매끄러움·성장·적분 가능성 가정 아래 유도되며, 필드도 상한과 하한, radial cutoff로 제약한다. 또한 8차원에서는 EDM이 forward/averaged KL과 모드 확률에서 더 나아 DK가 모든 지표에서 우월하지 않다. 저자는 score 기반 생성기로 모드와 확률을 잡은 뒤 DK로 이탈 샘플을 다듬는 조합, 사전학습된 score 모델로 드리프트를 초기화하는 방향, DK score identity를 이용해 더 일반적인 noising dynamics의 score 네트워크 학습 target을 만드는 방향을 후속 과제로 제시한다.