SOAR: 확산 모델의 최적 정렬 및 개선을 위한 자기 교정

SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models

현재 확산 모델의 추가 학습 파이프라인은 두 단계로 구성됩니다. 첫 번째는 선별된 데이터셋을 사용한 지도 미세 조정(SFT), 두 번째는 보상 모델을 활용한 강화 학습(RL)입니다. 이 두 단계 사이에는 근본적인 격차가 존재합니다. SFT는 노이즈 제거 모델을 순방향 노이징 과정에서 샘플링된 실제 상태에 대해서만 최적화합니다. 추론이 이러한 이상적인 상태에서 벗어나면, 후속 노이즈 제거는 학습된 수정보다는 일반화에 의존하게 되므로, 자기 회귀 모델에서 나타나는 동일한 노출 편향이 발생하지만, 토큰 시퀀스가 아닌 노이즈 제거 경로를 따라 누적됩니다. RL은 원칙적으로 이러한 불일치를 해결할 수 있지만, 종단 보상 신호는 희소하고, 신용 할당 문제가 있으며, 보상 악용의 위험이 있습니다. 우리는 이 격차를 해소하는 편향 수정 추가 학습 방법인 SOAR(Self-Correction for Optimal Alignment and Refinement)을 제안합니다. SOAR은 실제 샘플로부터 시작하여, 현재 모델을 사용하여 단일 단계의 역전파를 수행하고, 결과적으로 얻어진 경로 이탈 상태를 다시 노이징 처리한 후, 모델이 원래의 깨끗한 목표 상태로 회귀하도록 지도합니다. 이 방법은 온 정책이며, 보상이 필요 없으며, 신용 할당 문제가 없는 각 타임스텝에 대한 밀집적인 지도 신호를 제공합니다. SD3.5-Medium 모델에서 SOAR은 SFT를 기반으로 GenEval 지표를 0.70에서 0.78로, OCR 지표를 0.64에서 0.67로 향상시켰으며, 동시에 모든 모델 기반 선호도 점수를 높였습니다. 보상에 특화된 통제 실험에서, SOAR은 보상 모델에 접근할 수 없는 상황에서도 미적 감각 및 텍스트-이미지 정렬 작업 모두에서 Flow-GRPO보다 최종 지표 값에서 더 우수한 성능을 보였습니다. SOAR의 기본 손실 함수는 표준 SFT 목적 함수를 포함하므로, 사전 학습 후 더 강력한 첫 번째 추가 학습 단계로 SFT를 직접 대체할 수 있으며, 후속 RL 정렬과도 완전히 호환됩니다.

The post-training pipeline for diffusion models currently has two stages: supervised fine-tuning (SFT) on curated data and reinforcement learning (RL) with reward models. A fundamental gap separates them. SFT optimizes the denoiser only on ground-truth states sampled from the forward noising process; once inference deviates from these ideal states, subsequent denoising relies on out-of-distribution generalization rather than learned correction, exhibiting the same exposure bias that afflicts autoregressive models, but accumulated along the denoising trajectory instead of the token sequence. RL can in principle address this mismatch, yet its terminal reward signal is sparse, suffers from credit-assignment difficulty, and risks reward hacking. We propose SOAR (Self-Correction for Optimal Alignment and Refinement), a bias-correction post-training method that fills this gap. Starting from a real sample, SOAR performs a single stop-gradient rollout with the current model, re-noises the resulting off-trajectory state, and supervises the model to steer back toward the original clean target. The method is on-policy, reward-free, and provides dense per-timestep supervision with no credit-assignment problem. On SD3.5-Medium, SOAR improves GenEval from 0.70 to 0.78 and OCR from 0.64 to 0.67 over SFT, while simultaneously raising all model-based preference scores. In controlled reward-specific experiments, SOAR surpasses Flow-GRPO in final metric value on both aesthetic and text-image alignment tasks, despite having no access to a reward model. Since SOAR's base loss subsumes the standard SFT objective, it can directly replace SFT as a stronger first post-training stage after pretraining, while remaining fully compatible with subsequent RL alignment.