기술 요약: 다양하고 효율적인 확산 정렬을 위한 TMPO

1. 문제 제기

강화 학습 (RL) 은 확산 (diffusion) 및 흐름 매칭 (flow-matching) 모델을 하위 작업과 정렬시키는 데 상당한 잠재력을 보여주었습니다. 그러나 기존 방법들은 주로 스칼라 보상 최대화에 의존하는데, 저자들은 이것이 본질적으로 모드 추구 (mode-seeking) 성향을 가진다고 지적합니다. 이러한 최적화 목표는 모델이 소수의 고보상 궤적에 확률 질량을 집중시키게 하여, 두 가지 치명적인 실패 모드를 초래합니다:

1. 생성 다양성의 저하: 모델은 시각적 모드 붕괴 (visual mode collapse) 를 일으켜 제한된 구성과 스타일 집합으로 수렴하며, 사전 학습 중에 학습된 타당한 출력들의 이방성 분포를 무시합니다.
2. 불완전한 보상의 증폭: 특정 보상 (예: OCR 정확도 또는 특정 미적 점수) 을 선호하는 특정 특징들을 과도하게 강화함으로써, 모델은 전체 생성 품질을 희생하는 경우가 종종 있는 허위 속성들을 증폭시킵니다.

KL 정규화와 같은 현재의 완화 전략들은 목표의 근본적인 모드 추구 성향을 바꾸지 못하며, 기존 생성 흐름 네트워크 (GFlowNet) 접근법들은 종종 장기간의 탈노이즈 시간 범위에서 오류가 누적되는 비실용적인 상태 - 흐름 추정이나 국소적 제약에 의존합니다.

2. 방법론: 궤적 매칭 정책 최적화 (TMPO)

TMPO 는 최적화 패러다임을 스칼라 보상 최대화에서 궤적 수준 보상 분포 매칭으로 근본적으로 전환합니다.

2.1 소프트맥스 궤적 균형 (Softmax-TB)

기대 보상을 최대화하는 대신, TMPO 는 정책의 궤적 확률을 보상에 의해 유도된 목표 분포와 일치시키려 합니다.

• 목표: 동일한 프롬프트에서 생성된 $K$개의 궤적 그룹에 대해, 이 방법은 정책 확률 $P_\theta(\tau)$를 볼츠만 분포 $q(\tau) \propto \exp(\beta R(\tau))$와 일치시킵니다.
• 분할 함수 없는 공식화: $K$개의 궤적 그룹 내에서 확률과 보상을 정규화함으로써 비실용적인 전역 분할 함수를 제거합니다. 이 방법은 소프트맥스-TB 이점을 정의합니다:
  $$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$$
• 이론적 보장: 저자들은 이 목적 함수를 최소화하는 것이 그룹 내에서 순방향 KL 발산 ($D_{KL}(q \parallel p)$) 을 최소화하는 것과 동등함을 증명합니다. 표준 정책 경사에서 사용되는 역방향 KL 과 달리, 순방향 KL 은 모드 커버링 (mode-covering) 성질을 가집니다. 즉, 양의 목표 가중치 $q_i$를 가진 궤적 $i$에 대해 정책 확률 $p_i$가 0 에 가까워지면 무한한 패널티를 부과합니다. 이는 보조 KL 정규화 없이도 다양성 유지를 이론적으로 보장합니다.

2.2 동적 확률적 트리 샘플링

대규모 흐름 매칭 모델에 대해 여러 개의 독립적인 궤적을 샘플링하는 계산 비용을 해결하기 위해, TMPO 는 트리 구조 샘플링 전략을 도입합니다:

• 접두사 공유: 궤적들은 결정론적인 ODE 접두사 (공통 탈노이즈 단계) 를 공유하고 동적으로 계획된 확률적 지점에서 분기합니다.
• 동적 분기: 분기점은 고정되지 않으며, 훈련이 진행됨에 따라 초기 고잡음 영역 (거친 수준의 다양성) 에서 후기 저잡음 영역 (세밀한 변형) 으로 이동하는 커리큘럼을 따릅니다. 분기 위치는 고정된 트리 기하학에 대한 과적합을 방지하기 위해 베타 분포를 사용하여 확률적으로 교란됩니다.
• 효율성: 이 접근법은 공유 접두사에 대한 순방향 전달을 공유함으로써 불필요한 계산을 줄여 훈련 시간을 크게 단축하면서도 효과적인 탐색을 유지합니다.

3. 주요 기여

1. 근본 원인 식별: 이 논문은 보상 최대화의 모드 추구 성향을 확산 RL 에서의 보상 해킹 및 다양성 저하의 주된 원인으로 식별하고, 분포 매칭으로의 전환을 제안합니다.
2. 소프트맥스-TB 목적 함수: 순방향 KL 발산의 모드 커버링 특성을 계승하여 자연스럽게 생성 다양성을 보존하는, 분할 함수가 없는 궤적 수준 분포 매칭 목적 함수를 도입했습니다.
3. 동적 확률적 트리 샘플링: 탈노이즈 접두사를 공유하고 분기 일정을 동적으로 조정함으로써 대규모 모델의 훈련 시간을 단축하는 확장 가능한 샘플링 메커니즘입니다.

4. 실험 결과

저자들은 FLUX.1-dev 백본을 사용하여 세 가지 정렬 작업인 구성 이미지 생성 (GenEval), 시각적 텍스트 렌더링 (OCR), 인간 선호도 정렬 (PickScore) 에서 TMPO 를 평가했습니다.

• 다양성: TMPO 는 최첨단 방법 (Flow-GRPO, MixGRPO, TreeGRPO, GARDO) 에 비해 평균 9.1% 향상된 생성 다양성을 보였습니다. 이는 일관하게 양의 로그 기하 평균 거리 (LGMD) 점수를 유지한 유일한 방법이었으며, 이는 건강한 다양성을 나타내는 반면, 베이스라인들은 종종 음의 LGMD(모드 붕괴) 를 나타냈습니다.
• 성능: TMPO 는 작업별 지표 (GenEval 정확도, OCR 정확도) 와 인간 선호도 점수 (HPS-v2.1, ImageReward, PickScore) 에서 경쟁력 있거나 우수한 성과를 달성했습니다.
• 효율성: 접두사 공유 및 트리 샘플링을 활용함으로써 TMPO 는 기존 방법 대비 최대 27% 까지 반복 훈련 시간을 단축했습니다.
• 절충: 파레토 분석은 TMPO 가 보상 최대화, 다양성 보존, 훈련 효율성 사이의 최적 절충점을 달성함을 보여주었습니다.

5. 중요성 및 주장

이 논문은 TMPO 가 생성형 AI 정렬에서 보상 해킹 딜레마에 대한 근본적인 해결책을 제공한다고 주장합니다. 스칼라 보상을 최대화하는 대신 보상 유도 분포와 일치하도록 목적 함수를 재구성함으로써, TMPO 는 외부 정규화나 휴리스틱 재가중치에 의존하지 않고 구조적으로 다양성 보존을 강제합니다. 이 방법은 사전 학습 중에 학습된 풍부하고 다중 모달인 출력 분포를 유지하면서 높은 정렬 점수를 달성할 수 있음을 보여주며, 동시에 다중 궤적 RL 방법과 일반적으로 관련된 계산 오버헤드를 줄입니다. 저자들은 TMPO 를 대규모 흐름 매칭 모델에 적합한 확장 가능한 프레임워크로 위치시키며, 비디오 생성 및 로봇 행동 생성으로의 잠재적 확장을 제시합니다.