D-OPSD: 단계별 증류 확산 모델의 지속적인 튜닝을 위한 온폴리시 자기 증류

최근 FLUX.2와 같은 고성능 이미지 생성 모델들은 빠른 속도를 위해 적은 단계로 이미지를 생성하는 Step-distillation 기법을 사용하지만, 새로운 개념을 학습시키기 위해 추가 학습(Fine-tuning)을 하면 이 효율적인 생성 능력이 손상되는 문제가 있다. 이 논문은 모델이 스스로 생성한 데이터를 바탕으로 학습하는 On-policy 방식을 도입하여, 기존의 빠른 생성 속도를 유지하면서도 새로운 스타일이나 개념을 완벽하게 학습할 수 있는 해결책을 제시한다.

기존의 확산 모델 학습 방식인 SFT(Supervised Fine-tuning)는 정답 이미지에 노이즈를 섞어 입력하고 이를 복구하도록 학습시킨다. 하지만 적은 단계로 이미지를 만드는 모델(Step-distilled models)은 실제 추론 시에 자신이 직접 생성하며 거치는 경로가 정답 이미지의 노이즈 경로와 다르기 때문에, 학습과 실제 사용 환경 사이에 괴리가 발생하여 생성 품질이 급격히 떨어진다.

이 문제를 해결하기 위해 연구진은 강화학습의 On-policy 개념을 도입했다. 모델이 현재 실력으로 직접 이미지를 생성해보고(Student), 동시에 동일한 모델이 타겟 이미지라는 추가 정보를 참고하여 더 나은 결과물을 예측하도록(Teacher) 만든다. 학생 모델은 스승 모델이 제시한 정답을 따라가도록 학습되는데, 이때 학습 데이터가 모델이 직접 생성한 경로 위에 있으므로 추론 성능이 그대로 유지된다.

결과적으로 모델은 외부의 복잡한 보상 함수(Reward Function) 없이도, 자신이 가진 인컨텍스트 능력을 활용해 스스로를 가르치며 새로운 지식을 습득한다. 이는 마치 학생이 교과서(텍스트)만 보고 푼 문제와 해설지(이미지 포함 컨텍스트)를 보고 이해한 내용을 비교하며 스스로 오답 노트를 만드는 과정과 유사하다.

D-OPSD는 하나의 모델에 두 가지 역할을 부여하여 학습을 진행한다. 먼저 입력 프롬프트(y)만 사용하는 학생 조건(cs)과 프롬프트 및 타겟 이미지(x0)를 모두 사용하는 스승 조건(ct)을 생성한다. 이때 인코더로 Qwen3-VL과 같은 멀티모달 모델을 활용하여 이미지 정보를 컨텍스트로 주입한다.

학습 과정에서는 학생 모델이 현재의 파라미터로 몇 단계의 샘플링 궤적을 직접 생성한다. 각 단계(tk)에서 학생 모델의 예측 속도(uk)와 스승 모델의 예측 속도(u'k)를 계산한다. [학생/스승 조건 입력 → 확산 모델의 속도 필드 연산 → 각 시점의 속도 벡터 출력] 과정을 거친다.

최종 손실 함수(Loss Function)는 학생과 스승의 예측 속도 차이를 최소화하는 Mean Squared Error(MSE)로 정의된다. [두 속도 벡터의 차이 계산 → 제곱 연산 → 배치 평균 산출] 과정을 통해 오차를 구하고, 이를 줄이는 방향으로 학생 모델의 가중치를 업데이트한다. 스승 모델은 지수 이동 평균(EMA)을 통해 학생 모델의 성장을 부드럽게 추종하며 안정적인 가이드를 제공한다.

Z-Image-Turbo와 FLUX.2-klein 모델을 대상으로 한 실험에서 D-OPSD는 기존 SFT 및 PSO(Preference-based Step Optimization) 방식보다 월등한 성능을 보였다. 특히 LoRA 학습 설정에서 DINO-D(개념 보존)와 VLM-J(스타일 일치도) 지표에서 최고 수준을 기록했다.

풀 파인튜닝(Full Fine-tuning) 실험에서도 애니메이션 도메인으로의 전이 학습 시 기존 지식을 잊어버리는 '파괴적 망각' 현상을 효과적으로 억제했다. 수치적으로 FID 점수가 기존 모델 대비 약 15~20% 개선되었으며, 생성된 이미지의 미적 점수(Aesthetic Score) 또한 가장 높게 유지되었다.

Ablation Study 결과, EMA를 적용한 스승 모델을 사용하는 것이 학습의 안정성을 높이는 데 필수적임을 확인했다. EMA 계수를 0.9999로 설정했을 때 가장 높은 품질의 증류 결과가 나타났다.

D-OPSD는 Flow-matching 기반의 확산 모델 아키텍처를 활용한다. 핵심은 멀티모달 인코더(Qwen3-VL 등)의 가중치를 조정하여 훈련 시와 추론 시의 특징 공간(Feature Space) 불일치를 해결한 점이다. 단순히 인코더를 교체하는 대신, LLM 컴포넌트의 가중치를 재조정하여 고주파 아티팩트 발생을 억제했다.

수학적으로는 기존 LLM의 OPSD 수식을 확산 모델의 속도 예측 도메인으로 전이시켰다. 토큰 단위의 KL Divergence 대신 연속적인 속도 벡터 공간에서의 MSE Loss를 사용하여 온폴리시 궤적 정렬을 수행한다. 이는 모델이 스스로 생성한 샘플(Roll-outs)에 대해서만 그래디언트를 계산하므로 학습 효율이 높다.

학습 시 메모리 효율을 위해 LoRA(Rank 64, Alpha 128)를 주로 사용하며, 단일 H800 GPU에서 1,000회 반복 학습만으로도 충분한 개념 학습이 가능하다. 대규모 데이터셋의 경우 32개의 H800 GPU를 사용하여 10,000회 이상의 반복 학습을 수행한다.

D-OPSD의 성공 여부는 베이스 모델이 가진 인컨텍스트(In-context) 능력에 크게 의존한다. 만약 베이스 모델이 멀티모달 입력을 받았을 때 타겟 개념을 제대로 이해하지 못한다면, 스승 모델이 유효한 감독 신호를 생성할 수 없어 학습이 실패하게 된다. 또한 온폴리시 샘플링 과정으로 인해 일반 SFT 대비 약 2배의 학습 시간이 소요된다.