온폴리시 속도 회귀만으로 4스텝 확산 증류를 단순화하고 개선한다
Simplifying and Improving On-Policy Distillation for Few-Step Diffusion
무엇인가
확산 모델과 플로우 모델은 샘플 품질이 좋지만 추론할 때 함수 평가 횟수(NFE)가 많이 필요해 비용이 크다. few-step distillation은 이 비용을 줄이는 방법이지만, 기존 기법 상당수는 보조 score 모델을 따로 두거나 특수한 아키텍처를 요구해 적용 부담이 컸다. 이 논문은 그 부담을 걷어내면서 4-NFE 생성 품질을 끌어올리는 것을 목표로 한다.
어떻게 동작하나
저자들의 출발점은 단순하다. on-policy velocity regression만으로도 4-NFE 생성의 강한 기반이 된다는 것이다. 여기서 on-policy란 학생 모델이 실제로 만들어낸 궤적 위에서 teacher의 velocity를 회귀 대상으로 삼는다는 뜻으로, 별도의 온라인 score 추정 없이 순수 회귀 손실만으로 학습이 이뤄진다.
무엇과 다른가
핵심 구성 요소 하나는 Pivot이다. 하나의 거친 전이(coarse transition) 안에서도 teacher velocity는 계속 변하는데, 이를 표현하기 위해 학습 시에만 쓰이는 잔차 프록시를 도입한다. Pivot은 학생이 내놓는 구간 평균 예측(interval-mean prediction)과 국소 보정(local correction)을 분리해, 학생이 구간 평균만 흉내 내고 세부 변화를 놓치는 문제를 줄인다. 추론 시점에는 쓰이지 않는 학습 전용 장치라는 점이 중요하다.
어떻게 쓰나
또 하나는 약한 teacher 예측으로부터 학생을 밀어내는(repel) 경량 장치다. 적대적 학습이나 별도 판별기를 동원하지 않고도 바람직하지 않은 모드로 학생이 끌려가는 것을 막는 메커니즘이다. 이렇게 묶인 프레임워크가 Simplified Matching Distillation(SMD)이며, 전적으로 회귀 기반이고, 온라인 score 추정도 적대적 학습도 필요로 하지 않으며, 배포되는 생성기 구조 자체는 그대로 둔다.
전제와 한계
실험 범위는 이미지 생성과 편집, 3D 생성, 비디오 합성, 오디오와 비디오를 함께 생성하는 과제까지 넓다. 모델 규모는 330억 파라미터까지 확장했다. 4-NFE 조건에서 SMD는 강한 few-step 베이스라인들을 앞서고, 여러 벤치마크에서는 multi-step teacher까지 능가한다고 저자들은 보고한다. 또한 서로 다른 모델 간 증류(cross-model), 여러 teacher를 함께 쓰는 증류(multi-teacher), 인과적 비디오 증류(causal-video distillation)도 지원한다. 초록에는 벤치마크별 구체 수치나 데이터셋 이름이 제시되지 않아, 정량적 비교는 원문을 확인해야 한다.
개발자 입장에서 이 논문의 실용적 매력은 배포 파이프라인을 바꾸지 않아도 된다는 점이다. 생성기 아키텍처를 그대로 두고 학습 절차만 회귀 기반 증류로 교체하는 방식이므로, 이미 서빙 중인 모델의 추론 코드를 손대지 않고 4-NFE 버전을 만들 수 있다. 다만 도입 전에는 자신의 도메인에서 4-NFE 품질이 실제로 충분한지, teacher 선택이 결과에 얼마나 민감한지, 그리고 cross-model·multi-teacher 설정이 자기 파이프라인에 맞는지를 직접 확인하는 편이 좋다.
한계와 전제는 초록에 명시적으로 정리되어 있지 않다. 다만 방법 자체가 teacher의 velocity에 접근할 수 있는 증류 설정을 전제로 하고, Pivot이 학습 전용 잔차 프록시라는 점에서 학습 파이프라인 설계에 일정한 조건이 따른다. 평가도 4-NFE라는 특정 예산에 초점이 맞춰져 있으므로, 더 적거나 더 많은 스텝에서의 거동은 원문에서 별도로 확인해야 한다.