Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

기술 요약: 크로스 스페이스 디스틸레이션 (Cross-Space Distillation)

문제 정의
현대의 원스텝 확산 모델(one-step diffusion models)은 분포 기반 타임스텝 디스틸레이션을 통해 고품질 생성을 달sur합니다. 그러나 이러한 방법들은 교사(Teacher)와 학생(Student)이 동일한 잠재 공간을 공유하며, 동일한 잠재 해상도와 VAE 파라미터화를 가져야 한다는 제한적인 **공유 공간 제약(Shared-Space constraint)**에 의존합니다. 이러한 가정은 고해상도(예: 1024×1024)에서 작동하는 현대의 고용량 교사 모델(예: Stable Diffusion 3.5, Flux, Kolors)로부터 저해상도(예: 512×512)에서 작동하는 컴팩트하고 배포 친화적인 학생 모델(예: Stable Diffusion 1.5)로 지식을 전달하는 것을 방해합니다.

표준 디스틸레이션 목적 함수는 다음과 같은 이유로 인해 이 설정에서 실패합니다:

1. 교차 해상도(Cross-Resolution): 불일치하는 공간 그리드 형태 (예: $h \times w$ 대 $H \times W$).
2. 교차 VAE(Cross-VAE): 서로 다른 오토인코더 파라미터화로 인한 정렬되지 않은 잠재 공간.
3. 교차 아키텍처/메커니즘(Cross-Architecture/Mechanism): UNet과 DiT 백본 간의 차이, 또는 표준 확산 모델과 플로우 매칭(flow matching) 정식화 간의 차이.

네트워크 프루닝(pruning)이나 새로운 컴팩트 백본을 처음부터 구축하는 것과 같은 기존의 대안들은 엄격한 희소성 임계값, 광범한 재학습 요구 사항, 그리고 생태계 호환성 상실 등의 문제에 직면합니다.

방법론: 크로스 스페이스 디스틸레이션과 브릿지(Bridge)
저자들은 이 설정을 크로스 스페이스 디스틸레이션으로 공식화하고, 학생의 백본을 수정하지 않고도 학생의 잠재 공간을 교사의 잠재 공간으로 매핑하도록 설계된 경량 모듈인 **브릿지($B_\phi$)**를 도입합니다.

• 아키텍처: 브릿지는 불일치를 효율적으로 처리하기 위해 두 단계로 분해됩니다:

  1. 공간적 사전 지식(Spatial Prior, 고정): 학생 VAE 디코더의 처음 $n$개 디코딩 블록($D_S^{(n)}$)을 고정된 상태로 유지합니다. 이 블록들은 사전 학습된 공간 업샘플러 역할을 하여, 학생의 잠재 변수 $z_S$를 저해상도에서 교사의 고해상도($H \times W$)로 확장합니다. 이는 추가적인 학습 파라미터 없이 학생의 기존 공간 지식을 활용합니다.
  2. 학습 가능한 프로젝터($g_\phi$): SwinIR로 구현된 컴팩트하고 학습 가능한 네트워크가 업샘플링된 학생의 특징 공간을 교사의 잠재 표현 공간으로 매핑하여, 교사와 호환되는 잠재 변수 $\hat{z}_T$를 출력합니다.

• 학습 목적 함수: 브릿지는 동일한 이미지를 두 VAE로 각각 인코딩하여 얻은 쌍을 이룬 잠재 변수 데이터셋 $(z_S, z_T)$를 사용하여 학습됩니다. 최적화는 두 손실 함수의 가중 합을 최소화합니다:

  1. 잠재 재구성 손실($L_{rec}$): 예측된 잠재 변수 $\hat{z}_T$와 실제 교사 잠재 변수 $z_T$ 사이의 $\ell_1$ 손실입니다. 저자들은 이 문맥에서 $\ell_1$이 $\ell_2$보다 더 나은 안정성과 구조 보존력을 제공한다고 언급했습니다.
  2. 어텐션 충실도($L_{attn}$): $\hat{z}_T$와 $z_T$를 각각 입력받았을 때 고정된 교사 디노이저의 셀프 어텐션 분포 간의 역-KL 발산(reverse-KL divergence) 손실입니다. 이는 매핑된 잠재 변수가 실제 교사 잠재 변수와 동일한 내부 구조적 반응(장거리 의존성)을 유도하도록 보장하며, 픽셀 단위 손실이 놓칠 수 있는 미세한 디테일을 보존합니다.

• 디스틸레이션 및 추론: 학습이 완료되면 브릿지가 학생에게 부착됩니다. 디스틸레이션 과정에서 학생의 출력은 브릿지를 통과하여 $\hat{z}_T$를 생성하며, 이를 통해 표준 분포 기반 목적 함수(예: Variational Score Distillation, Adversarial Distillation)를 교사에 대해 직접 적용할 수 있습니다. 결정적으로, 학생의 백본은 변경되지 않습니다.

주요 기여

1. 크로스 스페이스 디스틸레이션의 공식화: 본 논문은 불일치하는 잠재 해상도 및 VAE 공간 간에 디스틸레이션이 발생하는 새로운 체계를 정의하여, 현대의 고해상도 모델로부터 컴팩트한 레거시 모델로 지식을 전달하는 데 있어 핵심적인 격차를 해결했습니다.
2. 브릿지 모듈: 학생의 아키텍처를 변경하지 않고 표준 디스틸레이션 목적 함수를 가능하게 하는 경량화된 모듈형 인터페이스입니다. 이는 고정된 공간적 사전 지식과 학습 가능한 프로젝터를 결합하고, 안정적인 정렬을 위해 어텐션 충실도 목적 함수를 활용합니다.
3. 실질적 효과성: 이 방법은 컴팩트한 학생 모델(예: SD 1.5)이 원스텝 추론, 낮은 지연 시간 및 폭넓은 생태계 호환성을 유지하면서도 다양한 현대적 교사 모델(SDXL, Kolors, PixArt-$\Sigma$, Flux, SD 3.5)의 시각적 사전 지식을 상속받을 수 있게 합니다.
4. 모델 병합 능력: 이 방법은 서로 다른 교사로부터 학생으로 다양한 지식을 디스틸레이션하므로, 결과물인 체크포인트는 파라미터 평균화를 통해 여러 교사 소스의 지식을 통합하는 단일 모델로 쉽게 병합될 수 있습니다.

결과
실험은 컴팩트한 원스텝 학생 모델의 생성 품질을 실질적으로 개선함을 보여줍니다:

• 정량적 이득: 브릿지를 통해 SD 3.5 Medium(교사)을 SD 1.5(학생)로 디스틸레이션하면 HPSv3 점수가 5.37에서 9.42로, ImageReward가 -0.29에서 0.62로 향상됩니다. 다섯 가지 서로 다른 교사로부터 온 체크포인트를 병합하면 SD 1.5의 HPSv3는 10.53까지 더욱 높아집니다.
• 정성적 개선: 시각적 비교 결과, 브릿지가 적용된 학생 모델은 초기화된 베이스라인이나 프루닝된 모델에 비해 더 선명한 질감, 더 강한 구조적 일관성, 그리고 더 적은 아티팩트를 생성합니다.
• 효율성: 브릿지는 무시할 만한 오버헤드(약 5M 파라미터)를 추가하며 원스텝 추론을 가능하게 합니다. 또한, 학생을 재학습시키지 않고도 512×512 학생 샘플을 브릿지를 통해 1024×1024로 매핑하여 교사의 디코더로 디코딩할 수 있는 추론 시 해상도 업그레이드를 가능하게 합니다.
• 프루닝과의 비교: 브릿지 기반 디스틸레이션은 네트워크 프루닝(예: SD 3.5를 30% 희소도로 프루닝)보다 우수한 성능을 보이며, 프루닝은 흐릿함과 아티팩트를 유발하는 경향이 있습니다.

의의
본 논문은 크로스 스페이스 디스틸레이션이 거대하고 현대적인 교사의 생성 능력을 새로운 아키텍처를 처음부터 학습시키는 비용 없이 효율적이고 배포 가능한 백본으로 전달할 수 있는 잠재력을 열어준다고 주장합니다. 이질적인 잠재 표현의 정렬을 재사용 가능한 인터페이스 문제로 다룸으로써, 본 연구는 다음을 가능하게 합니다:

• 생태계 호환성: 컴팩트한 모델인 SD 1.5가 기존의 도구 및 하드웨어와 호환성을 유지하면서도 최첨단 품질로 업그레이드될 수 있습니다.
• 지식 통합: 서로 다른 아키텍처를 가진 다양한 교사(예: Flux 및 SD 3.5)의 지식을 단일한 경량 학생 모델로 병합하는 능력.
• 모듈성: 브릿지는 디스틸레이션과 추론 시 해상도 향상 모두에 재사용될 수 있는 플러그 앤 플레이 구성 요소로서, 학생의 생성 능력과 잠재 표현의 제약을 분리합니다.

저자들은 이 연구를 텍스트-투-이미지를 넘어 비디오 및 멀티모달 생성에도 적용 가능한, 이질적인 잠재 공간 간의 더 넓은 디스틸레이션 탐구를 향한 단계로 포지셔닝합니다.