PaperlayerView exact revision on arXiv

Source-linked AI summary

FIRM: Flow-based Imaging via Regularized Minimization

Shirin Shoushtari, Edward P. Chandler, Xiao Shi, Ulugbek S. Kamilov

arXiv:2609.12953v2cs.CV

TL;DR

Measurement-conditioned flow fields typically use measurements as network inputs, while guidance-based approaches apply corrections during sampling; neither explicitly places the forward operator within each learned velocity evaluation. FIRM derives an operator-aware velocity from a variational posterior-mean characterization, achieving leading reconstruction results with fewer evaluations and allowing a test-time distortion-perception trade-off.

  • Problem

    Measurement-conditioned velocity fields typically use measurements as network inputs rather than explicitly applying the forward operator, while guidance-based methods make consistency corrections during inference.

    Source: §1 · p. 1

  • Method

    FIRM characterizes the measurement-conditioned posterior mean as a variational minimizer and uses explicit data-consistency updates with a learned estimator to define the velocity field.

    Source: §1 · p. 2, §3.1 · p. 3, §3.2 · p. 4

  • Results

    Across five restoration tasks and two datasets, FIRM achieves leading distortion metrics at a fraction of competing flow-based solvers’ inference cost, with up to 50× fewer network evaluations than competitive flow-based methods.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    Changing the number of sampling steps controls the distortion-perception trade-off without retraining.

    Source: §4.3 · p. 7

    The theoretical analysis assumes a known linear forward model with additive Gaussian noise, and experiments focus on two-dimensional image restoration.

    Source: §5 · p. 9

Abstract

from arXiv · show

Flow matching methods for imaging inverse problems typically incorporate measurements through network conditioning or guidance during sampling. Neither approach explicitly applies the forward operator within the learned conditional velocity field. We develop a principled measurement-conditional velocity parameterization that does. For a linear interpolation path, we express the optimal velocity through the posterior mean $E[x_1|x_t, y]$ and show that this mean is the unique minimizer of a variational objective with an explicit data-consistency term. The velocity defined by this minimizer provably transports the source distribution to the measurement-conditioned posterior. This result leads to a forward operator-aware velocity field that is trained end-to-end and requires no separate guidance during sampling. Across five imaging tasks, our method achieves leading reconstruction quality with up to $50\times$ fewer network evaluations than competitive flow-based methods. Varying the number of sampling steps also controls the distortion-perception trade-off without retraining.

1 Introduction

Imaging inverse problems need learned or designed priors because degraded measurements may not uniquely determine an image. FIRM incorporates explicit data consistency into a measurement-conditional velocity field, with strong reconstruction results and fewer network evaluations.

1 Introduction

  • Because imaging measurements may not uniquely determine an image, reconstruction requires prior information.Classical approaches combine data fidelity with hand-designed regularizers, while learning-based methods use learned priors.Source: §1 · p. 1
  • Flow matching and score-based diffusion provide learned priors, but inverse-problem solvers differ in how measurement information enters reconstruction..Source: §1 · p. 1
  • Prior-based solvers often apply repeated measurement-consistency corrections during inference, which can require many velocity evaluations.This allows prior reuse across forward models.Source: §1 · p. 1
  • Measurement-conditioned fields can avoid separate inference guidance, but typically treat measurements as network inputs rather than explicitly applying the forward operator.FIRM instead incorporates data consistency into each velocity evaluation and characterizes the optimal field through a variational minimizer.Source: §1 · pp. 1–2
  • Across five imaging tasks, FIRM achieves leading reconstruction quality with up to 50× fewer network evaluations than competitive flow-based methods.Varying sampling steps also controls the distortion-perception trade-off without retraining.Source: §1 · p. 2

2 Background

Imaging inverse problems combine degraded measurements with prior information to recover images, since the measurement operator may be underdetermined or ill-conditioned. Background methods include data-consistency updates paired with learned priors and flow-based transport models.

2 Background

  • An underdetermined or ill-conditioned measurement operator can make image recovery non-unique or unstable, so reconstruction requires prior information.A variational estimator combines measurement consistency with prior information.Source: §2 · p. 2
  • Plug-and-Play methods replace a hand-designed regularizer's proximal update with a learned denoiser, while HQS separates data fidelity and regularization using an auxiliary variable.In PnP/HQS updates, the x-step enforces data consistency and the learned denoiser supplies prior information.Source: §2 · p. 2
  • Flow matching learns continuous transport dynamics from a source distribution to a target distribution using intermediate probability paths.For a linear interpolation path, the optimal marginal velocity is a conditional expectation; conditional flow matching regresses to sample-wise path velocity.Source: §2 · p. 3
  • In standard flow matching, “conditional” refers to the probability path used for the training target, not conditioning on measurements.The straight sample-wise paths motivate numerical integration with relatively few steps.Source: §2 · p. 3

3 Method

FIRM characterizes a measurement-conditioned posterior mean as the variational solution defining an optimal velocity, then approximates it with an operator-aware learned update. The resulting field is trained end-to-end and sampled without separate guidance.

3.1 Theoretical Characterization

  • Under linear interpolation, the state and measurement are conditionally independent noisy observations of the clean image, and their posterior mean is the unique minimizer of a variational objective.The objective combines interpolation consistency, explicit measurement consistency, and an implicit regularizer; the resulting velocity transports the source to the measurement-conditioned posterior.Source: §3.1 · pp. 3–4

3.2 Operator-Aware Conditional Velocity

  • HQS separates the objective into an explicit measurement-consistency update and a learned update for interpolation consistency and implicit regularization.The K-iteration learned output estimates the measurement-conditioned posterior mean and defines the velocity; the forward operator enters through the data-consistency update at every iteration.Source: §3.2 · p. 4

3.3 Training the Conditional Velocity Field

  • The complete estimator is trained end-to-end with a weighted flow-matching loss using the operator-aware posterior-mean estimate.Training samples a clean image, a Gaussian source state, a time, and a noisy measurement generated with the forward operator.Source: §3.3 · p. 4

3.4 Sampling with the Conditional Velocity Field

  • Sampling integrates the learned conditional velocity with explicit Euler steps and requires no external guidance because the posterior-mean estimator already incorporates the forward operator.Changing the number of steps controls the distortion-perception trade-off without retraining.Source: §3.4 · p. 5

4 Experimental Results

Across five imaging tasks, the experiments evaluate reconstruction quality, inference cost, sampling-step and solver-iteration choices, and robustness to operator and noise mismatch. Results show strong reconstruction quality alongside test-time control of the distortion–perception trade-off.

4 Experimental Results

  • The evaluation covers denoising, deblurring, super-resolution, random-mask inpainting, and box inpainting on CelebA and AFHQ-Cat.Results report PSNR, SSIM, and LPIPS, averaged over 100 test images.Source: §4 · p. 5, §4.1 · p. 6

4.1 Main Results

  • FIRM ranks first or second in most settings on CelebA and AFHQ-Cat; AFHQ-Cat results show best reconstruction quality with best or competitive perceptual quality.Figure 1 shows fine-detail reconstructions on AFHQ-Cat using N=2 sampling steps.Source: §4.1 · p. 6, Figure 1 · p. 6, Table 2 · p. 8

4.2 Inference Cost and Scalability

  • 50× fewer evaluations than Flower5-OT; the best-distortion CelebA deblurring configuration takes 0.24 s versus 13.03 s for Flower5-OT.The N=2 configuration requires 10 NFEs; inference cost is O(NK), with K=5 inner HQS iterations.Source: §4.2 · p. 6, Table 3 · p. 8

4.3 Ablation Studies

  • Increasing N improves LPIPS but lowers PSNR across most tasks, letting one trained model control the distortion–perception trade-off at test time.For AFHQ-Cat deblurring, Figure 3 shows the displayed-image PSNR/LPIPS values and test-set averages.Source: Figure 3 · p. 7, §4.3 · p. 7
  • With N=2 on random inpainting, increasing K improves PSNR, while LPIPS is minimized at K=5; K=5 costs approximately 70% of K=7.Inference cost scales as O(NK), and K=5 is used throughout.Source: §4.3 · p. 8
  • Under operator mismatch, a model trained at mask ratio 0.7 transfers smoothly to ratios 0.5–0.8 using the actual test mask, with learned parameters fixed.The experiment does not quantify the mismatch penalty because no matched model was trained at each ratio; at two tested noise levels, the σ=0.01 model outperforms Flower5-OT on all three metrics.Source: §4.3 · p. 8

4.4 Additional Results

  • Additional evaluations cover larger test sets, error maps, inference cost, operator mismatch, solver iterates, posterior diversity, Fourier measurements, and toy-problem sampling and conditional-mean estimation.Tables 11 and 12 report results on 400 AFHQ-Cat and 1000 CelebA images, including standard deviations.Source: §4.4 · p. 8

5 Conclusion

FIRM incorporates explicit data-consistency updates into each evaluation of its learned measurement-conditional velocity field. The paper reports a variational characterization and posterior-transport guarantee, alongside reconstruction results and limitations on the theory and experimental scope.

5 Conclusion

  • FIRM incorporates explicit data-consistency updates into each evaluation of its learned measurement-conditional velocity field.Source: §5 · p. 8
  • The measurement-conditioned posterior mean minimizes a variational objective, and the velocity it defines transports the source distribution to the posterior.Across five restoration tasks and two datasets, FIRM achieves leading distortion metrics at a fraction of competing flow-based solvers’ inference cost; varying sampling steps controls the distortion-perception trade-off without retraining.Source: §5 · p. 9
  • The theory assumes a known linear forward model with additive Gaussian noise, while experiments focus on two-dimensional image restoration.Nonlinear measurements, more general noise, alternative paths or solvers, and other imaging modalities remain directions for further study.Source: §5 · p. 9

A Theoretical Results Proofs

The proof first reduces the joint observations to a Gaussian denoising problem, then uses the Gaussian MMSE characterization to express the measurement-conditioned posterior mean as a variational minimizer. The reduction preserves the information about the clean image, and the resulting objective includes interpolation and measurement consistency.

A.1 Variational characterization

  • The Gaussian MMSE characterization makes the posterior mean the unique minimizer of a quadratic data-fitting term plus an implicit regularizer.The proof roadmap applies this characterization to the equivalent denoising observation and rewrites the quadratic term using the original observations.Source: §A.1 · p. 13, A Theoretical Results Proofs · p. 13
  • The proof reduces the joint observation of the interpolation state and measurement to an equivalent Gaussian denoising observation without losing information about the clean image.The reduced statistic u_t is sufficient for x_1, so conditioning on it gives the same posterior mean as conditioning on (x_t, y).Source: §A.1 · p. 14
  • The resulting objective combines interpolation consistency, explicit measurement consistency, and an implicit regularizer; its posterior mean is the unique global minimizer and stationary point.The regularizer depends on the fixed forward operator and noise level, whose dependence is suppressed in the notation.Source: §A.1 · pp. 14–15

A.2 Conditional probability flow

Under linear interpolation, the optimal measurement-conditional velocity is determined by the posterior mean of the endpoint given the current state and measurement. The proof establishes regularity, a conditional continuity equation, and transport of the source distribution toward the measurement-conditioned posterior.

A.2 Conditional probability flow

  • The optimal measurement-conditional velocity is E[x_1 | x_t = x, y] − x.It follows by averaging the linear-interpolation trajectory velocity over endpoints conditioned on the current state and measurement.Source: §A.2 · p. 15
  • With bounded support for the posterior, the conditional mean is smooth and its covariance-form Jacobian yields a globally Lipschitz velocity on every interval ending before t = 1.The Lipschitz bound is uniform on each [0, τ] for τ < 1.Source: §A.2 · pp. 15–16
  • The conditional density satisfies the continuity equation associated with the measurement-conditional velocity.The proof obtains the equation by averaging the probability fluxes of the linear-interpolation paths over the posterior distribution of the endpoint.Source: §A.2 · pp. 16–17
  • The velocity’s regularity gives a unique ODE flow whose distribution agrees with the conditional path and converges to the measurement-conditioned posterior as t approaches one.The argument establishes the flow on intervals ending before t = 1; it does not require extending the flow map to t = 1.Source: §A.2 · p. 17

A.3 Proof of Proposition 1

The proof establishes that the measurement-conditioned posterior mean is the unique minimizer of the stated objective and that its velocity transports the source distribution to the measurement-conditioned posterior. The algorithm-details section identifies the HQS objective and its alternating subproblems.

A.3 Proof of Proposition 1

  • The proof links the unique posterior-mean minimizer to the conditional probability flow, which converges weakly to the measurement-conditioned posterior.Lemmas 5 and 6 establish the minimizer and velocity; Lemmas 9 and 10 establish the flow result.Source: §A.3 · p. 17

B Algorithm Details

  • The algorithm-details section identifies the HQS objective and its two alternating subproblems as equations (17) and (18).Source: B Algorithm Details · p. 17

B.1 Data-Consistency Update

The appendix describes operator-dependent data-consistency updates, learned regularization, and operator-specific solves used to parameterize the velocity field. It also positions the method against inference-time guidance, amortized conditional models, and operator-aware learned solvers.

B.1 Data-Consistency Update

  • The data-consistency update depends explicitly on the forward operator A, and the considered operators admit closed-form updates.The appendix gives operator-specific forms for inpainting, super-resolution, deblurring, and subsampled Fourier measurements; general operators can use conjugate gradients.Source: §B.1 · p. 18, §B.4 · pp. 19–20

B.2 Learned Regularization Update

  • Because ϕt is unavailable, the exact regularization update is replaced by a learned estimator; damping preserves a contribution from the data-consistent iterate.The estimator receives the interpolation state and data-consistent iterate separately, allowing their relative weighting to depend on t and the HQS parameters.Source: §B.2 · p. 18

B.3 Velocity Parameterization and Cost

  • Each velocity evaluation applies data-consistency solves that depend on A and y within every inner iteration of the learned velocity field.Each evaluation also includes K learned-estimator evaluations.Source: §B.3 · p. 19

B.4 Operator-Specific Data-Consistency Solutions

  • The method adapts data-consistency solves to each operator, including elementwise inpainting updates and Fourier-domain deblurring and subsampling updates.For super-resolution, the implementation uses stride subsampling without an anti-aliasing prefilter.Source: §B.4 · pp. 19–20

C Related Works

  • Inference-time flow solvers reuse priors but approximate measurement-conditional dynamics at inference, while amortized models typically encode measurements through y, A⊤y, or learned representations.FIRM instead applies data consistency inside the generative dynamics; operator-aware unrolled solvers provide a structural precedent but typically do not define posterior samplers.Source: C Related Works · pp. 20–21
  • The Euler step count N empirically controls a distortion–perception trade-off; with a learned velocity, N=1 is posterior-mean-like rather than exactly the posterior mean.Increasing N evaluates the learned estimator at progressively lower interpolation noise levels, where posterior modes become more resolvable.Source: C Related Works · pp. 20–21

D Implementation Details

The implementation shares one estimator across inner iterations and fixes solver parameters across tasks. Training uses a specified optimization schedule, and the paper describes its flow and reconstruction baselines.

D Implementation Details

  • The U-Net receives the data-consistent iterate and flow state, is conditioned on t, and shares weights across K inner iterations.The full model therefore has the parameter count of a single U-Net.Source: D Implementation Details · p. 21
  • The solver uses K=5 inner iterations, damping β=0.5, and coupling ρ=0.01 across operators and noise levels.The data-consistency solve is closed-form for all considered operators, and iterations start from z0=xt.Source: D Implementation Details · p. 21
  • Training uses AdamW and an exponential moving average; the first 20 epochs omit data consistency before the five-iteration solver is enabled.Training lasts 500 epochs on CelebA and 700 on AFHQ-Cat.Source: D Implementation Details · p. 21

D.1 Baselines

  • The flow baselines share a pretrained prior, while the paper also evaluates PnP-HQS and ISTA-Net+ with task-specific implementation choices.The PnP-HQS iteration count is tuned per task; box inpainting is not reported for the Gaussian-denoiser PnP-HQS or compact-convolutional ISTA-Net+ baselines.Source: §D.1 · pp. 21–22

E.1 Additional Ablation

The ablations examine damping and compare operator-aware HQS with measurement conditioning alone. At equal compute, the operator-aware model improves PSNR and SSIM while Flow-only attains lower LPIPS.

E.1 Additional Ablation

  • Among damped settings, quality differences are small and all outperform every baseline in Table 1; β=1.0 is worst on all three metrics.β=1.0 discards the data-consistent iterate from the update, although xk+1 remains an input to Rθ; β=0.5 is used across tasks.Source: §E.1 · p. 22, Table 8 · p. 23

E.2 Ablating the HQS

  • The HQS ablation replaces the solver with A⊤y conditioning in the same network and training setup, reducing cost by removing inner iterations.It evaluates random inpainting and deblurring across ODE step sizes.Source: §E.2 · p. 23
  • At equal compute, N=2 (10 NFEs) improves over Flow-only N=10 (10 NFEs) by 1.5 dB PSNR on random inpainting and 1.3 dB on deblurring, with higher SSIM, while Flow-only attains lower LPIPS.At matched step count, the model improves PSNR and SSIM in all but one setting; Table 9 compares rows matched by step count, with five times the function evaluations for the HQS entries.Source: §E.2 · p. 23, Table 9 · p. 23

E.3 Posterior recovery on a toy problem

A two-dimensional analytical-posterior test probes mode recovery, while image experiments report task-level reconstruction and sampling behavior.

E.3 Posterior recovery on a toy problem

  • At N=50, sampled mode weights are 0.49/0.51 versus the exact 0.50/0.50, with W1(v) ≈0.1 and the observed coordinate within ±0.02 of the measurement.The test uses a four-component Gaussian mixture whose conditional distribution over v is bimodal.Source: §E.3 · p. 23
  • Mode coverage reaches 77% at N=2 and 93% at N=4; for 6 ≤ N ≤ 25, recovered weights and mode locations are within 1% and 0.05 of analytical values.Increasing N moves the final estimator evaluation closer to t=1, where interpolation noise is smaller, and changes the states produced by preceding Euler steps.Source: §E.3 · pp. 23–24

E.4 Compressed-sensing Fourier sampling.

  • Ours (N=2) yields the lowest error on every task; on CelebA, it exceeds Flower5 at ≈50× lower cost.Figure 9 reports MAE / PSNR for five inverse problems; Figure 10 compares PSNR with per-image wall-clock cost.Source: Figure 9 · p. 31, Figure 10 · p. 32
  • With measurements fixed, posterior-sample variation is confined to masked regions and vanishes where measurements are available; N=25 widens the spread.The figure uses K=5 and eight reconstructions per test image.Source: Figure 14 · p. 34
  • The AFHQ-Cat Fourier evaluation uses a Cartesian mask retaining 21.88% of k-space and compares the N=2 reconstruction with Flower1-OT and the zero-filled adjoint.The measurements use complex Gaussian noise with σ = 0.002.Source: §E.4 · pp. 24–25, Figure 17 · p. 35

Related papers

Loading 2609.12953v2…