기술 요약: 소수 단계 확산 증류용 연속 시간 분포 매칭

문제 제기

확산 및 흐름 매칭 모델은 텍스트-이미지 생성을 혁신적으로 변화시켰으나, 반복적인 샘플링 과정으로 인해 수십에서 수백 번의 네트워크 평가가 필요한 심각한 계산 병목 현상을 겪습니다. 확산 증류 패러다임 (예: 일관성 증류, 분포 매칭 증류) 이 이 과정을 가속화하기 위해 등장했지만, 기존 분포 매칭 증류 (DMD) 방법은 구조적 한계에 직면해 있습니다. 구체적으로, 기존 DMD 는 훈련 및 추론 스케줄 간의 엄격한 정렬을 유지하기 위해 미리 정의된 몇 개의 이산 시간 단계에서 희소 감독에 의존합니다. 이 제한된 이산 시간 공식과 역 KL 발산의 모드 추구 성향이 결합되면, 시각적 아티팩트, 과도하게 평활화된 출력, 그리고 세밀한 디테일의 손실로 이어지는 경우가 많습니다. 결과적으로, 현재 최첨단 DMD 방법은 시각적 충실도를 회복하기 위해 생성적 적대 신경망 (GAN) 이나 보상 모델과 같은 복잡한 보조 모듈을 자주 필요로 합니다.

방법론: 연속 시간 분포 매칭 (CDM)

저자들은 DMD 패러다임을 이산 앵커링에서 연속 최적화로 전환하는 연속 시간 분포 매칭 (CDM) 프레임워크를 제안합니다. 이 방법은 두 가지 주요 연속 시간 설계로 구성됩니다:

1. 동적 연속 시간 스케줄

추론 스케줄과 일치하는 고정된 이산 앵커 집합으로 역방향 시뮬레이션을 제한하는 기존 DMD 와 달리, CDM 은 동적 연속 스케줄을 사용합니다.

• 메커니즘: 각 훈련 반복에서 시뮬레이션 길이 $N$은 균일 분포 $U\{1, N_{max}\}$에서 무작위로 샘플링됩니다. 시간 단계 $t_1, \dots, t_N$은 $(0, 1]$ 내에서 엄격하게 감소하는 연속 시퀀스로 생성됩니다.
• 효과: 이는 훈련 최적화를 고정된 추론 스케줄에서 분리합니다. 학생 모델은 몇 개의 고정된 점이 아닌 전체 연속 시간 공간에서 훈련되어 더 넓은 범위의 궤적에서 기울기 신호를 받습니다. 이론적 분석에 따르면 이는 학생이 희소 이산 앵커에 과적합하는 대신 연속 영역 전반에 걸쳐 더 매끄러운 속도장을 학습하도록 장려합니다.

2. 속도 기반 외삽을 통한 궤적 외 정렬

소수 단계 생성에 내재된 수치적 잘림 오차 (큰 적분 단계로 인해 궤적이 이상적 다양체에서 벗어나는 현상) 를 해결하기 위해, CDM 은 새로운 **CDM 손실 ($L_{CDM}$)**을 도입합니다.

• 메커니즘: 궤적 내 잠재 변수로 감독을 제한하는 대신, 이 방법은 능동적 궤적 외 매칭을 수행합니다. 학생이 예측한 속도장을 따라 1 차 오일러 단계를 취하여 궤적 외 잠재 변수 $x'_{t_i} = x_{t_i} + (t'_i - t_i)v_{t_i}$를 외삽합니다.
• 목적: 이 손실은 외삽된 점들에서 분포 매칭을 강제합니다. 외삽된 지점에서 로컬 클린 데이터 추정을 타겟으로 앵커링함으로써, 이 방법은 학생의 벡터장에 대한 자기 일관성 제약을 생성합니다. 이는 주 궤적 외부의 유효하지 않은 속도 예측을 효과적으로 패널티화하여 앵커 간 불일치를 완화하고 학생이 적분 오차를 스스로 수정할 수 있게 합니다.

훈련 목적 함수

전체 훈련 목적 함수는 세 가지 구성 요소를 통합합니다:
$$L = L_{CA} + L_{DM} + L_{CDM}$$

• $L_{CA}$ (CFG 증강): 텍스트 - 이미지 정렬을 보장합니다.
• $L_{DM}$ (분포 매칭): 학생의 주변 분포를 교사의 CFG-프리 분포와 정렬합니다.
• $L_{CDM}$ (연속 시간 매칭): 궤적 외 감독을 통해 연속 시간 영역 전반에 걸쳐 속도장을 정규화합니다.

주요 기여

1. 실증적 통찰: 저자들은 훈련을 고정된 이산 시간 단계에 앵커링하는 것이 불필요하며, 분포 매칭 (DM) 손실이 단순한 정규화제가 아니라 학생을 교사의 CFG-프리 분포에 정렬시키는 주된 동인임을 밝힙니다.
2. CDM 프레임워크: 동적 연속 스케줄링과 속도 기반 궤적 외 정렬 목적 함수를 결합한 통합 프레임워크를 제안합니다. 이는 적대적 훈련이나 외부 보상 모델에 의존하지 않고 최적화를 연속 시간 영역으로 전환합니다.
3. 최첨단 성능: 다양한 아키텍처 (SD3-Medium 및 Longcat-Image) 에서의 광범위한 실험을 통해 CDM 이 소수 단계 생성 (4 NFE) 에서 시각적 충실도와 의미적 준수를 달성하여 DMD2, D-DMD, Hyper-SD 와 같은 기존 베이스라인을 능가함을 입증했습니다.

실험 결과

이 논문은 1024×1024 해상도에서 4 개의 네트워크 함수 평가 (NFE) 를 사용하여 SD3-Medium과 Longcat-Image에 대한 CDM 을 평가합니다.

• 정량적 성능: SD3-Medium 에서 CDM 은 Aesthetic(6.075), DPGBench(85.26), PickScore(21.95), HPSv3(9.561) 에서 최우수 점수를 기록하여 강력한 베이스라인인 D-DMD 를 능가했습니다. 주목할 점은 4-NFE CDM 학생 모델이 여러 지표 (예: DPGBench 및 HPSv3) 에서 100-NFE 교사 모델을 일치시키거나 능가한다는 것입니다.
• 정성적 개선: 시각적 비교 결과, CDM 은 베이스라인에 비해 더 선명한 질감, 풍부한 세밀한 디테일, 복잡한 프롬프트에 대한 더 나은 준수를 보여주며, 베이스라인은 종종 과도한 평활화로 고통받는 것과 대조됩니다.
• 애블레이션 연구:
  • $L_{CA}$를 제거하면 구조적 붕괴가 발생합니다.
  • $L_{DM}$ 또는 $L_{CDM}$을 개별적으로 제거하면 프롬프트 준수 또는 시각적 충실도가 각각 저하됩니다.
  • 고정 시간 스케줄로 되돌리거나 속도 기반 외삽 대신 가우스 잡음 교란을 사용하면 성능이 저하되어 연속 및 속도 기반 설계의 필요성이 검증됩니다.

중요성 및 주장

이 논문은 CDM 이 이산 스케줄에서 연속 최적화 공간으로 DMD 증류 프레임워크를 이전한 첫 번째 작업이라고 주장합니다. 그 중요성은 일반적으로 이산 시간 증류의 아티팩트를 수정하기 위해 필요한 복잡한 보조 목적 함수 (GAN 이나 보상 모델 등) 에 의존하지 않고도 고충실도 소수 단계 생성을 달성할 수 있음을 입증한다는 점에 있습니다.

연속 감독을 통해 목표 데이터 다양체로 되돌아가 잘림 드리프트를 명시적으로 시뮬레이션하고 수정함으로써, CDM 은 4 단계 이미지 생성에 대한 새로운 최첨단 결과를 확립합니다. 저자들은 이 작업을 더 접근 가능한 확산 증류로의 한 걸음으로 위치시키며, 시각적 품질을 유지하면서 추론 비용을 한 자릿수 줄여 상용 하드웨어에서 고품질 생성을 가능하게 한다고 주장합니다.