Bias mitigation in graph diffusion models

1. 문제 정의 (Problem Statement)

기존의 그래프 확산 모델 (GDSS, GSDM, HGDM 등) 은 생성 품질 저하를 유발하는 두 가지 근본적인 편향을 가지고 있습니다.

• 역방향 시작 편향 (Reverse-Starting Bias):
  • 이상적인 확산 과정에서는 순방향 (Forward) 과정이 데이터를 표준 가우시안 분포까지 완전히 교란시켜야 하고, 역방향 (Reverse) 샘플링은 이 표준 가우시안 분포에서 시작해야 합니다.
  • 그러나 실제 그래프 학습에서는 데이터 규모와 네트워크 학습 능력의 한계로 인해 순방향 과정이 표준 가우시안 분포에 도달하기 전에 중단 (Truncation) 되는 경우가 많습니다. 즉, 최대 교란 분포가 표준 가우시안과 다릅니다.
  • 그럼에도 불구하고 역방향 샘플링은 여전히 표준 가우시안에서 시작하므로, 순방향의 실제 최대 교란 분포와 역방향 시작점 사이에 심각한 불일치가 발생합니다. 이는 초기 샘플링 단계에서 생성 품질을 크게 떨어뜨립니다.
• 노출 편향 (Exposure Bias):
  • 순방향 과정에서는 실제 데이터에 노이즈를 추가하지만, 역방향 과정에서는 모델이 예측한 샘플 ($\hat{x}_t$) 에 노이즈를 추가하며 반복적으로 샘플링합니다.
  • 스코어 네트워크의 예측 오차로 인해 순방향의 $x_t$와 역방향의 $\hat{x}_t$ 사이에 불일치가 발생하며, 이 오차는 샘플링이 진행됨에 따라 누적되어 최종 생성 결과의 품질을 저하시킵니다.

2. 제안 방법 (Methodology: S++)

저자들은 네트워크 구조를 수정하지 않고 기존 확산 모델을 활용하여 두 가지 편향을 동시에 완화하는 S++ 방법을 제안합니다.

A. 역방향 시작 편향 완화: 랑주빈 샘플링 (Langevin Sampling)

• 원리: 역방향 샘플링을 표준 가우시안에서 시작하는 대신, 순방향 과정의 실제 **최대 교란 분포 (Maximum Perturbation Distribution)**에 정렬된 분포에서 시작합니다.
• 구현: 사전 학습된 스코어 네트워크 $s_{\theta, T}(\cdot)$를 사용하여 랑주빈 샘플링 (Langevin Sampling) 을 수행합니다. 이를 통해 역방향 시작점을 실제 데이터가 도달한 저노이즈 (Low-noise) 상태 영역으로 이동시킵니다.
• 효과: 역방향 시작 편향을 제거할 뿐만 아니라, 모델이 예측 오차에 덜 민감한 저노이즈 상태에서 샘플링을 시작하게 함으로써 노출 편향의 누적 효과를 줄입니다.

B. 노출 편향 완화: 스코어 보정 메커니즘 (Score Correction)

• 원리: 스코어 네트워크의 예측 오차를 보정하기 위해 **스코어 차이 (Score Difference)**를 활용합니다.
• 구현:
  1. 사전 학습된 스코어 네트워크 ($s_{\theta, t}$) 로부터 샘플을 생성합니다.
  2. 생성된 샘플을 사용하여 가짜 (Pseudo) 스코어 네트워크 ($s_{\psi, t}$) 를 학습시킵니다. (가짜 네트워크는 원래 데이터의 편향을 학습하게 됩니다.)
  3. 두 네트워크의 스코어 차이 ($s_{\theta, t} - s_{\psi, t}$) 를 계산합니다. 이 차이는 실제 데이터 ($x_0$) 에 대한 정보를 포함하고 있습니다.
  4. 이 차이를 보정 항으로 사용하여 원래 스코어를 조정합니다:
     $$\hat{s}_{\theta, t} = s_{\theta, t} + \lambda (s_{\theta, t} - s_{\psi, t})$$
     여기서 $\lambda$는 보정 계수이며, $\omega$는 스코어 크기를 조절하는 계수입니다.
• 효과: 보정된 스코어는 실제 스코어에 더 가까워지며, 이는 역방향 시작 편향과 노출 편향을 동시에 완화하여 생성 품질을 향상시킵니다. 특히 역방향 시작 단계 ($T$) 에서의 보정이 가장 중요합니다.

3. 주요 기여 (Key Contributions)

1. 체계적인 편향 해결: 그래프 확산 모델의 역방향 시작 편향과 노출 편향을 체계적으로 분석하고, 랑주빈 샘플링과 스코어 보정을 통해 최초로 동시에 해결했습니다.
2. 네트워크 무변경 (Network-Agnostic): 기존 모델의 아키텍처를 변경하거나 생성기/판별자 같은 추가 학습 요소를 도입하지 않습니다. 기존 확산 모델에 플러그인 (Plug-and-play) 방식으로 통합 가능합니다.
3. 이론적 분석 및 실증: 스코어 차이가 실제 스코어에 더 가까운 방향을 가리킨다는 논리적 분석을 제공했으며, 다양한 모델 (GDSS, GSDM, HGDM, MOOD) 과 데이터셋 (일반 그래프, 분자 그래프) 에서 SOTA 성능을 입증했습니다.

4. 실험 결과 (Results)

저자들은 다양한 일반 그래프 데이터셋 (Community-small, Enzymes, Grid) 과 분자 생성 데이터셋 (QM9, ZINC250k) 에서 실험을 수행했습니다.

• 생성 품질 향상:
  • 일반 그래프 생성에서 MMD (Maximum Mean Discrepancy) 지표가 기존 베이스라인보다 크게 개선되었습니다. 예를 들어, GDSS-OC-S++ 는 Deg, Clus, Orbit 등 모든 지표에서 기존 GDSS-OC 보다 우수한 성능을 보였습니다.
  • 분자 생성 (QM9, ZINC250k) 에서 FCD (Fréchet ChemNet Distance) 와 NSPDK MMD 가 크게 감소하여 생성된 분자가 실제 분자 분포와 더 유사해졌습니다.
• 효율성 및 속도:
  • 기존 'Corrector'를 사용하는 방법 (예: GDSS-WC) 보다 계산 비용이 적게 들면서도 동등하거나 더 나은 성능을 달성했습니다.
  • 샘플링 단계 수 ($N$) 를 크게 줄여도 (예: 1000 단계 $\to$ 100 단계) S++ 를 적용한 모델은 성능 저하 없이 안정적인 생성을 보여주었습니다.
• 약물 발견 태스크: MOOD 모델에 S++ 를 적용한 결과, 새로운 약물 후보 물질의 발견률 (Novel hit ratio) 과 도킹 점수 (Docking score) 가 기존 모델보다 향상되었습니다.

5. 의의 및 결론 (Significance)

이 논문은 그래프 확산 모델이 직면한 역방향 시작 편향이라는 고유한 문제를 최초로 식별하고 해결책을 제시했다는 점에서 의의가 큽니다.

• 이미지 vs 그래프: 이미지 확산 모델에서는 신호 누출 (Signal Leakage) 이 미미하여 표준 가우시안 시작점이 유효하지만, 그래프 모델은 데이터의 희소성과 상호 의존성으로 인해 순방향 과정이 표준 가우시안에 도달하지 못합니다. 따라서 이미지용 기법 (예: DPM-Fixes) 을 직접 적용할 수 없으며, 본 논문에서 제안한 랑주빈 샘플링 기반의 정렬 방식이 필수적입니다.
• 실용성: 네트워크 재학습이나 복잡한 구조 변경 없이 기존 모델의 성능을 극대화할 수 있어, 대규모 그래프 데이터셋과 다양한 응용 분야 (신약 개발 등) 에 확산 모델을 적용하는 데 강력한 기반을 제공합니다.

결론적으로, 이 연구는 그래프 확산 모델의 생성 품질 한계를 극복하고, 더 안정적이고 고품질의 그래프 생성을 가능하게 하는 실용적이고 효과적인 프레임워크를 제시했습니다.