Inductive Moment Matching: 빠르고 안정적인 생성 모델 학습 기법

 

최근 이미지, 비디오, 오디오 생성 모델은 뛰어난 성능을 보이며 발전하고 있습니다. 특히 확산 모델(Diffusion Models)과 Flow Matching 방식은 고품질의 샘플을 생성하는 데 강력한 도구로 자리 잡고 있습니다. 하지만 이러한 모델들은 높은 연산량과 느린 추론 속도라는 문제점을 가지고 있습니다. 기존 연구에서는 이러한 속도를 개선하기 위해 일괄 증류(Distillation)나 일관성 모델(Consistency Models, CMs) 등의 방법을 제안했지만, 이들은 종종 불안정한 학습 과정이나 세밀한 하이퍼파라미터 튜닝이 필요하다는 단점이 있습니다.

 

이 논문에서는 이러한 문제를 해결하기 위해 새로운 생성 모델 학습 방식인 Inductive Moment Matching (IMM)을 제안합니다. IMM은 한 번의 학습 단계(single-stage training)로 고품질 이미지를 빠르게 생성할 수 있는 모델을 학습할 수 있도록 합니다. 특히 기존의 확산 모델이나 CMs와 비교하여

• 사전 학습된 모델을 필요로 하지 않으며,
• 하이퍼파라미터에 민감하지 않고 안정적인 학습이 가능하며,
• ImageNet-256×256에서 1.99 FID(Frechet Inception Distance)를 기록하며 기존 모델을 능가하는 성능을 보여줍니다.

그림 : 단 8 step 만으로 생성된 이미지 샘플

---

2. 본론

2.1. 기존 방식의 문제점

기존 생성 모델 학습 방식에는 몇 가지 주요 문제가 있습니다.

1. 확산 모델(Diffusion Models)
   • 높은 품질의 샘플을 생성할 수 있지만, 다수의 추론 단계가 필요하여 속도가 느림
   • 예를 들어, DDPM은 1000개의 스텝을 거쳐야 샘플을 생성할 수 있음
2. Flow Matching
   • 연속적인 흐름(trajectory)을 학습하는 방식으로 빠르게 샘플을 생성할 수 있지만, 학습이 복잡하고 불안정할 수 있음
3. 일괄 증류(Distillation)
   • 빠른 샘플링을 위해 확산 모델을 몇 개의 단계로 압축하는 방식이지만, 학습이 불안정하고 정교한 튜닝이 필요함
4. 일관성 모델(Consistency Models, CMs)
   • 사전 학습된 모델 없이 직접 학습할 수 있지만, 특정 하이퍼파라미터에 의존도가 높아 학습이 불안정함

2.2. 접근 방식

IMM은 위 문제를 해결하기 위해 시간 종속적인 확률 분포 간의 모멘트(moment)를 일치(match)시키는 방식을 사용합니다. 이를 통해 한 번의 학습 단계로 빠르고 안정적인 생성 모델을 학습할 수 있습니다.

• 확률 밀도 분포 간의 간접적 샘플링 방식 적용
  IMM은 확률 분포 간의 연결(interpolant)을 활용하여 특정 시점(t)에서 다른 시점(s)으로 바로 변환하는 방식을 학습함
  즉, 아래 그림에서 기존 방식이 t=1(노이즈)에서 시작해 많은 단계를 거쳐 t=0(이미지 데이터)에 도달하는데 반해 IMM은 t에 s로 직접 점프하는 방식을 사용하며 이때 특정한 중간 단계 r을 거쳐 빠르게 고품질 샘플을 생성함 

그림 : IMM이 데이터를 샘플링하는 방법

 

• 수학적 귀납법(Mathematical Induction) 기반 학습
  • 한 번의 샘플링 과정에서 여러 개의 중간 시점(r)을 설정하고, 서로 다른 시점의 분포가 일치하도록 학습함
  • 이를 통해 학습의 안정성을 보장하며, 점진적으로 원하는 분포에 수렴할 수 있도록 설계됨

그림 : IMM이 학습하는 과정에서 샘플이 어떻게 이동하는지 설명

 

빨간 선은 모델이 데이터를 변형하는 경로이고 점선은 경로를 따라 샘플이 이동하는 방식으로, 샘플이 다른 샘플과 비교되어 유사한 분포를 갖도록 학습되는 방식이라고 합니다. (영상 설명은 이곳)

 

 

2.3. 세부 적용 기술

1) 확률적 보간(Stochastic Interpolants)

• 데이터 분포와 노이즈 분포를 연속적으로 연결하는 확률 변수를 정의함
• 기존 Flow Matching과 달리, 특정 시점 ( t )에서 다른 시점 ( s )로 직접 이동할 수 있도록 설계됨 (점프방식)

2) 최대 평균 차이(MMD, Maximum Mean Discrepancy) 활용

• 확률 분포 간의 차이를 정량적으로 평가하기 위해 MMD 기반의 거리 척도를 활용함
• 기존의 GAN과 유사하지만, 신경망 판별자 없이 최적의 커널 함수를 사용하여 두 확률 분포 간 차이를 최소화하도록 설계됨

3) 귀납적 부트스트래핑(Inductive Bootstrapping)

• 모델이 자기 자신의 샘플을 통해 점진적으로 학습할 수 있도록 설계됨
• 초기에는 작은 시점 간 변화를 학습하고, 점차 더 큰 시간 차이를 학습하도록 만듦

4) 학습 안정성을 위한 설계

• 기존 CMs는 단일 샘플을 사용해 학습하지만, IMM은 다수의 샘플을 이용하여 분포 차이를 완화하고 안정적인 학습을 유도함
• 이를 위해 Laplace Kernel와 같은 수학적 기법을 활용하여 분포 차이를 보다 정밀하게 반영(적절히 분포되도록)하도록 설계됨

2.4. 제약사항

• IMM은 빠르고 안정적이지만, 완전히 무작위한 노이즈에서 시작하는 샘플링보다는 사전에 정의된 분포를 기반으로 생성됨
• 극단적으로 작은 데이터셋에서는 충분한 일반화가 어려울 수 있음

---

3. 결론

IMM은 기존 생성 모델들의 단점을 극복하면서도 높은 성능을 유지하는 새로운 방식입니다.

• 빠른 샘플링 : 기존 확산 모델보다 훨씬 적은 샘플링 스텝으로 고품질 이미지 생성 가능
• 안정적인 학습 : 하이퍼파라미터 튜닝 없이도 쉽게 학습 가능
• 고품질 성능 : CIFAR-10 및 ImageNet-256×256에서 최첨단 성능 달성

IMM은 추후 다양한 이미지 및 영상 생성 응용 분야에서 강력한 도구로 활용될 수 있을 것입니다.

---

참고자료

• 논문) Inductive Moment Matching (링크)
• 코드) Inductive Moment Matching Github (링크)

---

Q&A

Q. IMM이 기존 확산 모델과 가장 큰 차이점은 무엇인가요?

기존 확산 모델은 다단계 샘플링이 필수적이지만, IMM은 단일 또는 몇 단계 만에 샘플링이 가능합니다. 또한, IMM은 사전 학습된 모델 없이도 직접 학습할 수 있습니다.

 

Q. IMM이 Consistency Models(CMs)보다 안정적인 이유는 무엇인가요?

IMM은 MMD 기반의 다중 샘플 매칭을 사용하여 분포 차이를 최소화하며, 단일 샘플 방식의 CMs보다 학습 안정성이 높습니다.

 

Q. IMM은 GAN과 어떻게 다른가요?

GAN은 생성자-판별자 간의 경쟁을 통해 학습하지만, IMM은 분포 간의 MMD 최적화를 통해 보다 안정적으로 학습됩니다.