어텐션 라우팅을 탐험 행동으로 삼아 확산 모델 RL 훈련 비용을 4배 줄인다

무엇인가

확산(diffusion)·플로우 매칭 생성기를 인간 선호나 미학 점수 같은 보상 신호에 맞추는 온라인 강화학습은 이미지와 비디오 품질을 끌어올리는 유효한 수단이지만, 비용 구조에 근본적인 문제가 있다. DDPO·DPOK는 확률적 디노이징 전이를 정책 행동으로 보고, Flow-GRPO·DanceGRPO는 결정론적 ODE 샘플러를 SDE로 바꿔 확률성을 주입한다. 이 방식은 탐험을 위해 역방향 궤적 자체를 교란하므로 롤아웃 기반 최적화 비용이 크고, 학습된 변화가 추론 시 쓰는 샘플러와 어긋날 위험도 있다. 이 논문은 질문을 바꾼다. 생성기의 원래 샘플러를 보존한 채 생성기 내부에서 확률적 탐험을 만들 수 있는가.

어떻게 동작하나

제안 방법 Routing Policy Optimization(RPO)의 탐험 행동은 어텐션 라우팅이다. 디퓨전 트랜스포머는 잠재 표현을 시각 토큰으로 쪼개는데, RPO는 토큰을 B=128 크기의 연속 블록 Q개로 묶고 지역성을 보존하는 Gilbert 곡선 순서로 배치한다. 전체 노이즈 레벨 중 소수의 앞쪽 레벨만 routed level로 지정하고, 각 routed 레벨에서 무작위로 고른 일부 트랜스포머 블록(anchor layer)은 dense 어텐션을 유지한다. 각 질의 블록은 자기 블록(S_q)을 안전 집합으로 항상 유지하고, 나머지 후보 블록 중 어느 것을 볼지 결정한다. 이 결정의 로짓은 별도 정책 헤드가 아니라 생성기 자신의 dense 어텐션 질량에서 나온다. 질의 블록 q가 키 블록 j에 쏟는 어텐션 비중을 routed 레이어 전체에 대해 로그 평균한 값이 로짓이 되고, dense 어텐션을 유지하는 선택지도 후보에 포함된다. 온도 T를 둔 소프트맥스에서 Gumbel 노이즈로 순위를 매겨, dense가 1위면 그대로 dense로 가고 아니면 안전 블록과 상위 K_u개 후보를 남긴다. K_u는 keep ratio κ로 정해지며, K_u=1이면 정확히 π에서의 샘플링이 된다. 라우팅 제안 패스에는 그래디언트가 흐르지 않는다.

무엇과 다른가

라우팅은 해당 전이의 velocity를 바꾸므로, 정책 경사에는 그 변화의 우도가 필요하다. RPO는 dense velocity를 중심으로 한 가우시안 우도 log q_θ = -||ṽ - v_θ||²/(2D s_t²)로 라우팅된 전이를 채점한다. 여기서 폭 s_t²는 해당 레벨에서의 평균 제곱 라우팅 편차로 고정해, 라우팅이 얼마나 세게 흔들든 로그 우도가 상수 크기를 유지하게 만든다. 롤아웃 중 기록한 잠재값에서 ṽ를 복원하고, dense velocity는 그래디언트와 함께 다시 계산한다. 이 우도의 그래디언트는 dense velocity를 기록된 라우팅 velocity 쪽으로 당기므로, 갱신은 추론 시 샘플러가 실제로 평가하는 velocity field에 직접 작용한다. 어드밴티지는 통제된 2-롤아웃 추정기로 만든다. 같은 프롬프트, 같은 초기 잠재, 같은 routed 레벨, 같은 anchor layer를 공유하고 라우팅 시드만 다른 두 롤아웃을 짝지어, 보상 차이를 leave-one-out 베이스라인으로 쓴다. 짝 차이는 윈저화하고 사분위 범위로 나눠 차원별 가중합한 뒤 stop-gradient를 걸어 Â_i로 고정하며, 손실은 -Â_i S_i/N_i를 두 롤아웃에 평균한 형태다. 추론 시에는 모든 마스크를 제거하고 dense 어텐션과 원래의 결정론적 샘플러를 복원한다.

어떻게 쓰나

이미지 실험은 Stable Diffusion 3.5 Medium을 PickScore SFW 분할 프롬프트로 학습한다. 512×512, 결정론적 10스텝 플로우 샘플러, classifier-free guidance 없음, 보상은 HPSv3에 0.1을 곱한 값이다. 평가는 DrawBench에서 다섯 지표 모두 개선됐고, 학습 보상이 아닌 선호 지표에서 폭이 컸다. ImageReward는 -0.54에서 -0.18로, UnifiedReward는 2.10에서 2.32로 올랐고, 화질 지표인 DeQA-Score도 3.35에서 3.56으로 상승했다. 같은 LoRA 어댑터와 같은 보상을 쓰는 Flow-GRPO와의 학습 속도 비교에서는, 학습 롤아웃 보상만 보면 RPO가 30스텝에 9를 넘고 9.5 근처에 머무는 반면 Flow-GRPO는 8 아래에 머문다. 다만 이는 서로 다른 종류의 롤아웃 점수라 결정적 비교가 아니며, 실제 비교는 dense 모델 기준이다. dense 모델에서 Flow-GRPO가 5스텝에서 잠깐 앞서지만 RPO가 10스텝부터 역전해, 25스텝에서 8.90으로 Flow-GRPO의 50스텝 점수 8.81을 이미 넘는다. RPO는 40스텝에서 10.20으로 정점을 찍고 9.91로 끝나, 시작 대비 +5.97을 얻어 Flow-GRPO의 +4.86보다 1.1점 높다. 즉 절반의 학습 스텝으로 상대의 최종 dense 보상을 따라잡는다. 두 곡선 모두 단일 학습 시드 결과다.

전제와 한계

비디오 실험은 Wan2.2-TI2V-5B를 704×1280, 13프레임, 10스텝 샘플러, CFG 스케일 5로 학습하고 보상은 네 번째 프레임마다 HPSv3로 계산한다. 베이스라인은 Flow-GRPO와 DanceGRPO(그룹 G=8, 전체 모델 갱신)이며, RPO는 스텝당 64개의 짝(G=2)을 수집한다. VBench 16개 세부 차원에서 6,220개의 13프레임 비디오를 생성해 평가한 결과, RPO는 두 RL 베이스라인을 모두 앞서는 8개 차원을 가지며 베이스 모델보다 전체 점수가 높은 유일한 학습 체크포인트다. 의미 정렬에서 이득이 컸는데 object class +4.40, human action +4.80, scene +3.43, multiple objects +2.15다. 훈련 비용은 24.1 GPU-시간으로 Flow-GRPO의 75.4보다 3.1배, DanceGRPO의 102.3보다 4.2배 적다(단, 연산량을 맞춘 비교는 아니다). 반면 색상과 외형 스타일은 베이스 모델보다 낮아졌고, Flow-GRPO는 dynamic degree를 12.50점 올리는 대신 일관성·부드러움·미학·색상·객체 클래스 정확도를 잃는다. Wan2.2 전체 점수 네 개가 단일 시드에서 0.81점 범위에 불과하므로 순위는 방향성 정도로 해석해야 한다고 저자들은 밝힌다.

개발자 관점에서 이 논문의 실용적 핵심은 추론 파이프라인을 건드리지 않는다는 점이다. 라우팅 마스크는 학습 중에만 쓰이고, 배포 시에는 dense 어텐션과 기존 결정론적 샘플러를 그대로 쓰므로 기존 서빙 코드와 스텝 수, 스케줄을 유지한 채 가중치만 교체할 수 있다. 또한 별도 정책 헤드나 SDE 변환이 필요 없어, 어텐션을 계산하는 트랜스포머 기반 생성기라면 비교적 적은 수정으로 얹을 수 있다. 다만 routed level, anchor layer 비율, keep ratio κ, 온도 T 같은 하이퍼파라미터가 학습 안정성과 전이 품질을 좌우하므로, 자체 모델에서는 검증 프롬프트에 대한 dense 모델 보상을 주기적으로 측정하며 튜닝해야 한다. 학습 롤아웃 보상이 아니라 마스크를 제거한 dense 모델의 보상을 봐야 한다는 점이 이 논문이 명시적으로 강조하는 실무 포인트다.

저자들이 밝힌 한계는 세 가지다. 첫째, 프레임워크가 어텐션에 국한된다. 행동이 키-값 블록에 대한 마스크이므로 어텐션을 계산하는 백본이 필요하고, 어텐션 질량이 존재하지 않는 다른 아키텍처로의 일반화는 아직 하지 못했다. 둘째, 라우팅 제안과 우도가 거칠다. 각 라우트 로짓이 routed 레이어와 헤드 전체에 걸쳐 블록 단위 어텐션 질량을 평균하므로 토큰 수준·레이어별 구조가 가려지고, velocity 우도는 이산 라우트가 만든 편차를 가우시안으로 취급한다. 더 세밀한 제안 분포와 라우트 선택 자체를 반영하는 우도가 필요하다. 셋째, RPO는 현재 기존 생성기에 적응할 뿐 생성기를 설계하지는 않는다. 네트워크 아키텍처와 RPO를 함께 설계하는 것이 더 높은 효율과 더 큰 RL 이득을 얻는 다음 방향이라고 제안한다.