Residual Multiplicative Filter Networks for Multiscale Reconstruction
Abstract
Coordinate networks like Multiplicative Filter Networks (MFNs) and B ACON offer some control over the frequency spectrum used to represent continuous signals such as images or 3D volumes.
Yet, they are not readily applicable to problems for which coarse-to-fine estimation is required, including various inverse problems in which coarse-to-fine optimization plays a key role in avoiding poor local minima.
We introduce a new coordinate network architecture and training scheme that enables coarse-to-fine optimization with fine-grained control over the frequency support of learned reconstructions.
This is achieved with two key innovations.
First, we incorporate skip connections so that structure at one scale is preserved when fitting finer-scale structure.
Second, we propose a novel initialization scheme to provide control over the model frequency spectrum at each stage of optimization.
We demonstrate how these modifications enable multiscale optimization for coarseto-fine fitting to natural images.
We then evaluate our model on synthetically generated datasets for the the problem of single-particle cryo-EM reconstruction.
We learn high resolution multiscale structures, on par with the state-of-the art. Project webpage: https://shekshaa.github.io/ResidualMFN/ .
Introduction
Coordinate networks have emerged as a powerful way to represent and reconstruct images, videos, and 3D scenes [1, 2, 3] and for solving challenging inverse problems such as 3D molecular reconstruction for cryo-electron microscopy (cryo-EM) [4, 5, 6].
They typically take as input a point defined on a continuous, low-dimensional domain (e g., a 2D position for images) and they output the signal value at that point (e.g., the color).
In recently proposed architectures [7, 8], the frequency content of the represented signal can also be explicitly controlled.
Such networks are motivated in part by the effectiveness of multiscale methods in image processing and 3D reconstruction.
Nevertheless, while current scale-aware coordinate network architectures offer some control over scale [8, 9, 10], they are not readily compatible with classical multiscale methods used for coarse-tofine optimization, like pyramids [11, 12] or multigrid solvers [13].
State-of-the-art cryo-EM models, for instance, use frequency-marching to progressively estimate 3D density, beginning with a coarse structure, then adding finer structure at each iteration [14, 15].
Some networks use heuristics to empirically constrain the scale of network output for coarse-to-fine optimization, but they have no explicit constraints on the represented frequencies [9, 16, 17].
Others represent signals at multiple scales at inference time, but do not enable control of the frequency spectrum during training [8, 10].
Here, we introduce a new coordinate network architecture and training scheme that enables one to effectively control the frequency support of the learned representations during multiscale optimization, building on recent neural networks with an analytical and controllable Fourier spectra [7, 8].
For coarse-to-fine reconstruction, one can divide the training process into stages, each of which learns a single scale.
Training starts at the coarsest resolution and then progressively adds finer scales.
Naively employing this procedure with existing architectures will damage the signal reconstruction at coarser scales when fitting finer ones (e.g., see Fig. 3).
To mitigate this, we modify MFNs by adding skip connections [18], so reconstruction at one scale naturally incorporates the learned structure at coarser scales.
This alleviates the need to adapt previous layers in favor of fitting the finer grained reconstructions, effectively obviating the need to re-learn signal structure captured by coarser scales.
Second, we derive an initialization scheme which explicitly introduces hierarchical gaps in the frequency spectrum of a new layer so that, by construction, its frequency support has a controllable degree of overlap.
This requires the network to use the skip connections in order to fill the holes in the spectrum and ensures lower levels remain faithful to coarse scale representations of the signal.
We apply our technique to the inverse problem of single-particle cryo-EM reconstruction [6, 4, 14, 19] to determine the 3D structure of macromolecular complexes, like proteins, from collections of noisy 2D tomographic projections.
This is a challenging non-linear inverse problem, closely related to multiview reconstruction, requiring estimation of the pose and 3D structure of bio-molecules.
Successful 3D reconstruction relies on coarse-to-fine methods, sometimes called frequency marching [20], to reduce the risk of becoming trapped in poor local minima.
Our results demonstrate effective and efficient ab initio reconstruction of 3D structures with cryo-EM to high resolution.
This paper proposes a new architecture and training approach to enable powerful multiscale estimation techniques with coordinate-based networks.
Specifically, we make the following contributions:
• We develop residual multiplicative filter networks, a new coordinate network architecture with a tailored initialization scheme for mulitscale signal reconstruction and representation.
• We design fast and efficient coarse-to-fine training strategies for residual multiplicative filter networks that leverage the multiscale representation.
• We apply the proposed architecture and training strategy to cryo-EM reconstruction on two synthetic datasets and achieve final reconstructions competitive with the cryoSPARC [14], a state-of-the-art method.
Related Work
Coordinate-based Networks.
Coordinate-based networks (e.g., [2, 3]) offer a memory-efficient, continuous function parameterization that can be flexibly trained to reconstruct 3D appearance [1, 21, 22, 23, 24, 25] and structure [26, 27, 28, 29, 30, 31, 32], including applications in biomedical imaging [33, 34] and cryo-EM [5, 4, 6].
Although originally formulated as fully-connected
MLP architectures, new models have been proposed to improve performance and interpretability.
For instance, one can use multiple small fully-connected networks to improve representational capacity [35, 36, 37], or combine coordinate-based networks with explicit feature grids to improve efficiency [38, 39], or facilitate generalization across shapes or scenes [40, 41, 42, 43, 29, 44, 45].
Multiplicative Filter Networks (MFNs) replace the conventional MLP architecture with successive layers of Hadamard products and sine non-linearities [7].
Closest to our own work, MFNs have an analytical Fourier spectrum whose bandwidth can be explicitly constrained [8], improving the controllability and interpretability of the representation.
Inspired by this work, we establish Residual MFNs (rMFNs) with specialized control over the Fourier spectrum to enable coarse-to-fine optimization.
Multiscale Reconstruction.
Multiscale representations and reconstruction methods are fundamental concepts in signal and image processing.
For example, wavelets are a fast and efficient multiscale signal representation [46, 47] useful for denoising, compression, communications, and optical flow estimation [48].
Multigrid frameworks are used for many differential equation solvers for physics simulations [13].
Gaussian and Laplacian pyramids [49, 11] have broad application to image processing and have, in part, inspired modern deep learning architectures [50].
Recent coordinate-based network architectures build on these fundamental concepts of multiscale signal representation and reconstruction.
Some architectures use explicit feature grids at multiple scales to enable efficient training and inference for representing shape [39, 51] or rendering scenes [52, 38].
Multiscale fully-connected coordinate network architectures have been realized by progressively increasing the frequencies of positional encodings during optimization, for example, for bundle adjustment and neural scene representation [9, 17].
It is also possible to optimize separate networks with inductive biases towards low or high frequencies to improve shape representation [16].
Still, while these methods introduce multiscale training techniques, the architectures rely on inductive biases rather than explicit constraints on the scale or Fourier spectrum of the representation.
Other methods use coordinate networks for multiscale representation rather than efficient optimization.
For instance, scale-aware positional encodings are trained at all scales simultaneously so that different output scales can be queried at inference time [10].
Band-limited coordinate networks (B ACON ) extend MFNs to control the network bandwidth, but require training on multiple output scales simultaneously for multiscale representation at inference time [8].
Our work is inspired by B ACON , but significantly modifies the architecture, initialization scheme, and training techniques to make efficient multiscale optimization techniques compatible with coarse-to-fine reconstruction.
Cryo-EM Reconstruction.
While coordinate networks have previously been used for cryo-EM reconstruction [5, 4, 6], they differ from our method in that they explicitly represent signals in the
Fourier domain.
Our goal is to demonstrate coarse-to-fine reconstruction methods with coordinate networks in the real domain, and we demonstrate this for cryo-EM reconstruction.
The proposed method may also prove promising for advanced cryo-EM methods that resolve intrinsic motion of structures since spatial motion can be directly modeled in the primal domain [53].
Background
Coordinate neural networks have emerged as effective tools for approximating complex spatial data (e.g., see [1, 2, 7, 8, 10]).
In the case of images, for example, they provide a mapping from continuous positions on the image plane to RGB values.
Mapping 3D coordinates to volumetric density allows modeling 3D geometry.
In particular, these networks have been shown to be readily trained to accurately approximate low-dimensional, complex signals in a memory-efficient way.
Among this broad family of network architectures, in this paper we focus on Multiplicative Filter Networks (MFNs) [7] such as B ACON [8], as they provide explicit control over the Fourier spectra of the function approximations.
Multiplicative Filter Networks.
Most coordinate-based networks, like SIREN [2] and Random
Fourier Features [3], use an MLP architecture consisting of the successive composition of linear transformations and element-wise non-linearities.
In contrast, Multiplicative Filter Networks (MFNs) use a Hadamard product (i.e., element-wise multiplication) instead of composition [7].
Concretely, in an L -layer MFN, the d in dimensional input coordinates, x ∈ R d in , are transformed using L +non-linear filter modules, denoted by g ( i ) ( · ) : R d in → R d h , i = 0 , 1 , . . .
, L , and d h is the hidden layer dimension.
In [7], either sinusoidal or Gabor filters are used for this transformation.
The sinusoidal case is given by g ( i ) ( x ) = sin( ω ( i ) x
+ ϕ ( i ) ) , where ω ( i ) ∈
R d h × 3 and ϕ ( i ) ∈
R d h are referred to as frequencies and phases.
At layer i , the intermediate representation of previous layer, z
( i − 1) , after being linearly transformed, is multiplied by the non-linear filter of input g ( i ) ( x ) .
Formally,
[Equation]
The behaviour of MFNs can be understood by analyzing the frequency of intermediate layers.
Using the trigonometric identity
[Equation]
closed-form expressions for intermediate representations can be derived [7, 8].
Dropping the bias terms b ( i ) for simplicity, the individual components of the intermediate representation z ( i ) is equal to a weighted sum of an exponential number of sine terms [7]:
[Equation]
where n = ( n 0 , · · · , n i − 1 , n i ) is a tuple of indices with n j ∈ { 1 , . .
, d h } and s j ∈ {− 1 , 1 } .
Each term in the summation comprises an amplitude, frequency and phase shift, given by
[Equation]
B ACON .
This analysis of the frequency content of an MFN shows that the bandwidth of z ( i ) ( x ) is the sum of the input bandwidths.
This was leveraged in B ACON [8] to band-limit the representation of each individual layer.
Additionally, an output was created for each layer to produce intermediate reconstructions; i.e.,
[Equation]
where W ( i ) out ∈ R d out × d h and b ( i ) out ∈ R d out .
However, to encourage y ( i ) to preserve its signal representation at scale i , one requires extra scale-specific losses during training, without which, subsequent training will corrupt the lower resolution representation.
Motivating Example.
Although B ACON does provide control over spectral bandwidth, when adopted in the context of a coarse-to-fine optimization procedure, it fails to preserve and carry the learned coarse-scale representations when one moves to the next finer-scale stage of the optimization. Consequently, at each round, one must re-optimize the entire representations at all previous scales, hindering efficient coarse-to-fine optimization.
As an example, in Fig.
3 we examine the representation learned by B ACON within a coarse-to-fine training strategy, where successive output layers are trained to fit an image at finer scales.
We apply a staged training procedure such that a loss is applied to outputs at each scale in separate rounds of optimization.
Unfortunately, when fitting finer scales B ACON completely forgets the learned representations of coarse scale outputs from previous optimization rounds (see supplement).
For a fairer comparison, given any round, we keep all scale-specific losses, but let only the linear output layers for other scales to be updated.
Still, the results of B ACON (Fig. 3, top row) show that the representations at coarser scales yet become corrupted while training at finer scales.
In fact, reconstructions at different scales are highly coupled in B ACON , since new layers can result in new frequencies anywhere within the band limit.
In the following, we modify the architecture and its initialization and training scheme (Fig. 1) to make coordinate networks applicable to coarse-to-fine multiscale reconstruction.
Residual Multiplicative Filter Networks
From Eq.
4, we can express frequencies in layer i in terms of frequencies in previous layers with the recursion
[Equation]
where ˜ n and ˜ s are formed from n and s by removing last entry, e.g., ˜ n = ( n 0 , · · · , n i − 1 )
=
n <i .
Thus z ( i ) n i contains frequencies from the previous layer, shifted by ω ( i ) n i and − ω ( i ) n i , where ω ( i ) n i is the frequency of hidden unit n i in layer i .
More specifically, assuming that frequencies in the previous layer, i.e., ω ( i − 1) (˜ n , ˜ s ) , are band-limited to [ − B i − 1 , + B i − 1 ] , those in z ( i ) lie in the union of several of its shifted clones.
From this perspective, the frequencies at layer i , ω ( i ) , characterize two important concepts: 1) the growth in bandwidth from the previous layer, and 2) the extent of overlap between the previous layer frequency spectrum and each new individual clone.
For instance, given entries of ω ( i ) in [ − δ, + δ ] , the range of frequencies represented in z ( i ) expands to [ − B i − 1 − δ, B i − 1 + δ ] .
Also, if a specific ω ( i ) n i is set to a relatively large value of 2 B i − 1 this results in a clone with frequency support in the range [ B i − 1 , 3 B i − 1 ] , and it avoids overlap with the original bandwidth [ − B i − 1 , B i − 1 ] .
4.1 Frequency Initialization
Motivated by this interpretation of ω ( i ) , we propose a new initialization scheme which explicitly controls the overlap and expansion of frequencies in successive layers.
For simplicity, we describe this initialization for the 2D case, though it can be extended to higher dimensions straightforwardly.
First, we follow a procedure similar to B ACON [8] to initialize and fix frequencies of filters which allows limiting outputs to specified bandwidths.
Then, given the band-limited support [ − B i − 1 , + B i − 1 ] of frequencies in layer i − 1 , we initialize the 2D vectors ω ( i ) j as follows:
[Equation]
where R = { (0 , 1) , (1 , 0) , (1 , 1) , (1 , − 1) } , and U denotes either continuous or discrete uniform distributions.
Here, ω ( i ) j is composed of two terms.
The first term mainly specifies the direction ( r ( i ) j ) and the proportional amount ( λ 2 B i − 1 ) by which frequencies of layer i − 1 will be shifted.
The second term is a relatively small random perturbation to break the symmetry.
An intuitive illustration of these operations is provided in Fig.
In 2D, there are 8 different directions to expand, half of which are redundant since s i = ± 1 entails two shifts in opposite directions ± ( r j, 1 , r j, 2 ) .
Thus we only consider 4 possible values R .
Using this initialization z ( i ) is thus band-limited to ω ( i ) ∈ [ − (1 + λ 1 + λ 2 ) B i − 1 ,
(1 + λ 1 + λ 2 ) B i − 1 ] .
4.2 Residual Connections
Consider setting λ 2 = 2 + λ 1 during initialization.
In this case, by design, the frequency spectrum for sine terms in layer i would thus not overlap at all with that of layer i − 1 ; in at least one direction, the frequencies would either end up to left or right of interval [ − B i − 1 , B i − 1 ] .
Thus the output layer would include no sine terms with frequencies in the central interval [ − B i − 1 , B i − 1 ] .
To obtain the output of layer i , we sum it with the output of layer i − 1 .
That is,
[Equation]
This can be interpreted as a residual connection that ensures we keep an identical copy of frequencies from the previous layer; i.e., y ( i − 1) will fill the gap in the frequency space.
In practice, rather than the setting λ 2 = 2 + λ 1 , one can use other less strict combinations which allow for different amounts of overlap and expansion.
We found empirically setting λ 2 = 2 , λ 1 = 0 .
3 to be effective.
2D Image Fitting
We first evaluate each of the proposed modifications to band-limited networks in the context of 2D image fitting.
This task is useful for didactic purposes, allowing one to clearly see the benefits of skip connections and our new initialization scheme in the context of coarse-to-fine image reconstruction.
We adopt a staged training procedure and fit a series of band-limited coordinate networks to natural images provided by [1].
The dataset consists of 32 RGB images, each of size 512 × 512 .
We first downsample the images by a factor of 2 to obtain images of size 256 × 256 .
These images are fit during training using a Mean Squared Error (MSE) training objective, applied to the output at the scale corresponding to the stage of training.
As previously described, for the fairest comparison, the “Vanilla” B ACON baseline is optimized with losses at all scales, though, at the stage of a given scale, only the linear output layers of other scales can be updated.
Still, this keeps the optimization of the internal representations staged.
All networks use a four-layer architecture and are trained for 10 , 000 iterations, divided into chunks of 500 , 500 , 1000 , 8000 for optimizing output scales { 1 , 2 , 3 , 4 } , respectively.
The frequencies for all methods are initialized to have the same band-limit for consistency.
All methods were implemented in PyTorch [54] and trained on a single NVIDIA GeForce RTX 2080. [Table]
Figure 3 visualizes the outputs at each scale for B ACON , the residual MFN with the band-limited
B ACON initialization, and the residual MFN with our proposed initialization.
Since the coarse-to-fine training scheme can potentially disturb previously trained outputs when applying the loss at the finer scales, we visualize the differences between the Fourier spectra at each output scale after the corresponding training iterations and after the full training regiment (Fig. 3, right).
The proposed method shows the least disturbance after training.
For the entire image dataset, we report the Fourier spectrum mean absolute difference and standard deviation in Table 1.
Finally [Table] we examine all networks on the image generalization task by querying them on a finer grid of 512 × 512 and computing the Peak Signal-to-Noise Ratio (PSNR) between reconstructed signal and the original ground truth.
The average PSNR values over all images is summarized in the last column of Table 1.
We also benchmark SAPE [9] as a baseline for coarse-to-fine reconstruction, on the image generalization task.
SAPE gradually exposes higher frequencies as a function of time and space, while providing only a single reconstruction at a time, so only their full-scale reconstruction can be evaluated for comparison. While our approach improves interpretability of the reconstructions, by enabling reuse of previous learned representations, it achieves empirical performance competitive with baselines. Additional experiments on a textual dataset are provided in Appendix A.
Cryo-EM 3D Reconstruction
Single particle cryo-EM is an increasingly popular experimental technique to recover the 3D structure of macromolecular complexes, such as proteins and viruses.
Cryo-EM has gained widespread attention in the last decade with pioneering advances in hardware and data processing techniques.
This has enabled atomic or near-atomic resolution reconstruction of challenging structures [55].
Multiscale, coarse-to-fine reconstructions has long played a critical role in cryo-EM estimation, e.g., [20], making it an ideal testbed for our proposed methods.
Image Formation Model.
Based on a simplified model [56], the protein structure is represented as a function V : R 3 → R ≥ 0 , which maps points in 3D space to non-negative real values, expressing the
Coulomb potential induced by the atoms.
Images { I i } ni =1 are captured with a transmission electron microscope which are 2D orthogonal projections of randomly oriented potential maps.
Each image
I i has a corresponding latent pose (orientation)
R i ∈ SO (3) representing the 3D rotation of the molecule in the image.
In practice, to ensure sufficient contrast, the microscope is defocused when images are captured, resulting in a point-spread function (PSF)
g i .
To minimize radiation damage to delicate biological molecules, samples are exposed to low dosages of electrons resulting notoriously noisy images.
Taken together, the image formation model can be expressed as
[Equation] [Equation]
where the projection is, by convention, assumed to be along the z -direction after rotation, ⋆ corresponds to convolution and ϵ is Gaussian noise [14, 4].
The PSF can be estimated in a preprocessing step and is assumed here to be known; for more details see [57]. (A more detailed formulation for cryoEM reconstruction is given in Appendix B.)
[Figure] 3D Reconstruction Problem. The goal of a cryo-EM experiment is to estimate the unknown 3D potential map of the biomolecule present in the observed images { I i } ni =1 .
This task has some intrinsic challenges.
First, the images are extremely noisy (see Fig. 4c).
Second, the ground-truth poses are typically unknown, or only a rough estimate is available, so it is not possible to perform 3D reconstruction using the forward model of Eq.
directly.
Previous work [6, 4, 19] estimate pose parameters during reconstruction with various techniques including amortized inference, variational bounds or direct marginalization.
Here we use a stochastic gradient-based method, which optimizes the representation of the 3D structure V , parameterized by our proposed residual MFN, together with the unknown pose variables { R i } ni =1 .
The loss function is the squared error between the observed image I i and that generated by the model, i.e.,
[Equation]
where I ∗ i ( x, y ) is given above in Eq.
See Appendix D for more detail on the pose parameterization and optimization.
In practice, naively optimizing for both the poses and potential map suffers from poor local minima.
To avoid this, coarse-to-fine reconstruction approaches are commonly employed.
At early stages, low frequency content of the map is estimated as the SNR is higher so estimation is easier, which then allows for more reliable subsequent pose estimation as higher frequencies are included.
Previous approaches [4, 14] employ this technique, called Frequency Marching [20], as they model the map in the Fourier domain where individual frequencies are readily accessible.
In contrast, we address the reconstruction problem in a coarse-to-fine manner while expressing the potential map in real space.
6.1 Experiments
We evaluated our approach on two synthetic datasets which allow us to directly evaluate against ground truth structures.
The datasets are based on atomic models of 1OL5 ( https://www.rcsb. org/structure/1ol5 ) and 4AKE ( https://www.rcsb.org/structure/4ake ) from Protein
Data Bank (PDB).
We first convert atomic models to 64 × 64 × 64 discrete potential maps, with voxel sizes 1 Å and 1 . 2 Å respectively, using pdb2mrc.py module from EMAN2 [58].
Then we simulate the image formation model as in Equation 9, at 50 , 000 rotations uniformly sampled on SO (3) .
To generate each image, we rotate the map V , resample it on the original grid, project it along canonical z-direction and apply the PSF.
Zero-mean white Gaussian noise is added to these simulated clean images produce a realistic SNR of 10%.
The structures and sample images are visualized in Fig.
Experimental Settings.
We use a 3-layer network with 128 hidden units to represent the 3D structure.
The network outputs the structure in 3 hierarchical scales and, accordingly, the training process is divided into 3 stages, one for each scale.
We set the hyperparameters as λ 1 = 0 .
3 , λ 2 = 2 .and train the network for a total of 100 epochs.
For PDB1ol5, we spend respectively 15, 15, and 70 epochs optimizing the first to the third scales.
For PDB4ake, we find a better reconstruction if we dedicate more epochs on the first and second scales, namely 25 epochs on each of the first and second scales.
When computing the loss, images are band-limited to the corresponding scale using a Gaussian filter whose scale parameter is set to match the bandwidth of the current stage scale.
Mini-batches of images are used to update poses and weights of the network representing the structure.
For each mini-batch, we alternate between optimizing the structure and poses for 5 and 20 iterations, respectively.
We use Adam optimizer [59] with a learning rate 0 .
001 for the network parameters and a higher learning rate of 0 .
01 for pose parameters.
Each pose parameter is initialized to random value uniformly distribution over a geodesic distance of 45 to 90 degrees from the ground truth.
This can be considered as a very crude initial estimate for poses.
In Appendix E we consider even coarser initial pose values, showing that they achieve competitive qualitative and quantitative performance to those here.
Experiments are implemented in PyTorch [54], models are trained on NVIDIA TITAN
Xp GPUs, and ChimeraX [60] is used to visualize the 3D maps.
Results.
Qualitative comparison of reconstructions with those obtained by homogeneous refinement of cryoSPARC (as the baseline) are illustrated in Fig.
In the middle, we observe that, from low to high scales, increasingly more structural features are resolved.
The low-scale reconstructions capture only the global and low-resolution features, whereas the full-scale ones contain fine-grained details.
Our approach is able to well-resolve some peripheral parts of the structure, especially for the more complex structure of PDB4ake.
Also there are some structural features learned by our reconstruction yet totally missed by cryoSPARC.
All in all, our final reconstruction visually looks competitive or better when compared with the baseline.
As shown on the right side of Fig.
5, to quantitatively evaluate reconstructions, we compute the
Fourier Shell Correlations (FSCs) [61] between the reconstructed map and the ground-truth.
FSC is the gold-standard metric in cryo-EM area measuring the normalised cross-correlation coefficient between two 3D volumes in Fourier domain along shells of increasing radius (see Appendix C for more background on FSC).
We measure FSC between ground-truth and reconstructions at all scales, following the convention of the threshold 0 .
5 to compute final resolution, as ground-truth is known a priori in our experiments.
We also evaluate how well the poses are optimized and converge to the ground-truth.
To this end, the histogram showing number of images with various geodesic distances between estimated and ground-truth poses are visualized as in Fig.
For both datasets, there is a concentration of images on left hand side corresponding to slight error in estimation of pose parameters.
Thus, for most of the images, predicted pose parameters are within few degrees of the corresponding ground truth, verifying that poses are optimized well along with the reconstruction.
Furthermore, per dataset, we select two sample images, one with almost the correct pose estimate, and the other with one that is far from ground truth.
Then we demonstrate their reconstruction at all scales as in Fig.
Finally, we conduct an experiment in which we supervise the network only at the full-scale reconstruction throughout the entire training.
The output structures are depicted in the right panel of Fig. 6, realizing a situation where poses are trapped in local minima, and we thus observe no substantial improvement in the structure for many iterations.
Conclusions, Limitations and Future Work
In this paper, we build on Band-limited Coordinate Network (B ACON ) and propose a new architecture and training strategy providing fine-grained control over spectral characteristics of network outputs, yielding a multiscale representation suitable for coarse-to-fine reconstruction.
The combination of skip connections and a consistent initialization scheme allows the proposed network to ensure lower levels of the network represent the coarser scales of the signal which are then reused at finer scales.
We empirically show that these modifications enable coordinate networks to fit natural images in coarse-to-fine fashion.
They are also shown to be effective for cryo-EM reconstruction.
Limitations.
We find that the proposed multi-scale optimization procedure is effective in avoiding local minima in the Cryo-EM reconstruction problem. We have yet to explore whether residual MFNs are also effective on other problems with local minima due to the existence of nuisance variables. Although beyond the scope of this paper, it is also possible that minor modifications of the architectural and hyper-parameter settings would improve the robustness of the optimization; e.g., increasing the number of layers for more fine-grained scales, or a more diverse inspection of λ hyperparameters.
Moreover, the cryo-EM experiments in this paper used synthetic data, to enable more thorough the evaluation of the method.
Real datasets are likely to require the use of other bases as well, e.g., Gabor , which are better able to represent the spatial variability of resolution in cryo-EM reconstructions.
We have also focused on homogeneous and rigid structures.
An advantages of our approach is that the efficient multiscale representation of the signal in the real domain opens promising doors for modeling flexible structures, a major open problem in cryo-EM [4, 53].
Societal Impact.
Cryo-EM for macromolecular structure determination is one of the foremost exciting areas in molecular biology with enormous social impact.
Cryo-EM tools were key to determining the structure of the SARS-COV2 spike protein, the discovery of its pre-fusion conformation, and the assessment of potential medical countermeasures.
Advances in this field enable the discovery and exploration of large protein complexes and viruses more generally [62].
On the other hand, we strongly condemn any usage of our proposed coordinate networks for generating malicious representations, improperly modifying signals, or spreading misinformation. Acknowledgments and Disclosure of Funding
We thank Ali Punjani for numerous valuable discussions and for use of the cryoSPARC software package.
This research was supported in part by the Province of Ontario, the Government of Canada, through NSERC, CIFAR, and the Canada First Research Excellence Fund for the Vision:
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes] , [No] , or [N/A] . You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example: • Did you include the license to the code and datasets? [Yes] See Section.
• Did you include the license to the code and datasets? [No] The code and the data are proprietary. • Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below. For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] The arguments reflect the actual contribution and scope.
(b) Did you describe the limitations of your work?
[Yes] We have described the limitations of our work in 7.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] We discussed both positive and negative potentials of our work in 7.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We have reviewed the guidelines and the paper comply with it.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Code for generating synthetic data is provided along with the implementation of networks and training schemes in supplementary.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Training details are provided in experimental settings sections.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] When relevant, error bars are provided as in Table (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Number of GPUs is stated individually for each experiment.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] All used softwares and codes cited.
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] The code for data generation and visualizations are provided in supplementary.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Multiplicative Filter
Networks
for Multiscale Reconstruction A 2D Image Fitting.
We provide additional natural image fitting results in Figs. 8 and 9. We also include the case of training “unfair” vanilla B ACON in a staged fashion which, unlike the proposed fair variant, does not continue to update the linear output layer during the subsequent optimization stages of each scale. For the generalization task, we visualize the final reconstructions of SAPE [9], B ACON [8], and our proposed method and compare it with ground truth as in Fig.
We further investigate the proposed modifications on a new dataset of 32 RGB text images, each of size 512 × 512 , provided by [1].
Similar to natural images, we first obtain downsampled 256 ×images.
For all networks, we use the same architecture as above, in the main body of the paper.
The models are fit to images using Mean Squared Error (MSE) as the objective loss.
The only difference is that we also visualize learned reconstructions for the unfair case of staged training of vanilla B ACON .
We keep all the experimental settings same, including the number of iterations dedicated to each stage and the learning rates.
Further, the frequencies for all networks are initialized to have the same band-limit.
In Fig.
11, the outputs at each scale for both unfair and fair staged B ACON , the residual MFN with the band-limited B ACON initialization, and the residual MFN with our proposed initialization, are visualized.
The differences between the Fourier spectra at each output scale after the corresponding training iterations and after the full training regiment are illustrated too (Fig. 11, right).
We provide additional examples of fitting text images in Fig.
12, comparing reconstruction at coarser scales of two and three.
We also report the Fourier spectra mean absolute difference and standard deviation in
Table 2 for the text dataset.
Finally all networks are examined on the same image generalization task and the Peak Signal-to-Noise Ratio (PSNR) between reconstructed signal and the original ground truth is computed.
We include SAPE [9] as a baseline for coarse-to-fine reconstruction.
Average
PSNR values over all images is summarized in the last column of Table 2.
Similarly, visualization of the final reconstruction on this task is provided in Fig.
[Table]
B Cryo-EM
Single particle cryo-EM is a revolutionary imaging technique used by structural biologists to discover the 3D structure of macromolecular complexes, including proteins and viruses.
Recovering 3D structure is a major step towards understanding how these miniature biological machines perform actions in our body cells.
During sample preparation for a cryo-EM experiment, one first obtains a purified sample in solution, containing many instances of a specimen.
This sample is then plunged into a cryogenic liquid, like liquid ethane, which freezes the molecules at random orientations.
Once frozen in vitreous ice, the sample is loaded into a transmission electron microscope and exposed to parallel electron beams which are recorded underneath by special detectors.
The majority of these electrons pass through the structure without interaction, causing shot noise.
The rest are scattered either in an elastic or inelastic fashion.
Elastic scattering positively contributes to image formation by introducing a relatively weak phase contrast while the inelastic scattering damages the structure leading to garbage particles.
The image formation in cryo-EM can be well-approximated using the weak-phase object model [56].
This model assumes the 3D structure is a density map V : R 3 → R ≥ 0 represented under a canonical orientation.
Ideally, without any corruption, as depicted in Fig.
14, one can formalize clean images { I ∗ i } ni =1 as 2D orthogonal projections of the map under random orientations R i ∈ SO
[Equation]
Here, P is a linear operator performing the orthogonal projection and V R i corresponds to a structure rotated by R i .
However, in practice, images are intentionally captured under defocus in order to improve the contrast.
This is modeled by convolving the clean images with an image-specific pointspread functions (PSF) g
i .
We express all sources of noise with an additional term and, taken together, the observed image can be formally described as
[Equation]
Here, ⋆ denotes convolution, ϵ is the additive noise, and V R i ( x, y ) corresponds to clean (noiseless) projection of the structure rotated by R i , as in Fig.
Statistically, as several factors contribute to the noise, it is common practice to assume ϵ is white (or colored) Gaussian noise.
C Fourier Shell Correlation (FSC) and Resolution
Fourier shell Correlation (FSC) is the standard way to measure the resolution of a 3D density map.
In a nutshell is is the correlation between the Fourier coefficients of two aligned maps, as a function of frequency [61].
Let S r be a spherical shell with radius r centered at the origin of the Fourier domain (i.e., containing 3D frequencies whose wavelengths are close to 1 /r ).
Then, given two maps in the Fourier domain, F 1 and F 2 , FSC ( r ) is defined as the normalized cross-correlation between the
Fourier coefficients of F 1 and F 2 within S r ; i.e.,
[Equation]
One can show that FSC ( r ) is real-valued, and ranges from -1 to 1.
Once correlations are computed for all shells, we plot FSC as a function of r .
In practice, one defines D/ 2 shells for maps with Dvoxels.
Typpically, one of F 1 or F 2 is the grount truth and the other is an estimated map.
Or the dataset is randomly partitioned into two halves, used to independently estimate the two maps F 1 or
F 2 .
At small r (i.e., low frequencies), where observed images have strong signal, maps are higher quality and FSC is close to 1.
With higher frequencies the maps become noisier and FSC therefore decreases (see the right panel of Fig. 5).
When a reconstruction is compared to a ground-truth map, one usually defines the resolution of the estimated map as the r for which FSC ( r ) = 0 .
Under mild assumptions this corresponds to a SNR of 1 [63].
It is worth noting that FSC also plays a key role in a more advanced version of frequency marching.
In experimental settings, as the true map is not available, FSC is computed between two independent reconstructions.
During frequency marching, as we increase the frequency support, the FSC curve is used to cross-validate the refinement to each reconstruction.
In fact, due to high frequency noise in the data, refinements are likely to overfit independently, yielding inconsistencies in the subsequent high frequency regime.
This is often observed as a sudden drop in the FSC curve.
And when this is detected, we adaptively limit the frequency support [64].
In our work, we achieve high quality reconstructions without this adaptive regularization.
On the other hand, we only have few scales for computational efficiency and this adaptive version is justified when more scales are dedicated.
We leave further investigations of this issue to the future work.
D Pose Optimization
Here, we provide the details of learning pose parameters, describing the 3D orientation of the molecule in an observed given particle image.
Let us denote current estimate of pose parameters by ˆ
R i ∈ SO (3) .
By applying a small perturbation δR i ∈ SO (3) , one can obtain a new estimate for poses by multiplication ( δR i )
ˆ R i .
Since, SO (3) is a group closed under multiplication, the perturbed pose ( δR i )
ˆ R i remains a valid rotation.
To find a better estimate, we iteratively find the optimal perturbation, and accordingly update pose estimates at each iteration.
To parameterize the perturbation δR i , similar to [21], we adopt the exponential map which is a mapping from the tangent space to the points on the SO (3) manifold.
One way to realize the exponential map is through axis-angle representation ( θ i , u i ) , where θ i is a scalar angle between zero and π and u i ∈ R 3 is a unit vector.
One can simply use Rodrigues formula to get back to the matrix form
[Equation]
(14) where [ u i ] × denotes the skew-symmetric 3 × 3 matrix of u i .
We rely on gradient-based optimization to solve for optimal perturbation parameters ( θ i , u i )
[Equation]
(15) where the image-specific loss is conditioned on the fixed potential map V and the corresponding image
I i .
We iteratively compute gradients with respect to the pose parameters and then update them.
Note that, after such an update, the resulting pose angle may no longer lie in the valid interval [0 , π ] , or the new axis ˆ u i may no longer have unit length.
For this reason we re-establish these constraints by truncating angles back to [0 , π ] , and normalizing u i to be a unit vector.
E Coarser Initialization of Poses
We conduct further experiments where poses are initialized to be uniformly distribution with angles within [0 , π ] of the ground-truth.
In this more challenging coarser initialization of pose, we dedicate more epochs to optimize the first and second scales, namely 30 epochs each.
Qualitatively, we obtain reconstructions on par with those in the main paper.
In terms of our FSC resolution metric, we achieve full-scale resolutions of 2.91A and 2.65A for PDB1ol5 and PDB4ake, respectively, as shown in 7.
F λ hyperparameters
The λ hyperparameters control the growth and disjointness of frequency support in the residual MFN (see Sec. 4.1).
Our first intuition was to avoid settings that create gaps in coverage of the Fourier domain, i.e., when λ 2 > 2 + λ 1 .
Beyond this constraint, we investigated a range of hyper-parameter settings in our experiments.
In all cases, the network fits the signal well; but we found that for λ 2 = 2, λ 1 = 0.3, the model yields a lower test loss (based on held out data for image fitting experiments).
For cryoEM experiments (3D regression) we found the same hyperparameters perform well, based on FSC with ground truth 3D density.
We therefore used the same hyperparameters for cryo-EM experiments with 2D images, without further hyper-parameter search.
G Videos
For further visualization of cryo-EM experimental results, we use Chimera [60] to create a set of videos, provided under videos folder, showing the results obtained by our method.
In videos pdb1ol5_all_scales.mp4 and pdb4ake_all_scales.mp4 , we visualize final 3D structures learned at different scales, respectively for pdb1ol5 and pdb4ake.
Reconstructions are rotated in the visualization to show the resolved structural features from several views.
We further visualize the coarse-to-fine evolution of the structure for each dataset during the entire training in videos pdb1ol5_evolution.mp4 and pdb4ake_evolution.mp4 .
In both cases, we start from a random structure.
We then obtain a coarse reconstruction and later stages capture more fine-grained (i.e. high-resolution) details of the estimated structure.
Finally, we use Chimera to fit the available atomic model of pdb4ake into the corresponding final reconstruction.
As shown in video model_fit.mp4 , various secondary structures, such as α -helices, are well aligned with the reconstructed 3D density map.