기술 요약: 모델 예측 제어(MPC)를 통한 흐름 기반 모델의 역문제 해결

1. 문제 정의

역문제(Inverse problems)는 잠재적으로 비선형 연산자인 $A$와 가우시안 노이즈 $\epsilon$이 포함된 측정값 $y = A(x) + \epsilon$으로부터 신호 $x \in \mathbb{R}^d$를 복원하는 과정을 포함한다. 흐름 기반 생성 모델(특히 Flow Matching을 통해 학습된 Continuous Normalising Flows)은 이러한 과업을 위한 강력한 무조건부 사전 확률(unconditional priors)을 제공하지만, 이들의 동역학을 조건부 제약 조건(즉, 데이터 충실도)에 부합하도록 유도하는 것은 여전히 어려운 과제이다.

기존의 훈련 불필ệp(training-free) 가이던스 방법들은 이론적 일관성에 대한 보장이 부족한 휴리스틱 전략에 의존하는 경우가 많으며, 계산적 불안정성을 겪을 수 있다. 반대로, 조건부 생성을 결정론적 최적 제어 문제로 프레이밍하는 것은 원칙적인 접근법을 제공하지만, 과도한 계산 및 메모리 비용을 초래한다. 결과적인 궤적 최적화(trajectory optimization)를 해결하려면 전체 흐름 동역학에 대해 미분하거나 아드joint 방정식(adjoint equations)을 풀어야 하는데, 이는 모델의 크기와 궤적 길이에 따라 확장성이 떨어져 대규모 아키텍처에는 실용적이지 않다.

2. 방법론: MPC-Flow

저자들은 사전 학습된 흐름 기반 생성 모델을 가이드하기 위해 모델 예측 제어(Model Predictive Control, MPC)를 적용하는 프레임워크인 MPC-Flow를 제안한다. 전체 궤적을 한 번에 최적화하는 대신, MPC-Flow는 제어 문제를 일련의 짧은 호라이즌(short-horizon) 하위 문제들로 분해한다.

핵심 메커니즘

각 타임 스텝 $t$에서, 프레임워크는 계획 호라이즌 $[t, t+H]$에 대해 로컬 최적 제어 문제를 해결한다. 제어 $u(\tau)$는 흐름 동역학을 제약 조건으로 하여 제어 에너지와 터미널(또는 중간) 손실로 구성된 비용 함수를 최소화하도록 최적화된다:
$$\frac{dx(\tau)}{d\tau} = v_\theta(x(\tau), \tau) + u(\tau)$$
계산된 제어의 초기 부분(짧은 구간 $\Delta t$에 해당하는 부분)만이 시스템에 적용된다. 이후 상태가 업데이트되고, 호라이즌이 앞으로 이동하며, 이 과정이 반복된다(receding-horizon control).

두 가지 설계 레짐(Design Regimes)

본 논문은 이 프레임워크의 두 가지 구체적인 구현체를 분석한다:

1. 후퇴 호라이즌 제어 (Receding-Horizon Control, MPC-RHC):

   • 호라이즌: $H = 1 - t$ (전체 남은 궤적).
   • 메커니즘: 매 단계마다 컨트롤러는 남은 전체 경로를 다시 계획한다.
   • 이론적 보장: 하위 문제가 최적으로 해결된다면, MPC-R치 RHC는 원래의 궤적 최적화 문제에 의해 정의된 전역 최적 제어를 복구한다.
   • 구현: 실행 가능성을 확보하기 위해, 남은 호라이즌을 $K$개의 스텝으로 이산화한다.
     • 멀티 스텝 ($K > 1$): $K-1$번의 중첩된 흐름 모델 평가에 대한 역전파(backpropagation)가 필요하며, 이는 메모리 집약적이다.
     • 싱글 스텝 ($K = 1$): 호라이즌을 단일 스텝으로 이산화한다. 이는 흐름 경로를 효과적으로 선형화하여, 생성 모델 $v_\theta$에 대한 역전파 없이 제어를 해결할 수 있게 한다. 이 변형은 메모리 요구 사항을 획기적으로 줄여 대규모 모델로의 확장을 가능하게 한다.

2. $\Delta t$-호라이즌 제어 (MPC-$\Delta t$):

   • 호라이즌: $H = \Delta t$ (바로 다음 스텝만 고려).
   • 메커니즘: 즉각적인 다음 타임 스텝을 위해 제어를 최적화하는 탐욕적(greedy) 접근 방식이다.
   • 비용 함수: 중간 비용 함수 $\Phi_{MPC}$를 필요로 한다. 이론적으로 최적의 선택은 가치 함수(value function, cost-to-go)이다. 실제 가치 함수는 다루기 어렵기 때문에, 저자들은 전체 궤적에 대한 역전파를 피하면서 단일 Euler 스텝(확산 모델의 Tweedie 추정치와 유사함)을 사용하는 휴리스틱 근사치를 제안한다.

3. 주요 기여

• 프레임워크 제안: 훈련 불필ệp 가이던스를 통해 대규모 최첨단 아키텍처로 확장 가능한, 최적 제어 이론에 명시적으로 뿌리를 둔 최초의 조건부 생성 방법인 MPC-Flow를 도입하였다.
• 이론적 분석: MPC-Flow를 기저의 최적 제어 목적 함수와 연결하는 공식적인 증명을 제공한다. 구체적으로, Theorem 3.1은 MPC-RHC가 전역 최적 제어를 복구함을 입증하며, Theorem 3.2는 $\Delta t$-호라이즌 제어가 최적 정책을 산출하는 조건을 정의한다.
• 메모리 효율적 변형: 싱글 스텝 RHC 변형($K=1$)이 생성 궤적에 대한 역전파를 제거하여, 엄격한 메모리 제약이 있는 소비자용 하드웨어에서도 실행 가능하게 함을 입증하였다.
• 통합된 스펙트럼: 서로 다른 호라이즌 길이와 제어 파라미터화의 선택이 전역 최적(하지만 비용이 많이 드는) 방식부터 탐욕적(하지만 효율적인) 방식에 이르기까지, 가이던스 알고리즘의 스펙트럼을 형성함을 보여준다.

4. 실험 결과

저자들은 이미지 복원 작업(denoising, deblurring, super-resolution, inpainting)과 대규모 조건부 생성 작업에 대해 MPC-Flow를 평가하였다.

• 이미지 복원 (CelebA):
  • MPC-Flow 변형들은 성능(PSNR/SSIM)과 실행 시간 모두에서 전역 최적 제어 베이스라인(FlowGrad, OC-Flow)보다 일관되게 우수한 성능을 보였다.
  • MPC-$\Delta t$는 선형 작업 전반에서 최첨단 또는 이에 근접한 성능을 달eric했으며, 종종 PnP-Flow에 이어 2위를 기록했으나, 전역 제어 방법들에 비해 계산 비용이 현저히 낮았다.
  • **MPC-RHC ($K=1$)**는 가장 빠른 추론 시간을 제공했으나, 단일 스텝 근사가 누락된 정보를 채우는 데 어려움을 겪는 큰 영공간(large null spaces)을 가진 작업(예: 특정 인페인팅 또는 초해상도 시나리오)에서는 한계를 보였다.
• 대규모 비전 모델 (FLUX.2 및 SD3.0):
  • 본 프레임워크는 소비자용 하드웨어(NVIDIA RTX 3090, 24GB VRAM)의 4비트 양자화 설정에서 FLUX.2 (32B) 모델에 성공적으로 적용되었다.
  • 스타일 전이 (Style Transfer): MPC-Flow는 원래의 생성 궤적을 유지하면서 스타일 특성을 채택함으로써, FlowChef보다 더 현실적인 이미지를 생성하였다. 반면 FlowChef는 상당한 아티팩트를 유발했다.
  • 이미지 컬러라이제이션 (Image Colorization): MPC-Flow는 FlowChef 및 무조건부 베이스라인에 비해 더 높은 PSNR(23.8 dB)과 SSIM(0.891)을 달성하여, 학습된 매니폴드에서 벗어나지 않으면서 관측값과 일치하는 이미지를 향해 흐름을 유도하는 능력을 입증하였다.
  • 초해상도 (Super-Resolution): 기존 베이스라인들과 비교하여 개선된 지각-왜곡 트레이드오프(perception-distortion trade-offs)를 보여주었다.

5. 의의 및 주장

본 논문은 MPC-Flow가 원칙적인 최적 제어와 대규모 생성 모델을 위한 실용적인 추론 시점 가이던스 사이의 간극을 메운다고 주장한다. 글로벌 최적화를 관리 가능한 하위 문제로 분해함으로써 다음과 같은 성과를 가능하게 한다:

1. 확장성: 아드joint 방법이나 전체 궤적 역전파의 메모리 비용으로 인해 이전에는 어려웠던, 소비자용 하드웨어에서 거대 모델(예: 32B 파라미터)을 가이드하는 것을 가능하게 한다.
2. 훈련 불필ệp 동작: 베이스 모델의 재학습을 요구하지 않고, 사전 학습된 무조건부 사전 확률을 직접 활용한다.
3. 이론적 엄밀성: 휴리스틱 가이던스 전략과 달리, MPC-Flow는 최적 제어 이론에 근거하여 특정 이산화 및 호라이즌 조건 하에서의 최적성에 대한 공식적인 보장을 제공한다.

저자들은 싱글 스텝 RHC($K=1$)가 큰 영공간을 가진 선형 역문제에서 항상 최적의 결과를 내지는 못할 수 있다는 점과, 본 방법이 미분 가능한 목적 함수에 의존한다는 점 등 한계점을 인정하였다. 그러나 이들은 MPC-Flow가 실용적이고 제어 가능하며 확장 가능한 역문제를 위한 생성 모델링의 중요한 진전이라고 위치시킨다.