MAD: Manifold Attracted Diffusion

기술 요약: 매니폴드 어트랙티드 디퓨전 (Manifold Attracted Diffusion, MAD)

문제 정의
점수 기반 확산 모델(Score-based diffusion models)은 이미지 합성 분야에서 SOTA(State-of-the-art) 방법론으로 자리 잡았으며, 샘플의 품질과 다양성 측면에서 GAN 및 VAE를 종종 능가합니다. 그러나 측정 오차, 압축 아티팩트 또는 데이터 수집 과정에서 발생하는 노이즈로 인해 훈련 데이터가 오염된 경우 심각한 한계가 발생합니다. 노이즈가 섞인 데이터로 직접 훈련된 표준 확산 모델은 이러한 오염을 재현하도록 학습되어, 훈련 데이터의 노이즈 특성을 그대로 유지하는 샘플을 생성하게 됩니다. 노이즈가 섞인 훈련에 대응하는 기존 방식들은 대개 훈련 과정 자체를 수정해야 합니다(예: 손실 함수를 조정하거나 깨끗한 데이터와 노이즈가 섞인 데이터를 혼합하여 사용). 반면, 기존의 훈련 프레임워크와 사전 훈련된 네트워크를 그대로 사용하면서 노이즈 분포로부터 깨끗한 샘플을 생성할 수 있는 효율적인 방법에 대해서는 연구가 부족한 실정입니다.

방법론
제안된 방법인 **매니폴드 어트랙티드 디퓨전(Manifold Attracted Diffusion, MAD)**은 훈련 과정을 수정하는 대신 점수 기반 모델의 추론 과정을 수정함으로써 이 문제를 해결합니다. 이 접근법은 유의미한 데이터가 고차원 앰비언트 공간 내의 저차원 매니폴드에 존재한다는 **매니폴드 가설(Manifold hypothesis)**에 근거합니다. 저자들은 노이즈를 이 구조에 직교하는(off-manifold) 방향의 낮은 진폭 변화로, 유의미한 데이터의 변화는 매니폴드 상의(on-manifold) 변화로 해석합니다.

핵심적인 기술적 혁신은 **확장된 점수(Extended score)**의 도입입니다.

• 정의: 확률 측도 $p$와 파라미터 $\gamma > 0$에 대하여, 확장된 점수 $H_\gamma p$는 $(1 + \gamma)S(p * g_\gamma) + \gamma \frac{d}{d\gamma}S(p * g_\gamma)$로 정의됩니다. 여기서 $S$는 표준 점수 연산자($\nabla \log p$)이고 $g_\gamma$는 가우시안 커널입니다.
• 특성: 특정 방향으로 분산이 0인 디락 델타 분포(데이터가 저차원 매니폴드에 존재함을 나타냄)에 대해 정의되지 않는 표준 점수와 달리, 확장된 점수는 잘 정의됩니다. 결정적으로, 디락 델타에 대해 확장된 점수는 단위 분산을 가진 가우시안의 점수와 유사하게 동작합니다. 이를 통해 추론 과정은 저분산(off-manifold) 방향을 비제로 분산을 가진 것처럼 취급하여, 효과적인 "소프트 임계값 처리(soft thresholding)" 효과를 적용할 수 있습니다.
• 추론 알고리즘: 저자들은 확장된 점수의 근사치를 사용하는 반복적 추론 알고리즘을 제안합니다. 사전 훈련된 디노이저 네트워크 $D_\theta$(표준 점수를 근사함)가 주어졌을 때, 확장된 점수는 약간 편향된 노이즈 레벨에서 네트워크를 평가하고 $\gamma$에 대한 미분을 계산하는 유한 차분법을 통해 근사됩니다.
• 구현: 이 알고리즘은 "매니폴드 인력(manifold attraction)"의 강도를 조절하는 하이퍼파라미터($a, b, p$)를 도입합니다. 이들을 조정함으로써 추론 과정은 매니폴드 밖의 방향(노이즈)에 있는 작은 변화는 억제하는 동시에, 매니폴드 위의 방향(신호)에 있는 큰 변화는 보존합니다. 이 방법은 Karras 등(2022)의 확률 흐름 ODE(probability flow ODE) 정식화와 같은 기존 프레임워크와 호환되며, 사전 훈련된 네트워크를 활용할 수 있습니다.

주요 기여

1. 정의 및 분석: 본 논문은 확장된 점수에 대한 엄밀한 수학적 정의를 제공하며, 표준 분포와 디락 델타 측도 모두에서의 동작을 분석하여 퇴화된 분포(degenerate distributions)로 점수 연산자를 일반화할 수 있음을 입증했습니다.
2. 효율적인 추론 알고리즘: 저자들은 별도의 훈련 과정 없이 생성된 샘플의 노이즈를 줄이는 추론 알고리즘을 제시합니다. 이는 해당 방법이 기존의 훈련 파이프라인 및 사전 훈련된 모델과 호환됨을 의미합니다.
3. 실증적 검증: 제안된 방식의 효능은 토이 문제, 합성 데이터, 그리고 FFHQ, AFHQ, ImageNet, cryo-electron microscopy(Cryo-EM) 데이터를 포함한 실제 데이터셋을 통해 입증되었습니다.

결과

• 토이 및 합성 데이터: 2D 예시 실험에서, 확장된 점수는 고분산 성분의 퍼짐은 유지하면서 저분산 가우시안의 궤적을 평균(디락 델타 모사)으로 수렴시키는 데 성공했습니다. 가우시안 노이즈와 픽셀 제거가 추가된 합성 데이터셋에서도 MAD는 표준 점수 추론을 크게 능가하며 원래의 깨끗한 형태를 복구하는 샘플을 생성했습니다.
• 실제 이미지 데이터: FFHQ, AFHQv2, ImageNet 실험 결과, 매니폴드 인력 파라미터($b$)를 높일수록 배경은 더 깨끗해지고 주요 특징은 더 선명해졌으며, 이는 고주파 노이즈와 이차적 변동을 효과적으로 임계값 처리했음을 보여줍니다.
• 정량적 지표: 가우시안, 샷(shot) 노이즈 및 픽셀 제거로 오염된 CIFAR-10 및 FFHQ 데이터셋에 대해, MAD는 표준 추론에 비해 현저히 낮은 Fréchet Inception Distance(FID)를 달로했습니다. 특히, 표준 추론의 FID는 노이즈가 섞인 훈련 데이터의 FID와 거의 일치하는 반면, MAD는 성공적으로 디노이징을 수행하여 성능 차이를 증명했습니다.
• Cryo-EM 적용: 노이즈가 섞인 단일 입자 Cryo-EM 이미지(EMPIAR-11618 데이터셋)에 적용했을 때, MAD는 네트워크가 오직 노이즈 섞인 데이터로만 훈련되었음에도 불구하고 전통적인 방식으로 얻은 정밀한 형태와 매우 유사한 샘플을 생성했습니다.

의의 및 주장
본 논문은 MAD가 특정 열화 모델을 알 수 없고 노이즈가 섞인 데이터만 존재하는 상황에서의 "블라인드 생성 디노이징(blind generative denoising)"에 대한 실질적인 해결책을 제공한다고 주장합니다. 이 방법의 의의는 훈련 시점의 복잡한 적응에 의존하는 대신, 생성 단계에서 데이터 매니폴드의 기하학적 구조를 활용하여 노이즈를 필터링할 수 있다는 점에 있습니다.

저자들은 다음과 같이 연구 범위에 대해 신중한 입장을 취합니다:

• 이 방법이 단일 주어진 이미지를 디노이징하는 역문제(inverse problems)를 해결한다고 주장하지 않으며, 대신 노이즈가 있는 데이터셋으로부터 기저 분포의 깨끗한 샘플을 새롭게 생성하는 것을 목표로 합니다.
• Cryo-EM과 같은 특정 도메인에서 최첨단 성능에 도달하기 위해서는 도메인 특화된 사전 지식을 통합해야 하며, 이는 본 연구의 범위를 벗어난다는 점을 인정합니다.
• 방법론이 효과적이긴 하지만, (추가적인 네트워크 평가로 인해) 계산 비용이 표준 추론보다 최대 두 배 정도 높을 수 있음을 언급했습니다. 다만, 경험적 관찰에 따르면 추론 단계의 후반부에만 확장된 점수를 사용하는 것으로도 충분할 수 있습니다.

이 연구는 확장된 점수를 생성 품질을 향бо할 수 있는 기하학적으로 근거 있는 메커니즘으로 위치시키며, 이는 깨끗한 데이터로 훈련된 모델에도 적용 가능할 수 있고, 적응형 파라미터 선택 및 조건부 확산 모델과의 통합을 위한 미래 연구의 길을 열어줍니다.