⚡ electrical engineering

Amortized Guidance for Image Inpainting with Pretrained Diffusion Models

This paper introduces Amortized Inpainting with Diffusion (AID), a method that trains a small, reusable guidance module on a fixed pretrained diffusion backbone to achieve efficient, high-quality image inpainting without per-instance optimization, leveraging a novel Gaussian formulation and actor-critic algorithm to bridge deterministic guidance with learnable randomized objectives.

Original authors: Yilie Huang, Xun Yu Zhou

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yilie Huang, Xun Yu Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Fix-It" Dilemma

Imagine you have a master painter (a Pretrained Diffusion Model) who is incredibly talented at painting beautiful landscapes from scratch. Now, imagine someone hands you a photo of a landscape where a large chunk has been torn out (a Masked Image). Your goal is to fill in that missing chunk so it looks perfect and fits the rest of the picture.

Currently, there are two main ways to do this, and both have a big catch:

The "Custom Apprentice" Approach: You hire a new, specialized apprentice to learn only how to fix torn photos.
- The Catch: It takes a long time and a lot of resources to train this apprentice. Plus, if you want to fix a different type of tear later, you might need a different apprentice.
The "On-the-Fly" Approach: You ask the master painter to fix the photo right then and there. But because the painter wasn't trained specifically for "tearing," they have to stop, think really hard, and run a complex, step-by-step calculation for every single photo to figure out how to fill the hole.
- The Catch: It's slow. If you have 1,000 photos, the painter has to do this heavy thinking 1,000 separate times.

The Solution: AID (The "Smart Guide")

The authors propose a middle ground called AID (Amortized Inpainting with Diffusion).

Think of AID as hiring a Smart Guide who stands next to the Master Painter.

The Master Painter stays exactly the same. We don't retrain them or change their brain. They keep doing what they do best.
The Smart Guide is a tiny, lightweight assistant. We train this guide offline (in advance) on thousands of examples of torn photos.
How it works: When a new torn photo arrives, the Master Painter starts painting. The Smart Guide whispers instructions to the painter, saying, "Hey, look at the visible part of the photo; make sure the new paint matches the style of the existing trees."

Once the guide is trained, it can help fix any torn photo instantly. We don't need to stop and think for each new photo; the guide just does its job.

The Secret Sauce: The "Randomized Rehearsal"

The hardest part of this paper is the math behind how they trained the guide.

Usually, to teach a guide to give perfect instructions, you'd have to solve a massive, impossible math puzzle for every single photo. The authors realized that trying to solve the "perfect" puzzle directly is too hard.

So, they invented a clever trick:

The Analogy: Imagine you want to teach a guide to drive a car perfectly on a specific road. Instead of trying to calculate the exact steering angle for every second (which is hard), you let the guide practice driving the road while wearing blinders that make the steering wheel wobble randomly.
The Magic: The authors proved mathematically that if you train the guide to handle this "wobbly, random" driving, the average of all those wobbly movements will point exactly to the perfect, smooth path you wanted all along.
The Result: They used a method called Continuous-Time Actor-Critic (a type of Reinforcement Learning) to train the guide using this "random wobble" method. Once training is done, they simply take the "average" (the smooth path) and use that as the guide's instructions. This makes the training possible and the final result perfect.

What They Found (The Results)

The authors tested this on famous image datasets (AFHQv2, FFHQ, and ImageNet). Here is what happened:

Speed vs. Quality: AID was much faster than the "On-the-Fly" methods (like RePaint) because it didn't have to do heavy calculations for every new photo. In fact, it was about 10 times faster than the best existing method while producing better quality images.
Tiny Footprint: The "Smart Guide" (the trainable part) is incredibly small. It adds less than 1% to the size of the original Master Painter. It's like adding a sticky note to a library book rather than rewriting the whole book.
Versatility: The same guide trained on one type of tear worked perfectly on different types of tears (center holes, strip holes) without needing to be retrained.

Summary

The paper introduces a way to fix damaged images using a powerful AI model without slowing it down. They do this by training a tiny, reusable "guide" that whispers instructions to the main AI. They proved mathematically that they can train this guide using a "randomized rehearsal" technique, ensuring the final instructions are perfect. The result is a system that is faster, cheaper, and produces higher-quality fixes than previous methods.

Technical Summary: Amortized Guidance for Image Inpainting with Pretrained Diffusion Models (AID)

Problem Statement

Image inpainting is a conditional generation task where the goal is to reconstruct missing regions of an image while maintaining consistency with visible pixels. While pretrained diffusion models have significantly advanced this capability, existing approaches face a dichotomy in efficiency and adaptability:

Task-Specific Training: Methods like LaMa, MAT, and BrushNet train dedicated models or large adaptation modules. While this amortizes inference cost, it requires substantial training time and adds significant model complexity.
Per-Instance Optimization: Methods like RePaint, DPS, and MCG keep the pretrained diffusion backbone fixed but adapt it separately for each masked image at test time. This preserves the pretrained model but incurs repeated optimization costs during deployment, treating every image as a new inference problem.

There is a gap for a "middle-ground" solution that amortizes inference without training a large task-specific model or performing per-instance optimization at test time.

Methodology: Amortized Inpainting with Diffusion (AID)

The authors propose AID, a framework that keeps the pretrained diffusion backbone fixed and trains a small, reusable guidance module offline. This module is then reused across all masked images at test time without further optimization.

1. Control Formulation

The problem is formulated as a deterministic control problem.

Dynamics: The reverse diffusion process is governed by a probability-flow ODE driven by a fixed pretrained score function $\hat{S}$ .
Guidance: A task-dependent control term $u(t, x; \xi)$ is added to the drift, where $\xi = (M, y)$ represents the observable mask and visible pixels. The control does not depend on the hidden clean image $x^\dagger$ .
Objective: The goal is to minimize a terminal loss $\Psi$ (balancing fidelity to visible pixels and the missing region) plus a running cost penalizing the magnitude of the guidance field. This leads to a Hamilton-Jacobi-Bellman (HJB) equation where the optimal control $u^*$ is deterministic.

2. Theoretical Bridge: Deterministic to Stochastic

Solving the high-dimensional deterministic HJB equation directly is intractable. The authors introduce an auxiliary randomized control problem using Gaussian policies to make the problem learnable via Continuous-Time Reinforcement Learning (CT-RL).

Auxiliary Problem: Instead of a deterministic control, the system uses a Gaussian policy $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$ , where the mean $\mu$ is the candidate guidance law.
Optimizer-Preserving Bridge: The paper proves two critical theorems:
1. Sufficiency: The optimal deterministic guidance $u^*$ can be lifted to an optimal Gaussian policy by setting the policy mean to $u^*$ .
2. Necessity: Any optimal Gaussian policy for the auxiliary problem must have a mean that exactly recovers the optimal deterministic guidance $u^*$ .
Significance: This bridge ensures that learning a stochastic policy (which enables gradient-based updates) does not alter the target deterministic guidance used at deployment. The variance of the Gaussian policy is fixed and serves only as a learning device, not an entropy regularizer.

3. Amortized Actor-Critic Algorithm

Based on the theoretical bridge, the authors develop a fully data-driven actor-critic algorithm:

Parametrization: The actor (guidance mean $\hat{\mu}_\phi$ ) and critic (value function $\hat{V}_\theta$ ) are parameterized by neural networks.
Learning: The algorithm uses martingale orthogonality conditions derived from the auxiliary HJB equation. It updates the actor and critic using stochastic approximation based on sampled trajectories from the task distribution.
Deployment: Once trained, the guidance module is deployed deterministically. For a new masked image, the system uses the learned mean guidance $\hat{\mu}_\phi$ within the reverse-time solver, discarding the policy noise.

Key Contributions

Methodology: AID introduces a controlled-guidance framework that amortizes inference by training a lightweight guidance module offline, which is then reused across tasks without per-instance optimization.
Theory: The paper establishes a rigorous "policy-equivalence bridge" between the deterministic guidance problem and an auxiliary Gaussian formulation. It proves that the auxiliary problem is optimizer-preserving, justifying the use of stochastic actor-critic methods to learn a deterministic control field.
Efficiency: The approach adds less than 1% trainable parameters relative to the pretrained backbone while significantly improving the quality-speed trade-off.

Empirical Results

The authors evaluated AID on AFHQv2, FFHQ (using the pixel-space EDM pipeline), and ImageNet (using the latent-space EDM2 pipeline).

Performance vs. Fixed-Backbone Baselines: AID consistently outperforms strong fixed-backbone baselines (Unguided, Replacement, MCG, DPS) across multiple metrics (PSNR, SSIM, LPIPS) and mask types (free-form, center, strip).
Performance vs. RePaint: AID surpasses RePaint (a strong iterative optimization baseline) in quality while using only ~10% of the sampling budget (e.g., 35 steps vs. 350 steps for RePaint).
Performance vs. LatentPaint: On ImageNet, AID outperforms LatentPaint (a reusable latent-space adaptation method) on all metrics while using a guidance module with less than 45% of the trainable parameters.
Latency: AID supports a low-latency operating point (12 steps) that further reduces inference time while maintaining performance superior to non-RePaint baselines and remaining competitive with RePaint.
Generalization: The same trained guidance module transfers effectively to unseen mask types and different sampling budgets without retraining.

Significance and Claims

The paper claims that AID is the first amortized controlled-guidance framework for image inpainting that:

Keeps the pretrained diffusion backbone frozen.
Trains only a small, reusable guidance module.
Provides an exact optimizer-preserving bridge from stochastic actor-critic learning to the deterministic guidance field used in deployment.

The authors position this work as a step toward using control- and RL-based frameworks to provide theoretically grounded methods for studying and improving diffusion models. They emphasize that continuous-time RL underpins both the rigorous analysis and the practical algorithm design, offering a principled alternative to heuristic per-instance optimization or heavy task-specific retraining.