💻 computer science

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

본 논문은 분포 정렬과 인간 선호도 지표를 모두 최적화하여 최첨단 단단계 이미지 생성을 달성하기 위해 분포 정합과 보상 기반 강화 학습을 통합하는 2 단계 프레임워크인 보상 기울기 분포 정합 증류 (RTDMD) 를 제안합니다.

원저자: Yushi Huang, Xiangxin Zhou, Ruoyu Wang, Chi Zhang, Jun Zhang, Tianyu Pang

게시일 2026-05-26

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yushi Huang, Xiangxin Zhou, Ruoyu Wang, Chi Zhang, Jun Zhang, Tianyu Pang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

마스터 셰프 (교사) 가 완벽한 복잡한 요리를 만들 수 있지만, 그 작업에 50 시간이 걸린다고 상상해 보세요. 당신은 부셰프 (학생) 가 같은 요리를 단 4 시간 만에 만들 수 있도록 가르치고 싶습니다.

문제는 두 가지입니다:

속도 대 품질: 학생에게 단순히 "마스터를 따라 하라"고만 말하면, 학생은 모든 단계를 신중히 고려할 시간이 부족해 서두르다가 엉망이 될 수 있습니다.
맛: 마스터 셰프가 만든 음식은 기술적으로 완벽할지라도, 사람들이 실제로 선호하는 맛과는 다를 수 있습니다 (너무 짜거나 질감이 이상할 수 있음). 당신은 학생이 마스터의 기술을 배우되, 사람들이 좋아하는 음식을 만들 수 있도록 배우기를 원합니다.

이 논문인 RTDMD는 정확히 이 문제를 해결하기 위한 새로운 학습 방법입니다. 이 방법은 학생이 4 단계 만에 고품질 요리를 만드는 법을 가르치면서, 음식이 사람들에게 맛있게 느껴지도록 보장합니다.

다음은 이를 단순한 개념으로 나누어 설명한 것입니다:

1. "기울어진 메뉴" (보상 기울어진 분포, Reward-Tilted Distribution)

보통 학생에게 마스터를 따라 하라고 가르칠 때, "마스터의 출력을 정확히 맞추라"고 말합니다. 하지만 마스터는 "나쁜" 요리 (낮은 보상) 와 "훌륭한" 요리 (높은 보상) 를 동일한 빈도로 만들어낼 수 있습니다.

저자들은 기발한 트릭을 제안합니다: 메뉴를 기울이세요.
마스터의 레시피 책을 물리적으로 기울여 "훌륭한 요리"가 위로 미끄러지고 "나쁜 요리"가 아래로 미끄러지도록 상상해 보세요. 이제 학생이 마스터를 따라 할 때, 자연스럽게 이미 좋은 것을 선호하는 마스터의 버전을 따라 하게 됩니다.

수학: 그들은 마스터의 분포와 "보상 함수" (이미지의 품질 점수) 를 결합합니다. 이는 마스터의 스타일을 유지하면서도 인간이 좋아하는 부분에 가중치를 두는 새로운 목표를 만들어냅니다.

2. 2 단계 훈련 캠프

이 논문은 학생을 훈련시키기 위해 2 단계 프로세스를 사용합니다.

1 단계: "안정된 손" 워밍업 (AC-DMD)

첫 번째 단계에서 학생은 기초를 배우고 있습니다.

문제: 4 단계 프로세스에서 학생은 중간에 "노이즈가 섞인" 재료로 작업합니다. 누군가 테이블을 흔들면서 그림을 그리려는 것과 같습니다. 학생의 "가짜 점수" (최종 이미지가 어떻게 보여야 할지 추측하는 데 도움을 주는 도구) 는 목표가 계속 움직이기 때문에 혼란을 겪습니다.
해결책 (Ambient-Consistent): 저자들은 **일관성 정규화 (Consistency Regularizer)**를 도입합니다. 이는 "현실 점검"과 같습니다.
- 학생이 2 단계의 흐릿한 덩어리가 4 단계에서 고양이로 변할 것이라고 예측한다면, 일관성 규칙은 이렇게 말합니다: "잠깐, 그 흐릿한 덩어리를 가져와서 한 단계 앞으로 나아가면, 여전히 고양이로 향하고 있는 것처럼 보이나요?"
- 이는 학생의 내부 논리가 노이즈가 섞인 단계들 사이에서 일관되게 유지되도록 강제하여, 학습하는 동안 혼란에 빠지는 것을 방지합니다.

2 단계: "맛 평가" 강화 (하이브리드 정책 경사, Hybrid Policy Gradient)

이제 학생이 빠르게 요리할 수 있게 되었으니, 맛있게 요리하는 법을 배워야 합니다. 여기서 강화 학습 (RL) 이 등장합니다.

문제: 요리 과정은 두 가지 요소가 섞여 있습니다:
1. 확률적 (무작위) 단계: 처음 3 단계는 무작위 노이즈를 추가하는 과정입니다 (재료를 섞기 위해 공중으로 던지는 것과 같음).
2. 결정론적 (고정) 단계: 마지막 단계는 정밀하고 계산된 과정입니다 (요리를 완벽하게 접시에 담는 것과 같음).
실수: 기존 방법들은 무작위 단계만 보거나 마지막 단계만 보았습니다. 이는 샐러드를 던지는 방식만으로 셰프를 평가하거나, 디저트를 접시에 담는 방식만으로 평가하는 것과 같으며, 전체 요리를 무시하는 것입니다.
해결책 (하이브리드 정책 경사): 저자들은 둘 다를 고려하는 새로운 "점수" 계산 방식을 고안했습니다:
- 무작위 단계에는 SubGRPO라는 기법을 사용합니다. 같은 요리를 만드는 24 명의 학생 그룹이 있다고 상상해 보세요. 그들이 완전히 다른 무작위 재료를 모두 사용하게 하면 (누가 좋은지 판단하기 어렵게 만듦), 대부분의 단계에서는 재료를 공유하게 하되, 단 하나의 특정 단계에서만 재료를 바꾸도록 합니다. 이렇게 하면 왜 한 요리가 다른 요리보다 더 맛있게 느껴졌는지를 분리해 낼 수 있습니다.
- 마지막 단계에서는 단순히 보상을 **역전파 (backpropagate)**합니다. 마지막 단계는 무작위성이 없는 직선 경로이므로, 마지막 동작을 변경하는 것이 맛을 어떻게 개선하는지 정확히 계산하여 학생을 즉시 업데이트할 수 있습니다.

3. 결과

저자들은 SD3, SD3.5, FLUX.2 와 같은 가장 진보된 이미지 생성기들에서 이 방법을 테스트했습니다.

주장: 그들의 방법 (RTDMD) 은 거의 다른 어떤 방법보다도 4 단계로 더 나은 이미지를 생성하며, 인간의 선호도와 더 잘 부합합니다.
놀라운 점: 그들의 4 단계 모델 (40 억 파라미터 기반) 은 실제로 훨씬 더 큰 90 억 파라미터 모델의 원래 50 단계 버전보다 많은 카테고리에서 더 좋은 성과를 냈습니다. 그들은 느리고 거대한 슈퍼컴퓨터의 품질을 빠르고 작은 모델에 압축해 넣는 데 성공했습니다.

요약 비유

RTDMD 를 자율주행차를 위한 운전 학교로 생각해 보세요:

1 단계: 도로가 울퉁불퉁할 때도 차가 차선 안에 머물고 일정한 속도를 유지하도록 가르칩니다 (일관성 정규화).
2 단계: 수천 번의 주행을 시뮬레이션하여 차가 안전하고 효율적으로 운전하도록 가르칩니다. 단순히 마지막에 충돌했다고 처벌하는 것이 아니라, 중간에 발생한 무작위 핸들 조작과 최종 제동 동작을 함께 분석하여 최상의 피드백을 제공합니다 (하이브리드 정책 경사).

결과는 무엇일까요? 느리지만 최적화가 잘 안 된 차들보다 빠르지만 (4 단계), 승객에게 더 안전하고 편안한 차 (인간 선호도) 를 운전하는 것입니다.

기술 요약: 보상 기울어진 분포 매칭을 통한 소수 단계 생성기 강화

문제 제기

최근 확산 및 흐름 기반 생성 모델의 발전으로 고품질 텍스트-이미지 합성이 가능해졌으나, 반복적 샘플링 절차는 일반적으로 수십 개의 함수 평가 (NFE) 를 필요로 하여 상호작용 콘텐츠 제작 및 실시간 시각 시스템과 같은 지연 시간에 민감한 애플리케이션에 병목 현상을 초래합니다. 분별화 기술 (예: 분포 매칭 분별화, DMD) 이 이러한 모델을 소수 단계 생성기로 성공적으로 축소시켰음에도 불구하고, 이러한 효율적인 모델을 인간 선호도와 정렬시키는 것은 여전히 어렵습니다.

강화 학습 (RL) 을 소수 단계 생성에 통합하는 것을 방해하는 두 가지 주요 장애물이 있습니다:

불안정한 분별화 신호: 소수 단계 생성에서 중간 잠재 변수는 본질적으로 노이즈가 많습니다. DMD 의 "가짜 스코어" 모델은 제한된 연산 예산 하에서 이동하는 생성기 분포를 추적해야 하므로, 콜드 스타트 분별화 신호가 불안정해집니다.
하이브리드 샘플링 역학: 샘플링 궤적은 주입된 노이즈로 인한 확률적 중간 단계와 노이즈가 0 인 결정론적 최종 단계로 구성됩니다. 기존 RL 방법은 종종 확률적 단계만 또는 최종 매핑만 최적화하여 전체 궤적의 하이브리드 특성을 고려하지 못하며, 이로 인해 하위 최적 정책 기울기가 발생합니다.

방법론: RTDMD

저자들은 분포 매칭 분별화와 보상 유도 강화 학습을 통합하는 2 단계 프레임워크인 보상 기울어진 분포 매칭 분별화 (Reward-Tilted Distribution Matching Distillation, RTDMD) 를 제안합니다. 핵심 이론적 통찰은 보상 기울어진 (reward-tilted) 교사 분포에 대한 커널 - 라이블러 (KL) 발산을 최소화하는 것이 자연스럽게 분포 매칭 항과 보상 최대화 항으로 분해된다는 점입니다.

1 단계: 환경 일관 분포 매칭 분별화 (AC-DMD)

안정적인 콜드 스타트를 제공하기 위해 저자들은 노이즈가 있는 중간 잠재 변수에서 가짜 스코어 모델을 훈련하는 불안정성을 해결하는 AC-DMD 를 소개합니다.

환경 분포 매칭: 깨끗한 샘플을 가정하는 대신, 이 방법은 노이즈가 있는 중간 잠재 변수 $\hat{x}_{t_k}$ 에 조건부로 부분 구간 $[t_k, 1]$ 에서 분포 매칭을 수행합니다. 이를 통해 노이즈가 있는 입력에 유효하도록 DMD 목적 함수를 재도출합니다.
일관성 정규화: 생성기 분포가 이동함에 따라 가짜 스코어 모델 ( $s_\phi$ ) 을 안정화하기 위해 일관성 정규화자가 추가됩니다. 이는 $t'$ 타임스텝에서의 모델 예측이 $t''$ 로 역확산 단계 후의 기대 예측과 같아야 한다는 자기 일관성 속성을 강제합니다. 이는 타임스텝 간 예측을 결합하여 추정 분산을 줄입니다.

2 단계: 단계 서브셋 GRPO 를 통한 하이브리드 정책 기울기

두 번째 단계에서 프레임워크는 분포 매칭과 보상 최대화 항을 공동으로 최적화합니다.

하이브리드 정책 기울기: 소수 단계 생성기가 확률적 가우시안 전이 후 결정론적 최종 단계를 따르는 하이브리드 정책임을 인식하여, 저자들은 다음을 결합하는 기울기 추정기를 유도합니다:
1. 확률적 중간 전이에 대한 REINFORCE 스타일 추정기.
2. 결정론적 최종 단계를 통한 직접 보상 역전파.
단계 서브셋 GRPO (SubGRPO): 정책 기울기의 분산을 줄이기 위해 저자들은 SubGRPO 를 도입합니다. 모든 궤적의 모든 단계에 대해 독립적인 노이즈를 사용하는 단순한 GRPO 와 달리, SubGRPO 는 무작위성을 가진 단계의 부분 집합을 선택하여 독립적인 노이즈를 사용하고 나머지 단계에 대해서는 그룹 간 노이즈를 공유합니다. 이는 특정 단계가 보상에 미치는 영향을 격리하여 기울기 분산을 낮추는 Rao–Blackwellized 변형으로 작용합니다.

주요 기여

이론적 통합: 이 논문은 보상 기울어진 분포에 대한 KL 발산을 최소화하는 것이 분별화와 RL 의 원칙적인 분해를 제공하여 별도의 임시 목적 함수의 필요성을 제거함을 보여줍니다.
AC-DMD: 부분 구간 매칭과 일관성 정규화를 통해 노이즈가 있는 중간 잠재 변수에서 가짜 스코어 모델의 안정적인 훈련을 가능하게 하는 새로운 분별화 전략으로, 소수 단계 분별화의 "이동 목표" 문제를 해결합니다.
하이브리드 정책 기울기: REINFORCE 와 경로 미분 (pathwise derivatives) 을 모두 활용하여 소수 단계 흐름 생성기의 혼합 확률 - 결정론적 특성을 올바르게 처리하는 맞춤형 기울기 추정기입니다.
SubGRPO: 선택되지 않은 단계에 대해 궤적 간 노이즈를 공유하여 소수 단계 생성기의 분산 감소 기술로, 보상 최적화의 효율성을 향상시킵니다.

실험 결과

저자들은 4 개의 추론 단계를 사용하여 SD3-M, SD3.5-M, FLUX.2 4B에 대해 RTDMD 를 평가했습니다.

최첨단 성능: RTDMD 는 선호도 (HPSv2, ImageReward), 미적, 구성 (GenEval) 지표 전반에 걸쳐 새로운 최첨단 결과를 달성했습니다.
효율성 대 품질: FLUX.2 에서 분별화된 4B 모델 (4 단계) 은 대부분의 벤치마크에서 전체 FLUX.2 9B 모델 (50 단계) 을 능가하여, RL 유도 분별화가 모델 규모 축소로 인한 품질 격차를 효과적으로 해소할 수 있음을 보여주었습니다.
절대 연구: 실험은 AC-DMD 의 일관성 정규화자와 SubGRPO 전략이 각각의 기준선 (A-DMD 및 단순 GRPO) 대비 성능을 크게 향상시킨다는 것을 확인했습니다.

중요성 및 주장

이 논문은 RTDMD 가 인간 선호도와 정렬된 고품질 소수 단계 생성기를 훈련하기 위한 원칙적이고 효과적인 솔루션을 제공한다고 주장합니다. 분포 매칭과 보상 최적화를 통합하고 노이즈가 있는 중간 변수 및 하이브리드 역학의 특정 과제를 해결함으로써, 이 방법은 다단계 샘플링이나 더 큰 모델과 일반적으로 연관된 품질을 희생하지 않으면서 효율적인 이미지 생성 (4 NFE) 을 가능하게 합니다. 저자들은 이 작업을 실시간 및 상호작용 애플리케이션을 위한 고정밀 생성 모델을 실용화하는 한 걸음으로 위치시킵니다.

참고: 이 논문은 명시적으로 텍스트 - 이미지 생성에 범위를 제한하며, 보상 설계 및 시간적 일관성의 추가적인 어려움으로 인해 비디오 합성 또는 이미지 편집으로 프레임워크를 확장하는 것은 향후 과제로 남아 있다고 명시합니다.