Exploring the Design Space of Reward Backpropagation for Flow Matching

기술 요약: Flow Matching을 위한 보상 역전파(Reward Backpropagation)의 디자인 공간 탐색

1. 문제 정의

사전 학습된 텍text-to-image flow matching 모델은 고품질의 이미지를 생성하지만, 종종 인간의 선호도와 일치하지 않는 경우가 발생합니다. 정책 경사(policy-gradient) 방법들이 존재하지만, **직접 보상 역전파(direct reward backpropagation)**는 flow 샘플러의 미분 가능성을 활용하여 엔드투엔드로 보상 그래디언트를 전파함으로써 더 효율적인 샘플링 대안을 제공합니다. 그러나 이 접근 방식은 현대적 규모의 모델에서 발생하는 두 가지 잘 알려진 병목 현상으로 인해 어려움을 겪고 있습니다:

1. 활성화 메모리(Activation Memory): 전체 샘플링 궤적(trajectory)에 걸쳐 활성화를 저장하는 것은 불가능합니다.
2. 그래디언트 폭주(Gradient Explosion): 여러 단계에 걸친 체인된 자코비안 곱(Jacobian products)이 보상 그래디언트를 역방향으로 전달하면서 초기 타임스텝의 업데이트를 불안정하게 만드는 그래디언트 팽창을 유발합니다.

기존 솔루션들은 전체 역방향 궤적을 짧은 "고정된(pinned)" 경로(예: straight-through connectors 사용)로 대체하여 이러한 문제를 완화하려고 시도합니다. 그러나 이러한 방법들은 종 often 역방향 경로가 전체 롤아웃(rollout)을 어떻게 근사할 것인지에 대해 특정되고 고립된 디자인 선택에 의존합니다. 이는 특히 긴 구간에 대해 그래디언트 품질이 근사의 정확도에 크게 의존하게 만들며, 근본적인 디자인 공간을 탐구되지 않은 상태로 남겨둡니다.

2. 방법론: FlowBP 프레임워크

저자들은 역방향 궤적 자체를 주요 디자인 객체로 취급하는 통합된 대리 궤적(surrogate-trajectory) 프레임워크인 FlowBP를 제안합니다. FlowBP는 전체 $N$-단계 Euler 롤아웃을 통해 역전파하는 대신, 그래디언트가 없는 캐시된 롤아웃을 기반으로 구축된 가벼운 대리 궤적을 구성합니다.

이 프레임워크는 네 가지 독립적인 디자인 축으로 정의됩니다:

1. 보상 모델 입력(Reward-Model Input): 보상이 실제 캐시된 롤아웃 샘플($x_0$)에서 평가되는지, 아니면 사후 평균 추정치($\hat{x}_0$)에서 평가되는지를 결정합니다.
2. 활성 집합(Active Set): 어떤 속도(velocity) 평가들이 그래디언트와 함께 재순전파(re-forwarded)될지, 그리고 나머지는 상수로 취급될지를 지정합니다.
3. 적분 가중치(Integration Weights): 활성화된 속도들이 어떻게 결합되는지 정의합니다 (예: Euler 가중치 vs. Lagrange quadrature).
4. 브릿지 커플링(Bridge Coupling): 분할 인덱스 $j$가 전/후 세그먼트를 연결하는 데 사용되는지 여부와, 이 분할 지점에서의 중첩된 그래디언트 경로의 강도($\alpha$)를 제어합니다.

핵심 메커니즘:

• 메모리 효율성: 메모리는 롤아웃 길이($N$)가 아닌 활성 집합의 크기($|A|$)에 따라 스케일링됩니다.
• 그래디언트 안정성: 그래디언트 전파는 최대 하나의 자코비안 인자(브릿지 커플링을 통해)로 제한되어, 다단계 체인에 의한 지수적 팽창을 방지합니다.
• 대리 구조 구축: 프레임워크는 두 가지 유형의 대리 추정치를 지원합니다:
  • Leap 추정치: 커넥터를 사용하여 예측된 상태를 캐시된 상태에 고정하여(잔차를 도입함) 연결합니다.
  • 재구성(Reconstruction) 추정치: 캐시된 상태와 재순전파된 활성 속도를 사용하여 이산 Euler 구간을 재생함으로써, 포워드 값이 캐시된 롤아웃과 정확히 일치하도록 보장하며 커넥터 잔차를 발생시키지 않습니다.

3. 주요 기여

본 논문은 FlowBP 프레임워크의 세 가지 구체적인 구현체를 소개하며, 각각 디자인 공간의 서로 다른 영역을 탐구합니다:

1. FlowBP-Sparse:

   • 전략: 브릿지 없이($j=0$) 전체 롤아웃에 대해 희소(sparse) Euler 재구성을 사용합니다.
   • 메커니즘: 캐시된 Euler 업데이트를 재생하며, 중간 정도의 활성 집합만을 autograd에 노출합니다. 브릿지를 완전히 제거하여 교차 단계 자코비안이 없는 완전히 디커플된 그래디언트를 생성합니다.
   • 이점: 무조건적인 안정성을 보장하고, 실제 분포 내 샘플에 대해 보상을 평가하면서 메모리를 활성 집합 크기로 제한합니다.

2. FlowBP-Bridge:

   • 전략: 조절 가능한 커플링 강도($\alpha$)를 가진 브릿지로 구성된 전체 Euler-quadrature 궤적 재구성을 결합합니다.
   • 메커니즘: 롤아웃을 전/후 세그먼트로 나누고, 편향된 단일 속도 커넥터 대신 Euler 구성을 통해 형성된 브릿지로 연결합니다.
   • 이점: FlowBP-Sparse의 디커플링된 안정성과 브릿지 방식의 교차 단계 신용 할당(credit assignment) 사이를 보간하여, 제어 가능한 중첩 커플링을 허용합니다.

3. FlowBP-Lagrange:

   • 전략: Leap 단계의 quadrature 차수를 높여 커넥터 기반 접근 방식(LeapAlign과 유사)을 강화합니다.
   • 메커니一种: 두 개의 세그メント 구조 내에서 단일 속도 점프 대신 고차 Lagrange quadrature 규칙을 사용합니다.
   • 이점: 긴 구간에 대한 적분 정확도를 개선함으로써 커넥터 잔차($d_0$ 및 $d_j$)를 수치적으로 줄여, 기존 커넥터 기반 방법들에서 관찰되는 "엔드포인트 잔차 실패 모드"를 완화합니다.

4. 실험 결과

저자들은 세 가지 flow-matching 백본인 SD3.5-M, FLUX.1-dev, FLUX.2-Klein-base에 대해 세 가지 변형 모델을 평가했습니다.

• 성능: 선호도(HPSv2.1, PickScore, ImageReward), 품질(UR-Align, UR-IQ), 그리고 구성적 능력(GenEval) 지표 전반에서 FlowBP 변형 모델들은 직접 그래디언트 베이스라인(ReFL, DRaFT-LV, DRTune, LeapAlign)보다 일관되게 우수한 성능을 보였습니다.
• 가변성: 모든 지표와 모델을 압도하는 단일 변형 모델은 없었으며, 최적의 선택은 백본과 지표에 따라 달라졌습니다. 예를 들어, FlowBP-Lagrange는 SD3.5-M에 대해 HPSv2.1에서 뛰어난 성능을 보였고, FlowBP-Bridge는 FLUX.1-dev에 대해 PickScore에서 앞섰습니다.
• 구성적 생성: FLUX.1-dev에서 FlowBP-Lagrange는 GenEval 전체 점수를 63.25에서 69.88로 향상시켜, 선호도 정렬이 프롬프트 준수 능력을 저하시키지 않음을 입증했습니다.
• Ablation 통찰:
  • 엔드포인트 재구성: (재구성을 통해) 실제 롤아웃 샘플에서 보상을 평가하는 것이 ReFL과 같은 기존 방법들의 성능을 개선했습니다.
  • 커넥터 정확도: 고차 커넥터(Lagrange)는 단일 속도 leap에 비해 엔드포인트 예측 오류를 크게 줄여, 큰 잔차로 인한 학습 불안정성을 방지했습니다.
  • 브릿지 커플링: 중간 정도의 중첩 커플링($\alpha$)은 그래디언트 불안정성을 재도입하지 않으면서도 유용한 신호 전파를 제공했습니다.

5. 의의 및 주장

본 논문은 FlowBP가 이전의 직접 그래디언트 방법들을 특수 사례로 포함하는 동시에, 이전에 탐구되지 않았던 디자인 구성을 노출하는 통합된 관점을 제공한다고 주장합니다. 이 연구의 주요 의의는 다음과 같습니다:

• 샘플링과 최적화의 분리: 역방향 궤적을 디자인 객체로 취급함으로써, 샘플링 과정(캐시된 롤아웃)과 최적화 그래프(대리 모델)를 분리합니다.
• 확장성: 설계 단계에서부터 메모리와 그래디언트 폭주를 제한하여, 대규모 모델에서도 직접 보상 역전파를 가능하게 합니다.
• 실행 가능한 디자인: 네 가지 디자인 축(입력, 활성 집합, 가중치, 커플링)은 그래디언트 충실도, 메모리 사용량, 학습 안정성 사이의 트레이드오프를 제어할 수 있는 실질적인 조절 장치를 제공합니다.

저자들은 명시적인 대리 궤적 설계가 고립된 휴리스틱을 넘어, 역방향 경로에 대한 체계적인 탐구를 통해 flow 기반 생성 모델의 보상 정렬을 확장할 수 있는 실용적이고 해석 가능한 경로임을 결론짓습니다.