기술 요약: 강화 어드저인트 매칭 (RAM)

문제 제기

확산 (Diffusion) 및 흐름 매칭 (Flow-matching) 모델은 사전 학습이 지도 회귀에 의존하기 때문에 고충실도 생성을 달성합니다. 즉, 깨끗한 샘플을 분석적으로 노이즈 처리하고, 모델이 폐형식 (closed-form) 타겟에 대해 회귀합니다. 그러나 보상 함수 (예: 구성적 정확성, 텍스트 렌더링, 또는 인간 선호도) 와 이러한 모델을 정렬하기 위해 강화 학습 (RL) 후 학습을 수행하면 이러한 효율적인 구조가 붕괴됩니다.

기존 RL 후 학습 방법은 상당한 계산 및 구조적 장애물에 직면합니다:

1. 정책 경사 (Policy Gradient) 방법: 탈노이즈 과정을 마르코프 결정 과정 (MDP) 으로 취급하여 확률적 SDE 롤아웃을 요구하며 변분 하한을 최적화합니다. 이는 높은 분산을 도입하고 많은 샘플링 단계를 필요로 합니다.
2. 확률적 최적 제어 (어드저인트) 방법: KL 정규화된 보상 최대화 목표를 정확하게 타겟팅하지만, 저장된 SDE 궤적을 따라 역방향 어드저인트 스윕을 수행하고 샘플링 과정을 통해 역전파된 보상 기울기를 요구합니다. 이는 수치적으로 불안정하고 계산 비용이 많이 들며, 특히 노이즈 스케줄이 노이즈가 많은 끝부분에서 발산하는 흐름 매칭 모델에서 특히 그렇습니다.
3. 대리 손실 (Surrogate Losses): 최근 접근법은 처리 불가능한 가능도 (likelihoods) 를 대리물 (예: ELBO 또는 대비적 엔드포인트) 로 대체하여, 사전 학습의 확장성을 가능하게 하는 깨끗한 회귀 구조를 희생합니다.

핵심 과제는 보상을 효과적으로 최적화하면서도 사전 학습 회귀 프레임워크의 단순성과 확장성을 유지하는 RL 후 학습 목표를 도출하는 것입니다.

방법론: 강화 어드저인트 매칭 (RAM)

저자들은 SDE 롤아웃, 역방향 어드저인트 스윕, 또는 보상 기울기를 요구하지 않고 사전 학습 회귀 구조를 RL 후 학습으로 확장하는 강화 어드저인트 매칭 (RAM) 을 제안합니다.

이론적 기반

이 방법은 확률적 최적 제어 문제로 공식화된 KL 정규화된 보상 최대화에 기반합니다.

1. 최적 과정 구조: KL 정규화 하에서 최적 생성 과정은 깨끗한 엔드포인트 ($X_0$) 의 분포를 고보상 샘플 ($p^*_0 \propto p_{ref}^0 \exp(r)$) 로 기울이지만, 깨끗한 엔드포인트가 주어졌을 때의 노이즈 상태의 조건부 법칙은 변경되지 않습니다. 구체적으로, 고정된 $X_0$ 에서 노이즈 상태 $X_t$ 로 가는 궤적은 사전 학습에서 사용된 분석적 가우스 커널과 동일하게 유지됩니다.
2. 어드저인트 매칭 조건: 최적 제어 $u_t$ 는 제어값이 가치 함수 (어드저인트) 의 음수 스케일 기울기와 같아지는 고정점 조건을 만족합니다: $u_t = -\sigma_t \nabla_x V_t(x)$.
3. REINFORCE 분해: $\nabla r$ (비미분 가능할 수 있음) 을 계산하지 않기 위해, 저자들은 로그 도함수 (REINFORCE) 항등식을 가치 함수에 적용합니다. 이는 어드저인트를 역방향 브리지 점수 (bridge score) 로 가중된 보상 항과 경로 비용 보정 항으로 분해합니다.

RAM 목표 함수

RAM 은 경로 비용 보정을 제거하여 (낮은 분산과 규모에서의 안정성으로 정당화됨) 어드저인트를 근사하고, 노이즈 커널이 분석적으로 유지된다는 구조적 통찰을 활용합니다.

1. 엔드포인트 샘플링: 제어된 SDE 를 시뮬레이션하는 대신, RAM 은 기성 ODE 샘플러를 사용하여 현재 모델의 분포에서 깨끗한 엔드포인트 $X_0$ 를 추출합니다.
2. 분석적 노이즈 처리: 샘플링된 $X_0$ 를 분석적으로 노이즈 처리하여 훈련 상태 $X_t = (1-t)X_0 + t\epsilon$ 를 생성하며, 이는 사전 학습과 정확히 동일합니다.
3. 폐형식 타겟: 속도장 $v_\theta$ 에 대한 회귀 타겟은 보상 항과 베이즈 브리지 점수 (선형 노이즈 처리에 대해 폐형식 표현을 가짐) 를 어드저인트 조건에 대입하여 유도됩니다.

결과 손실 함수는 다음과 같습니다:
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
여기서 $X_0 \sim p_\theta^0$, $\epsilon \sim \mathcal{N}(0, I)$ 이며, $\text{sg}(\cdot)$ 는 그라디언트 정지 (stop-gradient) 연산자를 나타냅니다.

주요 효율성 메커니즘

• SDE 롤아웃 부재: 훈련 상태는 단일 엔드포인트를 $K$ 개의 독립적인 노이즈 벡터로 노이즈 처리하여 생성됩니다. 이는 모델 순전파당 $K$ 개의 조건부 독립 훈련 샘플을 생성하여, 단일 SDE 궤적에 의존하는 방법보다 단계당 더 많은 기울기 신호를 제공합니다.
• 보상 기울기 부재: 이 방법은 블랙박스 보상으로 작동합니다.
• 역방향 스윕 부재: 노이즈 커널의 분석적 특성은 역방향 어드저인트 적분의 필요성을 제거합니다.

주요 기여

1. 구조적 확장: 이 논문은 KL 정규화 하에서 사전 학습의 회귀 구조가 RL 후 학습으로 확장됨을 보여줍니다. 최적 과정은 엔드포인트 분포만 변경하고 조건부 노이즈 법칙은 변경하지 않습니다.
2. RAM 알고리즘: 기존 방법의 계산 병목 현상을 피하고 폐형식 베이즈 브리지 점수를 사용하여 보상 정보로 사전 학습 타겟을 보정하는 새로운 일관성 손실 (consistency loss) 입니다.
3. 이론적 정당성: RAM 의 고정점이 KL 정규화된 최적점에 해당함을 보여주는 유도 과정으로, 선형 보상 하에서 가우스 참조에 대해 정확성이 보장되며 일반적으로 1 차 일치성을 가집니다.
4. 분산 분석: 경로 비용 추정기에 대한 분석은 경로 비용 보정을 제거하는 것 (RAM 이 수행하는 방식) 이 전체 구간 또는 국소 어드저인트 추정기에 비해 분산을 크게 줄여주며, 이미지 규모 후 학습에 유일한 실현 가능한 추정기임을 보여줍니다.

실험 결과

저자들은 세 가지 작업에 걸쳐 Stable Diffusion 3.5 Medium (SD3.5M) 에서 RAM 을 평가했습니다:

1. 구성적 이미지 생성 (GenEval): 객체, 속성, 공간 관계의 정확성 측정.
2. 시각적 텍스트 렌더링 (OCR): 생성된 이미지 내 텍스트의 가독성 측정.
3. 인간 선호도 정렬 (PickScore): 인간 선호도와 정렬.

성능:

• 보상 최대화: RAM 은 Flow-GRPO, DiffusionNFT, AWM 을 포함한 베이스라인을 능가하며 세 가지 작업 모두에서 가장 높은 훈련 보상을 달성했습니다.
• 훈련 효율성: RAM 은 GenEval 의 경우 50 배 적은 단계로, OCR 은 48 배 적은 단계로, PickScore 는 34 배 적은 단계로 Flow-GRPO 의 최대 보상에 도달했습니다.
• 이미지 품질: 일부 베이스라인이 작업 지표를 위해 미적 품질을 저하시키는 "보상 해킹 (reward hacking)"을 겪은 것과 달리, RAM 은 작업 보상을 최대화하면서도 이미지 품질 지표 (Aesthetic, DeQA, ImageReward, HPSv2) 를 유지하거나 향상시켰습니다.
• CFG 호환성: RAM 은 후 학습 후에도 클래스프리 가이드 (Classifier-Free Guidance, CFG) 와 호환되지만, DiffusionNFT 및 AWM 과 같은 베이스라인은 CFG 가 적용될 때 작업 보상이 크게 감소했습니다.

중요성 및 주장

이 논문은 RAM 이 확산 및 흐름 매칭 모델의 RL 후 학습에서 패러다임 전환을 나타낸다고 주장합니다. 최적 제어 문제의 특정 구조적 속성을 활용함으로써 RAM, 사전 학습을 확장 가능하게 만든 "단순한 회귀" 패러다임을 복원합니다.

저자들은 다음과 같이 주장합니다:

• RAM 은 고해상도 이미지 생성으로 RL 을 확장하는 데 주요 병목 현상이었던 비용이 큰 SDE 롤아웃과 역방향 어드저인트 스윕을 제거합니다.
• 이 방법은 "간단하고 확장 가능하며", 기존 방법과 비교 가능한 단계당 계산 비용으로 훨씬 적은 훈련 단계로 최첨단 결과를 달성합니다.
• 이 접근법은 가시적인 품질 저하나 보상 해킹 없이 복잡한 보상 (구성성, 텍스트, 선호도) 과 모델을 성공적으로 정렬하며, 올바른 이론적 조건이 충족될 때 사전 학습 구조가 후 학습 정렬을 처리하기에 충분히 견고함을 보여줍니다.