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

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

왜 중요한가

핵심 기여

2단계 도약 궤적(Two-Step Leap Trajectory) 설계

수십 단계의 복잡한 생성 과정을 단 두 번의 예측 단계로 요약하는 궤적을 구축했다. 이를 통해 메모리 사용량을 일정하게 유지하면서도 생성 초기 단계까지 보상 그래디언트를 안정적으로 전달하여 모델 가중치를 업데이트한다.

그래디언트 할인(Gradient Discounting) 메커니즘

여러 단계를 거치며 발생하는 중첩된 그래디언트(Nested Gradient)의 폭주 문제를 해결하기 위해, 이를 완전히 제거하는 대신 특정 계수 α를 곱해 크기를 줄이는 방식을 도입했다. 이는 학습 신호를 보존하면서도 최적화 안정성을 극대화한다.

궤적 유사도 가중치(Trajectory-Similarity Weighting) 적용

압축된 2단계 궤적이 실제 다단계 생성 경로와 얼마나 일치하는지를 측정하여 손실 함수에 반영했다. 실제 경로와 유사한 샘플에 더 높은 가중치를 부여함으로써 정렬 학습의 신뢰도를 높였다.

관련 Figure

Screenshot
LeapAlign이 '자전거 위의 주차 요금기'나 '보트 오른쪽의 여행 가방'과 같이 복잡한 위치 관계와 객체 조합 프롬프트를 가장 정확하게 시각화하고 있음을 보여준다. 이는 초기 단계 튜닝을 통한 레이아웃 개선 효과를 입증한다.
GenEval 벤치마크에서 기본 Flux 모델 및 다른 정렬 기법들과 LeapAlign의 생성 결과를 비교한 이미지이다.

핵심 아이디어 이해하기

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

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

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

관련 Figure

Chart
α=1.0(할인 없음)일 때는 그래디언트가 폭주하여 성능이 저하되지만, α=0.3으로 적절히 조절하면 그래디언트 크기를 안정화하면서도 유용한 정보를 활용해 최적의 성능을 낼 수 있음을 증명한다.
중첩된 그래디언트(Nested Gradient)의 유무와 α 계수에 따른 성능 및 그래디언트 노름(Norm) 변화 분석 그래프이다.

방법론

전체 생성 궤적에서 무작위로 두 지점 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 값을 얻고] → [이 값이 클수록 실제 생성 과정과 일치하는 신뢰도 높은 데이터임을 나타낸다].

관련 Figure

Diagram
긴 ODE 샘플링 단계를 두 개의 큰 도약으로 압축하고, Latent Connector를 통해 실제 궤적과 연결하며 보상 그래디언트를 역전파하는 구조를 시각화했다. 이를 통해 모든 생성 단계의 가중치를 업데이트할 수 있음을 보여준다.
LeapAlign의 핵심 아이디어인 2단계 도약 궤적(Two-Step Leap Trajectory) 구축 과정을 보여주는 다이어그램이다.

주요 결과

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단계 이상으로 늘려도 성능 이득은 미미하면서 메모리 사용량만 증가한다는 점을 수치적으로 증명했다.

관련 Figure

Chart
LeapAlign이 기존 DRTune보다 더 빠르고 높은 보상 점수 상승을 보여줌을 확인할 수 있다. 특히 GenEval 레이더 차트에서 위치(Position), 개수(Counting) 등 구조적 정렬 지표에서 압도적인 우위를 점하고 있다.
LeapAlign의 학습 반복 횟수에 따른 보상 점수 변화 및 다양한 벤치마크에서의 성능 비교 차트이다.

기술 상세

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과 같은 다른 플로우 매칭 아키텍처에도 범용적으로 적용 가능하다는 일반성을 입증했다.

관련 Figure

Chart
α=0.3일 때 최적의 성능을 내며, 2단계 도약이 성능과 메모리 사용량 사이에서 가장 효율적인 선택임을 수치로 보여준다. 또한 최종 이미지 x0를 보상 모델의 입력으로 사용하는 것이 가장 효과적임을 나타낸다.
그래디언트 할인 계수, 도약 단계 수, 보상 모델 입력값 등 다양한 하이퍼파라미터에 대한 절제 실험(Ablation Study) 결과이다.

한계점

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

실무 활용

고해상도 이미지 생성 모델인 Flux 등을 적은 메모리로도 인간의 선호도에 맞게 미세 조정할 수 있는 실용적인 방법론을 제공한다.

특정 예술적 화풍이나 브랜드 가이드라인에 맞춘 이미지 생성 모델의 사후 학습
복잡한 문장 구조(위치 관계, 객체 개수 등)를 정확히 반영해야 하는 광고 및 디자인 에이전시용 모델 튜닝
제한된 GPU 자원 환경에서 대규모 플로우 매칭 모델을 효율적으로 정렬 학습하고자 하는 경우

코드 공개 여부: 공개

코드 저장소 보기

키워드

Flow Matching(플로우 매칭)Post-Training(사후 학습)Reward Alignment(보상 정렬)Gradient Discounting(그래디언트 할인)Flux Model(플럭스 모델)

용어 해설

플로우 매칭(Flow Matching): — 데이터 분포를 가우시안 노이즈로 변환하거나 그 반대로 변환하는 연속적인 벡터 필드를 학습하는 생성 모델 기법이다. 확산 모델과 유사하지만 결정론적인 궤적을 따라 샘플링을 수행하며, 미분 가능한 생성 과정을 제공하여 보상 그래디언트를 직접 전파할 수 있는 장점이 있다.
시간 기반 오차 역전파(Backpropagation Through Time): — 순차적인 데이터나 여러 단계를 거치는 생성 과정에서 각 단계의 연산을 거슬러 올라가며 그래디언트를 계산하는 방식이다. 생성 단계가 길어질수록 메모리 사용량이 급격히 증가하고 그래디언트가 폭주하거나 소실되는 문제가 발생하기 쉽다.
상미분 방정식 샘플링(ODE Sampling): — 학습된 벡터 필드를 따라 노이즈에서 이미지를 생성할 때 상미분 방정식을 수치적으로 풀어내는 과정이다. 여러 번의 반복적인 계산 단계를 거쳐 최종 결과물에 도달하며, 각 단계는 이전 단계의 출력값에 의존하는 연속적인 구조를 가진다.
그래디언트 할인(Gradient Discounting): — 역전파 과정에서 발생하는 그래디언트의 크기를 특정 계수를 통해 조절하여 학습의 안정성을 도모하는 기법이다. 특히 여러 단계를 거치며 증폭된 그래디언트 성분을 완전히 제거하는 대신 적절히 줄여줌으로써 유용한 학습 신호를 보존하면서도 수치적 불안정성을 해결한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환