기술 요약: LeapAlign

문제 정의

본 논문은 사후 학습을 통해 흐름 매칭 (flow matching) 모델을 인간 선호도와 정렬하는 과제를 다룹니다. 직접 기울기 (direct-gradient) 방법은 흐름 매칭 샘플링의 미분 가능성에 기반하여 정책 기울기 (policy-gradient) 접근법 (예: GRPO) 에 대한 유망한 대안을 제시하지만, 긴 생성 궤적에 적용될 때 상당한 장애물에 직면합니다. 구체적으로, 전체 ODE 샘플링 단계를 통해 보상 기울기를 역전파하는 것은 막대한 메모리 비용을 초래하고 기울기 폭발의 위험을 내포합니다. 결과적으로 기존 직접 기울기 방법들은 일반적으로 최종 이미지 근처의 단일 시간 단계로만 업데이트를 제한하거나 불안정성을 피하기 위해 중첩된 기울기 (nested gradients) 를 폐기합니다. 이러한 제한은 최종 이미지의 전역 레이아웃과 구조적 구성을 결정하는 데 중요한 초기 생성 단계의 최적화를 방해합니다.

방법론: LeapAlign

저자들은 과도한 메모리 오버헤드 없이 모든 생성 단계, 특히 초기 시간 단계까지 효율적이고 안정적인 직접 기울기 전파를 가능하게 하는 미세 조정 방법인 LeapAlign을 제안합니다. 핵심 방법론은 세 가지 주요 구성 요소를 포함합니다:

1. 이단계 점프 궤적 구성

전체 생성 궤적을 통해 역전파하는 대신, LeapAlign 은 표준 전체 실행 궤적에서 파생된 두 단계로만 구성된 단축된 **점프 궤적 (leap trajectory)**을 구성합니다.

• 과정: 각 반복 동안, 노이즈 ($x_1$) 에서 깨끗한 이미지 ($x_0$) 까지 전체 궤적이 샘플링됩니다. 두 시간 단계 $k$와 $j$ ($k > j$) 가 무작위로 선택됩니다.
• 점프: 이 방법은 정렬된 흐름 (rectified flow) 모델의 한 단계 점프 예측 속성을 활용하여 미래 잠재 변수를 한 단계로 추정합니다:
  • 점프 1: $x_k$에서의 속도장을 사용하여 $x_k$로부터 잠재 변수 $x_j$를 예측합니다: $\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$.
  • 점프 2: $x_j$에서의 속도장을 사용하여 $x_j$로부터 최종 잠재 변수 $x_0$를 예측합니다: $\hat{x}_{0|j} = x_j - j v_\theta(x_j)$.
• 잠재 연결자: 전체 궤적에서 샘플링된 실제 상태와 예측된 상태를 정렬하면서 미분 가능성을 유지하기 위해 stop_gradient 연산을 사용한 "잠재 연결자 (latent connectors)"가 도입됩니다 (예: $x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$).
• 이점: 이는 역전파 경로를 두 단계로 줄여 원래 궤적 길이에 관계없이 메모리 비용을 일정하게 유지하면서도, 생성 과정의 임의의 부분을 업데이트하기 위해 $(k, j)$를 무작위로 선택할 수 있게 합니다.

2. 기울기 할인

유용한 학습 신호를 보존하면서 기울기 폭발을 해결하기 위해 LeapAlign 은 기울기 할인 (gradient discounting) 메커니즘을 도입합니다.

• 문제: 두 단계를 통해 역전파하면 시간 단계 간 상호작용을 포착하는 "중첩 기울기 (nested gradient)" 항 ($\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$) 이 도입됩니다. DRTune 과 같은 이전 방법들은 모델 입력에서 기울기를 차단하여 이 항을 완전히 제거함으로써 가치 있는 의존성 정보를 폐기합니다.
• 해결책: LeapAlign 은 이 항을 제거하는 대신 $\alpha \in [0, 1]$ 인자로 중첩 기울기 항을 스케일링합니다. 이는 두 번째 점프에서 속도 예측 입력을 수정함으로써 달성됩니다: $\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$.
• 효과: 이는 큰 기울기 항의 크기를 줄여 훈련을 안정화시키면서도 중첩 기울기를 포함한 전체 기울기 구조를 유지합니다.

3. 궤적 유사성 가중치

단축된 점프 궤적이 신뢰할 수 있는 감독을 제공하도록 보장하기 위해, 이 방법은 실제 생성 경로와 유사한 점프에 더 높은 가중치를 부여합니다.

• 메커니즘: 손실은 연결점 ($x_j$ 및 $x_0$) 에서 예측된 상태 ($\hat{x}$) 와 실제 상태 ($x$) 사이의 평균 절대 차이의 역수로 가중치가 부여됩니다.
• 공식: $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$이며, 여기서 $d$는 거리를 나타내고 $\tau$는 거의 동일한 쌍에 대한 과도한 강조를 방지하기 위한 최소 임계값입니다.
• 목표: 이는 실제 역학에서 크게 벗어난 점프가 기울기 업데이트에 덜 기여하도록 하여 감독의 안정성과 충실도를 향상시킵니다.

4. 미세 조정 목적 함수

이 방법은 보상 해킹을 방지하면서 보상을 최대화하기 위해 힌지 스타일 손실을 사용합니다: $L_{raw} = \max(0, \lambda - r(x_0))$. 핵심적으로, 점프 예측에서 보상을 추정하는 방법들과 달리 LeapAlign 은 전체 궤적으로 생성된 **실제 최종 이미지 ($x_0$)**에 대해 보상을 평가하여 보상 모델이 실제 시각적 및 의미적 품질을 평가하도록 보장합니다.

주요 기여

1. LeapAlign 프레임워크: 일정한 메모리 비용으로 초기 생성 단계까지 직접 기울기 역전파를 가능하게 하는 두 단계 점프 궤적을 구성하는 새로운 사후 학습 방법.
2. 기울기 안정화 기술: 중첩 기울기를 제거하는 대신 스케일링하는 기울기 할인의 도입과 일관된 점프를 우선시하는 궤적 유사성 가중치를 통해 기존 직접 기울기 방법의 불안정성과 신호 손실 문제를 해결.
3. 종합 평가: 인간 선호도 점수 (HPSv2.1, HPSv3, PickScore) 와 구성 정렬 벤치마크 (GenEval) 를 포함한 여러 지표를 통해 LeapAlign 이 최첨단 정책 기울기 (GRPO 기반) 및 직접 기울기 방법보다 우월함을 입증.

결과

저자들은 Flux.1-Dev 모델을 LeapAlign 으로 미세 조정하고 DanceGRPO, MixGRPO, ReFL, DRaFT-LV, DRTune 등을 포함한 베이스라인과 비교 평가했습니다.

• 일반 선호도 정렬: LeapAlign 은 도메인 내 (HPSv2.1) 및 도메인 외 (HPSv3, PickScore, UnifiedReward) 지표에서 가장 높은 점수를 달성했습니다. 특히, 세 개의 보상 모델을 사용하는 MixGRPO 를 능가하면서도 단일 보상 모델 (HPSv2.1) 로 훈련되었습니다.
• 구성 정렬: GenEval 벤치마크에서 LeapAlign 은 "두 개의 객체", "색상", "위치", "속성 바인딩" 작업에서 특히 큰 개선을 보였습니다. 이러한 카테고리에서의 개선은 전역 레이아웃을 결정하는 초기 생성 단계의 최적화에서 이 방법의 성공을 강조합니다.
• 정성적 분석: 시각적 비교는 LeapAlign 이 텍스트 프롬프트와 더 잘 정렬되도록 전역 이미지 구조를 상당 부분 수정하는 반면, 후기 단계 업데이트로 제한된 방법들 (예: ReFL) 은 종종 사전 훈련된 모델의 레이아웃을 유지함을 보여주었습니다.
• 절차 분석 실험: 실험은 두 단계 점프 설계가 성능과 메모리 간 가장 좋은 균형을 제공함을 확인했습니다. $\alpha=0.3$인 기울기 할인이 최적의 결과를 낳았으며, 전체 범위 $[0, 1]$에 걸쳐 시간 단계를 무작위로 선택하는 것이 고정된 범위보다 우월하여 초기 단계 업데이트의 중요성을 강조했습니다.

의미 및 주장

본 논문은 LeapAlign 이 흐름 매칭 사후 학습에서 메모리 효율성과 기울기 범위 간의 트레이드오프를 성공적으로 해결했다고 주장합니다. prohibitive 한 계산 비용 없이 초기 생성 단계에 대한 업데이트를 가능하게 함으로써, 복잡한 구성 작업에 필수적인 전역 이미지 구조의 최적화를 허용합니다. 저자들은 LeapAlign 을 확률적 정책 기울기 방법과 불안정한 직접 기울기 접근법 모두에 대한 안정적이고 효과적인 대안으로 위치시키며, 우수한 이미지 - 텍스트 정렬 및 생성 품질을 달성했다고 밝힙니다. 이 작업은 중첩 기울기를 폐기하는 대신 보존하고 스케일링하는 것이 직접 기울기 정렬의 성능을 향상시키는 핵심 요소임을 시사합니다. 향후 작업은 이러한 원칙을 비디오 생성으로 확장하는 것으로 언급됩니다.