Learning Efficient and Effective Trajectories for Differential Equation-based Image Restoration

2024-10-07

arXiv.org

AI Summary

Balancing reconstruction quality and computational efficiency remains challenging in differential-equation-based image restoration. To address this, we propose a trajectory-level reinforcement learning (RL) guidance framework coupled with cost-aware variable-step trajectory distillation. We are the first to adapt the 12B-parameter FLUX large model to seven diverse restoration tasks—including denoising, super-resolution, and deblurring—enabling unified multi-task modeling. Our method jointly optimizes solution trajectories via RL, integrates diffusion model fine-tuning, and refines numerical ODE solvers to improve both sampling efficiency and accuracy. Experiments demonstrate up to a 2.1 dB PSNR gain over baselines, substantial visual quality enhancement, and efficient single-model execution across all seven tasks—eliminating the need for task-specific architectures. This work bridges large-scale foundation models and differential-equation-driven restoration, establishing a new paradigm for efficient, generalizable, and high-fidelity image recovery.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

The differential equation-based image restoration approach aims to establish learnable trajectories connecting high-quality images to a tractable distribution, e.g., low-quality images or a Gaussian distribution. In this paper, we reformulate the trajectory optimization of this kind of method, focusing on enhancing both reconstruction quality and efficiency. Initially, we navigate effective restoration paths through a reinforcement learning process, gradually steering potential trajectories toward the most precise options. Additionally, to mitigate the considerable computational burden associated with iterative sampling, we propose cost-aware trajectory distillation to streamline complex paths into several manageable steps with adaptable sizes. Moreover, we fine-tune a foundational diffusion model (FLUX) with 12B parameters by using our algorithms, producing a unified framework for handling 7 kinds of image restoration tasks. Extensive experiments showcase the significant superiority of the proposed method, achieving a maximum PSNR improvement of 2.1 dB over state-of-the-art methods, while also greatly enhancing visual perceptual quality. Project page: url{https://zhu-zhiyu.github.io/FLUX-IR/}.

Problem

• Optimizing trajectories for differential equation-based image restoration
• Reducing computational cost in iterative sampling processes
• Enhancing reconstruction quality and efficiency across multiple tasks

Innovation

• Reinforcement learning optimizes image restoration trajectories
• Cost-aware distillation reduces computational steps adaptively
• Fine-tuned 12B-parameter FLUX model unifies 7 tasks