Plug-and-Play Priors for Bright Field Electron Tomography and Sparse Interpolation
Suhas Sreehari, S. V. Venkatakrishnan, Brendt Wohlberg, Lawrence F. Drummy, Jeffrey P. Simmons, Charles A. Bouman
I Introduction
Transmission electron microscopes are widely used for characterization of material and biological samples at the nano-meter scale . In many cases, these electron microscopy samples contain many repeating structures that are similar or identical to each other. High quality reconstruction of these samples from tomographic projections is possible by exploiting the redundancy caused by repeating structures. As an important example, cryo-electron microscope (EM) tomography involves single particle reconstructions using several views of the same particle . However, in the more general area of 3D transmission electron microscopy (TEM) tomography, no solution currently exists to fully exploit the redundancy in images constituted by many similar or identical particles.
Another important imaging problem is that raster scanning an electron beam across a large field of view is time consuming and can damage the sample. For this reason, there is growing interest in reconstructing full resolution images from sparsely sampled pixels . The redundancy in material and biological samples suggests that it is possible to reconstruct such images with sufficient fidelity by acquiring only a few random samples in the image and using an advanced image reconstruction algorithm that exploits non-local redundancies.
Conventionally, model-based iterative reconstruction (MBIR) solves a single optimization problem that tightly couples the log likelihood term (based on the data) and the log of the prior probability . MBIR can, in principle, exploit redundancy in microscope images for tomographic reconstruction. This requires selection of the appropriate log prior probability, which is very challenging in practice. Patch-based denoising algorithms such as non-local means (NLM) and BM3D have been very successful in exploiting non-local redundancy in images. However, since NLM and BM3D are not explicitly formulated as cost functions, it is unclear how to use them as prior models in the MBIR framework. Venkatakrishnan et al. developed a semi-empirical framework termed plug-and-play priors to incorporate such algorithms into general inverse problems, but limited results were presented and the convergence of the algorithm was not discussed. Chen et al. proposed an MRF-style prior, but with non-local spatial dependencies, to perform Bayesian tomographic reconstruction. The authors adopted a two-step optimization involving non-local weight update, followed by the image update. However, the cost function changes every iteration, so that there is no single fixed cost function that is minimized. Chun et al. proposed non-local regularizers for emission tomography based on alternating direction method of multipliers (ADMM) , using Fair potential as the non-local regularizer, instead of non-local means. This model is restricted to convex potential functions, which in practice is a very strong constraint, and severely limits how expressive the model can be. Yang et al. proposed a unifying energy minimization framework for non-local regularization , resulting in a model that captures the intrinsically non-convex behavior required for modeling distant particles with similar structure. However, it is not clear under what conditions their method converges. Non-local regularizers using PDE-like evolutions and total variation are proposed to solve inverse problems .
Image interpolation is also a widely researched problem . The approaches can be broadly classified into two categories - those based on local regularization and those on non-local regularization. In local approaches, the missing pixels are reconstructed from an immediate neighborhood surrounding the unknown values to encourage similarity between spatially neighboring pixels . Spurred by the success of non-local means, there have been several efforts to solve the sparse interpolation problem using global patch based dictionary models . Li et al. adapted a two stage approach similar to to the problem of sparse image reconstruction using the BM3D denoising algorithm. However, this approach is not immediately applicable to denoising operators such as NLM and those formulated using a nonparametric point estimation framework . The simplicity and success of NLM and BM3D has also led to the question of how these algorithms can be used to solve other inverse problems. In fact, Danielyan et al. have adapted BM3D for image deblurring through the optimization of two cost functions balanced by the generalized Nash equilibrium.
In this paper, we present an algorithm for tomographic reconstruction and sparse image interpolation that exploits the non-local redundancies in microscope images. Our solution uses the plug-and-play (P&P) framework which is based on the alternating direction method of multipliers (ADMM) and decouples the forward model and the prior terms in the optimization procedure. This results in an algorithm that involves repeated application of two steps: an inversion step only dependent on the forward model, and a denoising step only dependent on the image prior model. The P&P takes ADMM one step further by replacing the prior model optimization by a denoising operator. However, while it is convenient to be able to use any denoising operator as a prior model, this new framework also begs the question as to whether P&P necessarily inherits the convergence properties of ADMM? We answer this important question by presenting a theorem that outlines the sufficiency conditions to be satisfied by the denoising operator in order to guarantee convergence of the P&P algorithm. We also present a proof for this convergence theorem partly based on the ideas presented by Moreau and Williamson et al. . Using this result, we then modify NLM to satisfy these sufficiency conditions and call it doubly-stochastic gradient NLM (DSG-NLM). We then apply DSG-NLM as a prior model to the tomographic reconstruction and sparse interpolation problems. This new DSG-NLM algorithm is based on symmetrizing the filter corresponding to the traditional NLM algorithm. Interestingly, Milanfar has also discussed the benefit of symmetrizing the denoising operator, albeit in the context of improving the performance of heuristic denoising algorithms.
The plug-and-play electron tomography solution presented in this paper builds on the existing MBIR framework for bright field electron tomography , which models Bragg scatter and anomaly detection. We demonstrate that our proposed algorithm produces high quality tomographic reconstructions and interpolation on both simulated and real electron microscope images. Additionally our method has improved convergence properties compared to using the standard NLM or the BM3D algorithm as a regularizer for the reconstruction. Due to the generality of the plug-and-play technique, this work results in an MBIR framework that is compatible with any denoising algorithm as a prior model, and thus opens up a huge opportunity to adopt a wide variety of spatial constraints to solve a wide variety of inverse problems.
II Plug-and-play framework
Splitting the variable of equation (1) results in an equivalent expression for the MAP estimate given by
This contained optimization problem can then be computed by solving the following unconstrained augmented Lagrangian cost function given by
where must be chosen to meet the constraint of , and is the augmented Lagrangian parameter The augmented Lagrangian parameter, , is related to the ADMM penalty parameter, , through a simple expression: ..
It is well known that the solution to equation (3) may be computed using the ADMM algorithm. For this particular problem, the ADMM algorithm consists of iteration over the following steps:
In fact, if and are both proper, closed, and convex functions, and a saddle point solution exists , then it is well known that the ADMM converges to the global minimum.
We can express the ADMM iterations more compactly by defining two operators. The first is an inversion operator defined by
and the second is a denoising operator given by
Using these two operators, we can easily derive the plug-and-play algorithm shown in Algorithm 1 as an alternative form of the ADMM iterations. This formulation has a number of practical and theoretical advantages. First, in this form we can now “plug in” denoising operators that are not in the explicit form of the optimization of equation (8). So for example, we will later see that popular and effective denoising operators such as non-local means (NLM) or BM3D , which are not easily represented in an optimization framework can be used in the plug-and-play iterations. Second, this framework allows for decomposition of the problem into separate software systems for the implementation of the inversion operator, , and the denoising operator, . In practice, as software systems for large inversion problems become more complex, the ability to decompose them into separate modules, while retaining the global optimality of the solution, can be extremely valuable.
The plug-and-play algorithm requires the selection of two parameters, and , and then the . The unit-less parameter can typically be chosen to be near 1, with larger or smaller values producing more or less regularization, respectively. In theory, the value of does not affect the reconstruction for a convex optimization problem, but in practice, a well-chosen value of can substantially speed up ADMM convergence ; so the careful choice of is important. Our approach is to choose the value of to be approximately equal to the amount of variation in the reconstruction. Formally stated, we choose
This choice for the value of is motivated by its role as the inverse regularizer in equation (7). In practice, this can be done by first computing an approximate reconstruction using some baseline algorithm, and then computing the sample variance in the approximate reconstruction.
Of course, for an arbitrary denoising algorithm, the question remains of whether the plug-and-play algorithm converges? The following section provides practical conditions for the denoising operator to meet that ensure convergence of the iterations.
III Convergence of the plug-and-play algorithm
In the following theorem, we give conditions on both the log likelihood function, , and the denoising operator, , that are sufficient to guarantee convergence of the plug-and-play algorithm to the global minimum of some implicitly defined MAP cost function. This is interesting because it does not ever require that one know or explicitly specify the function . Instead, is implicitly defined through the choice of .
There exist a in the range of such that ;
is a proper closed convex function which is lower bounded by a function such that is monotone increasing with
is a proximal mapping for some proper closed convex function ;
There exists a MAP estimate, , such that
The plug-and-play algorithm converges in the following sense,
where and denote the result of the iteration.
The first and second conditions of the theorem ensure that the conditions of Moreau’s theorem are met. This is because the doubly stochastic structure of ensures that is the gradient of some function , that is convex, and that is non-expansive.
The additional two conditions of Theorem III.1 ensure that the MAP estimate actually exists for the problem. Importantly, this is done without explicit reference to the prior function . More specifically, the third condition ensures that the set of feasible solutions is not empty, and the fourth condition ensures that the MAP cost function takes on its global minimum value, i.e., that the minimum is not achieved toward infinity.
Importantly, in the next section, we will show that real denoising operators can be modified to meet the conditions of this theorem. In particular, the symmetrized non-local means filters investigated by Milanfar are designed to create a symmetric gradient.
IV Non-Local Means Denoising with Doubly Stochastic Gradient
In order to satisfy the conditions for convergence, the gradient of the denoising operator must be a doubly-stochastic matrix. However, the standard NLM denoising algorithm does not satisfy this condition. So in this section, we introduce a simple modification of the NLM approach, which we refer to as the doubly stochastic gradient NLM (DSG-NLM), that satisfies the required convergence conditions. Interestingly, the symmetrized non-local means filters investigated by Milanfar also achieve a symmetric gradient, but requires the use of a more complex iterative algorithm to symmetrize the operator.
Using this notation, the NLM denoising method can be represented as
where is the denoised result, the coefficients are the NLM weights, and is the NLM search window defined by
Note that the integer controls the size of the NLM search window. In general, larger values of can yield better results but at the cost of higher computational cost.
Using this notation, the plug-and-play denoising operator is given by
Now if we fix the weights, then it is clear that
So the condition 2 of Theorem III.1, simply requires that be a doubly stochastic matrix.
Notice that all three steps of equations (11), (12), and (13) are symmetric in and , so they produce symmetric weights with the property that . While equation (12) results in rows and columns that are approximately normalized, and equation (13) guarantees normalization by adjusting the diagonal coefficient of the matrix . Theoretically, equation (13) could produce a negative coefficient, but in practice this does not occur in real data for two reasons. First, the diagonal coefficient, is always the largest value generated in step 1 of equation (11) because . Second, the normalization of equation (12) typically makes the subtracted quantity of equation (13) small.
Therefore, this algorithm generates a matrix which is symmetric with rows and columns that sum to 1, and in all practical cases, non-negative elements. This makes a doubly stochastic matrix, so it fulfills condition 2 of Theorem III.1 as is required for guaranteed convergence of the plug-and-play algorithm.
V 3D Bright Field EM Forward Model
In this section, we formulate the explicit form of the inversion operator, , for the application of 3D bright field EM tomography. For this problem, we adopted both the forward model and optimization algorithms described in . More specifically, the negative log likelihood function is given by
where is the number of tilts, is the electron counts corresponding to the -th measurement at the -th tilt, , is the blank scan value at the -th tilt, , is the tomographic forward projection matrix associated with the -th tilt, is the -th row of , is a proportionality constant, is a diagonal matrix whose entries are set such that is the variance of , is the offset parameter vector, is a constant, and is the generalized Huber function defined as,
The generalized Huber function is used to reject measurements with large errors. This is useful because measurement may vary from the assumed model for many practical reasons. For example, in bright field EM, Bragg scatter can cause highly attenuated measurements that otherwise would cause visible streaks on the reconstruction .
To compute the inversion operator of equation (7), we minimize the cost function below with respect to , , and .
The details of the optimization algorithm required for equation (16) are described in . The optimization algorithm is based on alternating minimization with respect the the three quantities and it uses a majorization based on a surrogate function to handle the minimization of the generalized Huber function .
For this complex problem, we note some practical deviations from the theory. First, the negative log likelihood function, , is not convex in this case, so the assumptions of the plug-and-play convergence do not hold. Moreover, with such a non-convex optimization, it is not possible to guarantee convergence to a global minimum, but in practice most optimization algorithms generate very good results. In addition, this cost function also violates condition 4 of Theorem III.1 because it only grows at a linear rate as . Again, this condition is used to guarantee that the plug-and-play algorithm does not drift off to a minimum tending to infinity. However, in practice, we have never observed this to happen with real data sets and useful denoising operators. Finally, the global optimization of equation (16) is approximated by three iterations of alternating minimization with respect to , , and . Nonetheless, in our experimental results section, we will illustrate our empirical observation that the plug-and-play algorithm consistently converges even with these approximations to the ideal case.
VI Sparse Interpolation Forward Model
For such a sparse sampling system, we can write the negative log likelihood function as
where is a constant. In order to enforce positivity, we also modify the negative likelihood function by setting for . We include positivity in rather than in the denoising operator so that remains continuously differentiable.
Using equation (7), the interpolation inversion operator is given by
Due to the simple structure of the matrix , we can also calculate an explicit pixel-wise expression for . Moreover, if we let , then reduces to the following form
where represents zeroing of any negative argument. In this case, the interpolation is forced to take on the measured values at the sample points.
VII Results
In this section, we present experimental results on both real and simulated data for the applications of bright-field EM tomography and sparse interpolation. For all experiments, we present convergence plots that compare both primal and dual residual convergence resulting from using different priors. The normalized primal and dual residues [25, p. 18], and respectively, at the -th iteration of the P&P algorithm are given by
where , , and are the values of , , and respectively after the -th iteration of the plug-and-play algorithm, respectively, and is the final value of the reconstruction, .
In this section, we present the results of bright field tomographic reconstruction of (1) a simulated dataset of aluminum spheres of varying radii, (2) a real dataset of aluminum spheres, and (3) a real dataset of silicon dioxide. We compare four reconstruction methods – filtered backprojection, MBIR with qGGMRF prior , plug-and-play reconstructions with 3D NLM and 3D DSG-NLM as prior models. We used qGGMRF, 3D NLM and 3D DSG-NLM as prior models within the plug-and-play framework. Filtered backprojection was used as the initialization for all MBIR-based reconstructions. All the reconstruction results shown below are - slices (i.e., slices parallel to the electron beam). The qGGMRF parameters used for all reconstructions are as follows: , , and . The NLM and DSG-NLM patch size used for all reconstructions is . In order to meet the conditions of convergence, we stopped adapting the DSG-NLM weights at 20 iterations of the plug-and-play algorithm. The P&P parameters used are given in Table II.
In all the experiments, we observe from Tables III and IV that the DSG-NLM ensures that the plug-and-play algorithm converges fully, while NLM achieves convergence to within a fraction of a percent.
The aluminum spheres simulated dataset contains 47 equally-spaced tilts about the -axis, spanning . The attenuation co-efficient of the spheres are assumed to be nm. The noise model is Gaussian, with variance set equal to the mean. The phantom also contains effects that resemble Bragg scatter. The dimensions of the phantom are 256 nm, 512 nm, and 512 nm – along , , and axes, respectively.
Fig. 1 shows a tilt projection of the simulated TEM data. Since this is a bright-field image, the aluminum spheres appear dark against a bright background. Fig. 2 shows the ground truth along with three reconstructions of slice 280 along the - plane. The NLM and DSG-NLM reconstructions have no shadow artifacts, and also have low RMSE values (see Table I). The edges are also sharper in the NLM and DSG-NLM reconstructions.
VII-A2 Aluminum spheres (real) dataset
The aluminum spheres dataset (see Fig. 4) has 67 equally-spaced tilts about the -axis, spanning . Fig. 4 shows a tilt projection of the real aluminum spheres TEM data. Fig. 5 shows three reconstructions along the - plane. The NLM-based reconstruction has less smear artifacts than the qGGMRF reconstruction, and more clarity than the filtered backprojection reconstruction. Also, the NLM and DSG-NLM reconstructions have visibly suppressed missing-wedge artifact.
VII-A3 Silicon dioxide (real) dataset
The silicon dioxide dataset (see Fig. 7) has 31 tilts about the -axis, spanning .
Fig. 7 shows a tilt projection of the real silicon dioxide TEM data. Fig. 8 shows three reconstructions along the - plane. The NLM and DSG-NLM reconstructions have less smear artifacts than the qGGMRF reconstruction, and far more clarity than the filtered backprojection reconstruction.
VII-B Sparse Interpolation
In this section, we present sparse interpolation results on both simulated and real microscope images. We show that a variety of denoising algorithms like NLM, DSG-NLM, and BM3D can be plugged in as prior models to reconstruct images from sparse samples. In all the sparse interpolation experiments, we stopped adapting the weights of the DSG-NLM after 12 iterations of the plug-and-play algorithm. The P&P parameters used are given in Table V.
Our first dataset is a set of simulated super ellipses that mimic the shapes of several material grains like Ni-Cr-Al alloy . The next dataset is a real microscope image of zinc oxide nano-rods . All the images are scaled to the range $$.
In all experiments, the plug-and-play sparse interpolation results are clearer than Shepard interpolation results. We observe from Table VIII that DSG-NLM typically results in the least RMS interpolation error. The RMSE values are normalized as , where is the interpolated image and is the ground truth image. Furthermore, we can see from Tables VI and VII that DSG-NLM makes plug-and-play converge fully.
VIII Conclusions
Microscope images of material and biological samples contain several repeating structures at distant locations. High quality reconstruction of these samples is possible by exploiting non-local repetitive structures. Though model-based iterative reconstruction (MBIR) could in principle exploit these repetitions, practically choosing the appropriate log probability term is very challenging. To solve this problem, we presented the “plug-and-play” (P&P) framework which is based on ADMM. ADMM is a popular method to decouple the log likelihood and the log prior probability terms in the MBIR cost function. Plug-and-play takes ADMM one step further by replacing the optimization step related to the prior model by a denoising operation. This approach has two major advantages: First, it allows the use of a variety of modern denoising operators as implicit prior models; and second, it allows for more modular implementation of software systems for the solution of complex inverse problems.
We next presented and proved theoretical conditions for convergence of the plug-and-play algorithm which depend on the gradient of the denoising operator being a doubly stochastic matrix. We also re-designed the non-local means (NLM) denoising algorithm to have a doubly stochastic gradient, thereby ensuring plug-and-play convergence.
In order to demonstrate the value of our method, we applied the plug-and-play algorithm to two important problems: bright field electron tomography and sparse image interpolation. The results indicate that the plug-and-play algorithm when used with the NLM and DSG-NLM priors were able to reduce artifacts, improve clarity, and reduce RMSE (for the simulated dataset) as compared to the filtered back-projection and qGGMRF reconstructions. Then we performed sparse interpolation on simulated and real microscope images with as little as 5% of the pixels sampled – using three denoising operators: NLM, doubly-stochastic gradient NLM (DSG-NLM), and BM3D. We then compared the results against Shepard’s interpolation as the baseline. In all experiments, DSG-NLM resulted in the least RMSE and also complete convergence of the plug-and-play algorithm, as predicted by theory.
Appendix A Proof of Plug and Play Convergence Theorem
Let and both be proper closed convex functions and let be proper. Then is proper, closed, and convex.
Using these results, we next provide a proof of Theorem III.1.
Proof: Without loss of generality, we will assume and in order to simplify the notation of the proof.
We next show result 2 of the theorem, that a MAP estimate exists. This is equivalent to saying that the function takes on its global minimum value for some .
First define the function . By condition 3 of Theorem III.1 there exists an and such that and . Since, we also know that . Therefore, and is proper. By Lemma A.3, must also be proper, closed, and convex.
Then since is proper, closed, and convex, we know that . Select any . So clearly, is nonempty.
By condition 4 of Theorem III.1, it is always possible to choose so that
In this case, it is easy to show that for all , we have that
So therefore, we know that , , and that is a nonempty bounded and therefore compact set. Consequently, must take on its global minimum value for some value in the compact set .
Finally, we show result 3 of the theorem, that the plug-and-play algorithm convergences. Since the plug-and-play algorithm is just an application of the ADMM algorithm, we can use standard ADMM convergence theorems. We use the standard theorem as stated in [25, p. 16]. This depends on two assumptions. The first assumption is that and must be a proper, closed, and convex functions, which we have already shown. The second assumption is that the standard (un-augmented) Lagrangian must have a saddle point.
The standard Lagrangian for this problem is given by,
and the associated dual function is denoted by
Now we have already proved that a solution to our optimization problem exists and is given by . So we know that the primal problem has a solution given by
So we have that . Furthermore since , we know that for all . So putting together these two results, we have that , thus proving the existence of a saddle point of the un-augmented Lagrangian, .
Adapting the theorem of [25, p. 16], we then have the stated convergence results of equation (3).
Acknowledgment
The authors thank Gregery Buzzard, professor and head of the Mathematics department at Purdue University, for many useful discussions regarding the conditions of convergence of the plug-and-play algorithm. They would also like to thank Marc DeGraef, professor of material science at Carnegie Mellon University, for providing simulated aluminum spheres tomography datasets.