Due: Friday, Oct. 23.
In Lab 4 you built a restoration that used a QGGMRF prior. Buried inside it, without being named, was a denoiser. In this lab you pull that denoiser out, hand it to a different algorithm, and then swap it for a denoiser that was never derived from a prior at all. The restoration keeps working.
That swap is the plug-and-play idea, and it is worth being careful about. Up to now every algorithm in this course minimized a cost function, and the answer was defined before the algorithm existed. After the swap there is no cost function. The algorithm still converges and the pictures still improve, but what it converges to is a new question. Module 4 asks you to answer it for your own code.
The background is Chapter 9.5 on variable splitting and ADMM, and Chapter 10 on plug-and-play. The notation here matches the book.
Plug-and-play was introduced in S. V. Venkatakrishnan, C. A. Bouman, and B. Wohlberg, “Plug-and-Play Priors for Model Based Reconstruction,” IEEE GlobalSIP, 2013. The consensus-equilibrium view used in Step 3 is G. T. Buzzard, S. H. Chan, S. Sreehari, and C. A. Bouman, “Plug-and-Play Unplugged: Optimization-Free Reconstruction using Consensus Equilibrium,” SIAM J. Imaging Sci., 2018. A recent overview is U. S. Kamilov, C. A. Bouman, G. T. Buzzard, and B. Wohlberg, “Plug-and-Play Methods for Integrating Physical and Learned Models in Computational Imaging,” IEEE Signal Processing Magazine, 2023.
The same image and blur kernel as Lab 3, so that your numbers can be compared with the ones you already have.
Source: image 23 of the Kodak Lossless True Color Image Suite, released for unrestricted use.
Source: A. Levin, Y. Weiss, F. Durand, and W. T. Freeman, "Understanding and Evaluating Blind Deconvolution Algorithms," IEEE CVPR 2009.
wget https://cabouman.github.io/grad_labs/labs/resources/kodim23.pgm wget https://cabouman.github.io/grad_labs/labs/resources/levin09_kernels.zip unzip levin09_kernels.zip
lab5.py, holding functions and nothing else. Run
your experiments from a separate script.forward, adjoint, and
load_pgm from your lab3.py. Everything
else on this page is new.float64. No for loop over
pixels.deepinv package for the pretrained
denoiser. Install it once with pip install deepinv.
Every reconstruction in this lab uses the same image and blur.
The image is kodim23.pgm on \([0,1]\), and \(A\) is
convolution with Levin kernel 1 and circular boundaries. The QGGMRF
potential always uses \( p = 2.0 \), \( q = 1.2 \), and \( T = 1.0 \).
The noise level \( \sigma_w = 0.02 \) matches the deblurring problem in
Lab 3 and Lab 4, so the three labs pose the same problem.
The table lists the settings for each hand-in item. Every experiment below points back to this table by its deliverable number, so you do not have to search the text for a number.
| Item | Experiment | Noise \( \sigma_w \) | \( \sigma_p \) | \( \sigma_n \) | \( \sigma_x \) | Iterations |
|---|---|---|---|---|---|---|
| D1 | DFT operator check (Task 7a) | none | — | — | — | — |
| D1 | \(F\) on its own (Task 7b) | 0.02 | 0.01 and 1.0 | — | — | — |
| D3 | PnP-ADMM, QGGMRF denoiser: match the Lab 4 MAP (Task 10a) | 0.02 | 0.05 | 0.05 | 0.020 † | 100 outer, 20 inner ‡ |
| D4 | PnP-ADMM, QGGMRF denoiser: \( \sigma_p \) sweep (Task 10b) | 0.02 | 0.025, 0.05, 0.10 | 0.05 | 0.020 † | 100 outer, 20 inner ‡ |
| D5 | PnP-ADMM, DRUNet denoiser: deblurring (Task 12a) | 0.02 | 0.05 | 0.03 | — | 50 |
| D6 | PnP-ADMM, DRUNet denoiser: \( \sigma_p \) sweep (Task 12b) | 0.02 | 0.025, 0.05, 0.10 | 0.03 | — | 50 |
| D7 | PnP-PG, DRUNet denoiser: unstable run (Task 16a), optional | 0.02 | 0.05 § | 0.03 | — | 50 |
| D8 | PnP-PG, DRUNet denoiser: stable rescaled run (Task 16a), optional | 0.02 | 0.025 § | 0.015 | — | 50 |
† \( \sigma_x = 0.020 \) is the QGGMRF value from Lab 4 (its run
R2). Feed the same \( \sigma_x \) to the PnP denoiser and to the direct
Lab 4 MAP, so the two minimize one cost.
‡ The QGGMRF proximal map is solved by 20 inner steepest-descent
iterations, warm-started from the previous outer iteration. The warm start
keeps the map accurate with only 20 inner iterations, so 100 outer PnP
iterations reproduce the Lab 4 MAP.
§ Proximal-gradient PnP (PnP-PG) drops the dual variable of PnP-ADMM and
replaces the forward proximal map \( F \) with one gradient step on the data
term, of size \( \sigma_p^2 \) (Step 13). D7 uses the D5 balance
\( \sigma_p = 0.05,\ \sigma_n = 0.03 \), which is unstable because
\( \sigma_p \) is above \( \sqrt 2\,\sigma_w \). D8 halves both to
\( \sigma_p = 0.025,\ \sigma_n = 0.015 \): the same balance
\( \sigma_n^2/\sigma_p^2 = 0.36 \), now stable (Step 15).
The MAP estimate you computed in Labs 3 and 4 minimizes a sum of two terms,
$$ \hat{x} = \arg\min_x \left\{ f(x) + h(x) \right\} \ , \qquad f(x) = \frac{1}{2\sigma_w^2} \| y - A x \|^2 \ , \tag{1} $$where \(f\) is the forward model term and \(h\) is the prior term. They come from different places. The forward model came from physics, from what the camera did. The prior came from a belief about what images look like.
In Labs 3 and 4 you minimized the sum with one optimizer, which meant the two models were welded together. Changing the prior meant rederiving the gradient of the whole cost. Changing the blur meant the same. This lab takes them apart, so that each model is handled by its own piece of code and the two pieces only exchange images.
The tool that separates them is the proximal map. For each of the two terms, define
$$ \begin{aligned} F(v) &= \arg\min_x \left\{ f(x) + \frac{1}{2\sigma_p^2} \| x - v \|^2 \right\} \ , \\[4pt] H(v) &= \arg\min_x \left\{ h(x) + \frac{1}{2\sigma_n^2} \| x - v \|^2 \right\} \ . \end{aligned} \tag{2} $$Read \(F(v)\) as a question: starting from the image \(v\), where should I move to reduce \(f\), without going far? The parameter \(\sigma_p\) sets how far. Small \(\sigma_p\) stays close to \(v\); large \(\sigma_p\) goes most of the way to the minimum of \(f\). \(H(v)\) asks the same of the prior term \(h\), with its own parameter \(\sigma_n\). In the plain ADMM of Chapter 9 these two are one shared penalty parameter; this lab keeps them separate, \(\sigma_p\) for the forward model and \(\sigma_n\) for the denoiser, so you can tune each on its own.
This is a change in kind, and it is the reason the chapter exists. A cost function is an objective: it scores images and says which are better. A proximal map is an action: give it an image and it hands you back another image. The proximal map converts one into the other.
Put the two proximal maps in a loop. Applying ADMM to the split cost (1) gives three lines, and this loop is the whole algorithm:
PnP-ADMM
Initialize u = 0 and v = y
Until converged:
x = F(v - u)
v = H(x + u)
u = u + (x - v)
Plug-and-play was introduced in S. V. Venkatakrishnan, C. A. Bouman, and B. Wohlberg, “Plug-and-Play Priors for Model Based Reconstruction,” IEEE GlobalSIP, 2013.
Read it as a negotiation between two parties who never speak the same language. \(F\) knows the camera and nothing about images; \(H\) knows images and nothing about the camera. The vector \(u\) is the running record of how far apart they are, and it is what forces them to agree.
The equilibrium. Ask what the loop settles on. At a fixed point the last line, \( u \leftarrow u + (x - v) \), leaves \(u\) unchanged, so \( x = v \). Write \(x^*\) for that common value and \(u^*\) for the final \(u\). The first two lines then read
$$ x^* = F( x^* - u^* ) \ , \qquad x^* = H( x^* + u^* ) \ , \tag{3} $$or, equivalently, \( F(x^* - u^*) = H(x^* + u^*) \): the forward-model agent and the denoiser agent return the same image when each is offset from \(x^*\) by \(\mp u^*\). These are the consensus equilibrium (CE) conditions of Chapter 10. They, not a cost function, are what defines the solution \(x^*\). The next two steps ask what this solution is (Step 4), and how the loop reaches it (Step 5).
When both maps are proximal maps, the equilibrium is the MAP estimate. Suppose \(F\) and \(H\) are the proximal maps of \(f\) and \(h\), and give them the same parameter, \( \sigma_p = \sigma_n = \sigma \). The first condition, \( x^* = F(x^* - u^*) \), says \(x^*\) is the image that minimizes \( f(x) + \frac{1}{2\sigma^2}\| x - (x^* - u^*) \|^2 \). Setting the gradient of that to zero at \(x^*\),
$$ \nabla f(x^*) + \frac{1}{\sigma^2}\bigl( x^* - (x^* - u^*) \bigr) = \nabla f(x^*) + \frac{1}{\sigma^2} u^* = 0 \ . $$The second condition, \( x^* = H(x^* + u^*) \), gives the same way
$$ \nabla h(x^*) + \frac{1}{\sigma^2}\bigl( x^* - (x^* + u^*) \bigr) = \nabla h(x^*) - \frac{1}{\sigma^2} u^* = 0 \ . $$Add the two equations. The \(u^*\) terms cancel:
$$ \nabla f(x^*) + \nabla h(x^*) = 0 \ . $$So \(x^*\) is a stationary point of \( f + h \) — the same image that minimizes the MAP cost (1). The equilibrium (3) and the optimization of (1) have the same solution, and the extra vector is just the scaled gradient, \( u^* = -\sigma^2 \nabla f(x^*) \). This is the ADMM of Chapter 9; nothing has been given up.
When \(H\) is a general denoiser, only an equilibrium is left. The point of plug-and-play is to replace \(H\) with an action — a denoiser designed or trained to remove white Gaussian noise — that is not the proximal map of any prior. Moreau's theorem says an operator is the proximal map of a proper closed convex cost only if it has two properties: it is nonexpansive, and it is a gradient, meaning \( H = \nabla \phi \) for some scalar function \(\phi\). A trained denoiser is generally not a gradient — its Jacobian is not symmetric — and its nonexpansiveness cannot be checked. So there is no cost (1) to minimize, and \(x^*\) is defined only by the equilibrium (3). This is the central idea of the lab: an ADMM-shaped loop whose \(H\) is not a proximal map, solving an equilibrium rather than an optimization.
Does the loop reach the equilibrium (3), and is the three-line form the best way to reach it? Reorganize (3) into a fixed-point equation. Change variables to \( w_1 = x - u \) and \( w_2 = x + u \). The two equilibrium conditions become
$$ ( 2F - I )\, w_1^* = w_2^* \ , \qquad ( 2H - I )\, w_2^* = w_1^* \ , $$where the reflected operator \( 2F - I \) sends the anchor to twice \(F\)'s output minus itself, and likewise \( 2H - I \). Substituting \(w_2^*\) into the second equation leaves a fixed point of the single operator \( T = ( 2H - I )( 2F - I ) \),
$$ T w_1^* = w_1^* \ , \qquad T = ( 2H - I )( 2F - I ) \ . \tag{4} $$There are robust algorithms for finding a fixed point of an operator. We use the Mann iteration,
$$ w_1 \leftarrow ( 1 - \rho )\, w_1 + \rho\, T w_1 \ , \qquad \rho \in ( 0, 1 ) \ . $$Once it converges, recover the reconstruction from \( w_2^* = ( 2F - I ) w_1^* \) and \( x^* = ( w_1^* + w_2^* ) / 2 \). In full:
PnP with Mann iteration
Initialize w1
Until converged:
w1_new = (2H - I)(2F - I) w1 # apply F, then H, with reflections
w1 = (1 - rho) * w1 + rho * w1_new
w2 = (2F - I) w1
x = (w1 + w2) / 2
Return x
When \( \rho = \tfrac{1}{2} \) this is exactly the PnP-ADMM loop of Step 3. Larger \(\rho\) takes a longer step toward the fixed point and usually converges faster; \( \rho = 0.8 \) is a common practical choice.
Convergence. The guarantee, from Chapter 10 and the proximal-operator appendix, is this: if \(T\) is nonexpansive and has a fixed point, the Mann iteration converges to a solution of the equilibrium (3). A proximal map is firmly nonexpansive, which makes its reflected operator \( 2H - I \) nonexpansive, so when \(F\) and \(H\) are proximal maps \(T\) is nonexpansive and convergence is assured.
The catch. For a trained denoiser you cannot check that \(T\) is nonexpansive; it would mean bounding the network's response over all image pairs. So the theorem does not apply, and the faster-with-larger-\(\rho\) behavior comes with no guarantee either — one reviewer of the original paper renamed the method “plug-and-pray.” In practice the loop almost always converges, and you confirm it empirically, which you do in Module 4.
\(f\) is quadratic, so the minimization in (2) is a linear system. Setting the gradient of \( \frac{1}{2\sigma_w^2}\|y-Ax\|^2 + \frac{1}{2\sigma_p^2}\|x-v\|^2 \) with respect to \(x\) to zero gives
$$ M x = \frac{1}{\sigma_w^2} A^t y + \frac{1}{\sigma_p^2} v \ , \qquad M = \frac{1}{\sigma_w^2} A^t A + \frac{1}{\sigma_p^2} I \ , \tag{5} $$and \( F(v) \) is the solution \(x\).
Take \(A\) to be any linear operator. From (5), \( F(v) = M^{-1}\!\left( \frac{1}{\sigma_w^2} A^t y + \frac{1}{\sigma_p^2} v \right) \). Rewrite it so the answer is the anchor \(v\) plus a correction. Add and subtract \( \frac{1}{\sigma_w^2} A^t A v \) on the right-hand side of (5):
$$ \frac{1}{\sigma_w^2} A^t y + \frac{1}{\sigma_p^2} v = \frac{1}{\sigma_w^2} A^t ( y - A v ) + \underbrace{\left( \frac{1}{\sigma_w^2} A^t A + \frac{1}{\sigma_p^2} I \right)}_{M} v \ . $$Multiplying through by \( M^{-1} \) gives the closed form for \(F\):
$$ F(v) = v + M^{-1}\, \frac{1}{\sigma_w^2}\, A^t ( y - A v ) \ . \tag{6} $$This is the form to remember. \(F\) keeps the anchor image \(v\) and adds a correction built from the data error \( y - A v \): back-project the error with \( A^t \), then filter it with \( M^{-1} \). In particular, if \( A v = y \) exactly, the data error is zero, the correction vanishes, and \( F(v) = v \): any image that already reproduces the data is a fixed point of the proximal map. More generally, the closer \(v\) is to fitting the data the less \(F\) moves it, and the further away, the larger the correction. A current estimate plus a filtered back-projection of the data residual is the shape of nearly every model-based update in this course. Nothing in (6) used any special structure of \(A\), so it holds for any linear forward model. For a general \(A\), the one hard part is applying \( M^{-1} \), which means solving the linear system (5).
Now let \(A\) be the blur: a linear shift-invariant filter, applied with the circular boundaries you chose in Lab 3. This one structural fact turns \( M^{-1} \) from a matrix inverse into a division, and \(F\) into a few FFTs. The rest of this step builds that, one piece at a time.
Circular convolution becomes multiplication. The convolution theorem says that circular convolution in the image domain is pointwise multiplication in the DFT domain. Let \(\hat{a}\) be the 2-D DFT of the blur kernel, placed in an array the size of the image with its center at index \((0,0)\) (the code below does that placement). Writing a hat for the 2-D DFT and \(\odot\) for the entrywise product, applying the blur is a multiply,
$$ \widehat{A u} = \hat{a} \odot \hat{u} \ . $$The adjoint \(A^t\) is the same filter run in reverse; in the DFT domain that is multiplication by the complex conjugate \(\overline{\hat{a}}\) (the bar), and applying \(A\) and then \(A^t\) multiplies by \( |\hat{a}|^2 \):
$$ \widehat{A^t u} = \overline{\hat{a}} \odot \hat{u} \ , \qquad \widehat{A^t A\, u} = |\hat{a}|^2 \odot \hat{u} \ . $$\(M\) becomes a real array, so \(M^{-1}\) is a division. Recall \( M = \frac{1}{\sigma_w^2} A^t A + \frac{1}{\sigma_p^2} I \). Each term is now a multiply in the DFT domain, so \(M\) itself multiplies by the real, positive array
$$ \hat{m} = \frac{|\hat{a}|^2}{\sigma_w^2} + \frac{1}{\sigma_p^2} \ . $$Since \(M\) is a pointwise multiply by \(\hat{m}\), \(M^{-1}\) is the pointwise divide by \(\hat{m}\) — no matrix inverse and no linear solve.
Assemble the update form (6) in the DFT domain. The update form is \( F(v) = v + M^{-1}\, \frac{1}{\sigma_w^2}\, A^t ( y - A v ) \). Take the DFT of each piece, with \( \hat{V} = \mathrm{DFT}(v) \) and \( \hat{Y} = \mathrm{DFT}(y) \): the residual \( y - A v \) transforms to \( \hat{Y} - \hat{a}\,\hat{V} \); back-projecting with \(A^t\) multiplies by \( \overline{\hat{a}} \); the \( 1/\sigma_w^2 \) is a scalar; and \(M^{-1}\) divides by \(\hat{m}\). Transform back with the inverse DFT, keeping the real part because \(v\) and \(y\) are real:
$$ F(v) = v + \operatorname{Re}\!\left( \mathrm{IDFT}\!\left[ \frac{ \overline{\hat{a}} \, ( \hat{Y} - \hat{a}\,\hat{V} ) / \sigma_w^2 } { \hat{m} } \right] \right) \ . $$No iteration, no step size, nothing to converge. Computing the small correction this way is also more stable than solving (5) as the single ratio \( ( \overline{\hat{a}}\hat{Y}/\sigma_w^2 + \hat{V}/\sigma_p^2 ) / \hat{m} \), which builds the answer from two large terms that nearly cancel and lose precision when \(\sigma_w\) is small.
Both functions are given here. In prox_data, denom
is \(\hat{m}\); Av is the blur \(A v\) (multiply by
\(\hat{a}\), transform back); resid_hat is
\( \hat{Y} - \hat{a}\,\hat{V} \); and corr is the bracket
above, transformed back to the image.
def kernel_dft(a, shape):
"""Return the DFT of the blur kernel, sized to the image.
Embed the odd-sized kernel in a zero array of the image shape and roll
it so the kernel center lands at index (0,0). Then
ifft2(fft2(u) * a_hat).real equals forward(u, a) from Lab 3, and
ifft2(a_hat.conj() * fft2(u)).real equals adjoint(u, a).
a (Ka, Kb) float64 tensor, odd sized
shape (H, W), the image shape
returns (H, W) complex128 tensor
"""
H, W = shape
Ka, Kb = a.shape
ac = torch.zeros(shape, dtype=a.dtype)
ac[:Ka, :Kb] = a
ac = torch.roll(ac, shifts=(-(Ka // 2), -(Kb // 2)), dims=(0, 1))
return torch.fft.fft2(ac)
def prox_data(v, y, a_hat, sigma_w, sigma_p):
"""Return F(v), the proximal map of the forward model term.
Uses the update form F(v) = v + M^{-1} (1/sigma_w^2) A^t (y - A v),
with M = (1/sigma_w^2) A^t A + (1/sigma_p^2) I, entirely in the DFT domain.
v (H, W) float64 tensor, the anchor point
y (H, W) float64 tensor, the measured data
a_hat (H, W) complex128 tensor from kernel_dft
sigma_w float, the noise standard deviation
sigma_p float, the forward-model proximal parameter
returns (H, W) float64 tensor
"""
denom = a_hat.abs() ** 2 / sigma_w ** 2 + 1.0 / sigma_p ** 2 # M in the DFT
Av = torch.fft.ifft2(a_hat * torch.fft.fft2(v)).real # A v
resid_hat = torch.fft.fft2(y - Av) # DFT(y - A v)
corr = torch.fft.ifft2(a_hat.conj() * resid_hat
/ sigma_w ** 2 / denom).real
return v + corr
roll above does. Place it wrong
and the blur shifts, and every image after it is shifted too. Task 7a
catches that in one number, which is easier than staring at pictures.
Task 7a. Check the DFT operator against Lab 3. For a random
image \(u\), compare
ifft2(fft2(u) * a_hat).real with
forward(u, a) from Lab 3. Report the largest absolute
difference. It should be at the level of round-off.
Task 7b. See what \(F\) does on its own. Form the blurred noisy data \( y = Ax + w \) with kernel 1 and \(\sigma_w = 0.02\). Compute \(F(y)\) for \(\sigma_p = 0.01\) and for \(\sigma_p = 1.0\), and display both next to \(y\).
The three images \(y\), \(F(y)\) at \(\sigma_p = 0.01\), and \(F(y)\) at \(\sigma_p = 1.0\), with the number from Task 7a. Add two sentences explaining what \(\sigma_p\) did, and why the large \(\sigma_p\) result looks the way it does even though it fits the data better.
By Step 2, \(H\) is a MAP denoiser for the prior. For the QGGMRF prior of Lab 4 that means
$$ H(v) = \arg\min_x \left\{ \sum_{ \{s,r\} \in \mathcal{P} } b_{s,r} \, \rho( x_s - x_r ) + \frac{1}{2\sigma_n^2} \| x - v \|^2 \right\} \ , \tag{7} $$with the QGGMRF potential from Lab 4, repeated here so you do not have to go looking:
$$ \rho( \Delta ) = \frac{ |\Delta|^p }{ p\,\sigma_x^p } \left( \frac{ \left| \frac{\Delta}{T \sigma_x} \right|^{q-p} } { 1 + \left| \frac{\Delta}{T \sigma_x} \right|^{q-p} } \right) \ . \tag{8} $$This is the Lab 4 problem with the blur removed. Solve it the same way, with majorization on the outside and gradient descent on the inside. If your Lab 4 code is organized into functions you can call, call them. If it is not, this is a good moment to fix that.
def denoise_qggmrf(v, sigma_n, sigma_x, p, q, T, num_iters):
"""Return H(v), the MAP denoiser of equation (7).
v (H, W) float64 tensor, the noisy image (the anchor)
sigma_n float, the noise level the denoiser assumes
sigma_x float, the prior scale
p, q, T floats, the QGGMRF parameters of equation (8)
num_iters int, outer majorization iterations
returns (H, W) float64 tensor
"""
Use \(p = 2.0\), \(q = 1.2\), and \(T = 1.0\) throughout this lab, the values from Lab 4.
Now that you have \(F\) (Step 7) and a denoiser to play \(H\) (Step 8), assemble the loop of Step 3 in code. The loop never names a prior and never sees one: it calls whatever denoiser you hand it, so the same function serves the QGGMRF denoiser here and the deep network of Module 4.
def pnp_admm(y, a_hat, denoiser, sigma_w, sigma_p, num_iters):
"""Run the PnP-ADMM loop of Step 3.
y (H, W) float64 tensor, the measured data
a_hat (H, W) complex128 tensor from kernel_dft
denoiser a function taking one (H, W) tensor and returning one, which
plays the role of H. Its own noise level sigma_n is baked in
(build it as, e.g., lambda z: denoise_qggmrf(z, sigma_n, ...)).
sigma_w float, the noise standard deviation
sigma_p float, the forward-model proximal parameter (used by F)
num_iters int
returns (x, history), the final image and a list of the values of
||x - v|| / ||x|| at each iteration
"""
pnp_admm takes the
denoiser as an argument, so nothing in the loop depends on which one you
use. That is the entire mechanism of plug-and-play, and both of the next
two modules rely on it.
The specification you gave your AI assistant for
denoise_qggmrf and pnp_admm, quoted, with
any corrections you had to make.
The QGGMRF denoiser of Step 8 is a true proximal map: it minimizes the prior plus a squared-distance penalty, equation (7). Step 4 showed that when both agents in the loop are proximal maps, PnP-ADMM is ordinary ADMM, and its equilibrium is the minimizer of the single MAP cost (1) — the Lab 4 reconstruction. So with the QGGMRF denoiser the loop should return your Lab 4 MAP result. Reproducing it is how you confirm the whole PnP machinery is correct, on a problem whose answer you already have.
Making the two agree takes three settings, each a direct consequence of “both agents are proximal maps of the one cost (1)”:
Because the warm start carries an image from one iteration to the next,
this case gets its own small routine rather than reusing the generic
pnp_admm.
def pnp_qggmrf_map(y, a_hat, sigma_w, sigma_p, sigma_n, sigma_x, p, q, T,
num_iters, inner_iters):
"""Run PnP-ADMM whose denoiser is the QGGMRF proximal map.
F uses sigma_p; the QGGMRF proximal map H uses sigma_n, and its inner solve
is warm-started from the previous iteration's output. When sigma_n = sigma_p
the two proximal maps share one penalty and the equilibrium is the Lab 4 MAP
estimate. When they differ, the equilibrium minimizes
f(x) + (sigma_n^2 / sigma_p^2) h(x), a reweighted balance.
y (H, W) float64 tensor, the measured data
a_hat (H, W) complex128 tensor from kernel_dft
sigma_w float, the noise standard deviation
sigma_p float, the forward-model proximal parameter (F)
sigma_n float, the denoiser proximal parameter (H)
sigma_x float, the QGGMRF prior scale (the Lab 4 value)
p, q, T floats, the QGGMRF parameters
num_iters int, outer PnP iterations
inner_iters int, steepest-descent iterations inside the QGGMRF prox
returns (x, history), the final image and the list of
||x - v|| / ||x|| at each iteration
"""
Task 10a. Reproduce the MAP. Form \(y = Ax + w\) with kernel 1 and
\(\sigma_w = 0.02\). Run pnp_qggmrf_map with
\(\sigma_p = \sigma_n = 0.05\), the \(\sigma_x = 0.020\) you used for your
Lab 4 MAP, 100 outer PnP iterations, and 20 inner steepest-descent
iterations in the warm-started QGGMRF proximal map. On the same data,
compute your direct Lab 4 MAP reconstruction. Compare the two.
The QGGMRF-PnP image and the direct Lab 4 MAP image, over \([0,1]\), each with a 128 by 128 crop (state the coordinates); the RMSE of each against the clean image; and the RMS difference between the two. One sentence: did the PnP loop reproduce the MAP?
Task 10b. \(\sigma_p\) as a regularization knob. Repeat Task 10a at \(\sigma_p \in \{0.025,\ 0.05,\ 0.10\}\), but now hold the denoiser fixed at \(\sigma_n = 0.05\) instead of tying it to \(\sigma_p\). With \(\sigma_n \ne \sigma_p\) the equilibrium minimizes \( f(x) + (\sigma_n^2/\sigma_p^2)\,h(x) \), so the prior weight is \( \sigma_n^2/\sigma_p^2 \): a small \(\sigma_p\) weights the prior more and smooths harder, a large \(\sigma_p\) weights it less and stays sharper, and \(\sigma_p = 0.05\) recovers the Lab 4 MAP of Task 10a. Report the RMSE of each and show the three crops.
The three \(\sigma_p\)-sweep crops with their RMSE, and one sentence on how \(\sigma_p\) sets the regularization: which direction is smoother and which is sharper, and why, in terms of the prior weight \( \sigma_n^2/\sigma_p^2 \).
A deep denoiser is not derived from a prior. It was trained on many pairs of clean and noisy images to remove noise, and it does that well. But nobody wrote down a \(p(x)\), took its logarithm, or minimized anything. It is a function fit to data, and you plug it in as \(H\) exactly as you did the QGGMRF denoiser.
Use DRUNet, a pretrained denoising network from the deepinv
library. It takes an image and a noise level \(\sigma_n\), the same
interface as your other denoiser. Install it once with
pip install deepinv; the pretrained grayscale weights
download on first use (about 125 MB). It runs on the CPU, a few seconds
per call on this image.
import deepinv as dinv
# Load the pretrained grayscale DRUNet once. The labs run in float64, so
# force the model to float32 with .float(); otherwise its weights load as
# float64 and the noise-level channel will not match the input.
_drunet = dinv.models.DRUNet(in_channels=1, out_channels=1,
pretrained="download").float().eval()
def denoise_dnn(v, sigma_n):
"""Return a DRUNet denoising of v, as an agent for pnp_admm.
v (H, W) float64 tensor in [0, 1], the noisy image
sigma_n float, the noise level the denoiser assumes, on the [0, 1] scale
returns (H, W) float64 tensor
"""
with torch.no_grad():
out = _drunet(v.to(torch.float32)[None, None], float(sigma_n))
return out[0, 0].to(torch.float64)
Now plug DRUNet into the same pnp_admm loop. It was not
derived from a prior: there is no cost (1) it minimizes and no
\(\sigma_x\). So the reasoning of Module 3 does not apply. The two
parameters \(\sigma_p\) and \(\sigma_n\) are free knobs, and the
equilibrium is whatever the loop settles on. Build the denoiser as a
one-argument closure that fixes \(\sigma_n\), for example
lambda z: denoise_dnn(z, 0.03).
Task 12a. On the data \( y = Ax + w \) with kernel 1 and
\(\sigma_w = 0.02\) (the same data as Module 3), run pnp_admm
for 50 iterations with the DRUNet denoiser at \(\sigma_p = 0.05\) and
\(\sigma_n = 0.03\). Compare to the blurred noisy \(y\) and to the Lab 4
MAP baseline.
The blurred noisy \(y\), the DRUNet-PnP result, and the Lab 4 MAP, over \([0,1]\). Show a 128 by 128 crop (state the coordinates) of the ground truth, \(y\), the DRUNet-PnP result, and the Lab 4 MAP together, so the four can be compared side by side. Give the RMSE of \(y\), DRUNet-PnP, and the MAP against the clean image, and the convergence history \( \|x - v\| / \|x\| \) versus iteration on a log vertical axis. One sentence: how does DRUNet-PnP compare to the MAP?
Reading off the equilibrium. At a fixed point of the loop the last line leaves \(u\) unchanged, so \( x = v \). Write \(x^*\) for that common image and \(u^*\) for the final \(u\). The PnP solution then satisfies the consensus equilibrium conditions (3),
$$ x^* = F(x^* - u^*) \ , \qquad x^* = H(x^* + u^*) \ . $$The quantity \( \|x - v\| / \|x\| \) from your convergence plot is the equilibrium residual: as it falls toward zero the iterates approach a solution of (3). This holds for both denoisers. It holds for the QGGMRF proximal map of Module 3, where the equilibrium is also the Lab 4 MAP, and it holds for DRUNet here, where it is the minimizer of no cost.
Task 12b. The \(\sigma_p\) sweep. Repeat the DRUNet run at \(\sigma_p \in \{0.025,\ 0.05,\ 0.10\}\) — the good value, halved, and doubled — holding \(\sigma_n = 0.03\) and everything else fixed. Report the RMSE of each and show the three crops. Unlike Module 3, the image changes: small \(\sigma_p\) over-regularizes (too smooth), and large \(\sigma_p\) under-regularizes (noise comes back).
The three \(\sigma_p\)-sweep crops with their RMSE, and two sentences on what \(\sigma_p\) controls for DRUNet: which direction over-regularizes and which under-regularizes.
PnP-ADMM leans on the proximal map \(F\), which in this lab is a few FFTs but for a general forward model is a full linear solve. This module builds a second algorithm, proximal-gradient PnP (PnP-PG), that needs only \(A\), \(A^t\), and the denoiser, and never solves a linear system. It is also a clean place to see why a denoiser is a principled stand-in for a prior. This module is optional.
Proximal-gradient PnP was introduced in U. S. Kamilov, H. Mansour, and B. Wohlberg, “A Plug-and-Play Priors Approach for Solving Nonlinear Imaging Inverse Problems,” IEEE Signal Processing Letters, 2017.
Recall the two proximal maps of the split cost (1), one for the data term \(f\) and one for the prior \(h\):
$$ F(v) = \arg\min_x \left\{ f(x) + \tfrac{1}{2\sigma_p^2}\|x - v\|^2 \right\} , \qquad H(v) = \arg\min_x \left\{ h(x) + \tfrac{1}{2\sigma_n^2}\|x - v\|^2 \right\} . $$The proximal-gradient method, also called forward-backward splitting, minimizes \( f + h \) by taking a gradient step on one term and a proximal step on the other. There are two symmetric forms:
$$ x \leftarrow H\!\left( x - \sigma_p^2\,\nabla f(x) \right) \ , \tag{9a} $$ $$ x \leftarrow F\!\left( x - \sigma_n^2\,\nabla h(x) \right) \ . \tag{9b} $$Each form takes the gradient step with the variance of the map it replaces: form (9a) replaces \(F\), so its step is \(\sigma_p^2\); form (9b) replaces \(H\), so its step is \(\sigma_n^2\). This is the step size that turns an exact proximal map into one gradient step. The forward map, for example, obeys \( F(v) = v - \sigma_p^2\,\nabla f(v) + O(\sigma_p^4) \) (set the gradient of its definition to zero and expand for small \(\sigma_p\)), so \( x - \sigma_p^2\,\nabla f(x) \) is the one-step approximation of \(F\).
The fixed point has zero gradient, exactly as in Step 4. Take form (9a). At a fixed point \(x^*\) the anchor is \( v = x^* - \sigma_p^2\,\nabla f(x^*) \), and \( x^* = H(v) \) minimizes \( h(x) + \tfrac{1}{2\sigma_n^2}\|x - v\|^2 \), so its gradient is zero there:
$$ \nabla h(x^*) + \tfrac{1}{\sigma_n^2}\left( x^* - v \right) = \nabla h(x^*) + \tfrac{\sigma_p^2}{\sigma_n^2}\,\nabla f(x^*) = 0 \ . $$Multiplying through by \( \sigma_n^2/\sigma_p^2 \), the fixed point satisfies \( \nabla f(x^*) + \tfrac{\sigma_n^2}{\sigma_p^2}\,\nabla h(x^*) = 0 \): it minimizes the weighted MAP cost $$ f(x) + \frac{\sigma_n^2}{\sigma_p^2}\, h(x) \ . $$ This is the same cost the PnP-ADMM proximal-map pair minimizes (Task 10b), with the same prior weight \( \sigma_n^2/\sigma_p^2 \). Form (9a) is the PnP-ADMM iteration with the exact forward map \(F\) replaced by the one gradient step above and the dual variable dropped, so it targets the same reconstruction without any linear solve.
The two forms ask different things of the prior. Form (9a) needs the proximal map \(H\); form (9b) needs the gradient \(\nabla h\). There is no explicit prior, only a trained denoiser, so we have a formula for neither. One identity, Tweedie's formula, bridges the gap and supplies both: the first box states it and handles form (9b), and the second draws the corollary for form (9a). Both readings are used in practice.
We use form (9a). Both forms reach the same MAP fixed point, but (9a) needs no linear solve: the denoiser plays \(H\), and the gradient \( \nabla f \) is one \(A\) and one \(A^t\). Form (9b) would still need the exact data proximal map \(F\), a linear solve, which is what this module set out to avoid. Form (9a) is also the standard proximal-gradient PnP.
Proximal-gradient PnP is one gradient step on the data term followed by one call to the denoiser:
$$ x \leftarrow \text{denoise}_{\sigma_n}\!\left( x - \sigma_p^2\,\nabla f(x) \right) \ , \tag{11} $$the denoiser at \( \sigma_n \) and the gradient step at \( \sigma_p^2 \) (Step 13). The data term is \( f(x) = \tfrac{1}{2\sigma_w^2}\|y - A x\|^2 \), so $$ \nabla f(x) = \frac{1}{\sigma_w^2}\, A^t\!\left( A x - y \right) \ , $$ with \(A\) and \(A^t\) the DFT multiplies of Step 2, \( \widehat{A x} = \hat a \odot \hat x \) and \( \widehat{A^t r} = \overline{\hat a} \odot \hat r \).
Unlike PnP-ADMM, this iteration can diverge. PnP-ADMM computes the exact proximal map \(F\), a linear solve that is stable for any \(\sigma_p\). PnP-PG replaces that solve with one explicit gradient step, and an explicit gradient step is stable only when its step size is small enough.
The condition. The denoiser is non-expansive: it never pushes two images further apart, so it cannot cause a blow-up. Any instability is in the gradient step alone. That step is affine, \( x - \sigma_p^2\,\nabla f(x) = \bigl( I - \tfrac{\sigma_p^2}{\sigma_w^2}\,A^t A \bigr) x + \text{const} \), so each iteration multiplies the error by $$ G = I - \frac{\sigma_p^2}{\sigma_w^2}\, A^t A \ . $$ The error stays bounded only if every eigenvalue of \(G\) has magnitude at most one. With \(\lambda\) an eigenvalue of \(A^t A\), the eigenvalues of \(G\) are \( 1 - \tfrac{\sigma_p^2}{\sigma_w^2}\,\lambda \), and requiring \( \bigl| 1 - \tfrac{\sigma_p^2}{\sigma_w^2}\,\lambda \bigr| \le 1 \) for every \(\lambda\) gives $$ \sigma_p^2 \le \frac{2\,\sigma_w^2}{\lambda_{\max}(A^t A)} \ . $$ In the DFT domain \(A^t A\) multiplies frequency \(\omega\) by \(|\hat a(\omega)|^2\), so \( \lambda_{\max}(A^t A) = \max_\omega |\hat a(\omega)|^2 \). The kernel is normalized to sum to one, so that maximum is one and the condition becomes $$ \sigma_p \le \sqrt{2}\,\sigma_w \ . $$
D5's parameters are unstable. With \( \sigma_w = 0.02 \) the ceiling is \( \sqrt 2\,\sigma_w \approx 0.028 \), and D5 uses \( \sigma_p = 0.05 \), which is above it. At \( \sigma_p = 0.05 \) the iteration diverges. This is the price of the explicit gradient step, not a defect of the denoiser.
Rescale to a stable point. The balance between data and prior is set by the ratio \( \sigma_n^2/\sigma_p^2 \) (Step 13), not by \( \sigma_p \) and \( \sigma_n \) separately. Divide both by the same factor and that ratio, and so the fixed point, is unchanged while the step \( \sigma_p^2 \) shrinks. Halving both, \( \sigma_p = 0.025 \) and \( \sigma_n = 0.015 \), keeps the ratio \( \sigma_n^2/\sigma_p^2 = 0.36 \) and now satisfies \( \sigma_p \le \sqrt 2\,\sigma_w \), so the iteration converges to the same balance. (In the idealized case of a true proximal-map denoiser the two share a fixed point exactly; DRUNet is trained, so lowering \(\sigma_n\) shifts its effective prior slightly.)
The denoiser. \( \text{denoise}_{\sigma_n} \) is an image denoiser: given a noisy image and an assumed noise standard deviation \(\sigma_n\), it returns an estimate of the clean image, treating its input as a clean image plus white Gaussian noise of standard deviation \(\sigma_n\). It is the same denoiser you plugged into PnP-ADMM as \(H\), the pretrained DRUNet of Module 4. It is not derived from any prior; equation (11) calls it in place of the proximal map \(H\), which Step 14 justifies through Tweedie's formula: the ideal such denoiser is the posterior mean \( E[X \mid V = v] \) of (10), and it agrees with \(H\) to leading order in \(\sigma_n\).
def pnp_prox_grad(y, a_hat, denoiser, sigma_w, sigma_p, num_iters):
"""Run proximal-gradient PnP: a gradient step on the data term, then denoise.
Each iteration applies equation (11),
grad_f = (1 / sigma_w**2) * A^t (A x - y) # gradient of the data term
x = denoiser( x - sigma_p**2 * grad_f ),
with A and A^t applied in the DFT domain through a_hat. There is no
proximal map for F and no linear solve.
y (H, W) float64 tensor, the measured data
a_hat (H, W) complex128 tensor from kernel_dft
denoiser denoise_{sigma_n}: a function taking one (H, W) tensor and
returning one, the denoiser at its assumed noise std sigma_n; it
plays the role of the proximal map H
sigma_w float, the noise level in the data term
sigma_p float, the forward-model proximal parameter; the gradient step
size is sigma_p**2
num_iters int
returns (x, history), the final image and the list of
||x - x_prev|| / ||x|| at each iteration
"""
Task 16a. Run both settings. On the same data as Module 4,
\( y = A x + w \) with kernel 1 and \( \sigma_w = 0.02 \), run
pnp_prox_grad with the DRUNet denoiser for 50 iterations at two
settings that share the balance \( \sigma_n^2/\sigma_p^2 = 0.36 \):
Optional. The proximal-gradient PnP run at \( \sigma_p = 0.05 \), \( \sigma_n = 0.03 \): its per-iteration change \( \|x - x_{\text{prev}}\| / \|x\| \) on a log axis, showing it grows rather than falls. One sentence naming the stability condition of Step 15 that \( \sigma_p = 0.05 \) violates.
Optional. The proximal-gradient PnP run at \( \sigma_p = 0.025 \), \( \sigma_n = 0.015 \) next to the PnP-ADMM result from D5, both over \([0,1]\) with the same crop; the RMSE of each against the clean image; and the convergence history of each on one log axis — the per-iteration change \( \|x - x_{\text{prev}}\| / \|x\| \) that proximal-gradient PnP returns, and the equilibrium residual \( \|x - v\| / \|x\| \) that PnP-ADMM returns from D5. One sentence comparing the two in image quality and in the number of iterations they take to settle.
Answer each question in a few sentences, using your own images and numbers as evidence.
Your answers to the five questions above.
Submit two files through Brightspace.
lab5.py, as a plain .py file. Do
not paste the code into the report and do not send a zip
file.
Also commit your code to your image-processing-labs
repository under lab5 and push it before you submit.
Concise and clear beats long.
You have now built a reconstruction whose prior can be replaced without touching the algorithm, and you have checked what that replacement cost you. Every method in the rest of the course fits into the same frame: an agent that knows the physics, one or more agents that know what the answer should look like, and a rule that makes them agree.