The Spacetime of Diffusion Models: An Information Geometry Perspective

이 논문은 확산 모델의 잠재 공간에 대한 정보 기하학적 관점을 제시하여, 기존 결정론적 ODE 디코더의 한계를 지적하고 '잠재 시공간'을 도입하여 피셔-라오 계량을 기반으로 한 새로운 기하학적 구조와 확산 편집 거리, 그리고 분자 시스템의 전이 경로 샘플링에 대한 효율적인 방법을 제안합니다.

핵심

확산 모델의 '시공간': 정보를 기하학적으로 바라보는 새로운 시각

이 논문은 최근 AI 이미지 생성의 핵심 기술인 **'확산 모델 (Diffusion Models)'**이 어떻게 작동하는지, 그 내부의 '잠재 공간 (Latent Space)'을 새로운 방식으로 해석한 연구입니다.

기존의 방식과 저자들이 제안한 새로운 방식을 쉽게 비유해서 설명해 드리겠습니다.

---

1. 기존 방식의 문제: "직선으로만 가는 착각"

확산 모델은 이미지를 만들 때, 완전히 잡음 (노이즈) 상태인 $x_T$에서 시작해 점점 선명해지는 이미지 $x_0$로 변해가는 과정을 거칩니다.

• 기존의 생각 (Pullback Geometry): 연구자들은 "잡음 공간에서 가장 짧은 길을 찾으면, 그 길은 이미지 공간에서도 가장 자연스러운 변환 경로일 거야"라고 생각했습니다. 마치 지도에서 두 도시 A 와 B 를 잇는 가장 짧은 직선을 그리는 것과 비슷합니다.
• 문제점: 이 논문은 "그건 틀렸습니다!"라고 지적합니다. 확산 모델의 수학적 구조상, 잡음 공간에서 아무리 복잡한 곡선을 그리더라도, 그것이 이미지로 변환되면 항상 직선이 되어버린다고 증명했습니다.
  • 비유: 마치 "우주선 (잡음 공간) 에서 가장 짧은 경로로 날아간다고 해서, 지구 (이미지 공간) 에 착륙할 때 지형의 구릉이나 산을 피할 수 있는 건 아니다"라는 뜻입니다. 데이터가 가진 복잡한 모양 (곡선, 구불구불한 길) 을 전혀 반영하지 못하고, 무조건 뻥 뚫린 직선으로만 이어져 버려서 실제 활용 가치가 떨어집니다.

2. 저자들의 해결책: "시간을 포함한 4 차원 시공간"

그렇다면 어떻게 해야 할까요? 저자들은 "잡음의 양 (시간)"까지 포함해서 생각하자고 제안합니다.

• 새로운 개념 (시공간 Spacetime): 단순히 "어떤 잡음 ($x_t$)"만 보는 게 아니라, **"어떤 잡음 ($x_t$) 을 가지고, 몇 번째 시간 ($t$) 에 있는지"**를 함께 고려합니다. 이를 $(x_t, t)$라고 부르는 '시공간 좌표'로 만듭니다.
• 비유:
  • 기존 방식: "이 사진이 흐릿할 때 (시간 $t$), 이 사진이 선명할 때 (시간 $0$) 로 가는 길"만 생각함.
  • 새로운 방식: "흐릿한 사진이면서 동시에 시간이 흐르는 과정" 전체를 하나의 지도로 봅니다.
  • 마치 영화 필름을 생각해보세요. 단순히 '한 장의 사진'만 보는 게 아니라, '흐릿한 장면에서 선명한 장면으로 넘어가는 모든 프레임의 연속'을 하나의 거대한 공간으로 보는 것입니다.

3. 핵심 발견: "이미지 변환은 '지우기'와 '다시 그리기'의 과정"

이 '시공간'을 이용하면 이미지 사이의 거리를 계산할 수 있는데, 이를 **'확산 편집 거리 (Diffusion Edit Distance)'**라고 부릅니다.

• 어떻게 작동할까?
  • 두 이미지 (예: '고양이'와 '강아지') 가 있을 때, 이 둘을 연결하는 가장 짧은 길은 단순히 고양이를 강아지로 변형시키는 게 아닙니다.
  • 과정:
    1. 잡음 추가 (지우기): 고양이 이미지에 노이즈를 섞어 '고양이'라는 정보를 잊어버리게 만듭니다. (완전한 잡음 상태에 가까워짐)
    2. 잡음 제거 (다시 그리기): 그 잡음 상태에서 '강아지'라는 정보를 다시 불러옵니다.
  • 결과: 이 '지우기'와 '다시 그리기'의 총비용이 두 이미지 사이의 거리입니다.
  • 비유: 두 도시를 연결할 때, A 도시를 완전히 부순 뒤 (지우기), 그 자리에 B 도시를 새로 짓는 (다시 그리기) 데 드는 비용과 시간을 계산하는 것과 같습니다. 이 거리가 짧을수록 두 이미지는 서로 변환하기 쉽다는 뜻입니다.

4. 실용적인 활용: "분자 세계의 안전한 길 찾기"

이 이론은 단순히 이미지 변환뿐만 아니라, 분자 (원자) 의 움직임을 예측하는 데도 쓰입니다.

• 상황: 분자가 에너지가 낮은 상태 A 에서 상태 B 로 이동할 때, 가장 에너지가 낮고 안전한 경로를 찾아야 합니다.
• 기존 방법: 무작위로 길을 찾아보거나, 이미 알려진 방법들을 쓰면 위험한 고에너지 지역 (분자가 터질 수 있는 곳) 을 지나가거나, 너무 단순한 경로만 찾습니다.
• 이 방법의 장점:
  • 시공간 지형도 활용: "여기는 위험한 고에너지 지역이야 (지형이 높음)"라고 표시된 지도를 가지고, 그 지역을 피하면서 A 에서 B 로 가는 가장 자연스러운 곡선을 찾습니다.
  • 제약 조건: "이 지역은 절대 지나가지 마" 혹은 "이 경로는 너무 흔들리면 안 돼" 같은 조건을 넣어서, 더 안전하고 현실적인 분자 이동 경로를 설계할 수 있습니다.

5. 요약: 왜 이 연구가 중요한가요?

1. 오해 깨기: 기존에 "잡음 공간의 직선이 이미지 공간의 직선이다"라고 믿었던 잘못된 상식을 깨뜨렸습니다.
2. 새로운 지도: 시간과 잡음의 양을 모두 포함한 '시공간 지도'를 만들어, 이미지나 분자 사이의 관계를 훨씬 더 정교하게 이해할 수 있게 했습니다.
3. 실제 적용: 이미지를 자연스럽게 변형시키거나, 복잡한 분자 구조를 설계할 때, 기존 방법보다 더 효율적이고 안전한 경로를 찾아낼 수 있습니다.

한 줄 요약:

"확산 모델은 잡음 속에서 이미지를 만드는 마법사인데, 이 논문은 그 마법사가 걸어가는 길이 단순한 직선이 아니라, 시간과 잡음이 섞인 복잡한 4 차원 지형임을 발견했고, 그 지형을 이용해 더 똑똑한 길 찾기를 가능하게 했습니다."

기술 요약

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 생성 모델링에서 뛰어난 성과를 보이고 있지만, 노이즈가 섞인 중간 상태 $x_t$ 를 통해 정보가 어떻게 진화하는지에 대한 잠재 공간 (latent space) 의 기하학적 구조를 이해하는 것은 여전히 열린 문제입니다.

기존 연구들은 데이터의 내재적 기하학을 이해하기 위해 잠재 공간에 유클리드 메트릭을 푸백 (pullback) 하는 방식을 사용했습니다. 그러나 이 논문은 확산 모델에서 **결정론적 확률 흐름 ODE (PF-ODE)**를 디코더로 사용할 때 발생하는 근본적인 한계를 지적합니다:

• 기하학적 붕괴 (Geometric Collapse): PF-ODE 기반의 푸백 메트릭을 사용하면, 잠재 공간에서의 측지선 (geodesic, 최단 경로) 이 데이터 공간에서 **항상 직선 (straight segments)**으로 디코딩됩니다.
• 이는 데이터가 저차원 매니폴드에 존재하더라도, 잠재 공간과 데이터 공간의 차원이 동일하고 디코더가 주변 유클리드 공간에서 직접 작동하기 때문에 데이터의 내재적 곡률 (curvature) 을 전혀 반영하지 못함을 의미합니다. 결과적으로 기존 푸백 접근법은 하류 작업에 유용한 기하학적 정보를 제공하지 못합니다.

2. 방법론 (Methodology)

저자들은 확산 모델을 **정보 기하학 (Information Geometry)**의 관점에서 재해석하고, 새로운 시공간 (Spacetime) 표현을 제안합니다.

A. 확률적 디코더와 피셔 - 라오 (Fisher-Rao) 메트릭

결정론적 ODE 대신 **확률적 역 SDE (Reverse SDE)**를 통해 정의된 제거 노이즈 분포 (Denoising Distribution) $p(x_0|x_t)$를 디코더로 사용합니다.

• 잠재 공간 $z$에 피셔 - 라오 (Fisher-Rao) 메트릭을 적용하여, 잠재 변수 $x_t$의 미세한 변화가 제거 노이즈 분포 전체를 어떻게 변화시키는지를 측정합니다.
• 하지만 단순히 $x_T$ (최종 노이즈) 를 잠재 변수로 사용하면, 확산 모델의 '기억 상실성 (memorylessness)'으로 인해 분포가 데이터 분포 $q(x_0)$에 수렴하여 메트릭이 0 으로 붕괴되는 문제가 발생합니다.

B. 잠재 시공간 (Latent Spacetime) 도입

이 문제를 해결하기 위해 저자들은 시공간 좌표 $z = (x_t, t)$를 도입합니다.

• 이는 모든 노이즈 스케일 $t \in (0, T]$에 걸친 제거 노이즈 분포의 가족 (family) 을 인덱싱하는 $(D+1)$차원 다양체입니다.
• 시간 $t$를 포함함으로써 피셔 - 라오 메트릭이 상태와 시간에 따라 변하게 되어 비자명한 (nontrivial) 기하학적 구조를 복원합니다.

C. 지수족 (Exponential Family) 성질과 계산 효율성

핵심적인 이론적 발견은 **모든 노이즈 스케일의 제거 노이즈 분포 $p(x_0|x_t)$가 지수족 (Exponential Family)**을 이룬다는 것입니다.

• 자연 매개변수 (Natural Parameter) $\eta$와 기대값 매개변수 (Expectation Parameter) $\mu$를 도출했습니다.
  • $\eta(x_t, t) = (\frac{\alpha_t}{\sigma_t^2}x_t, -\frac{\alpha_t^2}{2\sigma_t^2})$
  • $\mu(x_t, t) = (E[x_0|x_t], E[\|x_0\|^2|x_t])$
• 지수족의 성질을 이용하여 **측지선 에너지 (Geodesic Energy)**를 시뮬레이션 없이 (simulation-free) 추정할 수 있는 폐쇄형 근사식을 유도했습니다.
• 실제 계산에서는 Tweedie 공식과 Hutchinson's trick을 사용하여 $\mu$를 효율적으로 추정하며, 역 SDE 를 직접 실행하지 않고도 경로 길이를 계산할 수 있어 계산 비용이 크게 절감됩니다.

3. 주요 기여 (Key Contributions)

1. 푸백 기하학의 한계 증명: 확산 모델에서 PF-ODE 기반 푸백 메트릭이 데이터 공간의 직선 보간으로만 귀결되어 무의미함을 수학적으로 증명했습니다.
2. 확산 시공간 (Diffusion Spacetime) 프레임워크: $z=(x_t, t)$를 기반으로 한 피셔 - 라오 기하학을 제안하여, 노이즈 제거 분포의 변화를 측정하는 새로운 기하학적 구조를 확립했습니다.
3. 계산 가능한 측지선 추정: 제거 노이즈 분포가 지수족임을 증명하고, 이를 통해 역 SDE 시뮬레이션 없이도 측지선 길이와 에너지를 효율적으로 계산하는 알고리즘을 제시했습니다.
4. Diffusion Edit Distance (DiffED): 두 데이터 포인트 간의 최단 시공간 경로의 길이를 정의하여, 데이터를 변환하는 데 필요한 '최소 편집 비용'을 측정하는 새로운 거리 척도를 제안했습니다.
5. 응용 분야 확장: 분자 시스템의 전이 경로 샘플링 (Transition Path Sampling) 에 적용하여, 고에너지 영역을 회피하는 효율적인 경로 생성 방법을 제시했습니다.

4. 실험 결과 (Results)

• 이미지 보간 (Image Interpolation): ImageNet-512 모델에서 PF-ODE 샘플링 경로와 제안된 시공간 측지선을 비교했습니다. 시각적으로 큰 차이는 없었으나, 측지선이 정보를 조금 더 일찍 생성하는 경향을 보였습니다.
• DiffED 평가:
  • DiffED 는 인간의 지각적 유사성 (LPIPS) 과는 낮은 상관관계 (-7%) 를 보였으나, 구조적 유사성 (SSIM) 과는 53% 의 상관관계를 보였습니다. 이는 DiffED 가 픽셀 수준의 유사성이 아닌, 정보 제거 및 재구성의 '편집 비용'을 측정한다는 것을 시사합니다.
  • 이미지 쌍을 여러 메트릭으로 정렬했을 때, DiffED 는 직선 보간이나 LPIPS 와 다른 유사성 기준을 제공했습니다.
• 분자 전이 경로 샘플링 (Molecular Transition Path Sampling):
  • 알라닌 디펩타이드 (Alanine Dipeptide) 의 에너지 지형에서 두 저에너지 상태 간의 전이 경로를 생성했습니다.
  • 성능: 제안된 방법은 기존 MCMC 기반 방법 및 Doob's Lagrangian 방법보다 **최대 에너지 (MaxEnergy)**가 낮아 (37.36 vs 42.54 이상), 고에너지 장벽을 더 효과적으로 우회하는 경로를 생성했습니다.
  • 효율성: 기존 방법들에 비해 에너지 함수 평가 횟수가 수백 배 적게 소요되었습니다.
• 제약 조건付き 경로: 분산 감소 (Low-variance) 또는 특정 영역 회피 (Region avoidance) 와 같은 제약 조건을 페널티 함수로 추가하여 시공간 측지선을 최적화하는 데 성공했습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 모델의 잠재 공간을 단순한 노이즈 공간이 아닌, **통계적 다양체 (Statistical Manifold)**로 재정의하여 정보 기하학적 통찰을 제공합니다.

• 이론적 기여: 확산 모델의 역과정 (Reverse Process) 이 지수족을 이룬다는 사실을 밝혀, 복잡한 확률 분포 간의 기하학적 거리를 효율적으로 계산할 수 있는 토대를 마련했습니다.
• 실용적 가치:
  • DiffED는 데이터 간의 관계를 '편집 비용'이라는 새로운 관점에서 정의하여, 생성 모델의 내재적 구조를 이해하는 데 도움을 줍니다.
  • 분자 동역학 등 물리 기반 시뮬레이션 분야에서, 고에너지 영역을 피하는 효율적인 전이 경로를 찾는 데 기존 방법보다 우수한 성능을 입증했습니다.
  • 시뮬레이션 없이 측지선을 계산할 수 있는 효율적인 알고리즘은 고차원 이미지 생성 모델에서도 적용 가능하여, 향후 생성 모델의 제어 및 최적화 기법 개발에 중요한 기여를 할 것으로 기대됩니다.

요약하자면, 이 연구는 확산 모델의 '시공간'을 정보 기하학적으로 분석하여 기존 방법론의 한계를 극복하고, 분자 시뮬레이션 및 생성 모델 제어에 적용 가능한 새로운 이론적 프레임워크와 도구를 제시했습니다.