Active Deep Probabilistic Subsampling

1 INTRODUCTION . Present-day technologies produce and consume vast amounts of data , which is typically acquired using an analog-to-digital converter ( ADC ) . The amount of data digitized by an ADC is determined not only by the temporal sampling rate , but also by the manner in which spatial acquisitions are taken , e.g . by using a specific design of sensor arrays . Reducing the number of sample acquisitions needed , can lead to meaningful reductions in scanning time , e.g . in Magnetic Resonance Imaging ( MRI ) , radiation exposure , e.g . in Computed Tomography ( CT ) , battery drain , and bandwidth requirements . While the Nyquist theorem is traditionally used to provide theoretical bounds on the sampling rate , in recent years signal reconstruction from sub-Nyquist sampled data has been achieved through a framework called Compressive Sensing ( CS ) . First proposed by Donoho ( 2006 ) , and later applied for MRI by Lustig et al . ( 2007 ) , CS leverages structural signal priors , specifically sparsity under some known transform . By taking compressive measurements followed by iterative optimization of a linear system under said sparsity prior , reconstruction of the original signal is possible while sampling at sub-Nyquist rates . Researchers have employed CS with great success in a wide variety of applications , such as radar ( Baraniuk & Steeghs , 2007 ; Ender , 2010 ) , seismic surveying ( Herrmann et al. , 2012 ) , spectroscopy ( Sanders et al. , 2012 ) , and medical imaging ( Han et al. , 2016 ; Lai et al. , 2016 ) . However , both the need to know the sparsifying basis of the data , and the iterative nature of the reconstruction algorithms , still hamper practical applicability of CS in many situations . These limitations can be overcome by the use of deep learning reconstruction models that make the sparsity assumption implicit , and facilitate non-iterative inference once trained . Moreover , the ( typically random ) nature of the measurement matrix in CS does , despite adhering to the given assumptions , not necessarily result in an optimal measurement given the underlying data statistics and the downstream system task . This has recently been tackled by algorithms that learn the sampling scheme from a data distribution . In general , these data-driven sampling algorithms can be divided into two categories : algorithms that learn sampling schemes which are fixed once learned ( Huijben et al. , 2020a ; b ; c ; Ravishankar & Bresler , 2011 ; Sanchez et al. , 2020 ; Bahadir et al. , 2019 ; Bahadir et al. , 2020 ; Weiss et al. , 2019 ) , and algorithms that learn to actively sample ( Ji et al. , 2008 ; Zhang et al. , 2019 ; Jin et al. , 2019 ; Pineda et al. , 2020 ; Bakker et al. , 2020 ) ; selecting new samples based on sequential acquisition of the information . The former type of algorithms learn a sampling scheme that - on average - selects informative samples of all instances originating from the training distribution . However , when this distribution is multi-modal , using one globally optimized sampling scheme , can easily be sub-optimal on instance-level . Active acquisition algorithms deal with such shifts in underlying data statistics by conditioning sampling behavior on previously acquired information from the instance ( e.g . the image to be sampled ) . This results in a sampling sequence that varies across test instances , i.e . sampling is adapted to the new data . This adaptation as a result of conditioning , promises lower achievable sampling rates , or better downstream task performance for the same rate , compared to sampling schemes that operate equivalently on all data . In this work , we extend the Deep Probabilistic Subsampling ( DPS ) framework ( Huijben et al. , 2020a ) to an active acquisition framework by making the sampling procedure iterative and conditional on the samples already acquired , see Fig . 1 . We refer to our method as Active Deep Probabilistic Subsampling ( A-DPS ) . We show how A-DPS clearly exploits the ten different modalities ( i.e . the digits ) present in the MNIST dataset to adopts instance-adaptive sampling sequences . Moreover , we demonstrate both on MNIST ( LeCun et al. , 1998 ) and the real-world fast MRI knee dataset ( Zbontar et al. , 2018 ) , that A-DPS outperforms other state-of-the-art models for learned sub-Nyquist sampling . We make all code publicly available upon publication , in order to facilitate benchmarking to all provided baselines and A-DPS in future research . 2 RELATED WORK . Recently , several techniques for learning a fixed sampling pattern have been proposed , especially in the field of MR imaging , in which Ravishankar & Bresler ( 2011 ) were one of the firsts . In this work , the authors make use of non-overlapping cells in k-space , and move samples between these cells.During training Ravishankar & Bresler ( 2011 ) alternate between reconstruction and relocation of sampling positions . After a reconstruction step they sort the cells in terms of reconstructing error and an infinite-p norm . Selected samples from lower scoring cells are relocated to higher scoring cells in a greedy fashion . Sanchez et al . ( 2020 ) also propose a greedy approach , in which samples are not relocated between cells , but greedily chosen to optimize a reconstruction loss on a batch of examples . Both of the types of greedy optimization do however not allow for joint learning of sampling together with a downstream reconstruction/task model , as the reconstruction has to either be parameter-free or pretrained to work well with a variety of sampling schemes . Bahadir et al . ( 2019 ) on the other hand propose to learn the sampling pattern by thresholding pixelbased i.i.d . samples drawn from a uniform distribution , dubbed Learning-based Optimization of the Under-sampling PattErn ( LOUPE ) . The sample rate of LOUPE is indirectly controlled by promoting sparsity through the use of an ` 1 penalty on the thresholds . One of the first active sampling schemes was proposed by Ji et al . ( 2008 ) , who leverage CS reconstruction techniques that also give a measure of uncertainty of the reconstruction using Bayesian modeling . Ji et al . ( 2008 ) leveraged this uncertainty in the reconstruction to adaptivly select the next measurement that will reduce this uncertainty by the largest amount . However , this method - and other similar works from ( Carson et al. , 2012 ; Li et al. , 2013 ) - rely on linearly combined measurements , rather than discrete sampling , with which we concern ourselves here . In the field of MRI , Zhang et al . ( 2019 ) propose an active acquisition scheme by leveraging a reconstruction and adversarial neural network . Whereas the reconstruction network is trained to reconstruct MR images from the subsampled Fourier space ( k-space ) , the adversarial network is trained to distinguish between already sampled , and omitted lines in this space . The k-space line that is most believed to be ‘ fake ’ ( i.e . filled in by the reconstruction network ) by the adversarial network , is sampled next . However , This framework only works for undersampled Fourier to image reconstruction tasks , as the discriminator requires mappings of the image in k-space . Jin et al . ( 2019 ) put forth an active acquisition scheme for MRI by leveraging reinforcement learning ( RL ) . Two neural networks , one for sampling and one for reconstruction are trained jointly using a Monte-Carlo tree search , resulting in a sampling policy that is dependent on the current reconstruction of the image . Concurrently with our work , both Pineda et al . ( 2020 ) and Bakker et al . ( 2020 ) proposed RL-based active acquisition techniques . Pineda et al . ( 2020 ) leverages a Double Deep Q-Network . The model is trained using a modified -greedy policy , in which the best action is taken with probability 1− , and an exploratory action is taken with probability . Bakker et al . ( 2020 ) compare greedy with nongreedy training , finding that the greedy method leads to a higher degree of adaptability , especially for tasks with a long horizon ( i.e . more samples to be taken ) . Both of the frameworks proposed by ( Pineda et al. , 2020 ) and Bakker et al . ( 2020 ) make use of a pretrained reconstruction network , which differs from the proposed A-DPS method that enables joint training of both the reconstruction ( task ) network and sampling network . Even though subsampling is an extreme form of data compression , we differentiate from typical data compression architectures like deep encoder-decoder structures ( Theis et al. , 2017 ; Ballé et al. , 2017 ) , as these methods do not reduce data rates at the measurement stage . The feedback recurrent autoencoder proposed by Yang et al . ( 2020 ) is however related to A-DPS through its use of a recurrent context . But whereas Yang et al . ( 2020 ) learn a recurrent context to inform the encoder stage of the network , A-DPS uses this to inform the sampling pattern . 3 METHOD . 3.1 GENERAL FRAMEWORK . Given a prediction task s we are interested in learning to predict an optimal subsampling scheme A ∈ { 0 , 1 } M×N ( with M N ) on an input signal x ∈ RN , resulting in a measurement ỹ ∈ RM : ỹ = Ax . ( 1 ) Each row in A is constrained to have ` 0-norm of 1 , while each column in A is constrained to have an ` 0-norm of either 0 or 1 , i.e . each of the N candidate samples is selected at most once . In the rest of this paper we will index these candidate samples with n ∈ { 1 , . . . , N } , and the selected samples with m ∈ { 1 , . . . , M } . The percentage of selected samples from the candidate samples is called the sampling ratio r =M/N · 100 % . We also introduce a non-compressed form of the measurement ỹ , called y ∈ RN , that contains N −M zeros , and M non-zeros at the sampled indices specified by A , i.e . the masked input . This way , the location of samples from x is preserved , which is especially useful whenA changes during training . To acquire y from x , one seeks a subsampling mask d that can be applied on x via : y = d x = ATAx , ( 2 ) where denotes an element-wise multiplication . From the resulting measurement y we then aim at predicting the downstream task s through : ŝ = fθ ( y ) , ( 3 ) where fθ ( . ) is a function that is differentiable with respect to its input and parameters θ , e.g . a neural network . Normally , optimization of the task model fθ ( . ) is achieved through backpropagation of some loss function L ( s , ŝ ) . However , calculating gradients on the sampling matrix is blocked by its combinatorial nature , inhibiting joint training of the task with the sampling operation . The DPS framework provides a solution to this problem , on which we will elaborate in the next section . 3.2 DPS : DEEP PROBABILISTIC SUBSAMPLING . To enable joint training of the sampling operation with the downstream task model , Huijben et al . ( 2020a ) introduce DPS . Rather than optimizing A directly , they propose to optimize a generative sampling model P ( A|φ ) , whereφ are learned unnormalized logits of ( possibly multiple ) categorical distribution ( s ) . Each distribution expresses the probabilities for sampling any of the elements xn fromx through sampling matrixA . More specifically , φm , n is the log-probability for setting am , n = 1 , and thus sampling xn as mth sample . To generate a sampling pattern from these unnormalized logits , i.e . implementation of this conditional model , the Gumbel-max trick is leveraged ( Gumbel , 1954 ) . In the Gumbel-max trick the unnormalized logits are perturbed with i.i.d . Gumbel noise samples em , n ∼ Gumbel ( 0 , 1 ) . By selecting the maximum of this perturbation a realization of the sampling mask can be found using : Am , : = one-hotN { argmax n { wm−1 , n + φm , n + em , n } } , ( 4 ) where Am , : denotes the m-th row of A and one-hotN creates a one-hot vector of length N , with the one at the index specified by the argmax operator . Moreover , the cumulative mask wm−1 , n ∈ { −∞ , 0 } masks previously selected samples by adding minus infinity to those logits , thereby ensuring sampling without replacement . During backpropagation , gradients are computed by relaxing this sampling procedure using the Gumbel-softmax trick ( Jang et al. , 2016 ; Maddison et al. , 2017 ) , resulting in : ∇φmAm , : : = ∇φmEem [ softmaxτ { wm−1 , n + φm , n + em , n } ] , ( 5 ) where τ denotes the temperature parameter of the softmax operator . Setting τ > 0 results in a smoothed sampling matrix A ( i.e . elements can have values between 0 and 1 as well ) , allowing gradients to distribute over multiple logits during training . In the limit of τ → 0 the softmax operator approaches the one-hot argmax function of equation ( 4 ) . Although this approach – also known as straight-through Gumbel-softmax – leads to biased gradients , it has been shown to work well in practice , and Huijben et al . ( 2020a ) keep τ at a fixed value during training . Huijben et al . ( 2020a ) propose two regimes of DPS . First , Top-1 sampling , an expressive form of DPS where each of the M selected samples are separately conditioned on all N candidate samples , resulting in M ×N trainable logits φm , n . Second , Top-M sampling ( called Top-K in their paper ) , a constrained form where all M samples together are conditioned on all N candidate samples , i.e . the logits φn are shared between the M rows of A , resulting in only N trainable logits . While Top-1 sampling is more expressive , Huijben et al . ( 2020a ) noticed slightly better results for the Top-M regime , possibly thanks to the smaller number of trainable logits , therefore facilitating optimization . For scaleability reasons , we thus choose to continue with Top-M sampling in this work and refer to this regime as DPS in the rest of this paper . We refer the reader to Huijben et al . ( 2020a ) for more details regarding DPS .

Summary: This paper develops methods to perform active subsampling. That is, given some downstream task like classification or image reconstruction, it sequentially selects which elements of an image or signal to sample so as to perform said task. It does so by extending the Deep Probabilistic Subsampling (DPS) method developed by Huijben et al. The proposed method is applied to two problems as well as a simplified, low-resolution MRI reconstruction problem.