CoFlow가 프롬프트마다 스텝 수를 골라 플로우 모델 추론 비용을 줄인다

무엇인가

플로우 매칭(Flow Matching)은 연속시간 동역학을 적분해 이미지와 영상을 생성하는 방식으로, 확산 모델과 비슷하거나 더 나은 지각적 품질을 더 적은 함수 평가(NFE)로 얻을 수 있다는 점이 장점이다. 문제는 추론이다. 표본 하나를 만들려면 학습된 속도장을 시간축으로 여러 번 순차적으로 평가해야 하고, 모델이 커지고 해상도와 프레임 길이가 늘어날수록 이 직렬 계산이 병목이 된다. 기존 가속 방법은 증류·스텝 절단·일관성 학습처럼 추가 학습 비용이나 교사 모델, 대규모 데이터를 요구하거나, 학습이 필요 없는 휴리스틱 스텝 건너뛰기의 경우 스텝 중요도를 비싸게 추정해야 하고 그 추정 오차가 누적되면 품질이 되돌릴 수 없게 나빠진다. 저자들은 또 하나의 축을 지적한다. "고양이 사진"과 "UFO를 타고 달로 가는 고양이 사진"은 필요한 계산량이 다르다. 즉 고정 스텝 예산은 어떤 입력에는 과하고 어떤 입력에는 부족하다.

어떻게 동작하나

제안 방법 CoFlow는 이 문제를 추론 시점의 온라인 의사결정 문제로 바꾼다. 스텝 예산 선택을 신경망 컨텍스트 밴딧(Neural Contextual Bandit)으로 정식화하고, 프롬프트의 의미 임베딩(SentenceTransformer all-MiniLM-L6-v2)과 개체 수·평균 단어 길이 같은 8차원 수작업 복잡도 특징을 문맥 φ(x)로 쓴다. 행동 α는 [0,1] 구간의 이산화 스텝 수이고, 2층 MLP(ReLU, 은닉 차원 64)가 (문맥, 행동) 쌍의 기대 보상을 예측한다. 스텝 수가 정해지면 그 예산으로 전체 생성 궤적을 끝까지 실행한다. 중요한 점은 생성 모델 자체를 다시 학습하지 않는다는 것이다. 플러그 앤 플레이 방식이며 기존 솔버에도 그대로 얹힌다.

무엇과 다른가

보상 설계가 이 논문의 핵심 아이디어다. 저자들은 복잡한 프롬프트일수록 시간 스텝에 따라 속도장이 더 크게 변한다는 관찰에서 출발해, 궤적을 따라 인접 스텝의 속도 벡터 변화를 L1 상대 변화량으로 평균한 v̂(x,α)를 계산한다. 그리고 r(x,α) = max(0, (1 − v̂(x,α)) − μα)로 보상을 정의한다. 앞 항은 속도 변화가 작은 스텝 예산을 선호하고, μα 항은 추가 계산을 벌점으로 깎는다. μ는 1/max α로 정규화해 쓴다. 이 보상은 외부 품질 평가나 정답 라벨 없이 모델 내부 동역학만으로 계산되는 비지도 신호이며, 저자들 스스로 v̂가 이산화 오차나 생성 품질의 직접 측정치가 아니라 대리 지표(surrogate)라고 명시한다.

어떻게 쓰나

탐색과 학습 절차는 다음과 같다. 초반에는 각 행동을 한 번씩 뽑는 콜드 스타트와 ε-그리디로 충분히 탐색하고, 이후에는 예측 보상에 불확실성 보너스를 더한 NeuralUCB 규칙으로 행동을 고른다. 불확실성 σ는 신경망 출력의 파라미터에 대한 그래디언트 g와 공분산 행렬 V_t = λI + Σ g gᵀ로 계산하며, V_t의 역행렬은 랭크-1 Sherman–Morrison 업데이트로 점진적으로 갱신한다. (문맥, 행동, 보상) 튜플은 리플레이 버퍼에 쌓이고 P회마다 MSE 손실로 Adam(학습률 5×10⁻³, 20 에폭) 업데이트가 돌아간다. 밴딧 모듈은 약 3만 개 파라미터의 경량 신경망이라 생성 모델 평가에 비하면 오버헤드가 작고, 저자들은 이 지연 시간을 보고된 모든 추론 시간에 포함시켰다고 밝힌다.

전제와 한계

이론적 근거도 제시한다. 정규성 조건(속도장이 상태 인자에 대해 L-립시츠, 궤적을 따라 연속 미분 가능) 아래에서 전방 오일러 이산화 오차가 E_i(α) ≤ M_i(e^L − 1)/(2Lα)로 묶인다는 정리 1을 증명한다. 즉 오차가 O(1/α)로 줄어들되 상수 M_i가 입력마다 다르다는 것이고, 이것이 입력별 스텝 예산 선택이 정당화되는 이유다. 밴딧 쪽은 NeuralUCB의 Õ(√T) 의사-리그렛 결과를 조건부로 인용한다.

실험은 텍스트-이미지 생성(GenEval, 553개 프롬프트), 이미지 편집(GEdit, 606개 지시), 텍스트-투-비디오(VBench에서 16개 평가 차원마다 5개씩 뽑은 80개 프롬프트)에서 수행했다. 백본은 BAGEL, Flux-Kontext, PeRFlow, StableDiffusion, HunYuanVideo다. 비교 대상은 전체 스텝 생성, TeaCache, InstaFlow, PeRFlow, FlowCast, FastFlow, AdaDiff다. GenEval에서 BAGEL 기준 CoFlow는 17.7 NFE로 종합 0.77을 기록했고, 전체 추론은 50 NFE로 0.78이었다. FLUX에서는 19.5 NFE로 0.65, 전체 추론은 50 NFE로 0.66이었다. 저자들은 이를 두고 계산량을 크게 줄이면서 종합 품질 변화는 작다고 요약하며, 이미지·영상 전반에서 2.5배 이상 가속을 달성했다고 주장한다. 편집 품질(G_SC, G_PQ, G_O)은 GPT-4.1을 자동 판정자로 써서 측정했고, 영상은 VBench 지표와 프레임 단위 BRISQUE를 함께 보고했다. 절제 실험에서는 문맥을 조건으로 주는 정책이 비문맥 변형보다 품질-효율 절충이 낫고, 불확실성 기반 UCB 탐색이 그리디 정책보다 대리 목적과 계산 효율 모두 개선한다고 밝힌다. FastFlow, FlowCast 같은 기존 추론 가속 기법과 조합도 가능하다고 한다. 모든 실험은 NVIDIA A6000 2장에서 수행됐다.

개발자 관점에서 이 논문의 실용적 가치는 재학습이 필요 없다는 점이다. 이미 플로우 매칭 계열 모델을 서빙하고 있고 스텝 수가 지연 시간의 지배적 요인이라면, 프롬프트 임베딩과 몇 개의 복잡도 특징을 입력으로 받는 작은 정책망을 추론 파이프라인 앞단에 붙이는 형태로 시도할 수 있다. 다만 도입 전에 확인할 것이 있다. 첫째, 보상이 지각적·의미적 품질이 아니라 속도장 변화량이라는 대리 지표라는 점이다. 저자들도 v̂가 이산화 오차나 M_i를 직접 추정한다고 가정하지 않는다고 못 박는다. 따라서 자체 도메인에서는 고정 고스텝 레퍼런스 대비 품질 회귀를 별도로 검증해야 한다. 둘째, 행동 공간(스텝 예산 후보 집합)과 μ 값이 효율-품질 절충을 좌우하므로 워크로드에 맞춰 튜닝해야 한다. 셋째, 온라인 학습이므로 서빙 중에 정책이 갱신되는 구조를 감당할 수 있는지, 리플레이 버퍼와 공분산 역행렬 갱신 비용이 실제 지연 예산 안에 들어오는지 확인해야 한다.

저자들이 밝힌 한계는 두 가지다. CoFlow는 경량이고 완전 온라인이지만 문맥 정책을 학습하기 위한 짧은 탐색 단계가 필요하다. 다만 보고된 모든 결과에 탐색 표본이 이미 포함되어 있다고 강조한다. 또한 효과가 프롬프트 특징의 표현력에 의존하는데, 특징을 의도적으로 단순하고 효율적으로 설계했다. 더 표현력 있는 특징을 쓰면 적응 품질이 올라갈 수 있지만 계산 오버헤드가 늘어난다. 저자들은 실제로는 탐색 단계와 경량 특징 설계 모두 빠르게 안정화되어 전체 성능에 미치는 영향이 작다고 덧붙인다.