Programmable 3D snapshot microscopy with Fourier convolutional networks

1 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 neuronal activity across the whole brain of transparent model organisms such as larval zebrafish ( Danio rerio ) at higher volume rates , to enable richer studies of neural dynamics and signal propagation . Existing approaches to whole-brain imaging involve sequential acquisition of stacks of 2D images [ 2 ] , 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 a 3D sample volume using non-local optical encodings to combine information corresponding to different depth planes in the volume into a single 2D image , allowing for fast volume acquisition at 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 106 free parameters , e.g . pixels on a spatial light modulator ( SLM ) , enable the implementation of a 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-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 endto-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 either use far fewer optical parameters or operate on relatively compact or 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 : ( 1 ) the typical convolutional networks with small filters used by previous methods for computational reconstruction have a prior to decode locally , and therefore limit optical designs to 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 a 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 a parallel , multi-GPU optical simulation framework . Problem statement We define an optical encoder as Mφ parameterized by φ and a computational decoder as Rθ parameterized by θ . We wish to develop a reconstruction network architecture for decoding non-local encodings in 2D images c produced by Mφ and also to enable the end-to-end optimization of φ to produce non-local encodings . We ’ ll consider two concrete applications : ( 1 ) 3D snapshot microscopy using an SLM-based programmable microscope ( Figure 1A ) , where we wish to optimize both the reconstruction network Rθ and the optical encoding Mφ 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 to optimize only the reconstruction network Rθ 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 on a 4f optical model with a phase mask ( φ ) implemented using a spatial light modulator ( SLM ) , described further in Appendix A.1 . SLMs φ can have over 106 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 programmable 3D snapshot microscope . 3 . We introduce an efficient new reconstruction network architecture for decoding from nonlocal optical encoders using very large global convolutions implemented via Fourier convolutions . 4 . We show that our networks outperform the state-of-the-art 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 106 parameters . In simulation , we demonstrate the potential for significant improvements in imaging resulting from samplespecific 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 UNet-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 a non-end-to-end manner to reconstruct volumes from 3D snapshot microscopes designed using microlens arrays [ 33 , 34 ] . For photography and MRI , using deep learning 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-end learning of highly non-local optical encoders for 3D snapshot microscopy . Large convolution filters have been shown to be helpful for other computer vision applications such as 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 , 1 ] . A variety of microlens-array-based light field microscopes have been used to perform whole-brain imaging [ 16 , 27 , 37 , 36 , 30 , 26 , 7 , 12 , 38 ] . [ 38 ] optimizes the placement of microlenses , but not in an end-to-end manner . Despite variation in design , microlens-based 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 a coded detection strategy using a 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 a custom optical element or using a programmable SLM [ 10 , 14 , 35 , 13 , 38 , 23 ] . Using a programmable element allows different microscope parameters to be used for different samples . 2 METHODS . We show our network architecture and an overview of autoencoder training both the microscope parameters φ and reconstruction network parameters θ in Figure 1A . The programmable microscope is simulated by a differentiable implementation of a wave-optics model of light propagation . We have selected a programmable microscope design based on pupil-plane phase modulation with a programmable spatial light modulator , for which imaging is well-approximated by a computationallyefficient convolution [ 11 ] . 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 . 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 a 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 ] . FourierNet We propose a simple three layer convolutional network architecture with very large global convolutions at the very first layer , followed by two standard local convolutional layers ( Figure 1B ) . 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 a 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 a 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 down-sampling in image space corresponds to a simple cropping operation in the Fourier domain , resulting in a band-limited computation of a feature map . We efficiently implement multi-scale global Fourier convolutions ( Figure 1C ) to replace the encoding/ “ analysis ” pathway of a UNet . We then use the standard decoding/ “ synthesis ” 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 . 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 O ( N2 ) 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 kernel w , perform element-wise multiplication in Fourier space , and finally inverse Fourier transform , requiring onlyO ( N logN ) operations [ 19 , 32 ] . Following [ 28 ] , we store and optimize the weights in Fourier space W . This over-parameterization costs 8× 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 a Fourier convolution is defined : Re { F−1 { W F { x } } } ( 1 ) For image and kernel sizes of 256× 256 pixels , our implementation leads to nearly 500× speedup : standard PyTorch convolution takes 2860ms , while Fourier convolution takes 5.92ms on a TITAN X. Multi-scale Fourier domain convolutions It is well-known [ 24 ] that downsampling corresponds to cropping in the Fourier domain . Thus the Fourier convolution can be extended to efficiently produce multi-scale feature representations in one step ( Figure 1C ) . We define our multi-scale Fourier convolution as { Re { F−1 { W1 crop1 [ c ] } } , . . . , Re { F−1 { Wn cropn [ c ] } } } ( 2 ) where subscript denotes scale level ( higher subscript indicates lower spatial scale/more cropping in Fourier space ) and we precompute c : = F { c } once .

Summary: 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.