Teaching deep neural networks to localize single molecules for super-resolution microscopy
Artur Speiser Lucas-Raphael Müller Ulf Matti Christopher J. Obara Wesley R. Legant Jonas Ries Jakob H. Macke Srinivas C. Turaga


Frequency Fluctuations in Tunable and Nonlinear Microwave Cavities
    A. J. Rimberg
    December 30, 2023
==================================================================






 
Single-molecule localization fluorescence microscopy constructs super-resolution images by sequential imaging and computational localization of sparsely activated fluorophores. 
Accurate and efficient fluorophore localization algorithms are key to the success of this computational microscopy method.
We present a novel localization algorithm based on deep learning which significantly improves upon the state of the art. 
Our contributions are a novel network architecture for simultaneous detection and localization, and new loss function which phrases detection and localization as a Bayesian inference problem, and thus allows the network to provide uncertainty-estimates. In contrast to standard methods which independently process imaging frames, our network architecture uses temporal context from multiple sequentially imaged frames to detect and localize molecules. 

We demonstrate the power of our method across a variety of datasets, imaging modalities, signal to noise ratios, and fluorophore densities.  While existing localization algorithms can achieve optimal localization accuracy at low fluorophore densities, they are confounded by high densities. Our method is the first deep-learning based approach which achieves state-of-the-art on the SMLM2016 challenge. It achieves the best scores on 12 out of 12 data-sets when comparing both detection accuracy and precision, and excels at high densities. Finally, we investigate how unsupervised learning can be used to make the network robust against mismatch between simulated and real data.  The lessons learned here are more generally relevant for the training of deep networks to solve challenging Bayesian inverse problems on spatially extended domains in biology and physics.








§ INTRODUCTION
 

Super-resolution microscopy techniques such as stochastic optical reconstruction microscopy (STORM) <cit.> and photo-activated localization microscopy (PALM) <cit.> have made it possible to observe biological structures and processes that where not accessible through optical microscopy due to the 
 Abbe diffraction limit. 
These techniques, commonly referred to as Single Molecule Localization Microscopy (SMLM), critically rely on computational methods for accurately localizing sparsely activated fluorophores <cit.> (Fig. <ref>a).

  
  State-of-the-art localization algorithms typically operate in two steps: first, single fluorophore candidates are detected and extracted from the images, and second, fluorophores are localized by fitting a high resolution “generative” model of the point-spread function (PSF) to the image. To deal with overlapping fluorophores, peaks are either rejected based on a statistical test for the presence of multiple fluorophores (single emitter fitting <cit.>), or emitters are added throughout the fitting procedure until a predetermined threshold for the goodness of fit is met (multi-emitter fitting  <cit.>). 
More recently, deep learning approaches have been used to perform the localization step <cit.>.

This general approach can be highly effective under favourable conditions of high SNR and low fluorophore density <cit.>. However, even multi-emitter approaches produce sub-par results in datasets with high fluorophore densities.  As was noted in a systematic comparison of multiple algorithms on public benchmark data-sets (SMLM2016) <cit.>, they perform even worse than single-emitter algorithms for 3D data. These limitations imply that current fitting-based approaches to SMLM can not be applied to experiments with high emitter densities, which would be critical for the investigation of living or moving structures.

Furthermore, most previous algorithms base their predictions on a single observed image. Thus, they ignore potentially useful information in the sequence of imaging frames which can enable detecting and separating fluorophores in crowded high density data by taking into account their temporal dynamics.  Nevertheless, attempts at using information from multiple images during inference are rare <cit.>, and have not yet yielded state-of-the-art performance.

Deep learning methods have revolutionized computer vision, and biological image analysis is no different <cit.>. 
Many of these advances are the result of the supervised training of deep neural networks using large training datasets of pairs of example input images and desired output predictions.  While the analysis of SMLM data is not a standard supervised learning problem, ground truth localization data for training a deep network can be generated by simulating the imaging of fluorophores. The two first applications of deep learning to SMLM, DeepSTORM3D <cit.> and DeepLoco <cit.> took this approach and used simulated synthetic SMLM to train deep networks to localize single molecules, an approach we call “simulator learning” (SL) <cit.> in this paper. These two deep learning methods differ in the output representation used by the networks. DeepSTORM3D directly predicts a high resolution 3D volume which has the advantage of simplicity but the disadvantage that increasing the resolution of the predictions requires increasing the computation. In contrast, DeepLoco predicts continuous localizations for a fixed number of particles, and has the advantage that its computational complexity scales with the maximum number of possible particles, rather than size of the volume predicted. While both approaches produce detection uncertainty, neither predicts localization uncertainty.

We present a new method for fast, efficient, and accurate single-molecule localization based on a new deep neural network architecture we call DECODE (DEep COntext DEpendent) which achieves state of the art performance. DECODE also uses simulator learning, but is based on three main innovations:  First, we introduce a novel network architecture which uses temporal context for inferring fluorophore locations.  This single DECODE network is trained to produce accurate predictions at both low and high densities, alleviating the need for analysis methods which deal with these `single emitter' and `multi emitter' cases separately. Second, we phrase localization as a Bayesian inference problem, and provide a novel cost-function which makes it possible for the DECODE network to also predict uncertainty-estimates for each localized fluorophore. These uncertainty-estimators can, for example, be used for post-processing algorithms. Third, simulator learning depends on the faithfulness of the generative model, and might show reduced performance when there is a mismatch between the simulated and experimental data. We provide an alternative training approach, Combined Learning (CL) which combines Simulator Learning and Variational Auto Encoder learning (AEL)  <cit.>, and evaluate its performance on simulated data. 


We apply DECODE to data-sets from the public SMLM 2016 challenge <cit.>, and show that it outperforms all existing methods which have been evaluated on this challenge so far, on 12 out of 12 data-sets for which DECODE is applicable <cit.>. Our DECODE method leads to a improvement in performance which is 7× as big as the improvement of the second best algorithm over the third best algorithm. Performance benefits are particular pronounced on high-density data-sets, on where the advantage from using DECODE increases to 10× over the next best method . We also apply DECODE to four datasets where the same sample of labeled Tubulin-A647 protein was imaged with different densities of fluorophore activation, and demonstrate that we achieve high quality reconstructions with 10x less imaging time by accurately localizing fluorophores at high densities. To demonstrate the flexibility of our approach, we adapted it to reconstructing a large 3D volume of an entire COS7 cell with intracellular membranes densely labeled using PAINT, and imaged by lattice light sheet microscopy, imaged over several days. Our method significantly improved the reconstructions, but also enabled high quality reconstructions with only a fraction of the imaging time. Finally, we explore the performance benefits brought about by the use of local context, and the different training approaches.







§ RESULTS




 §.§ DECODE network for simultaneous detection and localization of fluorophores


We designed and trained a deep neural network to simultaneously detect and localize fluorophores in SMLM measurements. The input to the deep network is a sequence of image frames containing sparsely activated fluorophores, and the desired outputs are locations of an unknown number of active fluorophores in each frame. 



  
Deep network design for predicting detection, localization, and uncertainty using spatial and temporal context 
Previous deep learning approaches to SMLM have processed each frame independently, using one of two approaches, which we combine in our work. DeepSTORM3D <cit.> produces predictions on a super-resolved 3D voxel grid. For each super-resolution voxel, a detection probability is predicted for the presence of a fluorophore. DeepLoco <cit.> combines classification and regression by predicting a fixed sized 256× 4 matrix for each imaging frame, with each row representing the presence or absence of a fluorophore, followed by the 3D vector containing the x, y and z coordinate of the molecule. This has the advantage that the computational complexity of the output scales only with the maximum number of possible active fluorophores in any frame, but not the volume of imaged field of view. However, it requires the network to learn a highly non-local and non-linear transformation from images into 3D coordinates in an undetermined ordering.

The DECODE network architecture is a hybrid of these two approaches: for each image frame it predicts eight channels for each imaged pixel (Fig. <ref>c). The first two channels indicate the detection probability of a fluorophore near that pixel in an imaging frame p and its brightness α. The next three channels describe the continuous valued localization of the fluorophore with respect to the center of the pixel, Δ x, Δ y, Δ z. This hybrid approach allows DECODE to scale only with the number of imaged pixels (not super-resolution-pixels), and avoids a highly nonlinear and non-local mapping of pixels to coordinates.

DECODE is the first approach to provide fully probabilistic prediction of both fluorophore detection and localization. In addition to the first 5 output channels, three further  channels estimate the uncertainty of the localization along each coordinate given by σ_x, σ_y, σ_z. A final channel represents the uncertainty in the DECODE prediction of the brightness σ_α. Thus DECODE directly predicts and independently represents uncertainty about detection, localization, and particle brightness.














  
Using temporal context
We introduce a new mechanism to integrate information across frames, and show that it leads to improved detection and localization. The temporal dynamics of the fluorophores are such that a fluorophore can be active across multiple adjacent frames, inducing correlations which are local in time. We designed the DECODE network architecture (Fig. <ref>c) to infer the hidden states from three consecutive images and then use the combined information for the final localization. Using context has a substantial positive impact on performance.





  
Training the DECODE network using simulator learning 
We want to train the parameters of the DECODE network to simultaneously detect and localize fluorophore particles from images of sparsely activated fluorophores. As ground truth particle localizations are not easily available for real data, we can not directly use supervised learning. However, the forward image formation model of how a given set of fluorophores gives rise to the detected image is well understood. To detect and localize active fluorophores, we need to train a neural network to invert this forward model. We investigate two different approaches for training the DECODE network to do this: The first method simulates data from our forward model and uses the simulated data to train the deep network using supervised learning <cit.>. We call this method “simulator learning” (SL, Fig. <ref>a). The advantage of the method is its simplicity.  However, its accuracy depends crucially on the quality of the simulation and how well it matches the dataset being analyzed. We will describe a second method (called “auto encoder learning”, AEL, Fig. <ref>a) below. 

Simulator learning has been used by previous deep learning approaches to SMLM <cit.>. Since the physics describing how the camera image is generated by the imaging of a biological sample is well understood <cit.>, we can use a simulation of biological samples consisting of point source emitters representing active fluorophores distributed randomly across a small image patch. We then model the forward generative process of image formation as follows: We simulate the noise- and background-free image of the fluorophores by convolving the point emitters with a model of the point spread function. A random homogeneous background intensity is added to generate a mean intensity image, and finally the noisy measured camera image is then simulated by sampling from a gamma distribution. The density, brightness, activation and inactivation times of the simulated fluorophores, and the background intensity values are chosen randomly to generate a large diversity of simulated images.

We developed a specialized loss function for our representation of the final localizations by the discrete pixel positions and the in-pixel offset variables.
 We interpret the binary values p as the probability that an activation exists in that pixel while the outputs α,Δ x,Δ y,Δ z,σ_α,σ_x, σ_y,σ_z parametrize Gaussians which are components of a Gaussian mixture model (GMM) which describes the spatial distribution of emitter activations.
We then maximize the likelihood of the simulated continuous ground truth positions under this GMM. This allows us to optimize all the output variables jointly and to obtain uncertainty estimates which can be used to filter out localizations or to convolve the localizations with a Gaussian parametrized by the uncertainty for improved rendering <cit.>. 
To determine the correctness of our uncertainty estimates we compared them to parametric estimates of the Cramér–Rao bound obtained with the equation from <cit.>. Such estimates generally only take the brightness, the background and the PSF shape into account but not other important factors that increase uncertainty like other close-by PSFs or inhomogeneous background. We observe that under optimal conditions, with a single emitter per frame and a Gaussian PSF, our uncertainty estimates agree well with the parametric estimates. For denser data our method generally produces higher uncertainties, except when temporal context is used (see Fig. <ref>). 







 §.§ Quantitative evaluation on simulated datasets from the SMLM challenge show DECODE outperforming all algorithms across a variety of conditions




The 2016 SMLM challenge [http://bigwww.epfl.ch/smlm/challenge2016/index.html?p=datasets] is the second generation comprehensive benchmark evaluation developed for the objective, quantitative evaluations of the plethora of available localization algorithms <cit.>. The benchmark offers synthetic datasets for training and evaluation that were created to emulate various experimental conditions. 


A direct comparison of DECODE with other contenders (Fig. <ref>) in the SMLM 2016 challenge shows DECODE outperforming other approaches across datasets [Note that currently not all of our results are displayed in the plots on website, but all results can be downloaded from http://bigwww.epfl.ch/smlm/challenge2016/leaderboard.csv. Our results uploaded January 16 2020, results are current as of July 06 2020.].
DECODE outperforms all 39 currently ranked algorithms on 12 out of 12 datasets, and often by a substantial margin. The datasets include high (N1) and low (N2) signal to noise ratios (SNR), with low (LD) or high (HD) emitter densities, with 2D, Astigmatism (AS) and Double Helix (DH) point spread function based imaging modalities Fig. <ref>. We quantified performance using RMSE lateral or volume localization error, as applicable for 2D and 3D data respectively, and the Jaccard index JI which measures single molecule detection accuracy. The SMLM 2016 benchmark also reports a single score which combines particle localization and detection accuracy into a measure called efficiency.

DECODE achieves an average efficiency score of 66.61 out of the best possible score of 100 (achievable only by a hypothetical algorithm that accurately detects 100% particles with 0 nm localization error). This is compared to an average score 48.3, and 45.6 for all non-DECODE second and third place algorithms respectively. The improvement in performance by using DECODE is substantial, leading to 7 × the accuracy improvement gained by using the second best algorithm over the third best algorithm.
The difference is particularly large under difficult imaging conditions, when high emitter densities and low SNR can conspire to make detection and localization challenging, particularly so for the double helix point spread function. For example, in the Low SNR/high density/Double Helix condition, DECODE achieves an efficiency score of 44.23, whereas no other algorithms achieves a non-negative efficiency score. DECODE achieves an average efficiency of 57.29 on the six challenging high density datasets, while the average second best and third best algorithms achieve only scores of 30.76 and 27.16 respectively. This represents an average 10 × improvement relative to the improvement by the second best algorithm over the third best algorithm on these challenging datasets.


DECODE is the best algorithm on all 12 datasets, across a variety of imaging modalities, SNR and density conditions. In contrast no other algorithm previously achieved such universal superiority, instead specializing on a limited range of imaging conditions.
Qualitatively, DECODE improves super-resolution reconstructions by improving both the detection and the localization of single molecules. An example of this can be seen in Fig. <ref>a, where we compare the reconstructions obtained with DECODE to the multi-emitter fitting approach CSpline <cit.> on two 3D double-helix datasets with high fluorophore densities [We used settings provided by the authors: https://github.com/ZhuangLab/storm-analysis]. DECODE detects more fluorophores, and localizes them more accurately than CSpline for this dataset.








 §.§ DECODE enables accurate reconstructions with shorter imaging times at high emitter densities


In SMLM, there is a trade-off between the imaging time and the activated fluorophore density. Sparsely activating fluorophores leads to the best localization accuracy, but requires long imaging times in order to localize sufficient numbers of particles to reconstruct the sample faithfully. By enabling accurate particle localizations at higher densities, DECODE can yield accurate super-resolution reconstructions with significantly shorter imaging times. We demonstrate this by imaging and reconstructing the same sample of labeled microtubules at four different emitter densities using dSTORM (direct stochastic optical reconstruction microscopy) <cit.>.
For the first dataset the experimental conditions roughly correspond to high SNR and low density settings modelled in the challenge, with an average upper limit of emitter density of 0.14 fluorophores μ m^-2 per image [Measured by dividing the DECODE predictions into 1 /μ m^-2 bins and calculating the 99 percentile of densities.].
The other three datasets consisted of 2.5 ×, 4.5 × and 12 × fewer imaging frames, while the total number of active fluorophores is roughly constant across datasets. 

We trained and applied one common DECODE model to all four datasets (Fig. <ref> a: reconstruction on subsection of low density dataset). We compared DECODE reconstructions with those of CSpline across all four data sets (Fig. <ref> b) to investigate how the quality of reconstruction deteriorates for denser datasets. Similarly to the simulated challenge datasets we observe a sharper image and less spurious localizations for the original low density dataset. As the density increases, DECODE consistently yields reconstructions with similar accuracy, while the reconstructions produced by CSpline degrade for high densities.

To quantify the reconstruction performance across the different conditions, we calculated the resolution of the reconstructed image using Fourier Ring Correlation (FRC) <cit.>. The FRC estimates resolution by measuring the correlation of two different reconstructions of the same image across spatial frequencies. We split the localizations ordered in time into blocks of 10000 and created two different reconstructions of the same sample coming from even and odd blocks. We then used the spatial frequency at which the correlation drops below a threshold value of 0.143 <cit.> to estimate the resolution of the reconstruction. DECODE consistently improves resolution by 10 nm - 25 nm over CSpline across all imaging densities (Fig. <ref>c and d), and requires 10 × fewer imaging frames for the same quality of reconstruction. 





 §.§ DECODE enables high fidelity reconstructions of 3D lattice light sheet PAINT imaging with reduced imaging time


To illustrate the general applicability of DECODE, we applied it to 3D lattice light sheet (LLS) microscopy combined with the PAINT (point accumulation for imaging of nanoscale topography labeling) technique <cit.>. In PAINT microscopy, the fluorophore labeling a sample stochastically binds and unbinds from the sample, providing dense labeling. In lattice light sheet microscopy, thick volumes can be imaged at high resolution by scanning a thin (1.1 um)  light sheet, with axial localization within the sheet enabled by astigmatism. We reconstructed a previously reported dataset of a chemically fixed COS-7 cell with intracellular membranes preferentially labeled by azepanyl-rhodamine (AzepRh) <cit.> consisting of 147,500 3D volumes comprising more than 20 million 2D images acquired in 270nm steps.

For this dataset,  one complete scan of the volume involved moving the probe 141 times by 500 nm. The detection axis was oblique to the coverslip, resulting in emitters that are active in successive frames, and thus appear to move in the x- and z-direction by a fixed distance. We adjusted our algorithm to account for this movement, so that we could still employ local context. Furthermore, we used a modified noise model, as the images were made with a sCMOS camera (see methods for details).

We compare our reconstructions to the original reconstructions described in <cit.> which used a  custom-made iterative MLE fitter with a parametric PSF model <cit.>. DECODE detects 1.25 billion particles, compared to 400 million particles detected by the original algorithm.
While LLS-PAINT microscopy yields high resolution reconstructions over large 3D volumes, its usability is limited by the long imaging times required to localize a sufficient numbers of particles for reconstruction. For example, the dataset we analyzed was obtained in over 3 days of imaging time. We show that DECODE provides sharper images using a smaller number of frames, and could thus be used to obtain the same quality of reconstructions using only a fraction of recorded frames which is confirmed by FRC resolution estimates (Fig. <ref>, <ref>).  
We note that one challenge of long imaging times is that performance can be limited by nonlinear swelling of the sample over the time course of the imaging, which can only be partially corrected by non-rigid registration. Thus, reducing imaging time by improved reconstruction algorithms could also lead to better reconstructions with fewer artifacts. 








 §.§ Combining simulator learning with auto encoder learning to learn to simulate better

The effectiveness of simulator learning depends  on the availability of an accurate forward generative model at training time, as deviations of the true forward model from the simulated forward model can degrade performance. This problem can be solved by simultaneously estimating the parameters of the true forward model, and training the DECODE network using the real measurements, rather than a fictitious simulation. This is possible using the recently developed framework of variational autoencoders (VAEs) <cit.>. In the VAE framework, the stochastic forward generative model and the DECODE network are stacked to form a stochastic autoencoder. This autoencoder is then used to simultaneously optimize the parameters of the deep network and the forward model, with the goal of achieving image-reconstructions which are similar to the original measurement.
 
Formally, this can be achieved by maximizing a so-called `evidence-lower bound' via stochastic gradient optimization (see methods for details). VAEs have, e.g.,  previously been used on the related problem of inferring action potentials from calcium imaging data <cit.>. A drawback of VAE-based approaches are that gradients for training the DECODE network need to be approximated using Monte Carlo sampling, which can make optimization more challenging.

Simulator learning and autoencoder learning are two sides of the same autoencoder coin, as can seen by comparing Fig 4a to Fig 4b: In simulator learning, a known PSF is used to transform (encode) simulated emitter locations into a microscope image, and the DECODE network is used to recover (decode) the original simulated emitter locations. The network is optimized to minimze the discrepancy-measure computed by comparing simulated and inferred localizations.

In what we call autoencoder learning (due to the relationship to VAEs) the image measured by the microscope is stochastically transformed (encoded) by the DECODE network into predicted emitter locations, and these predicted emitters are then transformed (decoded) by the estimated PSF back into a reconstruction of the measured image. The objective function used to train the DECODE network is the difference between the measured and reconstructed images. Because AEL relies on autoencoding the real measured data, it enables learning of both the encoder (DECODE network) and the decoder (PSF). In contrast, SL only allows training of the decoder (DECODE network). Empirically, we found that combining SL and AEL in a manner we call combined learning (CL) often leads to the best performance. 

To highlight the difference between simulator and autoencoder learning we show how the two approaches behave for different degrees of mismatch in the assumed point spread function. 
To simulate PSF mismatch we generated datasets using a 2D elliptical Gaussian PSF with increasing ellipticity and then trained DECODE models that use a non-elliptical circular Gaussian (Left panel Fig. <ref>e, solid lines show performance of models with fixed generative parameters). For an ellipticity of zero the models have access to the true underlying generative model. In this case SL training sets an upper bound to the achievable performance with a given network as it is able to generate an infinite amount of labelled data with the correct simulation parameters, and so outperforms AEL. 
However, pure simulator learning is brittle and more sensitive to parameter mismatch as the DECODE network never `sees' elliptical PSFs during training. Autoencoder learning can  still try to infer the correct positions as placing the circular PSF into the middle of the elliptic one achieves the best reconstruction. Alternating between the two methods retains the advantages of both: Performance is virtually the same as simulator learning when there is no mismatch and performance degrades more gracefully when the mismatch is increased. 
When we add training of the generative model parameters (dashed lines), our PSF model learns to account for some of the PSF-mismatch, which further improves performance. 

To test these findings in a more realistic setting, we trained DECODE methods using simulator and combined learning and submitted them to the SMLM 2016 challenge (see Fig. <ref>). We also evaluated them on the training datasets to more precisely evaluate differences in performance (Fig. <ref>d). For astigmatism and double helix data we fit the PSF model to the provided bead-data. For 2D we instead used a more heuristic estimate, choosing sigma values and a z-dependent scaling that covers the PSF observed in the dataset.
Furthermore we used combined learning in two different settings: one for variational inference (VI) where autoencoder learning is performed on the same dataset on which performance is evaluated and one for amortized inference (AI) where training and testing takes place on different datasets.
These two models where trained using both local and global temporal context, while for simulator learning we only used local context.
Overall performance of the three submissions is very similar, which is to be expected given that we are able to approximate the generative model of the data very accurately. 
Adding autoencoder learning is especially helpful for 2D data where we did not fit the PSF model to beads-data and instead adjusted it during training of the model. As can be seen in Fig. <ref> the algorithm learned to add diffraction rings to the PSF model without any access to bead-data. 
Additionally we observe the combined learning (VI) performs better on the high density / low SNR datasets. As described below, this can be attributed to the fact that global context is especially helpful in these difficult conditions.








 §.§ Local temporal context is more informative than global context



Imaging frames from SMLM contain correlations in time across short and long time-scales. Once activated, fluorophores are usually active for more than one imaging frame and are therefore visible in multiple consecutive frames at the same position, leading to short time-scale correlations in the images. The spatial distribution of fluorophores in a sample is non-uniform and is concentrated around the biological structures labeled by the fluorophores. This spatial distribution of fluorophores leads to temporally global correlations. Our DECODE network is designed to exploit both kinds of temporal correlations via local and global context windows across imaging frames to improve the detection and localization of single molecules. We studied the contributions of local and global context to the performance of the DECODE network.

We trained DECODE models using CL on the 6 high density datasets of the challenge, using either no context, local context, global context or both forms of context.

Local and global context have different effects on the performance (Fig. <ref>e). While we find that both local and global context individually improve performance, local context is generally more helpful than global context. Together, global and local context give the best possible performance but at a minor improvement over local context alone.

It should be noted that a straight forward way of using local context called “grouping” is commonly used to improve localizations as a post-processing step in SMLM <cit.>. Localizations occurring in consecutive images that are closer to each other than a fixed threshold are assumed to belong to the same emitter and their localization is averaged, potentially weighted by the uncertainty of each localization. 
We applied grouping to the DECODE models trained with and without local context as well as CSpline (Fig. <ref>) . We observe that grouping is almost as effective as our method in reducing the localization error for the easiest condition (high SNR / low density) but performs worse for any of the more difficult datasets. Furthermore, our method for using context also improves detection accuracy while grouping only influences the localization error. Lastly, using grouping on top of DECODE with local context results in a small improvement which would further increase  performance on the challenge if one used it. 












§ DISCUSSION


We here described DECODE, a new deep-learning based method for single molecule localization for reconstructing super-resolution images. DECODE differs from existing localization algorithms by simultaneously performing detection and localization of particles. DECODE yields substantial improvements in performance over previously evaluated algorithms in a publicly available benchmark challenge: It achieves best performance in every condition, and often improves prediction performance by a large margin. When applied to high density dSTORM imaging of microtubules, and LLS-PAINT imaging of whole cells, it leads to reconstructions which have markedly improved resolution due to substantial improvements in particle detection. The performance benefits of DECDOE are especially pronounced in high-density imaging conditions, thereby opening up new opportunities for faster imaging of fixed samples, and even live imaging. 

DECODE leverages the flexibility of deep learning – for its predictions, the network can benefit both from temporal context (e.g. from fluorophores being active across multiple imaging frames), as well as spatial context (e.g. from clustering of fluorophores in space). DECODE can be used in a very flexible and general manner and can easily be applied to arbitrary PSFs and noise models – in this paper, we applied it to 4 different imaging modalities ranging from engineering point spread functions to 3D lattice light sheet microscopy.

The DECODE network is trained to produce probabilistic point process predictions – it predicts both the probability of detection and the uncertainty of localization for each detected particle. We showed that the localization uncertainties predicted by our network are superior to conventionally used CRLB uncertainties and are particularly useful for filtering particles to produce high resolution reconstructions.

We presented and evaluated two ways to train the DECODE network, using simulator learning and autoencoder learning. Simulator learning allows for fast an easy training of a DECODE network when the optical properties of the microscope are precisely known. Autoencoder learning uniquely enables the in situ estimation and refinement of imaging parameters such as the empirical point spread function and noise model directly from the experimental measurements, in principle enabling the tracking of drift in the optical system over the course of long imaging experiments, or the estimation local point spread functions across large fields of view in the presence of sample dependent optical aberrations.

One weakness of DECODE is that it currently requires the training of a new neural network whenever the optical properties of the microscope change. This training can currently take over 10 hours on a single GPU. However, it may be possible to train a single network to predict robustly across minor variations in the point spread function or noise distribution, i.e. to amortize simulator-learning across setups <cit.>. This can enable real-time reconstruction on a single GPU without the need for re-training, even under the most challenging conditions, since the computational complexity of the predictions depend only on the size of the image and not the number of particles in each imaging frame.















  §.§.§ Acknowledgments


This work was supported by the German Research Foundation (DFG) through SFB 1089 and  Germany’s Excellence Strategy – EXC-Number 2064/1 – Project number 390727645, the German Federal Ministry of Education and Research (BMBF, project `ADMIMEM', FKZ 01IS18052 A-D), the Howard Hughes Medical Institute and the  European Research Council (CoG-724489 to J.R.). We thank Daniel Sage for useful discussions,  David Greenberg and Poornima Ramesh for comments on the manuscript, and Eric Betzig and Jennifer Lippincott Schwartz for kindly sharing data with us. 




§ METHODS




  §.§.§ Software availability

All methods were implemented in Python and PyTorch <cit.>. Code and hyper-parameter settings for the challenge results are available at <https://github.com/mackelab/DECODE>/



  §.§.§ Approximate Bayesian inference with DECODE


DECODE is a Bayesian inference method which requires a formal probabilistic description of the entire SMLM measurement process. This description amounts to a stochastic simulator (known as a generative model) which can generate synthetic SMLM data, but which can also be optimized to fit the data. In our framework, at each time point, each active fluorophore i has a location in 3D at x_i, y_i, z_i, and a brightness α_i. The set of N_t single molecule locations and their brightness at each time point are the unknown hidden causes or latent variables h = {{x_i}_t,{y_i}_t,{z_i}_t,{α_i}_t} that give rise to the noisy low resolution images which constitute the measured data d = {I_t}.
A complete description of the generative model in Bayesian framework includes the prior distribution p(h) which describes the spatial distribution and temporal dynamics of fluorophores, and the likelihood distribution p(d|h) which describes the stochastic process describing the distribution of images generated by the microscope for a given configuration of fluorophores. The likelihood is formalized in terms of the point spread function describing the transfer function of the microscope, and the measurement noise at the camera.

The DECODE network is trained to perform approximate Bayesian inference, using simulator learning and autoencoder learning to approximate the true posterior distribution p(h|d) and predicts the hidden single molecule locations and brightness from the measured images.



  §.§.§ Spatial distribution and temporal dynamics of fluorophores.

We assume that on each imaging frame, a fluorophore can be activated in any given pixel with a constant probability of p_on. An active fluorophore on any given frame has a probability of p_off of turning off in the next frame. The location of an active fluorophores within a pixel is drawn from a uniform distribution in x and y. For z we chose a Gaussian distribution with a mean centered at the focal plane and a variance chosen to cover the range of the point spread function of the microscope instead, as the bulk of the recorded structure is usually located around the focal plane. The brightness of an active fluorophore is sampled from a uniform distribution U(0.1,1) times the maximum possible expected brightness of a single fluorophore. This describes the prior distribution over particle locations and brightness p(h).



  §.§.§ Point spread functions
 

We used the sum of a parametric function PSF_parametric  and a non-parametric interpolated pixel map PSF_pixmap to model arbitrarily complex point spread functions PSF(x,y,z)=PSF_parametric(x,y,z) + PSF_pixmap(x,y,z). In principle, the non-parametric interpolated pixel map is sufficient to represent any possible PSF within the support of the pixel map, we found that the parametric component helps with the learning when using AE training.

In this paper, we analyzed data imaged or simulated using three PSF models – one for 2D localization, and two for 3D localization. For 3D localization, we used the astigmatism (AS, <cit.>) and the double helix (DH, <cit.>) PSFs. The parametric components of these parametric PSFs as a function of the spatial coordinates x, y, and z is given below, along with the parameters of the PSF after the semi-colon.



    PSF_2D(x,y,z; a_1,a_2, b_1,b_2) =    ∑_n=1,2(a_n · e^-b_n(x^2+y^2)/(1+|z|)^2)
    
    
    PSF_AS(x,y,z;a,b_x,b_y,c) =     e^-ax^2/((z-b_x)^2+c)e^-ay^2/((z-b_y)^2+c)
    
    PSF_DH(x,y,z;a,b,c,d) =     e^-a(x-s_x(z))^2+a(y-s_y(z))^2 + e^-a(x+s_x(z))^2+a(y+s_y(z))^2
       s_x(z)= d·cos(b · z + c)    s_y(z)=d·sin(b · z + c).


PSF_2D is a weighted sum of two circular Gaussians whose variance increases as a function of the distance of the point source away from the focal plane. PSF_AS is described by an elliptical PSF whose eccentricity is a function of the defocus z <cit.>. PSF_DH is a double helical PSF modelled as two circular Gaussians that rotate around the fluorophore at a distance d as a function of z.

PSF_pixmap enables the approximation of arbitrarily complex PSFs and is a 3D image volume with the same pixel size as camera and with a z-spacing of 100 nm. This pixel map is interpolated by trilinear interpolation to evaluate the point spread function at any location within the support of the pixel map.



  §.§.§ Imaging fluorophores and camera noise
 
The number of photons emitted from a fluorophore follows a Poisson distribution. To model the noisy imaging by an EMCCD camera, the distribution of photon counts can be convolved with a gamma distribution that models the electron-multiplying (EM) gain and with a Gaussian distribution that accounts for read-out noise. 
The resulting distribution can not be expressed analytically. This is irrelevant for simulator learning, however autoencoder learning requires a differentiable expression for the probability distribution of the measured image I(x,y) as a function of the mean intensity image I̅(x,y). Therefore, we approximate the noise model using a single Gamma distribution with parameters that are given by the camera baseline BL, its electron-multiplying gain EM and electron conversion factor EC. 
Given N activated fluorophores, each located at x_n,y_n,z_n and with a brightness of α_n, and a constant background fluorescence of β, the intensity of a pixel I(x,y) located at x,y of the resulting imaging frame is  simulated as:

    I̅(x,y) =   ∑_i=1^N α_i PSF(x-x_i,y-y_i,-z_i) + β
    
    I(x,y) ∼   Gamma((I̅(x,y) - BL)/η, η) + BL
       η = 2· EM/EC.

This describes the likelihood function p(d|h).








  §.§.§ Fitting of 3D PSFs from bead stacks


For 3D inference, it is common to calibrate the PSF model on data with known axial offset as the exact relationship between its shape and the position cannot be estimated from unlabeled data. We estimated the AS and DH PSF's using calibration bead stacks, i.e. images of single fluorophores at different offsets with  high signal to noise ratios. We first obtain a rough estimate of the bead locations using a basic peak-finding routine. We then maximize the likelihood p_θ(d|h) by performing stochastic gradient descent on the exact x- and y- coordinates of each bead (which are constant across images), the shape parameters of the PSF model (<ref>) and the the pixel maps δ_xy. This simple method achieves localization errors of less then 0.3 nm on the challenge calibration stacks where the ground truth locations are available.
During training of the DECODE network we generally keep the PSF model fixed (see example fits in Fig. <ref>). For 2D datasets the PSF model can be learned simultaneously with the network parameters in a completely unsupervised way, i.e. without requiring access to bead stacks (See pixel maps δ_xy  in Fig. <ref>). We emphasize that neither the training algorithm, nor the network architecture, depends on the specifics of the generative model or the PSF model, and both could well be combined with more flexible functional forms of PSFs.




  §.§.§ DECODE network architecture for probabilistic single molecule detection and localization

Our frame analysis module as well as our temporal context module are U-nets with two up- and downsampling stages and 48 filters in the first stage. Each stage consists of three fully convolutional layers, where in each downsampling stage the resolution is halved, and the number of filters doubled, and vice versa in each upsampling stage.
Upsampling is performed using nearest neighbor interpolation to avoid checkerboard artifacts <cit.>. The final output representation is predicted after two additional convolutions layers.

For each camera pixel, the DECODE network predicts the probability that a fluorophore was detected near that pixel p, the location of the detected fluorophore relative to the center of the pixel, Δ x,Δ y,Δ z and the predicted fluorophore brightness α, and the uncertainties associated with each of these predictions σ_x,σ_y,σ_z,σ_α.
We used ELU nonlinear activation function <cit.> for all hidden units, the hyperbolic tangent nonlinearity for the coordinate outputs Δ x,Δ y,Δ z and the logistic sigmoid nonlinearity for the non-negative brightness and uncertainty outputs α,σ_x,σ_y,σ_z,σ_α.

The output channels of the DECODE network together represent a distribution of possible interpretations of a measured image q(h|d) and constitute an approximation to the true posterior distribution over possible detections and localizations p(h|d). Our approximate posterior q(h|d) is a type of Gaussian mixture model, with one Gaussian mixture component per camera pixel. It can represent at most as many particles as there are camera image pixels, and for each pixel represents the detection probability and localization mean and variance of one particle.
In the following sections we describe how this architecture is trained and how deterministic detections and localizations of particles can be obtained from the output of the DECODE network at test-time.



  §.§.§ Simulator learning


Given a set of simulated particles and the resulting images over time, simulated as described above such that h ∼ p(h) and d ∼ p(d|h)
we take samples from our generative model h∼ p(h), d∼ p_θ(d|h) and maximize the loglikelihood log q_ϕ(h|d).
This procedure amounts to minimizing the Kullback-Leibler divergence (D_KL) between the posterior of the generative model and the recognition network, averaged over the (simulated) data distribution:

    𝔼_p_θ(d) [-D_KL(p_θ(h|d)||q_ϕ(h|d)] =
    𝔼_p_θ(d,h) [log q_ϕ(h|d)] + const.



Given our representation of the final localizations by the discrete pixel positions and the in-pixel offset variables


    x = x_p + Δ x    y = y_p + Δ y    z = Δ z


we developed a loss function that allows us to jointly optimize the different output variables. We interpret the binary values p as the probability that an activation exists in that pixel while the outputs α, Δ x,Δ y,Δ z, σ_α, σ_x, σ_y,σ_z parametrize Gaussians 𝒩(μ,Σ) which are components of a Gaussian mixture model (GMM) which describes the distribution of emitter activations:


    q_ϕ(h|d) ∝∏_i∑_kp_k/∑_k p_k𝒩(X_i|μ_k,Σ_k)


where k indexes all pixels, and X_i are the ground truth location vectors.

The number of activations follows a Poisson Binomial distribution given that the binary probabilities vary strongly across pixels.
As the likelihood of this distribution is hard to evaluate we instead use its mean and variance to parametrize a Gaussian approximation to the likelihood of counts:


    q_ϕ(h|d) ∝𝒩(∑_k S_k | ∑_k p_k, ∑_k p_k - p_k^2)


where ∑_kS_k is the true number of emitters. 

As the GMM term scales linearly with the number of emitters, while the count term stays constant, we multiply it by the number of emitters to balance the two terms. 
The resulting total log-likelihood we use to train our inference network is:


    log q_ϕ(h|d) = ∑_i∑_klogp_k/∑_k p_k𝒩(X_i|μ_k,Σ_k) + ∑_kS_k·log𝒩(∑_k S_k | ∑_k p_k, ∑_k p_k - p_k^2)




  
Auto encoder learning: Optimizing a lower bound on p(d) 

For auto encoder learning we only treat the discrete outputs as stochastic latent variables, and use the deterministic mean values for the continuous outputs. 
log q_ϕ(h|d) is therefore calculated as the binary cross-entropy between the inferred probabilities and the discrete pixel activations.
We also performed experiments using the loss described in (<ref>) for autoencoder learning. In this case in addition to the discrete samples we also draw samples of our continuous offsets from 𝒩(μ_k,Σ_k) and optimize them using the reparametrization trick <cit.>. However this resulted in far higher gradient variance and overall reduced performance.  

Using Jensen's inequality we can derive a lower bound (ELBO) on the marginal likelihood p(d):

    log p(d) = log𝔼_q [p_θ(d,h)/ q_ϕ(h|d)] ≥𝔼_q [logp_θ(d,h)/ q_ϕ(h|d)] = ℒ(d)

by maximizing this ELBO with respect to θ we minimize the reverse D_KL averaged over the true data distribution

    𝔼_p(d) [D_KL(q_ϕ(h|d)||p_θ(h|d))]

Unlike simulator learning maximization with respect to ϕ also allows us to learn the parameters of the generative model. 
If we instead use an importance weighted average over j samples from our recognition model to estimate p(d), and again apply Jensen's inequality we obtain a tighter lower bound (which is identical to the ELBO for j=1):

    ℒ_IW^j(d) = 𝔼_ q(h_1:J|d)[ 1/J∑^J_j=1log[ p_θ(d,h_j)/q_ϕ(h_j|d)_ω_k(d,h_j)] ]

This objective is the basis for both the importance weighted autoencoder (IWAE) <cit.> and the reweighted wake-sleep algorithm (RWS) <cit.>.



  §.§.§ Updating θ


For a given value of ϕ, unbiased gradients for θ can be obtained by sampling the discrete variables h_1,...,h_j ∼ q_ϕ and calculating the gradients:

    ∇_θℒ_IW^j(d∼𝒟) = ∇_θlog (1/J∑^J_j=1ω_j) = ∑^J_j=1ω̃_j ∇_θlog p_θ(h_j|d) 
    ω̃_j = ω_j(d,h_j)/∑_j'=1^J ω(d,h_j')




  §.§.§ Updating ϕ


Obtaining gradients for ϕ is more involved, especially in the case of discrete latents when the reparametrization trick cannot be applied. 

The RWS algorithm includes two procedures to obtain gradients for ϕ. The sleep phase update matches simulator learning which minimizes the D_KL between p and q over data that is generated from the generative model p_θ(d|h)p(z). 

The wake phase update optimizes the same D_KL, but over the true data distribution p(d) (i.e. using samples from the data).

    ∇_ϕ D_KL(p_θ(h|d)||q_ϕ(h|d)) ≃∑^J_j=1ω̃_j ∇_ϕlog q_ϕ(h_j|d)

We also experimented with the VIMCO algorithm <cit.> and the The Thermodynamic Variational Objective <cit.> as alternative approaches to obtain low variance gradients for discrete latent variable models. Performance was comparable across methods, so we chose RWS for its easy implementation and low number of hyper parameters. 

Wake phase updates can be very noisy, especially during the first iterations when the network basically produces random samples. If the network predicts large numbers of detections training can fail to to memory restraints. Therefore, we start training with a warm up phase of 1000 iterations of simulator learning, where one iteration corresponds to the evaluation of the loss on one batch and a subsequent gradient update.



  §.§.§ Training details and hyper-parameters


Training is performed on 40×40 pixel sized regions that are simulated or randomly selected from recorded images at each iteration.
If the network is trained to make use of global context, we use a running average of the hidden states collected over the last 100 training batches. At test time we perform two passes over the dataset: the first one to collect the average hidden state h̃ = ∑_T h_t and the second one to obtain the inference results.

When training with local context we employ different strategies for SL and AE training steps. For simulator learning, when sampling data we align the spatial variables Δ x,Δ y,Δ z = μ⃗ (but not the intensity) to be identical when a fluorophore is active in consecutive frames.
For AE training, for each set of variables S_t, μ⃗_t which are inferred from the images I_t-1,I_t,I_t-1 we also infer offset variables and uncertainties σ_x, σ_y,σ_z = Σ for t+1 and t-1 to provide context.

We use these variables to calculate an error term that is the sum of log-likelihoods of the offset variables at each pixel under the Gaussian distribution given by the activations in consecutive images:

    δ_xyz = ∑_x,y    S_t· S_t-1 (𝒩(μ⃗_t|μ⃗_t-1,Σ_t-1) + S_t· S_t+1(𝒩(μ⃗_t|μ⃗_t+1,Σ_t+1)


This term is subtracted from the objective function during training.

A base value of 0.01 is added to the uncertainties to avoid instabilities.  We used the AdamW optimizer <cit.> with a learning rate of 6·10^-4 for the network parameters which is multiplied by 0.9 after every 1000 training iterations. When also learning the PSF we used a learning rate of 0.015 for PSF model parameters and 3·10^-6 for the pixel maps δ_xy. To stabilize training and ensure that gradient steps of simulator and autoencoder learning are roughly equal in size when performing combined learning we employ gradient norm clipping with a maximum norm of 0.03. We normalized the inputs to the network by first subtracting the mean of the dataset and then dividing by the maximum of the mean over the image dimension. To calculate the RWS objective we used 40 samples. 



  §.§.§ Obtaining localizations and post-processing



The DECODE network predicts the probabilities of a fluorophore being located at a specific pixel. While we sample from this distribution during training, we prefer to generate deterministic detections and localizations at test time. To get deterministic, fast and precise pseudo samples we instead use a variant of non-maximum suppression to obtain final localizations. 
To obtain a binary mask of fluorophore candidates for a given frame we identify probability peaks, i.e. pixels with values that are above 0.3 and higher than all values in a surrounding 3x3 patch. 
We then add the probability mass from the 4 directly adjacent pixels to the values at the candidate positions by convolving the probability map with a cross shaped filter and applying the mask. All candidates with added probability values above 0.7 are counted towards the localizations. 

The algorithm can be expressed purely in the form of pooling and convolution operations and therefore runs efficiently on a GPU.

For difficult imaging conditions, i.e. high densities, low SNR values and high offsets from the focal plane the lateral offset variables can be biased towards small absolute values. This effect scales with the uncertainty of the predictions and can produce artifacts in the reconstructed image as localizations are concentrated at the pixel centers. 
To counteract this we therefore divide all localizations into equally sized bins according to the total variance  var_tot = √(σ_x^2+σ_y^2+σ_z^2). Then we calculate an empirical CDF F̂_x,F̂_y from the histograms of the Δ x and Δ y variables in each bin. The variables Δx̂,Δŷ = F̂_x(Δ x)-0.5, F̂_y(Δ y)-0.5 have a uniform distribution as desired.
This transformation effectively removes image artifacts while having no impact on the performance metrics. 

As a final post-processing step our inferred uncertainties allow us to effectively filter bad localizations. As shown in Fig. <ref>c this is very effective in reducing the overall localization error. For the challenge data removing between 0 and 20 % results in the best performance. For real data this threshold should be individually chosen according to the amount of data collected.




  §.§.§ Evaluating localization accuracy and reconstruction resolution

To evaluate performance on the challenge datasets, as well as our own simulations we use the lateral or volume localization error in  and the Jaccard index J which quantifies how well an algorithm does at detecting all the fluorophores while avoiding false positives J = 100· TP/(FN+FP+TP). Localizations are matched to ground truth positions when they are withing a circle of 250 nm radius. As a single metric that evaluates the ability to reliably infer fluorophores with high precision we use the efficiency metric:


    E = 100 - √((100-J)^2 + α^2 RMSE^2))


Lateral and axial efficiency are calculated with alpha values of α=0.5 nm^-1 and α=1 nm^-1 respectively and then averaged to obtain the overall 3D efficiency. 
Super-resolution images were rendered by convolving inferred positions with a 2D Gaussian with a width of 5nm.

The Fourier ring correlation (FRC, <cit.>) in <ref> was calculated by constructing two super-resolution image volumes of the same sample (σ=8.5nm, pixelsize=10nm) by dividing the localizations into two sets. We did this by alternating blocks of 50k consecutive localizations.



  §.§.§ DECODE for LLS-PAINT microscopy

The LLS dataset differs from the the other datasets we analyzed in two respects. First, due to the movement of the light sheet, fluorophores that are active across multiple frames change their x and z position within a frame by a fixed amount. While usually this can be accounted for in post-processing, in order to use local context we also have to adjust our generative model. 
Therefore, when generating the image triplets for simulator learning, we move the emitters by the correct amount when they are active in multiple frames.
Second, the images were recorded with a sCMOS instead of an EMCCD camera, therefore requiring a different noise model.
We follow the description of the image generating process in <cit.> for our noise model. Given the mean intensity as calculated in (<ref>) it is given by: 


    I̅(x,y) =   ∑_i=1^N α_i PSF(x-x_i,y-y_i,-z_i) + β
    
    I(x,y) ∼   Gamma((I̅(x,y))/η, η)
       η = Var(x,y) + g·(I̅(x,y)-BL)/I̅(x,y).


Here, we assumed the camera gain g to be constant, while Var(x,y) is the pixel specific noise that can we estimated from dark images. For the results shown in Fig. <ref> we trained DECODE with CL and local context. The PSF was optimized on bead stacks, and the pixel maps δ_x,y where further optimized during the training of the network.

The reconstructions shown in Figures <ref> and <ref> are rendered by convolving each localization with a 3D Gaussian parametrized by the respective uncertainties to construct a 3D histogram of the volume with a voxel size of 10x10x20 nm. Histogram values are then clipped at 2.5 to remove most of the contribution from fiducials. We then plotted the color-coded maximum projection over the z-axis with the maximum intensity set to the 99.5 percentile of histogram values.

0.2in






PART: 
*Supplementary Information














