Programmable 3D snapshot microscopy with Fourier convolutional networks
Authors: Diptodip Deb
Zhenfei Jiao
Alex Bo-Yuan Chen
Misha Ahrens
Kaspar Podgorski
Srinivas C Turaga
Venue: ICLR 2022 Submitted (2022)
Keywords: computational microscopy
computational photography
computer vision
deep learning

Abstract: 3D snapshot microscopy enables fast volumetric imaging by capturing a 3D volume in a single 2D camera image and performing computational reconstruction. Fast volumetric imaging has a variety of biological applications such as whole brain imaging of rapid neural activity in larval zebrafish. The optimal microscope design for this optical 3D-to-2D encoding is both sample- and task-dependent, with no general solution known. Deep learning based decoders can be combined with a differentiable simulation of an optical encoder for end-to-end optimization of both the deep learning decoder and optical encoder. This technique has been used to engineer local optical encoders for other problems such as depth estimation, 3D particle localization, and lensless photography. However, 3D snapshot microscopy is known to require a highly non-local optical encoder which existing UNet-based decoders are not able to engineer. We show that a neural network architecture based on global kernel Fourier convolutional neural networks can efficiently decode information from multiple depths in a volume, globally encoded across a 3D snapshot image. We show in simulation that our proposed networks succeed in engineering and reconstructing optical encoders for 3D snapshot microscopy where the existing state-of-the-art UNet architecture fails. We also show that our networks outperform the state-of-the-art learned reconstruction algorithms for a computational photography dataset collected on a prototype lensless camera which also uses a highly non-local optical encoding.

Under review as a conference paper at ICLR 2022
PROGRAMMABLE 3D SNAPSHOT MICROSCOPY WITH FOURIER CONVOLUTIONAL NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
3D snapshot microscopy enables fast volumetric imaging by capturing 3D volume in a single 2D camera image and performing computational reconstruction. Fast volumetric imaging has variety of biological applications such as whole brain imaging of rapid neural activity in larval zebrafish. The optimal microscope design for this optical 3D-to-2D encoding is both sample- and task dusk-dependent, with no general solution known. Deep learning based decoders can be combined with a differentiable simulation an optical encoder for end-to-end optimization of both the deep learning decoder and optical encoder. This technique has been used engineer local optical encoders for other problems such as depth estimation, 3D particle localization, and lensless photography. However, 3D snapshot microscopy is known to require a highly non-local optical encoder which existing UNet-based decoders are not able to engineer. We show that neural network architecture based on global kernel Fourier convolutional neural networks can efficiently decode information from multiple depths in volume, globally encoded across 3D snapshot image. We show in simulation that our proposed networks succeed engineering and reconstructing optical encoders for 3D snapshot microscopy where the existing state-of-tee ofte-off-the-at UNet architecture fails. We also show that our networks outperform the state-of-the-art learned reconstruction algorithms for a computational photography dataset collected on a prototype lensless camera which also uses a highly non-local optical encoding.
INTRODUCTION
Volumetric imaging has been a valuable technique for measuring neural activity at single neuron resolution across the brain [2, 22, 36]. We are particularly motivated by the problem of imaging neu- ronal activity across the whole brain of transparent model organisms such as larval zebrafish (Danio rerio) at higher volume rates, enable richer studies of neural dynamics and signal propagation. Existing approaches to whole-brain imaging involve sequential acquisition of stacks of 2D images [21. typically resulting in 0.5Hz 2Hz volume capture rates, limited by the number of 2D depth plane acquisitions. However, since neural activity occurs at orders of magnitude faster timescales, faster volume imaging rates can provide understanding of neural computations with significantly more detail.
3D snapshot microscopy techniques can significantly speed up the imaging of 3D sample volume using non-local optical encodings to combine information corresponding to different depth planes in the volume into single 2D image, allowing for fast volume acquisition the frame rate of the camera sensor followed by computational decoding to reconstruct the 3D volume [5. 3, 4, 7, 13, 16, 18, 17, 20, 27, 26, 30, 38, 37, 33]. The optimal optical design for a 3D snapshot microscope is unknown in general [5], and likely depends on the specific sample structure. Newly developed programmable microscopes with up to 10 free parameters, e.g. pixels on a spatial light modulator (SLM), enable the implementation of rich space of microscope encodings, and present the possibility of direct optimization of snapshot microscope parameters specifically for particular classes of samples and imaging tasks.
Deep learning-basse based decoders combined with differentiable simulations of optical encoders have enabled end-to-end optimization of both the deep learning decoder and the optical encoder. This end- to-end optical engineering approach has recently been applied to depth from defocus [6. 10, 14, 35], lensless photography [13], or particle localization microscopy [23]. However, all such applications

---

Under review as a conference paper at ICLR 2022
either use far fewer optical parameters or operate relatively compact local encodings due to small filters in their deep networks [6. 23, 35, 10, 14, 34, 13]. There are two barriers to optimizing more parameters and producing global versus local encodings: (I) the typical convolutional networks with small filters used by previous methods for computational reconstruction have prior to decode locally, and therefore limit optical designs local encodings, unlike the global PSFs used in snapshot microscopy and (2) the optimization of optical parameters involves backpropagation through both the deep learning reconstruction and differentiable optical simulation, which is more expensive computationally for non-local optical encodings. We focus on two computational optics problems that demonstrate a superior deep network architecture that can successfully reconstruct from images encoded by non-local optical encoders and with orders of magnitude more optical parameters than previously attempted. To address the computational burden of simulating and optimizing non-local encoders, we also developed parallel, multi-GPU optical simulation framework.
Problem statement We define an optical encoder as MΦ parameterized by & and computational decoder as Re parameterized by 0. We wish to develop reconstruction network architecture for decoding non-local encodings in 2D images с produced by M4 and also to enable the end-to-end optimization of dp produce non-local encodings. We'll consider two concrete applications: (1) 3D snapshot microscopy using an SLM-based programmable microscope (Figure IA), where we wish to optimize both the reconstruction network R, and the optical encoding Ma based on the SLM parameters in order to image 3D samples v and reconstruct 3D volumes v, and (2) lensless photography (Figure 4), where we wish optimize only the reconstruction network Rg to produce images of natural scenes imaged by a lensless camera M implementing a fixed optical encoding.
1.1 OUR CONTRIBUTIONS
1. We developed a large-scale, parallel, multi-GPU differentiable wave optics simulation of a programmable microscope, based 4f optical model with phase mask (dp) implemented using a spatial light modulator (SLM), described further in Appendix A. SLMs &f can have over 10th optimizable parameters.
2. We collected a large dataset of high resolution 3D confocal volumes of zebrafish larvae for the purpose of end-to-end sample-specific optimization of the parameters of the pro- grammable 3D snapshot microscope.
3. We introduce an efficient new reconstruction network architecture for decoding from non- local optical encoders using very large global convolutions implemented via Fourier convo- lutions.
4. We show that our networks outperform the sfae-oo-thhe-rr deep networks for both volume reconstruction and microscope parameter optimization, and for image reconstruction from lensless photography [20].
5. Our method enables, for the first time, direct end-to-end optimization of highly non-local optical encoders in the space of SLM pixels with over 10 parameters. In simulation, we demonstrate the potential for significant improvements imaging resulting from sample- specific engineering of optical encoders.
1.2 PRIOR WORK
Neural network architectures for computational imaging have all used convolution layers with small filters. End-to-end optimization of optical encoders have largely been performed with bNet-based architectures [6, 10, 14, 35, 13], with one method using ResNet-based architectures [23]. Such optimization has always led to local optical encodings. End-to-end optimization has never been attempted for large-field of view 3D snapshot microscopy due to the difficulty of simulating and reconstructing from non-local encoders. However, small filter convolutional deep networks have been used in non-end-to-end manner to reconstruct volumes from 3D snapshot microscopes designed using microlens arrays [33, 34]. For photography and MRI, using deep lcarning combined with unrolling iterations of traditional deconvolution algorithms provides the benefits of fast amortized inference and higher quality reconstructions due to learning of structural priors [8. 9, 20, 34].
In our work, we demonstrate that convolution layers with large filters implemented efficiently in the Fourier domain enable the end-to-ent learning of highly non-local optical encoders for 3D snapshot microscopy. Large convolution filters have been shown to be helpful for other computer vision
2

---

Under review as a conference paper at ICLR 2022
applications such semantic segmentation and salient object detection [25, 21], and the Fourier domain parameterization of small filters has been described previously [28].
A pioneering strategy in 3D snapshot microscopy has been light field microscopy [16], which employs a microlens array at the microscope's ' image plane to create subimages encoding both amplitude and phase of light [16, I]. A variety of microlens-arran waweeen--rrra-baasaaea baaed light field microscopes have been used to perform whole-brait imaging [16, 27, 37, 36, 30, 26, 7, 12, 38]. [38] optimizes the placement of microlenses, but in an end-to-end manner. Despite variation in design, microlens- -ased microscopes have, to various degrees, three main limitations that can be improved: 1) blocking or scattering of light between microlenses, causing light inefficiency. 2) not making use of all pixels on the camera to encode a 3D sample, leading to inefficient compression and suboptimal reconstructions v, and 3) a fixed optical encoding scheme.
An alternative to using microlenses is to implement coded detection strategy using phase mask or diffuser to spread light broadly across the camera sensor [3, 4, 5, 13, 17, 18, 38). The designed phase masks can be implemented either by manufacturing custom optical element using programmable SLM [10, 14, 35, 13, 38, 23J. Using programmable element allows different microscope parameters to be used for different samples.
METHODS
A 3D snapshot microscopy
programmable microscope (optical encoder)
reconstruction network (computational decoder)
2D image c
3D sample
M
R.o
3D reconstruction
phase mask
Fourier convolution
B FourlerNet
F
W
convs
conv
RPN
multiscale Fourier convolution
C FourierUNet FUN
FF
W
convs
conv
crop Vix
upsample 2x and concatenate
W2
crop Vix
upsample 2x and concatenate
W
convs
Figure 1: Overview of our problem setup and our proposed network architectures. Top row (A) shows the problem of 3D snapshot microscopy, where we computationally reconstruct a 3D sample from 2D image. Middle row (B) shows our proposed FourierNet architecture, which includes a Fourier convolution layer that enables efficient computation of global features. Bottom row (C) shows an extension of our proposed architecture, the FourierUNet, which mimics the multiscale feature extraction of standard UNet efficiently and with global features using multiscale Fourier convolution.
We show our network architecture and an overview of autoencoder training both the microscope parameters dp and reconstruction network parameters in Figure 1A. The programmable microscope simulated by a differentiable implementation of wave-optics model of light propagation. We have selected programmable microscope design based on pupil-plane phase modulation with programmable spatial light modulator, for which imaging is apll-approximated by a computationally efficient convolution [1]]. A detailed description of our simulation is provided in Appendix A.1. For lensless photography, there is no optical simulation because the images have been collected on a real camera.

---

Under review as a conference paper at ICLR 2022
2.1 FOURIERNET AND FOURIERUNET ARCHITECTURES FOR RECONSTRUCTION FROM OPTICAL ENCODERS
Because images created by optical encoders can potentially encode signals from the incoming light field to any location in the camera image in a non-local manner, it is essential that a reconstruction network have global context. Existing multi-scale architectures such as the UNet [29] can achieve global context, but at the expense of many computation layers with small convolution kernels which we believe have local information bias which is inappropriate for computational optics. In this paper, we introduce relatively shallow architectures for computational imaging, which rely on convolution layers with very large filters, implemented efficiently in the Fourier domain [28].
NourierNet We propose simple three layer convolutional network architecture with very large global convolutions at the very first layer, followed by two standard local convolutional layers (Figure IB). We define a global convolution as a convolution with kernel size equal to the input image. Such a convolution achieves global context in a single step but is computationally expensive. Global convolutions are implemented more efficiently in the Fourier domain, yielding speed up of two orders of magnitude. Due to the use of Fourier convolutions to enable global context, we call our architecture FourierNet. In contrast to typical UNet which can contain many tens of convolution layers, the FourierNet is only three layers deep, which requires backpropagation through fewer layers compared to a typical UNet with the same receptive field.
FourierUNet We also propose a multi-scale variant of the FourierNet by bringing together elements of the multi-scale UNet and the single-scale FourierNet. Here, we take advantage of the fact that fown-sampling in image space corresponds to a simple cropping operation in the Fourier domain, resulting a band-limited computation of a feature map. We efficiently implement multi-scale global Fourier convolutions (Figure IC) to replace the incoding/"analyiiii use the standard sycoheing"synthhssss" pathway of the UNet to combine the multi-scale features into a single 3D volume reconstruction (Appendix A.3, A.4). Thus we can study whether multi-scale features or global context is more important for decoding non-local optical encoders.
pathway of a UNet. We then
Fourier domain convolutions It is well-known that large kernel convolutions can be implemented more efficiently in the Fourier domain [24, 19, 32]. A naive implementation of global convolution requires ((N²) operations, where N is number of pixels in both the image and the kernel. An alternative global convolution implementation is to Fourier transform the input x and convolution kemel w, perform element-wise multiplication in Fourier space, and finally inverse Fourier transform, requiring only O(N log NN operations [19, 32J. Following [28], we store and optimize the weights in Fourier space W. This over ave-paaramtterization costs 8x the memory of an equivalent real valued large filter but saves the computational cost of Fourier transforming the real-valued weights (Appendix A.3). Thus Fourier convolution is defined:
(I) For image and kernel sizes of 256 x 256 pixels, our implementation leads to nearly 500x speedup: standard PyTorch convolution takes 2860ms, while Fourier convolution takes 5.92ms TITAN X.
Re(F-1 (W @ [x}}}
Multi-scale Fourier domain convolutions is well-known [24] that downsampling corresponds to cropping in the Fourier domain. Thus the Fourier convolution can be extended to efficiently produce sulti-scale feature representations in one step (Figure IC). We define our multi-scale Fourier convolution
[Re(F-1 O crop; [c]}).... Re(F-1 (W,, c crop, [{}}]
(2)
where subseript denotes scale level (higher subscript indicates lower spatial scale/more cropping in Fourier space) and we precompute == JF {c} once.
2.2 PHYSICS-BASER AUTOENCODER FOR SIMULTANEOUS ENGINEERING OF MICROSCOPE ENCODER AND RECONSTRUCTION NETWORK DECODER
We describe the imaging process the following transformation from the 3D light intensity volume of the sample v to the 2D image formed on the camera c:
He M (v)
C = max ([H + /uEE), 0), 10,
(4)
where Mg denotes the microscope parameterized by a 2D phase mask, db This phase mask db describes the 3D-to-2D encoding of this microscope model completely. A Poisson distribution
4

---

Under review as a conference paper at ICLR 2022
with mean rate He describes the physies of photon detection at the camera, but sampling from this distribution not differentiable. We approximate the noise distribution with a rectified Gaussian. We include details on Mp in Appendix A.1 [11]. Jointly training reconstruction networks and microscope parameters involves image simulation, reconstruction, then gradient backpropagation to update the reconstruction network and mieroscope parameters. Details parallelization, planewise reconstruction networks, and planewise sparse gradients are provided in Appendix A.2.
Our loss function computes the normalized mean squared error (MSE) LANMSE between the high pass filtered sample and reconstruction weighted by the normalized MSE LNMSE between the original sample and reconstruction (to include low frequency content the objective). Formally, our loss function L(v, v) is defined:
L(v, v) = LHNMSE VV 4) BLNMSE (v, V)
(5)
LHNMSE (v, V) EE [(H(V) ((H(V))HHV))) H()))
..MMSEVV E [vV - + E(v²)
(6)
B(H(())))
where H.) denotes high pass filtering and B(-) denotes computing the mean. Both loss terms are normalized shown to reduce variance in L, which otherwise take on large magnitude values and cause training instability. For our experiments, we set the weight 3 for the LNMSE term to 0.1.
3 RESULTS
Larval Zebrafish Dataset High resolution volumes of transgenic larval zebrafish whole brains expressing nuclear-restricted GCaMP6 calcium indicator in all neurons were imaged using a confocal microscope. These images are representative of brain-wide activity imaging. We train on 58 different zebrafish volumes (which we augment heavily) and test on 10 held-out volumes. For all experiments, we downsample the high resolution confocal data to (1.0 um z, 1.625 μm y, 1.625 μm x). We created 4 datasets from these data called Type A, B, C, and D corresponding to imaging different fields of view. Full specifications for these datasets are in Table 4 (Appendix A.2). For Figure 2, we restrict the field of view to (200 μmm 416 μm 416 μm x) with a tall cylinder cutout of diameter 193 μm and height 200 μm and image with 256 x 256 pixels on the simulated camera sensor. We call this setting Type D. Figure 3 and Table 2 show our larger experiments with 512 x 512 simulated camera pixels, with a field of view (250 jmm z, 832 μm y. 832 um x) and sample types Type A, Type B, and Type C.
DiffuserCam Lensless Mirflickr Dataset We also test reconstruction performance on experimental computational photography data' from [20] (Figure 4). This is dataset constructed by displaying RGB color natural images from the MIRFlickr dataset on a monitor and then capturing diffused images by the DiffuserCam lensless camera. The dataset contains 24,000 pairs of DiffuserCam and ground truth images. The goal of the dataset is to Iearn to reconstruct the ground truth images from the diffused images. As [20], train on 23,000 paired diffused and ground truth images, and test on 999 held-out pairs of images.
3.1 FOURIERNETS OUTPERFORM UNETS FOR ENGINEERING NON-LOCAL OPTICAL ENCODERS
Table 1: Quality of reconstructed volumes after optimizing microscope parameters to image sample
Type D on 256 256 pixel camera (mean + s.e.m., 10)
Microscope
Reconstruction
LhnmSE 4
MS-SSIM t
PSNR 1
Time J (s)
FourierNet2D FourierNet2D
FourierNet3DD
0.6409 = 0.0213
0.955 + 0.004
34.78 + 0.88
0.71
FourierUNet3D
0.6325 = 0.0222
0.956 + 0.003
34.74 = 0.83
1.37
FourierNet2D
UNet3D
0.7659 + 0.0130
0.922 + 0.008
30.06 + 0.93
3.68
UNer2D
UNet3D
0.7120 + 0.0160
0.913 + 0.009
29.17 + 1.13
3.68
We compare optimizing microscope parameters with two neural networks: 1) using our FourierNet with 2D convolutions (FourierNet2I)) and 2) using a vanilla UNet with 2D convolutions (UNe12D).
Publicly available: https://wal 1ab. thub. 1o/lo////// less
lessl.ed lee ini ing/ http//daasee.htttt
5

---

Under review as a conference paper at ICLR 2022
FourierNet
UNet
light propagation
simulated camera image
ground truth
100 jmm
0.01
Xy max projection
um
3.00
reconstructed volume
max projection
xy max max.proppee projection
Figure 2: Comparing simulated camera images and reconstructions of Type D volume captured using our FourierNet (middle) versus UNet (right) microscopes. Top row shows simulated 256 x 256 pixel camera images; bottom right of camera image shows approximate acquisition time (given reasonable number of simulated photons per pixel, i.e. SNR). truth has no corresponding camera image, because the 3D volume is imaged directly via confocal microscopy (acquisition time in view). Colored arrows in left column show projection axis for each row. White arrows show individual neurons clearly visible for FourierNet, but for UNet.
Training is two step process in which planewise reconstruction network with fewer parameters and 2D convolutions used during optimization of cb, then once & is fixed a more powerful reconstruction network with 3D convolutions is used. We find this scheme achieves better reconstruction quality because the reconstruction network does not need constantly adapt changing optical encoding (Appendix A.3). We train these microscopes and reconstruction networks on Type D, where samples tall cylindrical cutouts of zebrafish with diameter 193 um and height 200 μum. Samples are imaged on camera with 256 x 256 pixels (Figure 2). FourierNer2D has 2 convolution layers (with 99.8% of kernel parameters in the initial Fourier convolution layer), while UNer2D has convolution layers (kernel parameters approximately uniformly distributed per layer). FourierNet2D and UNet2D both have ~ x 107 parameters; FourierNet3D has x 107 parameters vs. 10 for UNet3D. Architecture details are in Appendix A.3, A.4.
Simulated camera images (Figure 2) show that the UNet microscope does not make sufficient use of camera pixels, producing only single view of the sample. We believe this is due to local information prior in the small kernels of UNets. The FourierNet microscope uses more camera pixels and performs better than the UNet microscope for reconstruction (Figure 2, quantified in Table I).

---

Under review as a conference paper at ICLR 2022
Table 2: Sample specific microscope parameter optimization across 3 different zebrafish sample types imaged with 512 x 512 pixel camera (mean PSNR (top), MS-SSIN (bottom) = s.c.m., it 10). Green shaded regions show the regions of interest for each sample type, cropped from the full volume.
microscope parameters optimized for Type C
Type A
Type B
tested on
Type A
49.75 + 1.35
46.01 + 1.33
42.63 = 1.14
0.998 + 0.000
0.996 + 0.001
0.992 = 0.002
Type B
35,53 +: 1.41
37.28 + 0.96
35.34 + 1.16
0.965 + 0.004
0.972 + 0.003
0.967 0.003
Type C
30.87 + 1.15
31.48 + 0.93
33.79 + 0.90
0.912 + 0.007
0.920 + 0.006
0.935 = 0.006
Times in Table are for only the forward pass on a single GPU; one training iteration on 8 GPUs takes ~0.4 seconds FourierNet3D and ~0.8 seconds for UNet3D (Appendix A.3). Both reconstruction networks must reconstruct from images that have a compressed encoding of 3D information, but the FourierNet2D is clearly more effective than the UNet2D at optimizing this encoding.
3.2 FOURIERNETS OUTPERFORM UNETS FOR 3D SNAPSHOT MICROSCOPY VOLUME RECONSTRUCTION
We can determine which architecture is better for volume reconstruction by choosing fixed microscope parameters and varying the architecture. In Table 1, we compare results using a FourierNet with 3D convolutions (FourierNet3D), FourierUNet with 3D convolutions (FourierUNet3DD, and vanilla UNet with 3D convolutions (UNet3D). Architecture details are Appendix A.3, A.4.
Reconstruction results in Table 1 compare normalized MSE LHNMSE between the high pass filtered sample and high pass filtered reconstruction, the multiscale structural similarity MS SSIM between the sample and reconstruction, and finally the peak signal-to-noise ratio PSNR. We also visualize reconstruction results for a volume in the head of a zebrafish in Figure 2. The best reconstruction networks (equally) FourierNet3D and FourierUNet3D. The UNet3D reconstruction networks (using either microscope) fall significantly behind the FourienNer3D/FFourrrtUNNNNN networks in all metrics. 3.3
reconstruction
ENGINEERED OPTICAL ENCODING DEPENDS ON SAMPLE STRUCTURE
To explore the effect of sample structure and size on optimized & and the resulting reconstruction performance, we optimized dps for three different sample types: Type A, samples with short cylinder cutout of 386 μm diameter and 25 μm height. 2) Type B, samples with a tall cylinder cutout of 386 μm diameter and 250 μm height, and 3) Type C, samples without any cutout of dimension (250 μm x 832 μm y x 832 μm x) (Table 2). All sample types were imaged with 512 x 512 pixels on the simulated camera. We then tested reconstruction performance on all combinations of optimized microscopes and samples, shown in Table 2 and visualized in Figure 3 for Type B. We include architecture details in Appendix A.3, A.5.
We see in Table 2 that for all sample types, highest performance is achieved using the microscope optimized for that particular sample. These results show that our optimization process produces optical encoders that are more optimal for specific type of sample over others.
3.4 FOURIERNETS OUTPERFORM STATE-OF- OF-THE-ARE ART FOR RECONSTRUCTING NATURAL IMAGES CAPTURED BY DIFFUSERCAM LENSLESS CAMERA
We compare our FourierNet architecture to the best learned method from [20] using an unrolled ADMM and denoising UNet, well as to vanilla UNet from [20]. Architecture details are in Appendix A.3, A.6. We can see that both of our methods visually outperform the methods from [20] in Figure 4. Table shows that methods outperform the others on all metrics. In [20], combined loss using MSE and learned perceptual loss (LPIPS) suggested to improve visual reconstruction quality. Table shows that training FourierNet on both MSE and LPIPS results in the lowest
7

---

Under review as a conference paper at ICLR 2022
ground truth (Type
optimized for Type
optimized for Type
optimized for Type
jum
simulated camera image
pm
Figure 3: Reconstructed volumes resulting from imaging Type B samples by microscopes optimized for Type A, B, C. Imaging Type B samples with microscope parameters optimized for Type samples yields the best reconstructions. Top to bottom: xy max projection, projection, simulated 512 x 512 camera image.
Table 3: Quality of natural image reconstruction on the DiffuserCam Lensless Mirflickr Dataset
(mean + s.c.m. 999). Superscripts denote loss function: MSE, MSE+LPIPS.
Method
MSE (x10 LPIPS
MS-SSIM 1
PSNR
Time J (ms)
FourierNet
0.39 + 0.007
0.20 + 0.00
0.882 + 0.001
24.8 + 0.09
37.54
FourierNet
0.54 + 0.010
0.16 + 0.00
0,868 + 0.001
23.4 + 0.09
37.54
FourierUNet!
0,43 + 0.009
0.22 0.00
0.875 0.001
24.5 + 0.09
204.69
Le-ADMM-U [20]
0.75 + 0.021
0.19 + 0.00
0.865 + 0.002
22.1 + 0.09
59.01
UNet [20]
1.68 0.060
0.24 0.00
0.818 0.002
19.2 + 0.11
06.97
LPIPS loss across all methods, but does not offer better visual performance than our MSE-only model seen in Figure 4. FourierUNet does show improved results FourierNet.
3.5 GLOBAL RECEPTIVE FIELD IS MORE IMPORTANT THAN MULTISCALE FEATURES
UNets effective because: (1) features computed at multiple scales and (2) large receptive fields are achieved few layers. FourierUNets allowed us to decouple these two explanations because the receptive field is global in single layer. We see both our microscopy dataset (which does not have multiscale structure) and the lensless photography dataset (which does have multiscale structure) that the FourierUNet does not improve upon the FourierNet. Thus see that more important for decoding from non-local optical encoders to have global receptive field than multi-scale features.
4 DISCUSSION
Summary We have presented FourierNets for decoding from highly non-local optical encodings. We have demonstrated that FourierNets enable end-to-end optimization of non-local optical encodings
8

---

Under review as a conference paper at ICLR 2022
DiffuerCan
Imagee
FourierNee (ours) FounerNee lours) ground 0/:S LMMPS)
(MSE)
DASE LPIPS)
Le-
(NSE LPMPS)
UNet
Figure 4: Comparisons of our method (second and third rows) to state-of-the-art learned reconstruction methods lensless diffused images of natural scenes. Regions labeled ①1 indicate missing details, either resolution or textures in backgrounds. Regions labeled 2 indicate hallucinated textures. Note that the previous state-of-the-ar solutions [20] exhibit both issues more often compared to our models.
where UNets fail for simulated 3D snapshot microscopy, and also that FourierNets beat state-of-the-art decoders for lensless photography. FourierNets provide fast amortized inference and reconstruction times many orders of magnitude faster than traditional iterative reconstruction algorithms [20] and higher quality reconstructions due to effective learning of sample structural priors. Generally, our global kernel architecture using Fourier convolutions could be applied to other problems where global integration of features is necessary (though we have focused on computational microscopy and photography where such global mixing of information dominates).
Limitations Our engineered optical encoders have not yet been experimentally tested on pro- grammable microscope. This will require measuring and accounting for system aberrations and calibration to implement the engineered optical encodings. However, the parameters of even an imperfectly implemented optical encoder can be perfectly measured and used to train reconstruc- tion networks to near optimal performance. Our claim that sample type-specific optical encoders perform better on their respective samples has only been evaluated simulation. Further, imaging a different sample class (i.e. with different spatiotemporal statistics) will require retraining at least the reconstruction network, and ideally also the optical encoder. Our autoencoder training style requires simulation of imaging with gradients, which be memory expensive and limits the total field of view. Our training approach requires multi-GPU parallelization (described in Appendix A.2), and in practice we optimized microscopes using 8 Quadro RTX 8000 GPUs RTX 2080 GPUs for the larger and smaller experiments, respectively. Training a microscope and decoder requires at least two weeks of total training. Our Fourier convolution layers require fixed input size to truly be global, though this could be addressed by resampling.
Reproducibility We train on 8 Quadro RTX 8000 GPUs for the largest experiments, and have described our pre-pre pre-processing, training, and testing procedures in Appendix A.2, A.3, A.4. We will also release our simulation software, training scripts, and our datasets upon publication.

---

Under review as a conference paper at ICLR 2022
REFERENCES
[1] E.H. Adelson and J.Y.A. Wang. Single lens stereo with a plenoptic camera. IEEE Transactions on Pattern Analysis and Machine Intelligence, 14(2):99-106, Feb 1992. ISSN 1939-3539. doi: 10.1 10.1109/34.12173.
[2] Misha B. Ahrens, Michael B. Orger, Drew N. Robson, Jennifer M. Li, and Philipp J. Keller. Whole-brain functional imaging at cellular resolution using light-sheet microscopy. Nature Methods, 10(55):413-4000 May 2013. ISSN 1548-7105,
doi: 10.1038/mmeth.2444
[3] Nick Antipa, Grace Kuo, Reinhard Heckel, Ben Mildenhall, Emrah Bostan, Ren Ng. and Laura Waller. Diffusercam: lensless single-exposure 3d imaging. Optica, 5(1):1-9, Jan 2018. ISSN 2334-2536. doi: 10. 00/16///FT1CA...000000
[4] M. S. Asif, A. Ayremlou, A. Veeraraghavan, R. Baraniuk, and A. Sankaranarayanan. Flatcam: Replacing lenses with masks and computation. In 2015 IEEE International Conference Computer Vision Workshop (ICCVW), pp. 663-666, Dec 2015. doi: 10.1 (0.1109/CCYW.2215..2.
[5] Michael Broxton. Volume Reconstruction and Resolution Limits for Three Dimensional Snapshot Microscopy. PhD thesis, Stanford University, Aug 2017.
[6] Julie Chang and Gordon Wetzstein. Deep optics for monocular depth estimation and 3d object detection. arXiv:1904 08601 [cs, eess), Apr 2019. URL Stthttp:/ /ttp:///aaiiiii wwtpeiiiiirrrroooooooooeee org/abag /1904 08601. aiv: 1904.08601.
[7] Lin Cong, Zeguan Wang, Yuming Chai, Wei Hang. Chunfeng Shang, Wenbin Yang. Lu Bai, Jiulin Du, Kai Wang, and Quan Wen. Rapid whole brain imaging of neural activity in freely behaving larval zebrafish (danio rerio). eLife, 6:e28158, Sep 2017. ISSN 2050-084X. doi: 10.7554/eLife.28155
[8] Steven Diamond, Vincent Sitzmann, Felix Heide, and Gordon Wetzstein. Unrolled optimization with deep priors. arXiv e-prints, 770::arXiv1105.000 70000, May 2017.
[9] Jiangxin Dong, Stefan Roth, and Bernt Schiele. Deep wiener deconvolution: Wiener meets deep learning for image deblurring. Advances in Neural Information Processing Systems, 33: 1048-1059, 2020.
[10] Xiong Dun, Xiong Dun, Hayato Ikoma, Gordon Wetzstein, Zhanshan Wang. Zhanshan Wang, Xinbin Cheng, Xinbin Cheng, Xinbin Cheng, Yifan Peng, and et al. Learned rotationally symmetric diffractive achromat for full-spectrum computational imaging. Optica, 7(8):913-922. Aug 2020. ISSN 2334-2536. doi: 10. 10.1364/OPTICA.39441
[11] Goodman. Introduction Fourier Optics. Macmillan Learning, 4 edition, 2017. ISBN 978-1-319-1191644
[12] Logan Grosenick, Michael Broxton, Christina K. Kim, Conor Liston, Ben Poole, Samuel Yang, Aaron Andalman, Edward Scharff, Noy Cohen, Ofer Yizhar, and al. Identification of cellular-activity dynamics across large tissue volumes in the mammalian brain. bioRxiv, pp. 132688, May 2017. doi: 10. 1011/31183.888
[13] Y. Hua, S. Nakamura, M. S. Asif, and A. C. Sankaranarayanan. Sweepcam - depth-awarr aware lensless imaging using programmable masks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 4267):160-1161,, Jul 2020. ISSN 1939-3539. doi: 10. 2000 2986784.
[14] Hayato Ikoma, Cindy M. Nguyen, Christopher A. Metzler, Yifan Peng, and Gordon Wetzstein. Depth from defocus with leameed optics for imaging and occlusion- aware depth estimation. In 2 IEEE International Conference on Computational Photography (ICCP), pp. 1-12, May 2021. doi: 10. 1099/CCC515881 20119..662.....66..
[15] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. 1XX2:05002.0167 [cs], Mar 2015. URL http: : /arxiv. org/abs/155t 1102. 031 67, arviv: 1502.03167.
10

---

Under review as a conference paper at ICLR 2022
[16] Mare Levoy, Ren Ng, Andrew Adams, Matthew Footer, and Mark Horowitz. Light field microscopy. In ACM SIGGRAPH 2006 Papers, SIGGRAPH '06, pp. 924-934. Association for Computing Machinery, Jul 2006. ISBN 978-1-59993-334-5. doi: 10.1145/119952.1141997
URL hww.tttp//// /doioo..... orgg 1tt.//..111111111111
[17] Fanglin Linda Liu, Grace Kuo, Nick Antipa, Kyrollos Yanny, and Laura Waller. Fourier diffuserscope: single-shot 3d fourier light field microscopy with a diffuser. Optics Express, 28
(20):28969, Sep 2020. ISSN 1094-4087. doi: 10. 10.136400E.400876.
[18] Fanglin Linda Liu, Vaishnavi Madhavan, Nick Antipa, Grace Kuo, Saul Kato, and Laura Waller. Single-shot 3d fluorescence microscopy with fourier diffusercam. In Biophoton- ics Congress: Optics in the Life Sciences Congress 20/9 (BODA,BRAIN, BRAIN,NIM, (MP) (2019), paper NS2B.3, pp. NS2B.3. Optical Society of America, Apr 2019. doi: 10.1364) NTM.2019.NS2B.3. URL https: / /www. www.ussapplioooooo www.niiissiing.mmo org/ abstract. cfm? Uz1=NTM--0199-NS
[19] Michael Mathieu, Mikael Henaff, and Yann LeCun, Fast training of convolutional networks through ffts. arXiv:/312.5851 [cs], Mar 2014. URL 5851. arXiv: 1312.5851.
[20] Kristina Monakhova, Joshua Yurtsever, Grace Kuo, Nick Antipa, Kyrollos Yanny, and Laura Waller. Learned reconstructions for practical mask-based lensless imaging. Optics Express, 27 220):28075-28000, Sep 2019. ISSN 1094-4087. doi: 1..13440OE.7..02807
[21] Nan Mu, Xin Xu, and Xiaolong Zhang. Salient object detection in low contrast images via global convolution and boundary refinement. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pp. 743-751. IEEE, Jun 2019. ISBN 978-1-72812-506-0
doi: 10. 1109/////..220900000222 URL https: ieee orgg document 9025357/
[22] Yu Mu, Davis V. Bennett, Mikail Rubinov, Sujatha Narayan, Chao-Tsung Yang, Masashi Tanimoto, Brett D. Mensh, Loren L. Looger, and Misha Ahrens. Glia accumulate evidence that actions are futile and suppress unsuccessful behavior. Cell, (78(1):27-43.e) Jun 2019. ISSN 0092-8674. doi: 10. (0016/..cel.2019.05..oo
[23] Elias Nehme, Daniel Freedman, Racheli Gordon, Boris Ferdman, Lucien E. Weiss, Onit Alalouf, Tal Naor, Reut Orange, Tomer Michaeli, and Yoav Shechtman. Deepstorm3d: dense 3d localization microscopy and psf design by deep learning. Nature Methods, 17(77):734-440, Jul
2020. ISSN 1548-7105. doi: 1911388/41122-020--220
[24] Alan V. Oppenheim and Alan S. Willsky. Signals and Systems: Pearson New International Edition. Pearson Education Limited, Jul 2013. ISBN 978-1-292-025990 Google-Books-ID: UΩAAQQACAAI..
[25] Chao Peng, Xiangyu Zhang, Gang Yu, Guiming Luo, and Jian Sun. Large kernel matters improve semantic segmentation by global convolutional network. In 20/7 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1743-1751, IEEE, Jul 2017. ISBN 978-1-5388-0457-1. doi: 10. 10.100//VVR.2017..10 URL http:///eeeeeepeeeeeeeeee ieee. (te.ddccmetttt999950
[26) Robert Prevedel, Young-Gyu Yoon, Maximilian Hoffmann, Nikita Pak, Gordon Wetzstein, Saul Kato, Tina Schrödel, Ramesh Raskar, Manuel Zimmer, Edward S. Boyden, and al. Simultaneous whole-animal 3d imaging of neuronal activity using light-field microscopy. Nature Methods, 11(77):727-73, Jul 2014. ISSN 1548-7105. doi: 10. 00.1038/mmeth.966
[27] Nicolas C. Pégard, Hsiou-Yuan Liu, Nick Antipa, Maximillian Gerlock, Hillel Adesnik, and Laura Waller. Compressive light-field microscopy for 3d neural activity recording. Optica, 3(5): 517-524, May 2016. ISSN 2334-2536. doi: 10. Stt./OPTICA..0000555
[28] Oren Rippel, Jasper Snoek, and Ryan P Adams. Spectral representations for convolu- tional neural networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Asso- ciates, Inc., 2015. URL https: httpee////cpcpppeeeeeeeeeeeeeeee 43577 wtwpeeeiiiinnnnnnnnineeeeeeeeeee nwwrips..r.ccmm http:/pp..ccppppp////33///////// .ddf.

---

Under review as a conference paper at ICLR 2022
[29] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net:
Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015, Lecture Notes in Computer Science, pp. 234-241. Springer International Publishing, 2015. ISBN 978-3-319-24574-4 doi: 10.1007/978-3-3111 2-319-2457444_28.
[30] Oliver Skocek, Tobias Nöbauer, Lukas Weilguny, Francisca Martínez Traub, Chuying Naomi Xia, Maxim I. Molodtsov, Abhinav Grama, Masahito Yamagata, Daniel Aharoni, David D. Cox, and et al. High-speed volumetric imaging neuronal activity in freely moving rodents. Nature doi:
Methods, (5(66))429-432 Jun 2018. ISSN 1548-7105.
[31] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv:1607 08022 [cs]. Nov 2017. URL http://arxiv: erg/abs / 607. 08022. arXiv: 1607.08022.
[32] Nicolas Vasilache, Jeff Johnson, Michael Mathieu, Soumith Chintala, Serkan Piantino, and Yann LeCun. Fast convolutional nets with fbfft: A gpu performance evaluation. arXiv:1412.7588 [cs), Apr 2015. URL http://ttaaa /twpvv//a/vvv//aa//////ee./aiav. org/ abs 1412. 7580. arXiv: 1412.7580.
[33] Nils Wagner, Fynn Beuttenmueller, Nils Norlin, Jakob Gierten, Juan Carlos Boffi, Joachim Wittbrodt, Martin Weigert, Lars Hufnagel, Robert Prevedel, and Anna Kreshuk, Deep learning- enhanced light-fieli field imaging with continuous validation. Nature Methods, 18(5)5557-563, May
2021. ISSN 1548-7105. doi: 00.1038/41192--21---1 136--0
[34] Zhaoqiang Wang, Lanxin Zhu, Hao Zhang, Guo Li, Chengqiang Yi, Yi Li. Yicong Yang. Yichen Ding, Mei Zhen, Shangbang Gao, and et al. Real-time volumetric reconstruction of biological dynamics with light-field microscopy and deep learning. Nature Methods, pp. 1-6, Feb 2021.
ISSN 1548-7105. doi: 10.1038/s41592-111 -2-01108-.
[35] Yicheng Wu, Vivek Boominathan, Huaijin Chen, Aswin Sankaranarayanan, and Ashok Veer- araghavan. Phasecam3d - learning phase masks for passive single view depth estimation. In 2019 IEEE International Conference on Compulational Photography (ICCP), pp. 1-12, May
2019. doi: 101110/CCCHH00::0111177777
[36] Weijian Yang and Rafael Yuste. In vivo imaging neural activity. Nature Methods, 14(44): 349-359, Apr 2017. ISSN 1548-7105, doi: 10. n0.103//mmeth.42.0 4230.
[37] Kyrollos Yanny, Nick Antipa, Ren Ng, Laura Waller, and Laura Waller. Miniature 3d fluores- cence microscope using random microlenses. In Biophotonics Congress: Optics in the Life Sciences Congress 2019 (BODA,BRAIN.NE OMA, OMP) (2019), paper BT3A.4, pp. BT3А.4. URL https: Optical Society of America, Apr 2019, doi: 10. /336//BAAIN.2019BBTT3 BT3A.4. // //www... www.obsapblisse webbbisshing..mmm wwtg/abbtractt. wtt.::t.ccccmmrrrrraaaaa----
[38] Kyrollos Yanny, Nick Antipa, William Liberti, Sam Dehaeck, Kristina Monakhova, Fan- glin Linda Liu, Konlin Shen, Ren Ng, and Laura Waller. Miniscope3d: optimized single-shot miniature 3d fluorescence microscopy. Light: Science & Applications, 9111:117 Oct 2020. ISSN 2047-7538. doi: 10. 177-0700000000----1
12

---

Under review as a conference paper at ICLR 2022
A
APPENDIX
A.1
FORWARD SIMULATION OF PROGRAMMABLE 3D SNAPSHOT MICROSCOPE
lens
V
lens
f
2D camera image
3D sample
phase mask
(pupil plane)
Figure 5: Diagram of a 4f optical model that is the basis for our simulated microscope Map showing the Fourier plane in which we have the programmable and trainable 2D phase mask dp.
Here we describe our wave optics simulation of the mieroscope Mø. which we model as 4f system [11]. The 4f optical system consists of two lenses, the first spaced one focal length from the object plane and the second spaced one focal length away from one focal length beyond the first lens (Figure 5). In between these two lenses, we can place phase mask to manipulate the light field before passing through the second lens and forming an image on the camera sensor.
We are concerned here with fluorescence microscopy, meaning that the sources of light that we image are individual fluorescent molecules, which we can model as point emitters. Because these molecules emit incoherent light, the camera sensor in effect sums the contributions of each point emitter. In order to model such an imaging system, we first need to address modeling a single point emitter's image the camera.
We can analytically calculate the complex-valued light field one focal length after the first lens (which call the pupil plane) due to point source centered at plane 2 (where 2 is a distance from the object plane 0). If the point source were centered (æ 0, = 0) in the object focal plane 2 = we would have plane wave at the pupil plane, but for the more general case of point source an arbitrary plane z relative to the object plane = O, we can analytically calculate the complex-values light field entering the pupil plane:
Upoint (k; exp 12
n (7)
||
where Mpoint is the incoming light field entering the pupil due to a point source centered in the plane at z, k E 2R22 denotes frequency space coordinates of the light field in the pupil plane, n is the refractive index, and is the wavelength of light [1]].
In this pupil plane, we can then apply a phase mask dp to the light field, which is modeled as multiplication of upaine (k;
(k;z) z) and e (k)( the complex phase of the pupil function. The light field exiting the pupil is therefore described by
Upupil (k; 2) Upoine (k; =)p(k)
(8)
where p(k) is the pupil function, composed of an amplitude a akk) and phase (k):
p(k) a = (kk
a(k) = {
k
2 S
NA
(10)
k
NA
where NA is the numerical aperture of the lens [11].
The light field the camera plane can then be described by Fourier transform [11]:
Ugamera F (Upppi).
(11)
13

---

Under review as a conference paper at ICLR 2022
The camera measures the intensity of this complex field:
s(x; 2) |eaamera (x: 2)
(12)
where x E R2 denotes spatial coordinates in the camera plane [11].
We can call this intensity s the point response function (PRF). If the shape of the PRF is translationally equivariant in x, meaning that moving a point source in-plane creates the same field at the camera, just shifted by the corresponding amount, then we call this PRF a point spread function (PSF). Note that moving the point source in will not give the same shape, which allows our system to encode depth information through the PSF [5].
In order avoid edge effects during imaging, we simulate the PSF at larger field of view, then crop and taper the edges of the PSF:
Staper = crop[s] + t
(13)
where t is a taper function created by taking the sigmoid of a distance transform divided by a width factor controlling how quickly the taper goes 0 at the edges and @ denotes elementwise multiplication. We intentionally simulate a larger field of view than the sample in order to avoid edge artifacts. The purpose of the crop[-] is to cut the PSF to the correct field of view. The purpose of the tapering is to remove artifacts at the edges of the cropped PSF After we compute this cropped and tapered PSF, we also downsample Staper the size of the data in order save memory.
Imaging is equivalent to the convolution of the incoming light field volume intensity v and the cropped and tapered PSF Staper for given plane. At the camera plane, the light field intensity is measured by the camera sensor. Therefore, we can describe the forward model as the following convolution and integral over planes:
/ʃ ((Txx Txi z) dTx da
Me
(14)
We then model shot noise of the camera sensor to produce the final image c, for which the appropriate model is sampling from Poisson distribution with a mean of He [11]:
~ Poisson (He)
(15)
However, because we cannot use the reparameterization trick to take pathwise derivatives through the discrete Poisson distribution, instead approximate the noise model with rectified Gaussian distribution:
6 e~N(0,1) = max (1H + v))
(16)
(17)
We now turn our attention to selecting the number of pixels used in the phase mask, i.e. the number of parameters for M We first need to determine the pixel size for Nyquist sampling the image plane with an objective of a given NA (numerical aperture). For a given pixel size Ax, we know that in frequency space coordinates we will have bandwidth of J., spanning we must have
2Az to 2Ar Because
NA A
|kkl|2 S
(18)
we know that the Nyquist sampling pixel size is given by A
Ar
(19)
2NA
Therefore, in the image plane, for a desired field of view L we must have least
L
N
(20)
Ar
pixels. The discretization in the pupil plane will be the same, which means we will need to have at least N* pixels in the pupil plane to achieve the appropriate light field in the image plane. For our settings of NA = 0.8, À = 0.532 and L 823 we have Ar" 0.3325 and N = 2476 pixels. Thus, a reasonable choice is Az = 0.3255um and N = 2560 pixels. Note that these simulation parameters are independent of the camera pixels: we have determined only how many pixels must be used in the phase mask in order to ensure our PSF can occupy the full field of view. The camera sensor can sample the field at the image plane an independent pixel size.
14

---

Under review as a conference paper at ICLR 2022
A.2 TRAINING PSFS AND VOLUME RECONSTRUCTION NETWORKS
FourierNet2D microscope 256 x 256 pixels, Type D
UNet2D microscope
256 x 256 pixels, Type
jm
Figure 6: FourierNet successfully optimizes PSF to image and reconstruct Type D where UNet fails. The FourierNet learned to produce multiple pencils in its PSF, which create multiple views of the sample in the camera image. UNet learned only a single pencil and fails to utilize the majority of pixels in the camera image to encode views of the sample. Top row shows simulated camera image of a Type D example, middle row shows xy max projection of the PSF, and bottom row shows xz max projection of the PSF
Given simulation of imaging, we can define two modes of autoencoder training: jointtly training the phase mask parameters and weak reconstruction networks in order learn good PSF for particular class of samples (i.e. samples with the same spatiotemporal statistics), and (2) training a stronger reconstruction network only with a fixed, pre-trained dp.
Definition of terms For both cases of training, the general framework is to simulate imaging using confocal volumes of nan-neuronal labeled larval zebrafish, reconstruct from the simulated image, then update the reconstruction network and, if desired, the microscope parameters. We will define the microscope parameters as dp and the reconstruction network parameters as 0 for any reconstruction network R (c) where R. maps 2D images to 3D volume reconstructions. For our training algorithms listed below, we also define: D our dataset, v ground truth volume, reconstructed volume, L computed loss, 2, list of plane indices that will be imaged/reconstruuted,, ag the learning rate for the microscope parameters, dg the learning rate for the reconstruction network parameters, and the weight of the non-high pass filtered component of the loss. When selecting a random ground truth volume, we also perform random shift, rotation, flip, and brightness augmentations.
Microscope simulation parameters When simulating the zebrafish imaging, of 0.532 um for all simulations. The NA of our microscope 0.8. The refractive index n is 1.33. We downsample all volumes to (1.0 μm 1.625 μm y, 1.625 μm x). We taper width of 5 for all
wavelength
15

---

Under review as a conference paper at ICLR 2022
optimized and imaged 512 x 512 pixels, Type A
optimized and imaged on: 512 x 512 pixels, Type B
optimized and imaged on:
512 x 512 pixels, Type C
200 umm
Figure 7: Optimizing PSFs for different samples result in PSFs tailored to each sample. Note that PSF optimized for Type A (left) has pencils with a span in that matches Type A. PSF optimized for Type (middle) has pencils that span the entire z depth. PSF optimized for Type C (right) has pencils spread farther apart to account for the larger sample. Top row shows simulated camera image of a Type A, B. or C example respectively, middle row shows xy max projection the PSF, and bottom row shows max projection of the PSF.
simulations, and simulate the PSF 50% larger dimensions in and y. The resolution of the camera (for all zebrafish datasets) also (1.625 μm y, 1.625 μmm x).
Initialization of dp For Type A, B, and D, initialize to produce PSF consisting of 6 pencil beams at different locations throughout the depth of the volume, with the centers of these beams arranged in hexagonal pattern in x and y. Because our optimizations generally find PSFs with many pencils, we find that initializing with such pattern helps to converge to a more optimal PSF (data shown).
For Type C, we instead initialize with a single helix spanning the depth of the volume (the "Potato Chip" from [5]), which seems find local minimum for dp that produces a PSF with more pencils (and therefore views in the camera image).
Data settings and augmentation for Type A, B, C, Using our total 58 training zebrafish volumes and 10 testing zebrafish volumes (imaged through confocal microscopy), we crop in four different ways to create four different datasets. For training volumes, we crop from random locations from each volume form of augmentation. For testing, crop from the same location. Physically, these crops correspond to either placing a circular aperture before light hits the 4 system or changing the illumination thickness in because samples would be illuminated from the side in real implementation of this microscope. We model these by cropping cylinders (or cubes if there is no aperture) of different diameters and heights. We show details for all types Type A, B, C, D in Table
16

---

Under review as a conference paper at ICLR 2022
where the diameter of the cylinder is labeled "Aperture Diameter" and the illumination thickness is labeled "Height".
We augment our volumes during training by taking random locations from these volumes, randomly flipping the volumes in both z and y, and also randomly rotating in pitch, yaw, and roll. Most importantly, also randomly scale the brightness of our samples and add random background levels which serve to adjust the signal-to-noise ratio (SNR) of the resulting simulated images. The only exception these augmentations is Type C, where we all the volumes the same in-plane vertical orientation (while still applying rotation augmentations in pitch and roll).
Table 4: Specifications of all zebrafish datasets Type A, B, C, D for reconstruction
Dataset
Camera (px)
Height (planes)
Span (z, y, x) (μm) Aperture Diameter (μım)
Type A
512 x 512
12
(25, 832, 832)
386
Type B
512 x 512
128
(250, 832, 832)
386
Type C
512 x 512
128
(250, 832, 832)
Type D
256 x 256
96
(200, 416, 416)
193
Parallelizing imaging and reconstruction Furthermore, because this simulation can become too expensive in memory to fit single device, we generally perform the simulation, reconstruction, and loss calculation in parallel for both training modes. Therefore, any variable that has a, subscript refers to list of chunks of that variable that will be run on each device. A 1 superscript indicates a particular chunk for GPU j. For example 2 is a list of plane indices to be imaged/reconstrruted and 2f is the jth chunk of plane indices that will be imaged/reconssrruceed on GPU j. We denote parallel for any operations that are performed in parallel and scatte e for splitting data into chunks and spreading across multiple GPUs Imaging can be cleanly parallelized: chunks of a PSF and sample can be partially imaged on multiple GPUs independently because the convolution occurs per plane, then finally all partial images can be summed together onto single GPU. The reconstructions can similarly take the final image and reconstruct partial chunks (as well as calculate losses on partial chunks) of the volume independently per device. We implicitly gather data to the same GPU when computing sums ( or means (E). The functions paralle image and compute PSF follow the definitions above in equations 14 and 12. In the algorithms shown here, parallel image applies the same convolution described above in equation 14.
Sparse gradients and downsampling We additionally support training and reconstructing only some of the planes for imaging and potentially a different partial group of planes during reconstruction, as a way to sparsely compute gradients for optimization of On, and save memory. The planes not imaged with gradients can still contribute to the image (without their gradients being tracked) in order to make the problem more difficult for the reconstruction network. Over multiple iterations, this can become equivalent to the more expensive densely computed gradient method, essentially trading training time for memory. An additional memory saving measure not written in the algorithms is to compute the PSF at a high resolution, then downsample the PSF using 2D sum pool to preserve total energy in order reduce memory usage when performing the imaging and reconstruction. We denote with no gradient tracking to show an operation without gradients.
A.3 IMPLEMENTATION DETAILS
Fourier convolution details Our Fourier convolution uses complex number weights, implemented as two channels of real numbers. Furthermore, in order to prevent the convolution from wrapping around the edges, we have to pad the input to double the size. The size of the weight must match the size of this padded input. This means that the number of parameters for our Fourier convolution implementation is 8x the number of parameters required for a global kernel in a spatial convolution (though the Fourier convolution is significantly faster). We do this to save an extra padding and Fourier operation, trading memory for speed. Because the simulation of imaging requires more memory than the reconstruction network, we found this to be acceptable tradeoff.
Common network details All convolutions (including Fourier convolutions) use "same" padding. For FourierUNets and vanilla UNets, downsampling and upsampling is performed only in the and y
17

---

Under review as a conference paper at ICLR 2022
Algorithm 1: Parallel PSF engineering by joint training of reconstruction network and phase mask. Microscope M6 parameters are Φp reconstruction network R. parameters are ê, dataset D, learning rates for and 0 are a and og respectively, plane indices to image and reconstruct from 2s, and weight for LNMSE is B.
Input :M a, ag 2.
for v do
// select plane indices to be imaged with and without gradients
Ra.gradient, 2,,n0 gradiont select planes (z,)
// move sample planes to be imaged with gradients to multiple GPUs
V V.gradlent + scatter (v, (agadiet))))))
// move sample planes to imaged without gradients to multiple GPUs
V, gradient scatter (v, 2, gradient)
// compute PSF with gradients
multiple GPUs
Se.gradient parallel (compute PSF (Mo, 2 for z in &egradient)
// compute partial image with gradients on multiple GPUs Cgradient f parallel image gradiint V.ggadieent)
// compute PSF without gradients on multiple GPUs with gradient tracking
sono gradient. + parallel (compute PSF (M 2 for 2 in ssno gradient)
end
// compute partial image without gradients multiple GPUs
10
with no gradient tracking
Cno gradient parallell image ns, gradient Vs, gradient)
12
end
// compute full image by summing partial images onto GPU
13
c Eggradient, Cno gradiens]
// se eeeect plane indi ces be reconstructed
14
Re.roconttruct + select planes (z,)
// move sample planes that will be reconstructed to multiple GPUs
15
V..roonstructt +- scatter (V, (s
// compute mean of high passed sample for loss normalization
(H((v) + (
// compute mean of sample for loss normalization
Pv + E(v? seconsttruul)
// move reconstruction networks to multiple GPUs
18
Ro. + scatter (Rş)
// compute reconstruction and loss multiple GPUs
19
+ parallel reconstruct/lost (c, Vereconsttuut R8 MHlyy:
// compute gradients for all parameters
ge + VoL
21
I + V
update all parameters
22
ADAM (oo, 0,go)
Φ ADAM (0 4194)
end
dimensions (we do not downsample or upsample in a because there could potentially not be enough planes do so). We train all networks using the ADAM optimizer with all default Py Torch parameters except the learning rate, which we always set to 10 for the reconstruction network parameters and 10-2 for the phase mask parameters db
18

---

Under review as a conference paper at ICLR 2022
Algorithm 2: Parallel training a reconstruction network given pre-trained phase mask. Mi- croscope Ms parameters are & (phase mask), reconstruction network R. parameters are dataset is D, learning rates for & and 0 are a and ag respectively, plane indices to image and reconstruct from are Rx, and weight for LNMSE is 3.
Input :Mo. Φ, a Re, 0 ag, 2a. // compute PSF without gradients on multiple GPUs with no gradient tracking
gradient paralllll (compute PSF (M4: 2/) for of in
end
for v E D do
// select plane indi ces to imaged without gradients
Za gradient f- select planes (2,)
// move sample planes to be imaged without gradients to multiple CPUs
V gradient + scatter (v, 2,,00 gradient)
// move necessary PSF planes to multiple GPUs
8s,no gradient + scatter (Sno gradient, &s,o gradient)
// compute image without gradients on multiple GPUs with no gradient tracking
t --rallllll image (S,,no gradients Va,o gradient) end
10
// select plane indices to reconstructed
2 re.onnnntt + select planes (z,)
// move sample planes that will be reconstructed to multiple GPUs
12 Ve.reconstruct + scatter (v, recoossuutt)
// compute mean of high passed sample for loss normalization
(H(() v.ccooscct))))
// compute mean of sample for loss normalization
Mv E[v? rrecomssuttt
// move reconstruction networks to muitiple GPUs
15
Re, + scatter (Re)
// compute reconstruction and loss on multiple GPUs
16
+ parallel teconstruct/lose (c, V.rooosttmnct, Pe,, // compute gradients for reconstruction networks only
ge Vol.
// update reconstruction network parameters only
18
ADAM (aa, 0,
end
Algorithm 3: Parallel imaging, PSF planes multiple GPUs are sample planes multiple GPUs to be imaged are v,.
Input :S,; V.
Output :c // compute images in parallel on multiple GPUs, then sum to single GPU
f return
E[parallel (convolve (s}, v for (s), v!) in (s., v&))]
Normalization We use input scaling during both training and inference in order to normalize out differences in the brightness of the image and prevent instabilities in gradients. This means we divide out the median value of the input (scaled by some factor in order to bring the loss to reasonable range) and then undo this scaling after the output of the network. This effectively lincarizes our reconstruction networks, meaning a scaling of the image sent to the network will exactly scale the output by that value. We also find this is more effective and simpler alternatiee
19

---

Under review as a conference paper at ICLR 2022
Algorithm 4: Parallel reconstruction/los calculation. Camera image is c, sample planes on multiple GPUs are va reconstruction networks on multiple GPUs are Rg mean for LHNMSE normalization is MH((), mean for LNMSE normalization is Av, and weight for LNMSE is B.
Input :c, Va, Ro,,s H
Output :L
// compute reconstruction and loss in parallel on multiple GPUs
V. concatenate (parallel (c) for R in Re,)))
L parallel (oM((()))(()))))) 8 Vv for (V) vg) in (f., v.))
An(v)
// compute mean of scattered losses on single GPU
+ + return L
using a BatchNorm on our inputs. We continue to use BatchNorm between our convolution layers within the reconstruction network [15], which is effectively InstanceNorm in our case where batch size [31].
Planewise network training logic When we train PSFs by optimizing Φ. we train separate recon- struction networks per plane. This allows us to flexibly compute sparse gradients across different planes from iteration to iteration, described in Appendix A.2. In order to do this, create placeholder networks on any number of GPUs, then copy the parameters stored on CPU for each plane's reconstruction network network the GPU needed during forward pass. After calculating an update with the optimizer, we copy the parameter values back to the corresponding parameter CPU.
Training times We optimize our Type D microscopy experiments on RTX 2080 Ti GPUs. For these, compare training times for the different network architectures. One training iteration (including microscope simulation, reconstruction, backpropagation, and parameter update) takes ~0.6 seconds for FourierNer2D and ~133 seconds for UNet2D when optimizing both di and O. One training iteration takes ~0.4 seconds for FourierNet3D, ~0.7 seconds for FourierUNet3D, and ~0.8 seconds for UNet3D when only optimizing O. More details are found in Tables 5 and 6.
Table 5: Type D experiment training times
Network
Optimizing
# *parameters
# train steps
Train step time (s)
Total time (h)
FourierNer2D
0, Φ
4.2 107
106
0.8
222
FoDD
e
6.3 107
106
~ 0.4
~ 111
FouDierUNet3D
8.4 X 107
106
~ 0.7
~ 194
UNet2D
0. 0.
4.0 X 107
106
1.3
~ 361
UNet3D
1.0 X 10
106
~ 0.8
~ 222
A.4 DETAILS FOR FOURIERNETS OUTPERFORM UNETS FOR ENGINEERING NON-LOCAL OPTICAL ENCODERS AND 3D SNAPSHOT MICROSCOPY VOLUME RECONSTRUCTION
For our experiments in Sections 3.1 and 3.2, we use 40 planes at 5ym resolution in and therefore 40 reconstruction networks to train PSFs. When training reconstruction networks only to produce the higher quality reconstructions, actually span 200 μm in z). We train in both settings without any sparse planewise gradients, meaning image and reconstruct all 40 all 96 planes, respectively. We show details of all datasets used for training reconstructions in Table 4.
96 planes at Iam resolution in z (chosen so that the planes
We show the details of our FourienNet2D architecture for training PSFs in Table 7 and our Fourier- Net3D architecture for training reconstruction networks in Table We also show details for training times for both training PSFs and for training more powerful reconstruction networks in Table 5. We
20

---

Under review as a conference paper at ICLR 2022
Table 6: Type A, B, C experiment training times
Network
Optimizing
# parameters Type # train steps Train step time (s)
Total time (h)
FourierNet2D) 0, 3.
1.7 10
A
5.8
10
1.1
177
FourierNet3DD
0 (fixed <p for A)
~.4 x 10a
A
2.6 x 10
~
1.6
116
FourierNe(3D
(fixed <p for A)
~.4 x 10a
B
1.3 x 10a
1.6
FourierNet3D (fixed Φ for A)
~.4 10
C
1.3 10a
1.6
58
FourienNet2D 0, +
~1.7 10a
B
5.8 x 10
1.1
177
FourierNet3D (fixed op for B)
~34 10
A
1.2
105
1.6
53
FourierNet3D e (fixed <p for B)
~ 3.4 x 10
B
106
1.6
444
FourierNe(3D (fixed cp for B)
3.4 x 10a
C
5.0 X 10a
1.6
222
FourierNet2D e, dp
~ 1..7 x 108
C
5.8 X 10
1.1
177
FoD3 (fixed dp for C)
3.4 x 10a
A
3.4 X 105
1.6
~ 151
FourierNet3D (fixed dp for C)
~ 3.4 10
B
~
3.4 105
1.6
151
FourierNet3D 0 (fixed dp for C)
~ 3.4 x 10
C
3.7 X 10
1.6
~ 164
Table 7: FourierNet2D detailed architecture (1 per plane)
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
InputScaling
scale: 0.01
1,, 256)
FourierConv2D
(256, 256)
(2, 2)
(8, 1,66, 256)
LeakyReLU
slope: -0.01
(8, 256, 256)
BatchNorm2D
(8, 256,, 256)
Conv2D
(11, 11)
(1, 1)
1, 256)
ReLU
(1, 1, 256, 256)
InputRescaling
scale: 0.01
(1, 256, 256)
Table 8: FourierNet3D detailed architecture (8 GPUs)
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
InputScaling
scale: 0.01
(1, 1, 256, 256)
FourierConv2D
(256, 256)
(2, 2)
(60, 1, 256, 256)
LeakyReLU
slope: -0.01
(60, 256, 256)
BatchNorm2D
(60, 256, 256)
Reshape2D3D
(5, 12, 256, 256)
Conv3D
(,1,7,7)
(1,
(5, 12, 256, 256)
LeakyReLU
slope: -0.01
(5, 12, 256, 256)
BatchNorm3D
(5, 12, 256, 256)
Conv3D
(,1,7,7)
(L. 1,
(1, 12, 256, 256)
ReLU
(1, 12, 256, 256)
InputRescaling
scale: 0.01
(1, 12, 256, 256)
trained all networks for Type D for the same number of iterations (more than necessary for PSFs to meaningfully converge)?.
The architecture FourierUNet3D is scales, with cropping factor of 2 per scale in the encoding path and an upsampling factor of 2 in the decoding path. For each scale, we perform a Fourier convolution in the encoding path producing 480 feature maps, which are concatenated with the incoming feature maps of the decoding convolutions at the corresponding scale (just in normal UNet). In the decoding path, we use convolutions with kernel size (3, 5, 5), producing 12 3D feature maps each. There two such convolutions per scale. Note that this requires reshape the 2D feature maps from the Fourier convolutions to 3D. This is followed by Ixl convolution
TTraining times approximate, and actual total time longer due to checkpoint- ing/snapshottinn//allaaaiio of data and/or differences in load the clusters being used.
21

---

Under review as a conference paper at ICLR 2022
Table 9: FourierUNet3D detailed architecture (8 GPUs)
Scale
Repeat
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
1
InputScaling
scale: 0.01
(1, 1, 256, 256)
Multiscale
(256, 256)
(2, 2)
(60, 1, 256, 256)
FourierConv2D
+ ReLU
BatchNorm22
2
(128, 128)
(2, 2)
(60, 1, 128, 128)
3
(64, 64)
(2, 2)
(60, 1, 64, 64)
4
(32, 32)
(2, 2)
(60, 32, 32)
4
1
Reshape2D3D
(5, 12, 32, 32) (5, 12, 64, 64)
3
Upsample2D
3
2
Conv3D
7,1,7,7)
1,,11)
(5, 12, 64, 64)
ReLU
BatchNorm33
2
Upsample2D
(5, 12, 128, 128) (5, 12, 128, 128)
2
2
Conv3D
(11,7,7)
(1,1,1)
ReLU
BatchNorm33
1
1
Upsample2D
(5, 12, 256, 256) (5, 12, 256, 256)
1
2
Conv3D
(11,7,7)
11,1,1
ReLU
BatthhNorm3D
1
1
Conv3D
1,,
(1,1,1)
(1, 12, 256, 256)
ReLU
InputRescaling
scale: 0.01 (1, 12, 256, 256)
producing the 3D reconstruction output. We show diagram of this architecture in Figure IC, and details of this architecture in Table 9.
For our UNet2D, each encoding convolution produced 24 feature maps (except the first scale, for which the first convolution produced 12 feature maps and the second convolution produced 24 feature maps). Each decoding convolution produced 24 feature maps, but took input of 48 feature maps where feature maps were concatenated from the corresponding encoding convolution at that scale. At the end of the UNet2D, (1, 1) convolution reduced the 24 final feature maps to I feature map. This single feature map is interpreted as the final output of the network, i.e. the reconstructed plane. UNet2D requires many more feature maps per plane and more layers than FourierNet, because these are necessary in order for the network to be able to integrate information from larger field of view. We show the details of our UNet2D architecture in Table
The architecture of the vanilla UNet3D is also 4 scales, with max pooling factor of 2 per scale in the encoding path and an upsampling factor of 2 in the decoding path. Each scale the encoding path produces 480 2D feature maps. These are concatenated to the incoming feature maps of the decoding convolutions at the corresponding scale, again with reshape from 2D to 3D. Each scale of the decoding path produces 48 3D feature maps. Again, this is followed by 1x1 convolution producing the 3D reconstruction output. All convolutions are in 3D with kernel size (5, 7, 7), with the 2 dimension being ignored for the encoding path because the input is 2D. We show the details of our UNet3l architecture in Table 11.
A.5 DETAILS FOR ENGINEERED OPTICAL ENCODING DEPENDS ON SAMPLE STRUCTURE
For our experiments in Section 3.3, we use 64 planes at Ium resolution in Z and therefore 64 reconstruction networks to train PSFs. When training reconstruction networks only to produce the higher quality reconstructions, we use 128 planes at I resolution actually span 250 μm in z). We train in the reconstruction only setting without any sparse planewise gradients, meaning we image and reconstruct all 128 planes. However, when training a PSF we image and reconstruct 40 planes time with gradient per iteration (spread across GPUs). These 40
(chosen so that the planes
22

---

Under review as a conference paper at ICLR 2022
Table 10: UNet2D detailed architecture (1 per plane)
Scale
Repeat
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
1
1
InputScaling
scale: 0.01
(1. I. 256, 256)
1
1
Conv2D
(7,
(1, 1)
(12, 1, 256, 256)
+ ReLU
+ Narmh2
1
Conv2D
(7,
(1,
(24, 1, 256, 256)
+ ReLU
+ BatchNNrmm2D
2
1
MaxPool2D
(2,
(2,2)
(24, 1, 128, 128)
2
2
Conv2D
(7,
(1,
(24, 1, 128, 128)
+ ReLU
+ BatchNorm22
n
1
MaxPool2D
(2 2)
(2
(24, 1, 256 256
2
Conv2D
(7,
(1,
(24, 1, 256 256
+ ReLU
+ BatchNNomm2D
8
1
MaxPooI2D
(2
(2,2)
(24, 1, (24, 1, 2,
8
2
Conv2D
(7.
(1,
+ ReLU
+ BatchNNrm22
7
1
Upsample2D
(24, 1, 4, 4)
7
2
Conv2D
(7,7)
(I.
(24, 1, 4,
+ ReLU
+ BatchNorm22
n
1
Upsample2D
(24, 1, 256
24
2u
256
n
2
Conv2D
(7.7)
(1,
(24, 1, 2a-1x 2n-1
256
256
+ ReLU
+ Batchhorm22
1
1
Upsample2D
(24, 1, 256, 256) (24, 1, 256, 256)
I
2
Conv2D
(7,7)
(1,
+ ReLU
+ Batchh
Conv2D
1)
(1
(1, 1, 256, 256)
+ ReLU
I
1
InputRescaling
scale: 0.01 (1, 1, 256, 256)
planes chosen randomly every iteration from the 64 total possible planes, making potentially separate draws of planes for imaging and reconstruction. We show details of all datasets used for training reconstructions in Table 4.
We show the details of our FourierNe12D architecture for training PSFs at the larger field of view in Type A. B. C in Table 12 and our FourierNet3D architecture for training reconstruction networks the larger field of view in Type A, B, in Table 13. There are no other networks used for these larger field of view experiments. We also show details for training times for both training PSFs and for training more powerful reconstruction networks in Table 6. All PSFs in these networks were trained for the same number of iterations. However, reconstruction networks for some of these experiments were only trained for as long necessary to converge (with some exceptions where attempted longer training to check for performance gains with long training periods). Generally, we observed that performance for such reconstruction networks does not meaningfully change with many more iterations training..
*Training times approximate, and actual total time longer due to checkpoint- ing/snapshoting//vliaaaan of data and/or differences in load the clusters being used.
23

---

Under review as a conference paper at ICLR 2022
Table 11: UNet3D detailed architecture (8 GPUs)
Scale
Repeat
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
1
InputScaling
scale: 0.01
(1, 1, 256, 256)
Conv2D
(7.7)
(1,
(30, 1, 256, 256)
ReLU
+ BatchNorm22
Conv2D
(7.
(1, 1)
(60, 1, 256, 256)
ReLU
BatchNorm2D
2
MaxPooI2D
(2,
(2,2)
(60, 1, 128, 128) (60, 1, 128, 128)
2
2
Conv2D
(7,
(1,
ReLU
BatchNorm22
3
1
MaxPool2D
(2
(2,2)
(60, 1, 64, 64) (60, 1, 64, 64)
3
Conv2D
(7,
11,
ReLU
BatchhNrm22
4
1
MaxPool2D
(2
(2,2)
(60, 32, 32) (60, 1, 32, 32)
2
Conv2D
(7,7)
(1,
ReLU
BatchNorm22
4
Reshape2D3D
(5, 12, 32, 32) (5, 12, 64, 64) (5, 12, 64, 64)
3
Upsample2D
3
2
Conv3D
(11,7,7)
(1,1,1
ReLU
BatchNorm33
Upsample2D
(5, 12, 128, 128) (5, 12, 128, 128)
2
2
Conv3D
(11,7,7)
(1,1,1)
ReLU
BatchNorm33
1
Upsample2D
(5, 12, 256, 256) (5, 12, 256, 256)
1
2
Conv3D
(11,7,7)
(1,
+ ReLU
BatchNorm3D
1
1
Conv3D
(1, 1,
(1,
(1, 12, 256, 256)
ReLU
1
1
InputRescaling
scale: 0.01 (1, 12, 256, 256)
Table 12: FourierNet2D detailed architecture (1 per plane)
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
InputScaling
scale: 0.01
(1, 1,,2, 512)
FourierConv2D
512,,122
(2, 2)
1,12,
LeakyReLUU
slope: -0.01
(5, 1,11,,2,,
BatchNorm22
1,,12
Conv2D
(11, 11)
(1,
1,,1, 512)
ReLU
(1, 1,,1, 512)
InputRescaling
scale: 0.01
1,,1,
A.6 DETAILS FOR FOURIERNETS OUTPERFORM STATE-OF-THE-ART FOR RECONSTRUCTING NATURAL IMAGES CAPTURED BY DIFFUSERCAM LENSLESS CAMBRA
We performed no augmentations for this of trainings reconstructing RGB color images of natural scenes from RGB diffused images taken through a DiffuserCam [20]. We modified our FourienNet2D architecture create the FourierNeIRGBB architecture and Nou2ierUNet2D architecture create
24

---

Under review as a conference paper at ICLR 2022
Table 13: FourierNet3DD detailed architecture (8 GPUs)
Layer type
Kernel size
Stride
Notes
Shape (C, D, H, W)
InputScaling
scale: 0.01
(I, 112,,
FourierConv2D
(512, 512)
(2, 2)
(80, 1,512,
LeakyReLU
slope: -0.01
(80, 1,512,
BatchNorm22
(80, 1,5122,512
Reshape2D3D
16, 512)
Conv3D
(11,7.7)
11
16,51, 512)
LeakyReLU
slope: -0.01
16,
BatchNorm3D
16,1,,25512
Conv3D
(11,7.7)
11,
(6, 512)
ReLU
16, 512, 512)
InputRescaling
scale: 0.01
(1, 16, 512, 512)
Table 14: DLMD experiment training times.
Superscripts denote
loss function:
MSE,
MSE+LPIPS.
Network
Optimizing
parameters
train steps
Train step time (s)
Total time (h)
FourierNetRGB
0
1.6 x 107
2.2 x 105
0.43
26
0
1.6 x 107
x 10a
0.47
14
FourierUNeeRGiBI
0
7.1 x 107
2.5 x 105
~ 3.3
229
Le-ADMM-U [20]
0
4.0 x 107
UNet? [20]
0
1.0 x 10
Table 15: FourierNeiRGB detailed architecture
Layer type
Kernel size Stride
Notes
Shape (N, C, H, W)
FourierConv2D
(270, 480)
(2,2)
(4, 3,,0,27,, 480)
LeakyReLU
slope: -0.01
(4, 20, 270, 480)
BatchNorm2D
(4, 20, 270, 480)
Conv2D
(11,
11,
(4, 27, 480)
BatchNorm2D
(4, 64, 480)
LeakyRel.U
slope: -0.01
(4, 270, 480)
Conv2D
(11,
(1,
(4, 64, 270, 480)
BatchNorm22
(4, 64, 270, 480)
LeakyReLU
slope: -0.01
64, 480)
Conv2D
(11, 11)
(1,1)
3,7 480)
ReLU
2, 480)
the FourierUNetRGH architecture, outlined in Table 15 and Table 16 respectively. Training details are shown in Table 14. Because these reconstructions are of 2D images only and required no microscope simulation, we were able use batch size of 4 images per iteration.
25

---

Under review as a conference paper at ICLR 2022
Table 16: FourierUNetRGiB detailed architecture
Scale
Repeat
Layer type
Kernel size
Stride
Notes
Shape (N, C, H, W)
Multiscale
(270, 480)
(2, 2)
(4, 64, 270, 480)
FourierConv2D
+ ReLU
BatchNorm2D
2
(135, 240)
(2, 2)
(4, 64, 135, 240)
3
(67, 120)
(2, 2)
(4, 64, 67, 120)
4
(33, 60)
(2, 2)
(4, 64, 33, 60)
3
1
Upsample2D
(4, 64, 67, 120)
3
2
Conv2D
(11, I1)
1I,
(4, 64, 67, 120)
ReLU
+ BatchNorm22
2
1
Upsample2D
(4, 64, 135, 240) (4, 64, 135, 240)
2
2
Conv2D
(I),
(I,
ReLU
+ BatchNNorm2
1
1
Upsample2D
(4, 64, 270, 480) (4, 64, 270, 480)
1
2
Conv2D
(11,
(1,
ReLU
BatchNorm22
Conv2D
(1, I)
(1, 1)
(4, 3, 270, 480)
+ ReLU
26

---

Under review as a conference paper at ICLR 2022
ground truth
M FourierNet2D
Mg FourierNet2D
Mg FourierNet2D
M4 UNet2D R. UNet3D
Rs FoureerNet3D
Rs FourjerUNet3D
R. UNat3D
50 um
256 256 px camera
Figure 8: Slab views of a Type D example volume reconstruction, showing our methods (Fouri- erNes/FurrirUUNet) do the best job of reconstructing throughout the volume. Note that the UNet reconstructions are blurry across all slabs, with few exceptions. Colored boxes show which sample planes a particular slab comes from, corresponding to boxes in projection view at top. Annotation M shows which network architecture was used for microscope optimization; annotation R shows which architecture used for reconstruction.
27

---

Under review as a conference paper at ICLR 2022
ground truth (Type A)
optimized for Type A
optimized for Type
optimized for Type
50 um
512 512 px camera
Figure 9: Slab views of an example Type A volume show that the microscope optimized for Type A results in the best reconstructions. Note that the reconstruction with a microscope optimized for Type A is almost identical to the ground truth, while the other microscopes create blurrier reconstructions. Slabs are xy max projections in thinner chunks as opposed to projecting through the entire volume. Colored boxes show which sample planes particular slab comes from, corresponding to boxes in xz projection view at top.
28

---

Under review as a conference paper at ICLR 2022
ground truth (Type B)
optimized for Type A
optimized for Type B
optimized for Type
50 um
512 512 px camera
Figure 10: Slab views of an example Type B volume show that the microscope optimized for Type results in the best reconstructions: other microscopes result in blurrier reconstructions. Colored boxes show which sample planes a particular slab comes from, corresponding to boxes in xz projection view top.
29

---

Under review as a conference paper at ICLR 2022
ground truth (Type C)
optimized for Type A
optimized for Type
optimized for Type
50 um
512 512 px camera
Figure 11: Slab views of example Type C volume show that microscope optimized for Type C
provides most consistent reconstruction. Colored boxes have same meaning as Figures 9, 10.
30


Reviews:
[{"review_id": "s6RAooELydv", "reviewer": "Reviewer_cUSr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n-------------\n- The paper is well written and easy to understand.\n- The experimental results are promising. \n- I really appreciate the thorough limitations section.\n\n\nWeaknesses\n-----------------\n- The paper has limited technical novelty when compared to a previous paper which implemented Fourier-domain networks, and does not currently contain experiments showing that the differences in approaches provide a meaningful improvement (see Detailed Comments).\n- The paper is missing some implementation details which affect the reader’s understanding of the Fourier structure of the neural network (see Detailed Comments).\n\nDetailed Comments\n--------------------------\n- The following paper contains very closely related work and should be cited and discussed (I am not affiliated with this paper):   \n\n    Rippel, Oren, Jasper Snoek, and Ryan P. Adams. \"Spectral representations for convolutional neural networks.\" Proceedings of the 28th International Conference on Neural Information Processing Systems-Volume 2. 2015.  \n\n    In particular, that paper proposes to learn multiplicative weights in Fourier space instead of learning image space convolutional kernels. The learned weights are parameterized to correspond to convolutional kernels with limited extent in the image space, which outperform image-space convolutions on a CIFAR classification task (Fig. 5a). This paper does differ from Rippel et al. in that the the learned convolutions have a kernel the size of the full image. It thus seems necessary to show that it is the global extent of the filters in this paper, and not the fact that the optimization is being performed in the space of the multiplicative Fourier weights, is driving the performance improvements. Further, that paper previously introduced the spectral cropping approach described in this paper for doing multi-scale learning in the Fourier space (although this is used for a classification task and not in the context of a U-net architecture). \n\n- I believe that the input images and output images of the networks are real-valued, and thus have a complex-valued Fourier transform with conjugate symmetry. (I do not have experience with this type of imaging and am going off Eq. 12-17). I see that the learned weights are complex valued.  How is the real valued-ness of the output of each layer maintained? Are the multiplicative weights chosen such that the output in Fourier space has conjugate symmetry? Since the core contribution of this paper is the Fourier network, these details would be useful to include in the main text, along with the fact that the multiplicative weights are complex-valued -- I think this is currently only mentioned in the Appendix.\n\n- Could you provide a more concrete sense of the memory requirements of the FourierNetRGB Architecture as compared to those of Le-ADMM-U and/or a comparison of the number of parameters? I see these for the other Fourier architectures but not FourierNetRGB in the Appendix. If the memory requirements/number of parameters is significantly smaller for Le-ADMM-U, a comparison may be needed with a version of that architecture scaled to match for fair comparison. I understand for the microscopy experiments that the imaging simulation requires significant memory, which is what motivates the multi-GPU implementation. However, it’s hard to know concretely what that means for the lensless imaging comparison, since only reconstruction is being performed in this case, and it would be good to know that the comparison in Table 3 is fair.  \n\n    Relatedly, could you provide the total training times for the architectures in the lensless imaging example? I assume the times in Tables 2 and 3 are the total inference time for reconstructing a single image; it is good to see that FourierNet is faster than Le-ADMM-U^2 there, but I don’t have a sense of whether the 20ms difference is relevant for a practical application. It would be great for readers to know whether there are significant differences in training time when choosing between these methods.", "summary_of_the_paper": "This paper introduces the use of Fourier convolutional neural networks to model both image encoding and reconstruction processes. These neural networks learn multiplicative weights in the Fourier domain, enabling efficient global convolutions (i.e., convolutions with kernels the size of the entire image), which are particularly important for imaging applications with a wide PSF. These networks compare favorably to UNets for the task of joint 3D snapshot microscope parameter optimization and image reconstruction. These networks also outperform the state-of-the-art architectures for a lensless imaging reconstruction task.", "summary_of_the_review": "This is a well-written paper with good reconstruction results. In its current form, the paper lacks important comparisons/contextualization with a very closely related work (Rippel et al.), which previously introduced many of the main technical ideas in this in this work. Also, for the lensless imaging task in this paper, more details are needed to understand whether the resources required for this method and the state-of-the-art are comparable. \n\nOverall, I would be inclined to improve my rating if the paper is extended with (1) a more careful comparison with Rippel et al., demonstrating which specific aspects of this work provide innovation and improvement over previous approaches and (2) more details about the resources required for the lensless imaging experiments, to better understand whether the comparison provided is fair.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "r9lcbJFyvDX", "reviewer": "Reviewer_3cwP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Positive:\n* The paper is nicely written, illustrated. Nice results.\n* On the method side, the integration of Fourier-domain layers, while not necessarily 100% novel, seems to do its job quite well and could inspire follow-up work in computational imaging.\n* The method is validated against multiple scenarios including an existing database and newly captured data, as well as simpler, more obvious architectures..\n\n\nQuestions:\n* I would assume that Fourier transforming the signal is not a new concept in the machine learning community. Among the things that come to mind are papers like Rippel et al., Spectral Representations for Convolutional Neural Networks, NIPS 2015 - but also other interpretations like \"positional encoding\" (Vasvani et al., Gehring et al., both 2017). Efficient large-kernel convolutions by virtue of the convolution theorem comes up as the main application scenario when doing a web search on \"Fourier CNN\". I would appreciate if the authors could clarify the relation of this work to those.\n\nDoubt:\n* Doing the Fourier convolution in various frequency bands (by cropping) seems somewhat redundant, as the downsampling amounts to cropping in Fourier domain and so one and the same frequency shows up in different layers. Could the desired effect (giving more weight to low frequencies) not be achieved just as well by appropriate frequency-dependent weighting or feature depth?", "summary_of_the_paper": "This paper introduces a convolutional neural network architecture for designing phase patterns to display on a spatial light modulator in the imaging path of a microscope, as well as decoding the resulting image into a 3D volume. The authors propose to use Fourier features in the decoding stage to enable features that are more strongly nonlocal than in regular (primal-domain) convolution than what they claim would be possible to achieve using \"vanilla\" UNet/ResNet-like architectures. They demonstrate their approach for snapshot 3D microscopy as well as the reconstruction of light fields from input captured using lensless system (with randomized optical elements).\n", "summary_of_the_review": "In my opinion, the foremost contributions of this paper are in the application and the pipeline, and not quite so much in machine learning methodology as the main idea (Fourier convolution) has been around for quite some time. The proposed pipeline is thoroughly motivated, described and evaluated. To me, this looks like an acceptable paper.\n\nThat said, I'm not an expert on the application (though it does feel like a bit of a niche). My relation to machine learning is that of a casual user rather than a contributor of hardcore methodology. So I'm looking forward to reading the other opinions on the work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "c-jB4VzHtJU", "reviewer": "Reviewer_jREz", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well written with a clear description. I like the idea of incorporating Fourier convolutional neural networks to more efficiently propagate gradients to the global structures, which allow for handling large PSFs often desired in high-dimensional imaging. Although the proposed architectures are simple modifications to the existing CNNs, their motivation is clear and demonstrates the advantages in synthetic experiments. In addition, all the evaluations are pretty well designed and executed. \n\nNevertheless, I have some questions and concerns about the claimed contributions shown in Section 1.2. \n\nFirst, the authors highlight that the parallel differentiable wave optics simulation achieved the optimization of 10^6 phase-modulating pixels. It is not clear whether achieving 10^6 pixels as optimization variables can be only achieved using the parallel simulator. Also, having 10^6 pixels is necessary to achieve the presented performance? In other words, I wonder how the performance of the proposed method would vary as the pixel count changes. Even though this is often regarded as a training detail in most end-to-end optimization papers, validating this would make the main contribution of this paper clear to readers because it could explain the motivation for the main contribution. \n\nThe second concern is, probably as expected by the authors, is the missing real-world experiments. The domain gap between the synthetic and real experiments is generally large for the end-to-end optimization of optics and reconstruction, making it challenging to incorporate many optimization variables for the optics encoder. This means that directly applying the proposed high pixel-count optics optimization to real-world experiments may impose challenges partially conflicted with the core idea of increasing optimization parameters in this work. I do not believe that this question comes as a surprise to the authors, but a more detailed discussion on this potential problem is worth to be given in the paper.", "summary_of_the_paper": "This paper proposes a neural network architecture that applies a series of convolution kernels on the Fourier domain to focus on global information of the captured images. The authors make use of the Fourier-plane phase modulation and apply this idea to single-shot 3D microscopic imaging using the Fourier-plane phase modulation for 3D microscope and lensless imaging. Also, this is further extended to the end-to-end optimization of the phase modulation and the reconstruction network. An extensive synthetic evaluation is conducted to validate the effectiveness of the proposed method over existing non-Fourier neural networks in the context of computational imaging.", "summary_of_the_review": "In summary, there are some major issues better be addressed in the paper. However, I don't see them as a critical barrier for blocking this paper from being accepted, given its high-quality execution and potential impact on computational imaging research. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PTejv45s2fn", "reviewer": "Reviewer_nest", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#### Strengths\n* The idea of using Fourier transform to achieve large-size PSF is interesting.\n* The design of the Fourier network is well explained: There are some features to reduce the computation in the Fourier network (e.g., save the parameters in Fourier space to save a padding operation) and they are explained in detail. The explanation of the multi-scale features in the Fourier domain is also helpful to understand the effectiveness of the proposed architecture.\n* The paper has a nice section on limitations.\n\n#### Weaknesses\n* Lack of evaluation: While this paper introduces an interesting architecture, the results are not sufficient to demonstrate the effectiveness of the architecture. First, the 3D snapshot microscope is evaluated only on simulated data. As the author mentioned in the limitations paragraph, the real optical encoder will have some artifacts because of aberrations and calibration error, and these effects need to be evaluated. Second, the evaluation is only on one type of volume (i.e., Larval Zebrafish). While the paper shows additional results on lensless imaging, it cannot support 3D snapshot microscopy, which is the title of this paper.\n* The result from UNet is unclear to me. Although it is computationally inefficient, UNet should be able to achieve global context if it has enough layers and thereby produce similar quality of results. However, in Figure 2, UNet provides much worse results compared to FourierNet. Why does UNet provide such bad results? Is it because the parameters of UNet are not sufficient to handle the global context?\n* Table 3: Why is UNet much faster compared to the FourierNet results? This is the opposite of what is claimed in the paper. Also, why is the FourierUNet result missing in this experiment?\n* I understand Lensless imaging results are included to demonstrate the effectiveness of Fourier network, but it is somewhat unnatural to have this new application abruptly considering the title of the paper starts with “3D snapshot microscopy”\n\n#### Minor comments\n* Figure 2: What is the 3.00s in the ground truth of the second row?\n\n##### Typos\n* “field field” in the 3rd paragraph of A.1\n* Table 3: unnecessary line break in the caption\n* Table 4: unnecessary line break in the second row.\n", "summary_of_the_paper": "This paper introduces an end-to-end optical encoder and deep learning decoder optimization for 3D snapshot microscopy. The main challenge is that the decoder needs to be able to handle the global PSFs used in the optical encoder in 3D snapshot microscope. This is difficult to be achieved in conventional UNet architectures as the global context requires many layers that is computationally expensive. The authors use a Fourier convolutional network that achieves the global context in a single layer with a much lower computational cost. The method is demonstrated on simulated data of one volume with four different fields of view. The authors also applied the technique to lensless imaging.\n", "summary_of_the_review": "This paper proposes an interesting architecture for the end-to-end optimization of the optical encoder and deep learning decoder for 3D snapshot microscopy. In particular, the authors introduce the Fourier network to handle the global PSF in the encoder with a lower computational cost compared to conventional UNet architecture. However, the method is evaluated only on simulated data, with a single type of object, and the explanation of conventional UNet result is insufficient to demonstrate the effectiveness of the proposed architecture. So the decision is weak reject.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes an application of the CNN to the microscopy problem of constructing 3D volumes from 2D captured images. The four reviewers thought the paper was a straightforward application of existing techniques to a new problem, while they were borderline towards accept the overall sentiment was that the technical novelty was very low from an ML perspective, and the ICLR community would only find the application potentially of interest. (Two reviewers changed from borderline reject, while the other two chose not to change their scores following the authors’ request.)
Decision: Reject