[논문 리뷰] Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

Alexander H. Denker, Moshe Eliasof외 2명|arXiv (Cornell University)|2026. 02. 09.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

MS-Flow는 흐름 궤적을 중간 잠재 상태로 표현하는 다중 샷(multiple-shooting) 접근법을 도입하여, 메모리 효율적이고 안정적인 역문제 해결을 가능하게 하는 flow 기반 모델.

ABSTRACT

Flow-based generative models have emerged as powerful priors for solving inverse problems. One option is to directly optimize the initial latent code (noise), such that the flow output solves the inverse problem. However, this requires backpropagating through the entire generative trajectory, incurring high memory costs and numerical instability. We propose MS-Flow, which represents the trajectory as a sequence of intermediate latent states rather than a single initial code. By enforcing the flow dynamics locally and coupling segments through trajectory-matching penalties, MS-Flow alternates between updating intermediate latent states and enforcing consistency with observed data. This reduces memory consumption while improving reconstruction quality. We demonstrate the effectiveness of MS-Flow over existing methods on image recovery and inverse problems, including inpainting, super-resolution, and computed tomography.

연구 동기 및 목표

• 이미징에서 흐름 기반 생성 선험(priors)을 사용하여 역문제를 해결하도록 동기를 부여한다.
• 플로우 모델에서 단일 샷 잠재 최적화의 높은 메모리 요구와 조건화 문제를 다룬다.
• 다이나믹스, 데이터 일관성, 정규화를 분리하는 다중 샷 프레임워크를 제안한다.
• 로컬 궤적 이음(스티칭) 페널티를 통해 확장 가능하고 병렬화 가능한 최적화를 가능하게 한다.

제안 방법

• 흐름 궤গ를 단일 초기 코드 대신 K개의 중간 잠재 상태(샷 포인트)로 표현한다.
• 데이터 일관성, 방사형 가우시안 잠재 정규화, 그리고 궤적 일관성 페널티를 포함하는 MS-Flow 목적식을 형식화한다.
• 교대 최소화(교대 최소화): (i) 자코비 행렬 없이 경사 근사치를 사용하여 좌표 하강으로 shooting points를 업데이트하고, (ii) 근방/정규화된 데이터 일관성 단계로 말단 이미지 추정치를 업데이트한다.
• 동역학을 명시적 Euler으로 이산화하고, 전체 시간 구간을 역전파하지 않도록 로컬 스티칭 제약을 적용한다.
• 정규성 가정 하에 궤적 업데이트의 수렴 보장을 제공하고 시간 스텝에 독립적인 메모리 복잡도를 보인다.
• shooting 간격에 걸친 배치 평가를 가능하게 하여 확장성과 병렬성을 시연한다.

Figure 1 : Peak GPU memory for D-Flow vs MS-Flow (Ours) using Euler discretization of the ODE on CelebA. D-Flow scales linearly with the number of timesteps, while MS-Flow is constant.

실험 결과

연구 질문

• RQ1 MS-Flow가 메모리 사용량 및 재구성 품질 측면에서 단일 샷 플로우 방법과 비교하여 어떤 차이가 있는가?
• RQ2 다중 샷 포뮬레이션이 플로우 기반 역문제에서 잠재 공간 최적화의 조건화와 수렴을 개선할 수 있는가?
• RQ3 자코비안-프리 트레이젝토리 업데이트와 정확한 그래디언트의 계산 비용 및 메모리 차이는 무엇인가?
• RQ4 MS-Flow는 다양한 이미징 역문제(인페인팅, 초해상도, CT)와 잠재-플로우 모델에서 어떻게 성능을 보이는가?
• RQ5 MS-Flow가 대규모 잠재-플로우 아키텍처(예: Stable Diffusion)에서도 재구성 품질을 보존하며 확장 가능한가?

주요 결과

• MS-Flow는 ODE 이산화에 대해 일정한 메모리를 유지하는 반면, 단일 샷 접근 방식은 시간 스텝에 따라 메모리 크기가 증가한다.
• 좌표 하강 기반의 궤적 업데이트는 궤적 목적함수를 단조롭게 감소시키고 정합성을 가지는 가정 하에서 정상(point)에 수렴한다.
• Jacobian-free 궤적 업데이트는 충분한 감소 조건 하에서 수렴 보장을 유지하면서 상당한 속도 향상을 제공한다.
• MS-Flow는 CelebA 기반 작업(블러 제거, 인페인팅, 초해상도)에서 PSNR/SSIM 면에서 최첨단 플로우 기반 역해결기에 비해 동등하거나 더 우수하게 성능을 달성하고, 하이퍼파라미터에 대해서도 강건함을 보여준다.
• 잠재-플로우 모델(예: Stable Diffusion)에 적용 시, Gaussian 디블러링에서 MS-Flow가 FlowDPS보다 PSNR/SSIM 면에서 더 나은 성능을 보였다.
• 이 접근법은 shooting 간격에 걸친 병렬 처리를 가능하게 하여 최신 하드웨어에서의 확장 가능한 계산을 촉진한다.

(a) Changing the regularization terms.