기술 요약: Any-OPD

문제 정의

현재의 온-폴리시 증류(On-Policy Distillation, OPD)는 플로우 매칭(flow-matching) 및 확산 모델(diffusion models)에 있어 **동질성 가정(homogeneity assumption)**에 의존합니다. 즉, 교사(teacher)와 학생(student) 모델이 동일한 VAE 잠재 공간(latent space), 아키텍처, 그리고 노이즈 스케줄(timestep grid)을 공유해야 한다는 것입니다. 이러한 가정은 증류를 동일한 모델 제품군 내의 쌍으로 제한합니다.

실제 상황에서 가장 바람직한 증류 시나리오는 **이종(heterogeneous)**적인 경우입니다. 예를 들어, 사용자는 (FLUX.1-dev와 같은) 서로 다른 제품군의 거대하고 최첨단인 교사의 우수한 품질을 활용하면서도, (SD3.5-Medium과 같이) 작고 효율적인 모델을 배포하고자 할 수 있습니다. 기존의 OPD는 다음과 같은 두 가지 근본적인 불일치로 인해 이러한 설정에서 실패합니다:

1. 공유된 공간의 부재: 교사의 잠재 변수(latents)는 학생의 좌표계에서 유효한 타겟이 될 수 없습니다. 또한, 픽셀 단위 회귀(pixel-wise regression)는 실패하는데, 이는 교사의 정교화(refinement) 과정이 의미론적 내용은 유지하되 국소적 디테일은 다시 그리는 확률적 재합성(stochastic re-synthesis)이기 때문입니다. 이 확률적 평균을 향해 회귀하는 것은 학습 붕괴나 심각한 블러링(blurring)을 초래합니다.
2. 공유된 시계(Clock)의 부재: 서로 다른 모델은 고유한 노이즈 시프트 함수(noise shift functions)를 사용합니다. 결과적으로, 동일한 솔버 인덱스(solver indices, 단계 수)가 서로 다른 노이즈 레벨에 대응하게 되어, 인덱스가 정렬된 감독(supervision) 쌍이 서로 비교 불가능한 상태가 됩니다.

기존 방식들은 자신의 잠재 공간과 노이즈 스케줄이 다른 교사를 사용하여 학생의 궤적(trajectory)을 감독할 수 없습니다.

방법론: Any-OPD

저자들은 임의의 래턴트 플로우 매칭 생성기 쌍 사이에서 온-폴리시 증류를 가능하게 하는 최초의 프레임워크인 Any-OPD(Any Teacher, Any Student, On-Policy Distillation)를 제안합니다. 이 방법은 두 모델의 내부 파라미터를 참조하지 않고, 대신 제3의 외부 표현(representation)과 노이즈 레벨 정렬에 의존함으로써 교사와 학생을 분리합니다.

1. 앵커링 단계 (오프라인)

분포 수준의 격차를 해결하기 위해, 학생은 먼저 교사의 매니폴드(manifold)에 "앵커링(anchoring)"됩니다.

• 교사가 프롬프트로부터 이미지를 생성하면, 이를 학생의 VAE를 사용하여 **재인코딩(re-encode)**합니다.
• 이를 통해 교사의 출력 분포를 학생의 잠재 좌표계로 투영합니다.
• 학생은 이러한 재인코딩된 타겟에 대해 표준 플로우 매칭 손실(flow-matching loss)을 통해 학습됩니다. 이는 학생이 온-폴리시 정교화를 시도하기 전에 "올바른 종류"의 이미지를 생성하는 법을 배우도록 보장합니다.

2. 온-폴리시 증류 단계

핵심 증류 루프는 공유된 잠재 변수나 동기화된 타임스텝 없이 작동합니다:

• 블랙박스 샘플러로서의 교사: 교사는 순수하게 노이즈 및 디노이즈 투영 연산자($\Pi^\sigma_T$)로 기능합니다. 학생의 샘플이 주어지면, 교사는 이를 특정 노이즈 레벨 $\sigma$로 다시 노이즈를 추가하고 정교화된 버전을 재생성합니다.
• 노이즈 레벨 매칭 (시계 격차 해결): 인덱스 단위로 정렬하는 대신, 이 방법은 노이즈 레벨에 따라 정렬합니다. 교사가 노이즈 레벨 $\sigma_T^k$에서 샘플을 정교화한 후, 그래디언트는 학생의 궤적 단계 중 노이즈 레벨 $\sigma_S^j \ge \sigma_T^k$인 단계에 대해서만 역전파됩니다. 이는 학생이 정교화에 해당하는 특정 노이즈 스케일에 적합한 수정을 학습하도록 보장합니다.
• 표현 공간 브릿징 (공간 격차 해결): 학생의 정교화된 출력과 교사의 정교화된 출력은 각각의 VAE를 사용하여 픽셀로 디코딩됩니다. 이후 이 이미지들은 동결된 모델 불가지론적 특징 공간(구체적으로, 동결된 DINOv2의 CLS 토큰)에서 비교됩니다.
  • 손실은 $1 - \cos(f(\hat{x}_S), f(x_{ref}))$로 정의되며, 여기서 $f$는 동결된 추출기입니다.
  • 이 전역적이고 좌표 불변적인 비교는 픽셀이나 잠재 회귀를 무력화시키는 국소적 확률적 재합성에 불변합니다.

주요 기여

1. 이종 OPD의 공식화: 논문은 기존의 온-폴리시 방식이 서로 다른 모델 제품군 간에 작동하는 것을 방해하는 두 가지 주요 장애물(공유된 공간 및 공유된 시계의 부재)을 식별하고 공식화했습니다.
2. Any-OPD 프레임워크: 오직 동결된 외부 표현 공간을 통해서만 모델을 결합하고, 솔버 인덱스가 아닌 노이즈 레벨 매칭을 통해 궤적을 정렬하는 새로운 아키텍처 불가지론적 증류 절차를 제시합니다.
3. 실증적 검증: 서로 다른 모델 제품군 간의 첫 성공적인 온-폴리시 증류(FLUX.1-dev $\to$ SD3.5-Medium)를 보여주며, 2.5B 규모의 학생 모델이 선호도 지표에서 12B 규모의 교사 모델의 성능에 근접하거나 이를 능가할 수 있음을 입증했습니다.

실험 결과

저자들은 **FLUX.1-dev (12B 파라미터)**를 **SD3.5-Medium (2.5B 파라미터)**로 증류함으로써 Any-OPD를 구현했습니다.

• 성능 향상: 증류된 학생 모델은 베이스라인인 SD3.5-Medium을 크게 상회하며, 여러 지표에서 12B 교사 모델에 근접하거나 이를 능가했습니다:
  • PickScore: 0.846 (베이스라인)에서 0.884로 증가 (교사의 0.878을 상회).
  • HPSv3: 9.12에서 10.97로 증가 (교사의 11.19에 근접).
  • ImageReward: 0.922에서 1.116으로 증가 (교사의 0.971을 상회).
• 절제 연구(Ablation) 결과:
  • 잠재 회귀(Latent Regression): VAE 경계를 가로지르는 직접적인 잠재 회귀는 즉각적인 학습 붕괴를 일으킵니다.
  • 픽셀/MSE 손실: 교사 정교화의 확률적 특성 때문에 실패합니다.
  • DINOv2 CLS: 유일하게 안정적이고 효과적인 목적 함수이며, 이는 전역적 의미론적 정렬이 국소적 픽셀/패치 정렬보다 필요함을 확인시켜 줍니다.
  • 앵커링: 안정성을 위해 필수적입니다. 앵커링 단계 없이 OPD를 시작하면 열등한 성능을 보입니다.
  • 일반화: 동일한 하이퍼파라미터를 사용하여 교사를 또 다른 별개의 제품군인 Z-Image로 교체해도 일관된 개선 효과를 얻었으며, 이는 방법론의 상호 교환성을 검증합니다.

의의 및 주장

본 논문은 Any-OPD가 기존의 온-폴리시 방식이 가진 구조적 장벽을 제거하여, 가장 강력한 가용 생성 모델을 더 작은 배포 가능 모델의 교사로 사용할 수 있게 한다고 주장합니다. 교사를 순수하게 샘플러로 취급하고 출력을 동결된 외부 특징 공간에서 비교함으로써, 이 방법은 완전히 다른 아키텍처, VAE, 노이즈 스케줄을 가진 모델 간에도 고품질 생성 능력을 전이할 수 있게 합니다.

저자들은 이러한 디커플링(decoupling)을 통해 "누가 만들든 상관없이 가장 강력한 가용 생성기"를 "사용자가 배포하고자 하는 어떤 모델"을 위한 교사로 만들 수 있으며, 이는 구조적 동질성을 요구하지 않고도 최첨단 생성 품질에 대한 접근을 민주화한다고 강조합니다.