LeapAlign: 2단계 궤적 구축을 통한 모든 생성 단계에서의 사후 학습 플로우 매칭 모델 정렬

기존의 이미지 생성 모델 정렬 방식은 메모리 한계로 인해 이미지의 전체 구도를 결정하는 초기 생성 단계를 직접 수정하기 어려웠다. 이 논문은 긴 생성 과정을 단 2단계의 '도약(Leap)'으로 압축하여 모든 단계에서 효율적으로 학습할 수 있는 방법을 제시함으로써 이미지 품질과 텍스트 일치도를 동시에 개선했다.

이미지 생성 모델은 노이즈에서 시작해 점진적으로 이미지를 만들어가는데, 초기 단계는 전체적인 구도를 잡고 후기 단계는 세부 묘사를 담당한다. 하지만 기존의 직접 그래디언트 방식은 역전파 경로가 너무 길어지면 메모리가 부족해지거나 숫자가 너무 커져서(Gradient Explosion) 초기 단계까지 학습 신호를 보내지 못하고 마지막 단계만 겨우 수정하는 한계가 있었다.

LeapAlign은 이 긴 경로를 '지름길'로 연결하는 아이디어에서 출발한다. 현재 상태에서 여러 단계를 건너뛰어 다음 상태를 한 번에 예측하는 'Leap Prediction'을 활용해, 전체 과정을 단 두 번의 큰 도약으로 압축한다. 이렇게 하면 아무리 초기 단계라도 단 두 번의 미분 연산만으로 보상 신호를 전달할 수 있게 된다.

단순히 지름길만 만드는 것이 아니라, 이 지름길이 실제 구불구불한 생성 경로와 너무 동떨어지지 않도록 관리한다. 실제 경로와 비슷한 지름길에서 나온 학습 신호를 더 중요하게 취급하고, 지름길 연산 중에 숫자가 너무 커지는 부분만 살짝 깎아내어 학습이 튕겨 나가지 않도록 정교하게 제어한다. 결과적으로 모델은 이미지의 세부 사항뿐만 아니라 초기 구도부터 사용자의 의도에 맞게 생성하는 법을 배우게 된다.

전체 생성 궤적에서 무작위로 두 지점 k와 j(k > j)를 선택하고, 현재 모델의 속도 예측값 vθ를 사용하여 xk에서 xj를 예측하는 첫 번째 도약과 xj에서 최종 이미지 x0를 예측하는 두 번째 도약을 구성한다. [잠재 변수 xk와 시간 k를 입력으로] → [vθ를 통한 선형 근사 연산을 수행해] → [예측된 잠재 변수 x_hat_j를 얻고] → [이는 여러 ODE 단계를 건너뛴 중간 상태를 의미한다].

예측된 상태와 실제 샘플링된 상태 사이의 불일치를 해결하기 위해 Latent Connector를 도입한다. [예측값 x_hat_j와 실제값 xj의 차이를 계산하고] → [stop_gradient 연산을 통해 미분값은 차단한 채 실제값 xj로 보정하여] → [다음 단계의 입력으로 전달하고] → [미분 가능성을 유지하면서도 실제 생성 궤적의 흐름을 따르게 한다].

최종 손실 함수는 힌지 로스(Hinge Loss) 형태의 보상 함수와 궤적 유사도 가중치 w_sim을 결합하여 정의한다. [예측 궤적과 실제 궤적의 유클리드 거리를 입력으로] → [역수 관계의 가중치 연산을 수행해] → [w_sim 값을 얻고] → [이 값이 클수록 실제 생성 과정과 일치하는 신뢰도 높은 데이터임을 나타낸다].

Flux.1-dev 모델을 기반으로 한 실험에서 LeapAlign은 기존의 직접 그래디언트 방식인 DRTune 및 정책 그래디언트 방식인 MixGRPO를 모든 지표에서 능가했다. 특히 일반적인 선호도 지표인 HPSv2.1에서 0.4092를 기록하여 기본 모델(0.3078) 대비 비약적인 향상을 보였으며, 경쟁 모델인 DRTune(0.3882)보다 우수한 성능을 입증했다.

복합적인 프롬프트 이해 능력을 측정하는 GenEval 벤치마크에서도 전체 점수 0.7420을 달성하여 MixGRPO(0.7232)와 DRTune(0.7101)을 앞섰다. 특히 '두 객체 간의 관계', '색상', '위치' 등 이미지의 전역적인 구조 이해가 필요한 항목에서 큰 폭의 개선이 확인되었는데, 이는 초기 생성 단계를 효과적으로 업데이트한 결과로 분석된다.

Ablation Study를 통해 그래디언트 할인 계수 α가 0.3일 때 가장 안정적이고 높은 성능을 보임을 확인했다. 또한 1단계 도약보다 2단계 도약이 성능과 메모리 효율 사이의 최적의 균형점(Trade-off)을 제공하며, 3단계 이상으로 늘려도 성능 이득은 미미하면서 메모리 사용량만 증가한다는 점을 수치적으로 증명했다.

LeapAlign은 Rectified Flow Matching의 결정론적 특성을 활용하여 다단계 ODE 샘플링 과정을 2단계의 선형 근사 궤적으로 치환한다. 핵심은 Nested Gradient 항인 j(k-j) * (∂vθ(xj)/∂xj) * (∂vθ(xk)/∂θ)를 완전히 제거하지 않고 α 계수를 통해 스케일링하여 보존한다는 점이다. 이는 DRTune이 해당 항을 완전히 제거하여 학습 신호를 손실했던 문제를 해결한다.

학습 시 전체 궤적 [0, 1] 범위에서 k와 j를 무작위로 샘플링함으로써 모델이 생성의 모든 타임스텝에 대해 보상 신호를 받을 수 있도록 설계되었다. 실험 결과, 초기 단계(타임스텝 1에 가까운 영역)를 포함하여 학습하는 것이 이미지의 레이아웃과 구성을 정확하게 정렬하는 데 필수적임이 밝혀졌다.

구현 측면에서는 16개의 GPU에서 300회 반복 학습만으로도 SOTA 성능에 도달할 정도로 수렴 속도가 빠르다. 또한 CLIP 기반의 HPSv2.1뿐만 아니라 VLM 기반의 HPSv3 등 다양한 보상 모델과 호환되며, Stable Diffusion 3.5 Medium과 같은 다른 플로우 매칭 아키텍처에도 범용적으로 적용 가능하다는 일반성을 입증했다.

현재 LeapAlign은 미분 가능한 보상 모델(Differentiable Reward Model)이 존재할 때만 직접적으로 그래디언트를 전파할 수 있다. 미분 불가능한 보상 모델이나 인간의 직접적인 피드백을 활용하기 위해서는 미분 가능한 가치 모델(Value Model)을 추가로 도입하는 등의 확장이 필요하다.