Image AID via continuous-time reinforcement learning
Abstract
We study image inpainting with generative diffusion models. Existing methods typically either train dedicated task-specific models, or adapt a pretrained diffusion model separately for each masked image at deployment. We introduce a middle-ground model, termed Amortized Inpainting with Diffusion (AID), which keeps a pretrained diffusion backbone fixed, trains a small reusable guidance module offline, and then reuses it across masked images without per-instance optimization. We formulate it as a deterministic guidance problem with a supervised terminal objective. To make this problem learnable in high dimensions, we derive an auxiliary Gaussian formulation and prove that solving this randomized problem recovers the optimal deterministic guidance field. This bridge yields a principled continuous-time actor--critic algorithm for learning the guidance module in a fully data-driven manner. Empirically, on AFHQv2 and FFHQ under the pixel EDM pipeline and on ImageNet under the latent EDM2 pipeline, AID consistently improves the quality--speed trade-off over strong fixed-backbone and amortized inpainting baselines across multiple mask types, while adding less than one percent trainable overhead.
Related
- Continuous Delayed-Memory Stochastic Gradient Descent and Continuous-Time Reinforcement Learning from History of Astrophysical Time Series Studies
- Continuous-Time Reinforcement Learning for N-Player Stochastic Differential Games with Exploratory Policies
- Reinforcement learning for constrained stochastic linear--quadratic control in continuous time
- PhiBE-Q-Learning: Bridging Off-Policy Reinforcement Learning and Continuous-Time Control
- Model-free Reinforcement Learning for Continuous Time and State: A Stochastic Maximum Principle Approach
BibTeX
@misc{huang2026image,
title = {{Image AID via continuous-time reinforcement learning}},
author = {Yilie Huang and Xun Yu Zhou},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2605.13010},
eprint = {2605.13010},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.13010},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Image AID via continuous-time reinforcement learning
AU - Yilie Huang
AU - Xun Yu Zhou
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2605.13010
UR - https://arxiv.org/abs/2605.13010
AB - We study image inpainting with generative diffusion models. Existing methods typically either train dedicated task-specific models, or adapt a pretrained diffusion model separately for each masked image at deployment. We introduce a middle-ground model, termed Amortized Inpainting with Diffusion (AID), which keeps a pretrained diffusion backbone fixed, trains a small reusable guidance module offline, and then reuses it across masked images without per-instance optimization. We formulate it as a deterministic guidance problem with a supervised terminal objective. To make this problem learnable in high dimensions, we derive an auxiliary Gaussian formulation and prove that solving this randomized problem recovers the optimal deterministic guidance field. This bridge yields a principled continuous-time actor--critic algorithm for learning the guidance module in a fully data-driven manner. Empirically, on AFHQv2 and FFHQ under the pixel EDM pipeline and on ImageNet under the latent EDM2 pipeline, AID consistently improves the quality--speed trade-off over strong fixed-backbone and amortized inpainting baselines across multiple mask types, while adding less than one percent trainable overhead.
ER -