worldbench/DiffusionOPSD
🔥 On-Policy Self-Distillation in Diffusion Models
📚 DiffusionOPSD는 무엇인가요?
DiffusionOPSD(On‑Policy Self‑Distillation)는 논문 "DiffusionOPSD: On‑Policy Self‑Distillation in Diffusion Models" (arXiv 2608.24646)에 설명된 알고리즘을 구현한 연구용 코드베이스입니다. 이는 대규모 확산 이미지 생성기(예: Stable Diffusion 3.5‑M 및 Z‑Image‑Turbo)를 다양한 학습된 이미지 품질 또는 선호도 측정 지표(PickScore, CLIPScore, HPS 등)에 따라 이미지 생성 품질을 향상시킬 수 있도록 세부 조정(fine‑tune)할 수 있는 방법을 제공합니다.
핵심 아이디어는 보상(reward)이 전체 노이즈 제거 경로 후에만 관측될 때 발생하는 감독 갭(supervision gap)을 해결하는 것입니다. DiffusionOPSD는 반복적으로 다음을 수행합니다:
- 현재(고정된) "행동" 정책에서 낮은 노이즈 상태를 수집합니다.
- 각 상태의 클린 출력 기준점 주변에서 몇 단계의 정규화된 보상-기울기(gradient)를 취하여 명시적인 긍정 및 부정 타겟을 생성합니다.
- 제한된 수의 최적화기 업데이트 내에서 이러한 분리된 타겟을 따라가도록 새로운 "학습 가능한" 정책을 훈련합니다.
- 학습 가능한 정책의 지수이동평균(EMA)으로 행동 정책을 갱신한 후 반복합니다.
타겟 생성과 모델 적합이 분리되어 있으므로, 이 방법은 DiffusionNFT 및 FlowGRPO와 같은 기존 접근법과 공정하게 검증하고 비교할 수 있습니다.
🚀 주요 기능
| 기능 | 왜 중요한가요 |
|---|---|
| 온-정책 쿼리 수집 | 타겟은 현재 모델이 실제로 방문한 상태에서 생성되므로, 오프라인 또는 인위적으로 노이즈가 추가된 상태를 사용할 때 발생하는 불일치를 피합니다. |
| 명시적인 보상-지향 타겟 | 보상 기울기 방향으로 제한된 단계를 취하여 긍정/부정 기준점을 생성함으로써 모델이 개선할 방향을 명확히 제공합니다. |
| 분리된 유한 적합 | 정책 업데이트 전에 보상/디코더 그래프를 분리하므로, 타겟의 품질은 모델이 얼마나 잘 학습하는지와 독립적으로 측정할 수 있습니다. |
| EMA 기반 행동 갱신 | 각 외부 루프 후 행동 정책은 지수이동평균으로 갱신되어 훈련 전반에 걸쳐 새로운 감독을 보장합니다. |
| 두 가지 백본 지원 | Stable Diffusion 3.5‑M (512²) 및 Z‑Image‑Turbo (1024²)와 즉시 호환되며, 표준 및 소수 단계 확산 루틴 모두를 커버합니다. |
| 단일 및 혼합 보상 훈련 | 논문에서 사용된 7개의 공개 평가자에 대한 사전 정의된 프리셋과 함께, 임의의 양의 가중합(예: PickScore/26 + CLIPScore + HPSv2.1)을 결합할 수 있습니다. |
| 효율성 향상 | DiffusionNFT보다 GPU 시간을 40 %–63 % 절감하면서 19/20의 보상 매칭 설정에서 더 높은 검증 점수를 달성했습니다. |
| 풍부한 스크립트 | 공개 훈련, 베이스라인, 스모크 테스트, 혼합 보상 실험을 위한 준비된 실행 스크립트와 보상 모델 설정 확인용 유틸리티를 제공합니다. |
🛠️ 설치 및 빠른 시작
사전 요구 사항 – CUDA를 지원하는 GPU가 있는 Linux, Python 3.10‑3.11, 일치하는 PyTorch 빌드.
# 저장소 복제
git clone https://github.com/worldbench/DiffusionOPSD.git
cd DiffusionOPSD
# 새 conda 환경 생성 (선택 사항이지만 권장)
conda create -n diffusionopsd python=3.11 -y
conda activate diffusionopsd
# 핵심 패키지 설치 (보상 스택 확장 포함)
pip install -e "[rewards]"
# ImageReward는 오래된 timm 버전을 고정하므로 종속성 없이 설치
pip install --no-deps 'image-reward==1.5'
# 공개 프리셋에서 사용되는 보상 모델 체크포인트 다운로드
export REWARD_CKPT_PATH="$PWD/reward_ckpts"
bash scripts/download_reward_weights.sh
# Pick‑a‑Pic 프롬프트 매니페스트 준비 (논문 실험에서 사용)
python scripts/prepare_pickapic_prompts.py
Z‑Image‑Turbo에서 훈련하려면 ZImagePipeline를 포함하는 Diffusers 빌드가 필요합니다:
git clone https://github.com/huggingface/diffusers.git
pip install -e "./diffusers[torch]"
선택적 환경 설정
export HF_HOME=/path/to/huggingface-cache # 모델 가중치 캐시 위치
export WANDB_MODE=offline # 기본적으로 온라인 로깅 비활성화
▶️ 작은 스모크 테스트 실행
이 저장소는 단일 GPU 노드에서 단일 최적화기 업데이트를 실행하는 작고 간단한 "스모크" 스크립트를 제공하며, 보상 기울기 파이프라인의 작동 여부를 기계에서 확인하는 데 유용합니다.
# SD3.5‑M 백본, HPSv2.1 보상 (빠른 점검)
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh sd35 hpsv2
# Z‑Image‑Turbo 백본, CLIPScore 보상
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh zimage clipscore
또한 HPSv3, DeQA와 같은 무거운 보상 모델을 테스트할 수 있으며, 이들은 큰 7‑/8‑B 모델에 의존하므로 별도의 환경이 필요합니다:
SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage hpsv3
SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage deqa
📊 전체 규모 훈련 예제
아래는 논문에서 사용된 전체 100개 업데이트 예산(약 100개 최적화기 단계 ≈ 수백 번의 롤아웃 라운드)을 사용하여 공개 보상 전용 모델을 훈련하는 일반적인 명령어입니다.
# PickScore/26 + CLIPScore + HPSv2.1 혼합 목표로 SD3.5‑M 훈련
NPROC=8 UPDATES=100 OUTPUT_DIR=outputs/sd35_mixed_opsd \
bash scripts/train_public.sh sd35 mixed
# HPSv3 평가자로 Z‑Image‑Turbo 훈련 (6개 정책 랭크 + 1개 보상 서버 랭크)
NPROC=7 UPDATES=100 OUTPUT_DIR=outputs/zimage_hpsv3 \
bash scripts/train_public.sh zimage hpsv3
체크포인트는 10개 업데이트마다 저장되며(사용자 설정 가능) 최종 모델은 자동으로 기록됩니다.
📦 저장소 내부 내용
| 디렉터리/파일 | 목적 |
|---|---|
scripts/ |
보상 다운로드, 프롬프트 준비, 스모크 테스트 실행, 분산 훈련 시작을 위한 보조 스크립트 |
config/ |
각 백본, 각 보상, 혼합 보상 설정에 대한 YAML 형식의 구성 파일 |
scripts/train_*.py |
분산 실행을 설정하고, 행동 및 학습 가능한 정책을 로드하며, 궤적을 수집하고, 타겟을 생성하고, 최적화기 단계를 수행하는 진입점 |
scripts/check_reward_setup.py |
긴 작업 전에 필요한 보상 모델 체크포인트 및 라이브러리 버전이 존재하는지 확인 |
scripts/smoke_reward_gradient.py |
단일 보상 모델을 로드하고 이미지 공간 기울기가 0이 아닌지 확인 |
assets/ |
README에서 사용된 그림(질적 갤러리, 훈련 곡선, 아블레이션) |
LICENSE |
Apache 2.0 – 자유로운 오픈소스 라이선스 |
📜 라이선스
코드는 Apache 2.0 라이선스 하에 배포되며, 자유로운 사용, 수정, 배포가 가능합니다(기여 표시 필요). 모델 가중치 자체는 상위 확산 및 보상 모델의 원래 라이선스에 따라 관리되며, 이 저장소에서는 재배포되지 않습니다.
🎯 누구에게 적합한가요?
- 보상 지향 확산 세부 조정, 자기-분산(self-distillation), 또는 생성 모델을 위한 온-정책 강화학습 스타일 방법을 탐구하는 연구자.
- 특정 미적 취향이나 선호도 지표에 맞춰 고품질, 제어 가능한 확산 모델이 필요하고 재현 가능한 베이스라인을 원하는 실무자.
- 다양한 미분 가능한 보상 모델(대규모 언어-시각 평가자 포함)을 확산 훈련 파이프라인에 통합하는 방법을 구체적이고 종단 간(end-to-end) 예제로 찾는 학생.
📚 추가 읽기
- 전체 논문(arXiv 2608.24646)은 수학적 공식, 이론적 동기, 광범위한 아블레이션을 설명합니다.
- 프로젝트 페이지(https://diffusionopsd.github.io/)에는 추가 시각적 결과와 FAQ가 있습니다.
- 경쟁 방법에 대한 배경은 인용된 논문 DiffusionNFT 및 FlowGRPO를 참조하세요.
TL;DR: DiffusionOPSD는 여러 학습된 보상 함수에 따라 확산 이미지 생성을 향상시키는 새로운 온-정책 자기-분산 기법의 잘 문서화된 오픈소스 구현입니다. 기존 방법보다 훈련 비용을 줄이며, 실행 가능한 스크립트, 두 가지 주요 확산 백본 지원, 유연한 보상 혼합 시스템을 포함합니다.
관련
- 프로젝트
- 프로젝트
- 프로젝트
- 프로젝트