MEND는 보상 이득이 이동 거리 값을 치를 때만 샘플을 움직인다

무엇인가

텍스트-이미지 플로우 모델을 학습된 보상 모델로 후학습하는 일에서 열린 문제는 스칼라 점수를 어떻게 "무엇을 얼마나 바꿀지"에 대한 업데이트로 바꾸느냐다. 기존 접근은 두 갈래다. Flow-GRPO 같은 정책 경사 방법은 KL 페널티 아래 그룹 상대 어드밴티지로 확률적 궤적을 재가중하고, DiffusionNFT 같은 가중 회귀 방법은 회귀 손실 안에서 모델 자신의 샘플 가중치를 조정한다. 둘 다 이미 있는 샘플의 기여도를 바꿀 뿐이다. ReFL·DRaFT 같은 보상 역전파는 방향을 주지만 모든 샘플을 움직이며, 그 이동이 이동 크기만 한 가치가 있는지 확인하지 않는다. KL 페널티나 참조 항은 분포 전체의 드리프트만 제한할 뿐 샘플 단위 결정을 하지 않는다. 논문은 3모드 예시에서 두 실패를 지적한다. 재가중은 어려운 모드를 억제하고, 검증 없는 상승은 샘플을 저밀도 영역으로 밀어 넣는다.

어떻게 동작하나

MEND는 이 결정을 샘플마다 명시적으로 내리는 방법이다. 원리는 하나다. 샘플은 자신이 얻는 보상이 이동 거리의 값을 치를 때만 움직인다. 한 라운드는 롤아웃, 타깃 선택(캡·제안·검증), 회귀 업데이트 한 번으로 구성된다. 표기는 rectified flow의 보간 z_t=(1-t)x0+tε와 clean 예측 x̂=z_t-t·v를 쓰고, 학습 어댑터 θ와 행동 어댑터 θ_old가 베이스를 공유한다. 프롬프트마다 θ_old가 결정론적 10스텝 궤적 G개를 생성하고, 저장 상태 z_q는 t_q=0.278에 가장 가까운 지점을 쓴다. 노름은 RMS 노름이다.

무엇과 다른가

절차는 세 단계다. 첫째, 캡. 이미 잘 나온 샘플을 더 밀면 보상은 조금 오르고 다양성만 잃는다. 그래서 프롬프트 그룹 안에서 보상을 분위수 q=0.75에서 자르고, 라운드가 진행되며 천천히 올라가는 전역 하한 κ_glob와 비교해 더 큰 값을 캡 κ(c)로 삼는다. 캡 이상 샘플은 이동량이 0이고, 캡을 넘는 초과분에는 아무 크레딧도 주지 않는다. 둘째, 제안. 캡 아래 샘플은 보상 그래디언트를 정규화한 방향으로 K=3개의 이동을 제안받는다. 이동 길이 η는 0.1, 0.2, 0.4 세 값이며 보상 스케일과 무관하다. 각 제안은 실제로 디코드되어 한 번씩 채점된다. 그래디언트가 0인 샘플은 그대로 유지된다. 셋째, 검증. 후보 집합에 원래 샘플 x를 항상 포함시킨 뒤 J(y)=R_κ(y,c)-‖y-x‖²/(2τ)를 최대화하는 후보를 고른다. 이는 근위점(proximal-point) 목적함수를 유한 후보 집합 위에서 최대화한 것이고, τ가 거리의 가격을 정한다. 캡 적용 보상 이득이 ‖y-x‖²/(2τ)를 넘을 때만 이동이 수용된다. 동점이면 x, 그다음 짧은 제안을 택한다. τ는 라운드 안에서는 고정이고 라운드 사이에 조정해 제안 샘플 중 30~60% 수용률을 목표로 한다.

어떻게 쓰나

회귀 단계에서 수용된 샘플은 d=y*-x, 그 외는 d=0으로 두고, 저장 상태 z_q에서 행동 예측에 d를 더한 값에 매칭한다. 엔드포인트 y* 자체에 회귀하면 롤아웃 나머지 부분에서 생기는 보상 무관 격차가 업데이트에 섞이기 때문이다. 손실은 이동 샘플 항과 유지 샘플 항(λ_keep=10)의 합이며, 속도 형태로 쓰면 목표는 v*=v_old-d/t_q다. 유지 샘플의 목표는 v_old 그대로다. 행동 어댑터는 고정 참조가 아니라 학습을 따라가는 지연 어댑터로, EMA 계수 μ_u=min(0.001u, 0.5)를 쓴다. 라운드당 AdamW 1스텝, 학습률 3e-4, ε=1e-12. 업데이트마다 샘플당 보상 역전파 1회, 캡 아래 샘플당 최대 3회 디코드·보상 평가, z_q에서 네트워크 순전파·역전파 1회가 필요하고, 샘플러를 통과하는 역전파는 하지 않는다.

전제와 한계

이론은 두 가지를 보인다. 명제 1은 수용된 타깃이 ‖y*-x‖²/(2τ) ≤ R_κ(y*)-R_κ(x) ≤ κ-R_κ(x)를 만족한다는 것, 따라서 이동 크기가 min{η_K, √(2τ(κ-R_κ(x)))}로 제한된다는 것이다. 캡에 가까운 샘플은 남은 여유 κ-R_κ(x)가 0으로 가면서 받을 수 있는 최대 이동도 0으로 줄어든다. 명제 2는 θ_old에서 회귀 손실의 음의 그래디언트가 (2/|A|)Σ η_i ∇_θ⟨ĥ_i, x̂_θ(z_q^(i))⟩, 즉 수용된 시드에만 제한되고 선택된 스텝으로 스케일된 보상 역전파와 같다는 것이다. 유지된 시드는 θ_old에서 그래디언트 기여가 0이다.

실험은 SD3.5-M을 PickScore로 학습하며 Zhou 등(2026)의 동일 예산 프로토콜을 따른다. 48개 Pick-a-Pic 프롬프트 × 업데이트당 24 이미지, rank-32 LoRA, 512픽셀 10스텝 롤아웃, 100 업데이트다. 베이스라인은 같은 프로토콜의 ReFL과 DiffusionNFT이고, 추가 비교 대상은 Flow-GRPO PickScore 어댑터(약 4천 업데이트)와 5보상 DiffusionNFT 모델(1.7천 업데이트)이다. 평가는 DrawBench 200 프롬프트 × 5 시드, 40 오일러 스텝으로 PickScore, HPSv2.1, HPSv3, ImageReward, CLIPScore, 미학 품질을 채점하고, 충실도는 같은 프롬프트·노이즈로 만든 베이스 이미지와의 DreamSim 거리로 잰다.

결과 수치는 이렇다. 베이스 거리 0.313이 같다는 조건에서 MEND는 6개 평가자 중 5개에서 Flow-GRPO를 앞서며 업데이트 수는 약 40분의 1이다. 미학 점수만 5.88 대 5.90으로 낮다. 동일 예산에서 update 100에 PickScore 24.03으로 ReFL 23.92, DiffusionNFT 23.43을 넘고, update 25 시점에 이미 약 4천·1.7천 업데이트를 쓴 두 모델의 수준을 추월한다. 학습 비용은 GB200 GPU 3장으로 10.0 GPU-시간(평가 제외)이고, 같은 설정을 다시 돌린 두 번째 런은 update 50에서 PickScore가 0.01 차이 났다. update 25에서 100 사이 DrawBench PickScore는 23.36에서 23.70으로 오르는 동안 베이스 거리는 0.294에서 0.313으로 거의 변하지 않는다. 백본 일반성도 확인된다. SD3-M에서 PickScore 20.36→23.70(Linear-DPO 20.96보다 높음), held-out HPSv2.1 0.215→0.298. 9스텝 1024픽셀 증류 모델 Z-Image-Turbo에서 PickScore 22.86→24.10, held-out HPSv2.1 0.295→0.315이고, HPSv2.1로 학습하면 0.357에 held-out PickScore 23.27이다. 절제 실험에서 전체 규칙은 캡 아래 모든 시드를 움직이는 것보다 다양성을 0.037 더 지키고 PickScore는 0.35 손해 본다. MEND의 라운드당 수용률은 34%로 PickScore 23.10, 다양성 0.297이었고, 캡을 빼면 45%, 가격을 빼면 78%로 수용률이 오르며 다양성은 0.287, 0.291로 떨어졌다. 캡 아래 전부를 움직이면(81%) PickScore는 23.45로 가장 높지만 다양성은 0.260으로 가장 낮다. 초기 가격 변경, 후보 1개 사용, 행동 EMA 제거는 update-50 PickScore를 최대 0.06 바꾸고, 후보 5개는 0.25 낮춘다.

개발자 입장에서 이 방법의 실질적 매력은 채택 비용이다. 차별화 가능한 보상만 있으면 어떤 플로우 백본에도 붙고, KL 항·고정 참조 모델·어드밴티지 가중치가 필요 없으며 샘플러를 통과하는 역전파도 없다. SD3.5-M 런이 10 GPU-시간이라는 점은 새 보상이나 새 백본마다 후학습을 반복해 볼 수 있는 수준이다. 도입 전에 확인할 것은 보상의 미분 가능성, τ와 수용률(30~60%) 튜닝, 캡 분위수 q=0.75와 κ_glob 스케줄이 자신의 프롬프트 분포에 맞는지다. 또 학습 보상이 오르는 동안 held-out 지표가 같이 오르는지 반드시 따로 추적해야 한다.

저자가 밝힌 한계도 분명하다. MEND는 미분 가능한 보상을 요구하고, 이론적 보장은 각 라운드의 타깃에 대해서만 성립한다. 더 길게 학습하면 보상을 과최적화할 수 있고, 일부 단일 보상 런은 100 업데이트 예산 안에서도 held-out 점수가 떨어진다. 실제로 CLIPScore와 ImageReward로 학습한 런은 held-out PickScore가 22.27, 22.33으로 가장 낮았고, 4개 런 모두 held-out 평가자에서 베이스보다 높게 끝나지만 그중 3개는 update 100 이전에 held-out 하락 구간을 보인다.