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 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).