PnP-Flow: Plug-and-Play Image Restoration with Flow Matching

2024-10-03

arXiv.org

AI Summary

This work addresses the bottleneck of generative image restoration in imaging inverse problems—including denoising, super-resolution, deblurring, and inpainting—by proposing the first paradigm that embeds Flow Matching (FM) into a Plug-and-Play (PnP) optimization framework. Our method leverages a pre-trained FM model to construct a time-varying denoiser and alternates between data-fidelity gradient descent, flow-path projection, and denoising steps—bypassing ODE backpropagation and trace computation, thereby significantly improving memory efficiency and inference speed. Key contributions include: (i) the first end-to-end co-modeling of FM and PnP without fine-tuning the generative model; and (ii) consistent state-of-the-art performance across all four inverse problems, outperforming existing PnP methods and FM-based SOTA approaches.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

In this paper, we introduce Plug-and-Play (PnP) Flow Matching, an algorithm for solving imaging inverse problems. PnP methods leverage the strength of pre-trained denoisers, often deep neural networks, by integrating them in optimization schemes. While they achieve state-of-the-art performance on various inverse problems in imaging, PnP approaches face inherent limitations on more generative tasks like inpainting. On the other hand, generative models such as Flow Matching pushed the boundary in image sampling yet lack a clear method for efficient use in image restoration. We propose to combine the PnP framework with Flow Matching (FM) by defining a time-dependent denoiser using a pre-trained FM model. Our algorithm alternates between gradient descent steps on the data-fidelity term, reprojections onto the learned FM path, and denoising. Notably, our method is computationally efficient and memory-friendly, as it avoids backpropagation through ODEs and trace computations. We evaluate its performance on denoising, super-resolution, deblurring, and inpainting tasks, demonstrating superior results compared to existing PnP algorithms and Flow Matching based state-of-the-art methods.

Problem

• Combines PnP framework with Flow Matching for image restoration
• Addresses limitations of PnP methods in generative tasks like inpainting
• Improves computational efficiency and memory usage in image restoration

Innovation

• Combines PnP framework with Flow Matching
• Uses time-dependent denoiser from FM model
• Avoids backpropagation through ODEs