Flow Factory 흐름 모델 강화학습 통합 프레임워크

초록

Flow‑Factory는 흐름‑매칭(Flow‑Matching) 및 확산 모델에 강화학습(RL)을 적용할 때 발생하는 코드 파편화·메모리·보상 설계 문제를 해결하기 위해 레지스트리 기반 모듈식 아키텍처를 제시한다. 모델, 알고리즘, 보상, 스케줄러를 각각 독립적인 컴포넌트로 분리하고, 사전 인코딩 캐시와 다중 보상 로더를 통해 GPU 메모리를 13 % 절감하고 학습 속도를 1.7배 가속한다. GRPO, DiffusionNFT, AWM 등 최신 RL‑for‑Flow 알고리즘을 동일한 인터페이스로 구현·비교할 수 있으며, Flux, Qwen‑Image, WAN 비디오 등 다양한 대형 모델에 손쉽게 적용 가능하도록 설계되었다.

상세 분석

본 논문은 흐름‑매칭 모델에 강화학습을 적용하는 최신 연구들의 급격한 성장에도 불구하고, 실무에서 마주치는 세 가지 핵심 장애물을 명확히 짚어낸다. 첫째, 알고리즘마다 고유한 SDE/ODE 구현과 학습 파이프라인이 존재해 코드베이스가 파편화되고, 모델‑특화 로직과 얽혀 있어 새로운 모델에 기존 알고리즘을 옮기려면 대규모 리팩터링이 필요했다. 둘째, 대형 멀티모달 모델은 텍스트 인코더·VAE 등 고정 파라미터가 차지하는 메모리가 압도적이라, 실제 파인튜닝 단계에서 GPU 메모리 한계에 부딪힌다. 셋째, 기존 프레임워크는 점별 보상만을 지원해, 최근 등장한 그룹‑와이즈 보상이나 다중 보상 조합을 구현하기가 어려웠다.

Flow‑Factory는 이 세 문제를 레지스트리 기반 설계, 사전 처리 메모리 최적화, 다중 보상 시스템이라는 세 축으로 해결한다. 레지스트리 메커니즘은 BaseAdapter, BaseTrainer, BaseRewardModel, SDESchedulerMixin 네 가지 추상 클래스를 정의하고, 각 구현체를 전역 레지스트리에 등록한다. YAML 설정만으로 모델‑알고리즘‑보상을 자유롭게 조합할 수 있게 함으로써, 통합 복잡도는 O(M + N)으로 감소한다(여기서 M은 지원 모델 수, N은 알고리즘 수).

메모리 최적화는 두 단계로 이루어진다. 학습 전 모든 조건(프롬프트 임베딩, VAE 라텐트 등)을 디스크에 캐시하고, 학습 중에는 트랜스포머 본체만 GPU에 올린다. 이 방식은 고정 인코더를 완전히 오프로드함으로써 평균 13 %(≈8 GB) 메모리를 절감하고, 중복 인코딩을 없애 학습 스텝당 시간을 144 s에서 82 s로 1.74배 단축한다.

다중 보상 시스템은 PointwiseRewardModel과 GroupwiseRewardModel을 통합 인터페이스로 제공하고, MultiRewardLoader가 동일 모델을 중복 로드하지 않도록 자동 디듀플리케이션한다. 또한, 가중합 방식과 GDPO‑스타일 정규화를 포함한 다양한 어드밴티지 집계 전략을 플러그인 형태로 교체 가능하게 설계했다. 이는 Pref‑GRPO와 같은 순위 기반 보상부터, 텍스트‑렌더링·PickScore와 같은 복합 보상까지 폭넓게 지원한다.

알고리즘 구현 측면에서는 Flow‑GRPO, MixGRPO, GRPO‑Guard, DiffusionNFT, AWM 등 최신 방법을 모두 동일한 BaseTrainer와 SDESchedulerMixin 위에 올려 구현했다. 특히 SDE 스케줄러는 Noise Schedule와 σₜ 파라미터를 표준화해 Table 1에 제시된 네 가지 동적 모델(Flow‑SDE, Dance‑SDE, CPS, ODE)을 설정 파일 하나로 전환 가능하게 만든다. 이는 연구자가 샘플링 동역학이 성능에 미치는 영향을 체계적으로 탐색할 수 있게 한다.

실험에서는 Flux‑1‑dev 모델을 백본으로 삼아 세 알고리즘을 동일한 보상(PickScore) 하에 재현했으며, 원 논문과 거의 동일한 보상 성장 곡선을 얻었다. 또한, 메모리 최적화와 속도 향상이 실제 배치 사이즈 확대와 학습 효율성에 기여함을 정량적으로 입증했다. 정성적 이미지 비교에서도 RL‑파인튜닝된 모델이 원본보다 시각적 품질과 인간 선호도 정렬이 현저히 개선된 것을 확인했다.

전체적으로 Flow‑Factory는 “알고리즘·모델·보상”이라는 세 축을 완전하게 분리함으로써, 새로운 흐름‑매칭 모델에 RL을 적용하는 엔지니어링 장벽을 크게 낮춘다. 레지스트리 기반 플러그인 구조와 사전 인코딩 캐시, 다중 보상 로더는 현재와 미래의 대규모 멀티모달 생성 모델 연구에 필수적인 인프라가 될 것으로 기대된다.