노이즈 주입으로 개선된 1단계 평균 흐름 모델

초록

RMFlow는 1‑NFE(단일 함수 평가) 평균 흐름(MeanFlow) 전송 뒤에 맞춤형 노이즈 주입 단계를 추가함으로써 텍스트‑이미지, 분자 구조, 시계열 등 다중모달 생성 작업에서 품질 저하를 크게 완화한다. 평균 속도장을 신경망으로 근사하고, Wasserstein 거리 최소화와 샘플 로그우도 최대화를 동시에 고려한 새로운 손실 함수를 도입해 학습한다. 실험 결과, 기존 MeanFlow 대비 1‑NFE만 사용해도 거의 최첨단 수준의 성능을 달성하면서 계산 비용은 비슷하게 유지한다.

상세 분석

본 논문은 기존 MeanFlow가 1‑NFE 설정에서 발생하는 품질 저하 문제를 근본적으로 해결하고자 두 단계의 생성 파이프라인을 제안한다. 첫 번째 단계는 기존 MeanFlow와 동일하게 평균 속도장 ˆu₀,₁(x₀;θ)을 한 번 평가해 prior 샘플 x₀를 중간 노이즈 상태 x₁으로 이동한다. 여기서 평균 속도장은 순간 속도 uₜ(xₜ|z)의 시간 평균을 근사한 것으로, 기존 MeanFlow와 동일하게 CMFM 손실 L_CMFM을 최소화해 학습한다. 두 번째 단계는 x₁에 추가적인 가우시안 노이즈를 주입해 최종 목표 분포 p_tgt와의 KL 발산을 감소시키는 역할을 한다. 이 노이즈 주입은 VAE의 ELBO와 유사한 형태의 NLL 손실 L_NLL을 도입함으로써 로그우도 최대화 목표를 명시적으로 포함한다.

핵심 이론적 기여는 두 손실을 결합한 공동 목표식 L_RMFlow = L_CMFM + λ₁L_NLL + λ₂‖ϕ_ω(c)‖²이다. 첫 항은 평균 흐름의 Wasserstein 거리 상한을 보장하고, 두 번째 항은 목표 분포에 대한 정확한 확률 밀도 추정을 촉진한다. 세 번째 항은 멀티모달 조건부 생성 시 텍스트 등 컨텍스트 c를 인코더 ϕ_ω에 매핑해 prior에 반영하도록 하는 정규화 항이다. 논문은 또한 Theorem 4.1을 통해 L_NLL이 기대 로그우도의 하한을 제공하고, 이를 최소화하면 KL 발산이 감소함을 증명한다.

학습 전략에서는 대규모 데이터셋에 대한 메모리 효율성을 위해 초기에는 기존 MeanFlow 손실만으로 사전 학습하고, 이후 파라미터 효율적 파인튜닝(PEFT)을 적용해 L_RMFlow를 미세 조정한다. 특히 분자 생성 실험에서는 물리적 피드백을 정책 그라디언트 형태로 결합해 구조적 유효성을 강화한다.

실험에서는 1‑D 혼합 가우시안, 2‑D 체커보드, 텍스트‑이미지(예: MS‑COCO), 텍스트‑분자(QM9), 시계열 데이터 등 다양한 베치마크에서 RMFlow가 1‑NFE만 사용했음에도 불구하고 8‑NFE 혹은 32‑NFE MeanFlow에 근접하거나 이를 초과하는 FID, IS, 유효 분자 비율 등의 지표를 기록했다. 특히 QM9에서 1‑NFE MeanFlow가 생성한 무효 구조가 다수였던 반면, RMFlow는 거의 모든 샘플이 화학적으로 유효한 그래프를 형성했다. 이는 노이즈 주입 단계가 미세한 확률 질량을 보정해 고차원 복합 구조의 일관성을 회복시킨다는 실증적 증거다.

전체적으로 RMFlow는 “1‑NFE 고품질 생성”이라는 목표를 달성하기 위해 평균 흐름의 효율성과 확률적 정밀성을 동시에 만족시키는 설계·학습 프레임워크를 제시한다. 이는 향후 확산·흐름 기반 모델의 경량화와 실시간 응용에 중요한 방향성을 제공한다.