흐름 기반 모델로 역문제 해결: 모델 예측 제어 접근법

초록

본 논문은 사전 학습된 연속 정규화 흐름(continuous normalizing flow) 모델을 이용해 이미지 복원·초해상도 등 다양한 역문제에 대한 조건부 생성 방법을 제안한다. 최적 제어 문제를 모델 예측 제어(MPC) 프레임워크로 분할하여 메모리와 연산 부담을 크게 낮추면서도 이론적 최적성 보장을 제공한다. 특히 대규모 사전 학습 모델인 FLUX.2(32 B)를 양자화된 상태에서도 소비자 하드웨어에서 실시간으로 활용할 수 있음을 실험으로 입증한다.

상세 분석

본 연구는 흐름 기반 생성 모델, 특히 Flow Matching으로 학습된 연속 정규화 흐름(CNF)을 역문제 해결에 활용하는 새로운 최적 제어 기반 프레임워크인 MPC‑Flow를 제시한다. 기존 작업들은 사전 학습된 흐름에 데이터 적합도 항이나 손실 함수를 직접 삽입해 조건부 샘플링을 수행했으나, 이들 방법은 이론적 일관성 부족과 메모리·연산 불안정성을 안고 있었다. Liu et al.와 Wang et al.이 제시한 최적 제어 형태는 원칙적으로는 최적성을 보장하지만, 전체 시간 구간에 대한 제어 변수 u(t)를 동시에 최적화해야 하므로 자동 미분을 통한 역전파 혹은 adjoint 방식이 필요하고, 이는 특히 대규모 모델에서는 메모리 사용량이 O(N)·시간 복잡도가 O(N²)로 급증한다는 치명적 한계를 가진다.

MPC‑Flow는 이러한 문제를 “재귀적·단계적” 접근으로 전환한다. 시간 t에서 현재 상태 (\hat{x}_t)를 초기값으로 삼아, 제한된 계획 호라이즌 H를 설정하고, 그 구간 내에서만 최적 제어 문제를 풀어 얻은 제어 신호를 짧은 실행 구간 (\Delta t)에만 적용한다. 이후 상태를 업데이트하고 호라이즌을 앞당겨 다시 최적화를 수행한다는 전형적인 모델 예측 제어(MPC) 절차를 흐름 모델에 그대로 적용한 것이다.

두 가지 주요 설계가 제시된다. 첫 번째는 재귀‑호라이즌 제어(RHC) 로, H를 현재 시점부터 최종 시간 1까지 전체 구간으로 잡고, 이를 K개의 서브스텝으로 이산화한다. 이 경우 전체 최적 제어와 동일한 최적성을 보장한다(정리 3.1). K가 크면 메모리 요구가 증가하지만, K를 작게(예: K=1) 설정하면 흐름 모델에 대한 역전파 없이도 제어를 근사화할 수 있다. 여기서는 제어 에너지 (|u|^2)와 최종 비용 (\Phi(x(1)))만을 고려하므로, 단일 스텝 선형화가 가능해 GPU 메모리 사용량을 크게 절감한다.

두 번째는 Δt‑호라이즌 제어 로, H를 바로 다음 스텝 (\Delta t)로 제한한다. 이 경우에는 단순히 최종 비용 (\Phi)를 사용하면 앞을 내다보지 못하는 비효율적인 제어가 된다. 논문은 중간 비용 (\Phi_{\text{MPC}}(x,t+\Delta t))를 값 함수 (V(t+\Delta t, x)) 로 정의하면, 베일리 원칙에 따라 전역 최적 정책과 동일함을 정리 3.2에서 증명한다. 실질적으로는 현재 상태에서 한 스텝 앞의 예측 상태를 흐름 벡터장으로 이동시킨 뒤, 그 예측 상태에 대한 값 함수를 근사(예: 라그랑주 승수 또는 신경망)하여 제어를 결정한다. 이 방식은 매 스텝마다 단일 변수 최적화만 필요하므로, 대규모 모델에 대한 실시간 적용이 가능하다.

이론적 분석 외에도, 저자는 2‑차원 toy 예시와 OrgancMNIST 기반 CT 복원 실험을 통해 RHC와 Δt‑호라이즌이 각각 전역 최적 해와 근사 최적 해에 얼마나 근접하는지를 정량적으로 검증한다. 이어서 CelebA 데이터셋을 이용한 이미지 인페인팅, 디블러링, 초해상도 등 다양한 선형·비선형 역문제에 대해 기존의 FlowGrad, OC‑Flow, 그리고 최근의 확산 기반 가이드와 비교한다. 특히 FLUX.2(32 B) 모델을 4‑bit 양자화 후에도 소비자 수준 GPU(예: RTX 4090)에서 실시간(≈1 fps)으로 실행할 수 있음을 보여, 메모리 절감 효과가 실제 대규모 모델에 적용 가능함을 입증한다.

요약하면, MPC‑Flow는 (1) 최적 제어 관점에서 조건부 생성의 이론적 근거를 제공하고, (2) 호라이즌·스텝 수 조절을 통해 메모리·연산 비용을 선형 혹은 상수 수준으로 낮추며, (3) 대규모 사전 학습 흐름 모델에 대한 훈련‑프리 가이드를 가능하게 한다는 세 가지 핵심 기여를 한다.