기술 요약: 확산 모델 미세 조정을 위한 효율적 어드저인트 매칭

문제 정의

보상 기반 미세 조정은 텍스트-이미지 (T2I) 생성에서 사전 훈련된 확산 및 흐름 모델을 인간 선호도와 정렬시키는 표준 접근법이 되었습니다. 이 중 **어드저인트 매칭 (Adjoint Matching, AM)**은 보상 미세 조정을 확률적 최적 제어 (SOC) 문제로 재구성함으로써 원칙적인 정의를 제공합니다. AM 은 사전 훈련된 모델을 $r(x)$를 보상 함수로 하는 목표 분포 $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$로 변환하는 것을 목표로 합니다.

그러나 AM 은 실제 적용을 방해하는 상당한 계산 비효율성으로 고통받고 있습니다:

1. 전방 시뮬레이션 병목 현상: AM 은 "메모리리스" 동역학 하에서 확률 미분 방정식 (SDE) 솔버를 사용하여 완전한 생성 궤적을 시뮬레이션해야 합니다. 이는 많은 수의 함수 평가 (NFEs) 를 필요로 하며, 전체 궤적을 메모리에 저장해야 합니다.
2. 후방 시뮬레이션 병목 현상: 제어 신호를 계산하기 위해 AM 은 샘플링된 궤적을 따라 어드저인트 상태의 후방 상미분 방정식 (ODE) 시뮬레이션을 요구합니다. 이는 각 적분 단계에서 사전 훈련된 네트워크를 통해 자코비안 - 벡터 곱 (JVPs) 을 계산해야 하므로 계산 비용이 매우 큽니다.

저자들은 이러한 병목 현상이 표준 AM 공식화에서 사전 훈련된 모델로부터 물려받은 비자명한 기본 드리프트와 본질적으로 연결되어 있음을 관찰합니다.

방법론: 효율적 어드저인트 매칭 (EAM)

본 논문은 이론적 보장을 유지하면서 두 가지 계산 병목 현상을 제거하기 위해 SOC 문제를 재구성한 **효율적 어드저인트 매칭 (Efficient Adjoint Matching, EAM)**을 제안합니다.

1. 기본 동역학 재설계

EAM 은 기본 동역학에 대한 두 가지 구조적 조건을 도입합니다:

• 선형성: 기본 드리프트 $b(x,t)$는 $x$에 대해 선형으로 재설계됩니다 (즉, $b(x,t) = D(t)x$).
• 메모리리스: 기본 동역학은 메모리리스로 유지되어 중간 상태를 엔드포인트에서 직접 샘플링할 수 있게 합니다.

확산 계수 $\sigma(t)$를 표준 설정과 일치하도록 고정함으로써, 저자들은 이러한 조건을 만족하는 허용 가능한 선형 드리프트의 특정 계열 (식 15) 을 유도합니다. 이 선형 드리프트는 어드저인트 상태에 대한 **폐쇄형 해 (closed-form solution)**를 가능하게 하여 후방 ODE 시뮬레이션과 JVP 평가의 필요성을 제거합니다.

2. 터미널 비용 보정

기본 드리프트를 변경하면 제어되지 않은 시스템의 터미널 분포가 변합니다. 최적 제어된 분포가 여전히 원하는 보상-기울어진 목표 분포 $p^\star_1$과 일치하도록 보장하기 위해 EAM 은 수정된 터미널 비용 $g(x)$를 유도합니다. 이 비용의 기울기 $\nabla g(x)$에는 사전 훈련된 스코어, 새로운 기본 터미널 스코어, 그리고 보상 기울기에 대한 항이 포함됩니다 (식 16).

3. 훈련 알고리즘

EAM 훈련 루프 (알고리즘 1) 는 다음과 같이 진행됩니다:

1. 엔드포인트 생성: 전체 SDE 롤아웃 대신 엔드포인트 이미지 $X_1$을 몇 단계의 결정론적 ODE 솔버 (예: 10 NFEs) 를 사용하여 생성합니다.
2. 중간 상태 구성: 중간 상태 $X_t$는 생성된 엔드포인트에 조건부인 원래 노이즈 커널 $q_t(\cdot | X_1)$에서 직접 샘플링되어 전방 SDE 시뮬레이션의 필요성을 우회합니다.
3. 폐쇄형 어드저인트: 어드저인트 상태는 엔드포인트 기울기 $\nabla g(X_1)$의 폐쇄형 스케일링을 통해 계산되어 후방 시뮬레이션을 완전히 제거합니다.
4. 제어 파라미터화: 제어 $u(x,t)$는 사전 훈련된 모델에서 초기화된 학습 가능한 속도 모델 (예: LoRA) 을 통해 파라미터화되어, 미세 조정이 사전 훈련된 동역학에서 시작되도록 합니다.

주요 기여

• 알고리즘 제안: 동일한 이론적 목표를 달성하지만 계산 비용을 현저히 줄인 보상 기울기 기반 미세 조정 알고리즘인 EAM 의 도입.
• 이론적 특성화: 효율적인 궤적 구성 (메모리리스) 과 폐쇄형 어드저인트 계산을 위한 요구 사항을 만족하는 선형 기본 드리프트에 대한 특성화.
• 실증적 검증: 표준 T2I 벤치마크에서 EAM 이 AM 보다 최대 4 배 빠르게 수렴하며 여러 지표에서 AM 과 동등하거나 더 나은 성능을 보임을 입증.

결과

저자들은 Pick-a-Pic 데이터셋을 사용하여 **Stable Diffusion 3.5-Medium (SD3.5-M)**에서 EAM 을 평가하고, 이를 AM 및 베이스라인 SD3.5-M 과 비교했습니다.

• 효율성: EAM 은 40 단계 SDE 시뮬레이션을 10 단계 ODE 엔드포인트 샘플링으로 대체하고 후방 어드저인트 시뮬레이션을 제거함으로써 반복당 훈련 비용을 최대 4 배까지 줄입니다.
• 성능: PickScore, ImageReward, HPSv2.1, CLIPScore, Aesthetics를 포함한 벤치마크에서 EAM 은 일관되게 AM 과 동등하거나 더 나은 성과를 냅니다.
  • 단일 보상 (PickScore) 설정에서 EAM 은 AM 의 23.00 대비 PickScore 23.16 을 달성했습니다.
  • 다중 보상 설정 (PickScore + HPSv2.1 + Aesthetics) 에서 EAM 은 PickScore 22.93 과 Aesthetics 점수 6.56 을 달성하여 AM(6.21) 을 능가했습니다.
• 정성적 개선: 시각적 비교에 따르면 EAM 은 AM 이 때때로 세부 사항을 부드럽게 만드는 것과 달리 미세한 세부 사항과 복잡한 구도를 더 잘 보존합니다. 이는 훈련 (엔드포인트 조건부 노이즈) 과 추론 궤적 간의 EAM 의 일관성 때문입니다.

중요성과 주장

본 논문은 어드저인트 매칭의 주요 한계가 SOC 공식화 자체에 내재된 것이 아니라 특정 기본 드리프트 선택의 산물이라고 주장합니다. 기본 드리프트를 선형이고 메모리리스가 되도록 재설계하고 이에 따라 터미널 비용을 조정함으로써, EAM 은 이론적 정렬 보장을 희생하지 않으면서 두 가지 주요 계산 병목 현상 (전방 SDE 시뮬레이션 및 후방 어드저인트 시뮬레이션) 을 제거합니다.

저자들은 EAM 을 샘플 효율성을 유지하면서 계산 오버헤드를 대폭 줄여 확산 모델의 고품질 정렬을 더 접근 가능하게 만드는 실용적이고 확장 가능한 보상 미세 조정 솔루션으로 위치시킵니다. 이 연구는 모든 최적화 문제를 해결한다고 주장하지는 않습니다 (예: 버퍼 재생 또는 적응형 스케줄은 향후 과제로 언급됨) 만, 이 분야의 효율성에 대한 새로운 기준을 제시합니다.