Deep probabilistic subsampling for task-adaptive compressed sensing

1 INTRODUCTION . In many real-world prediction problems , acquiring data is expensive and often bandwidthconstrained . Such is the case in regimes as medical imaging ( Lustig et al. , 2007 ; Choi et al. , 2010 ; Chernyakova & Eldar , 2014 ) , radar ( Baraniuk , 2007 ) , and seismic surveying ( Herrmann et al. , 2012 ) . By carefully reducing the number of samples acquired over time , in pixel-coordinate space or in kspace , efficient subsampling schemes lead to meaningful reductions in acquisition time , radiation exposure , battery drain , and data transfer . Subsampling is traditionally approached by exploiting expert knowledge on the signal of interest . Famously , the Nyquist theorem states that when the maximum frequency of a continuous signal is known , perfect reconstruction is possible when sampled at twice this frequency . More recently , it has been shown that if the signal is sparse in a certain domain , sub-Nyquist rate sampling can be achieved through compressive measurements and subsequent optimization of a linear system under said sparsity prior ; a framework known as compressed sensing ( CS ) ( Donoho et al. , 2006 ; Eldar & Kutyniok , 2012 ; Baraniuk , 2007 ) . CS methods however lack in the sense that they do not ( under a given data distribution ) focus solely on the information required to solve the downstream task of interest , such as disease prediction or semantic segmentation . Formalizing such knowledge is challenging in its own right and would require careful analysis for each modality and downstream task . In this work , we propose to explore ∗Equal contribution . 1https : //github.com/IamHuijben/Deep-Probabilistic-Subsampling.git the deep learning hypothesis as a promising alternative : reducing the need for expert knowledge in lieu of large datasets and end-to-end optimization of neural networks . As subsampling is non-differentiable , its integration into an end-to-end optimized deep learning model is non-trivial . Here we take a probabilistic approach : rather than learning a subsampling scheme directly , we pose a probability distribution that expresses belief over effective subsampling patterns and optimize the distribution ’ s parameters instead . To enable differentiable sampling from this distribution , we leverage recent advancements in a continuous relaxation of this sampling process , known as Gumbel-softmax sampling or sampling from a concrete distribution ( Jang et al. , 2017 ; Maddison et al. , 2016 ) . This enables end-to-end training of both the subsampling scheme and the downstream model . Naively , the number of parameters of a distribution over an n-choose-k problem scales factorially , which is intractable for all practical purposes . We propose a novel , expressive yet tractable , parameterization for the subsampling distribution that conditions on the output sample index . We hypothesize that such conditioning prevents redundant sampling : it enables modeling the scenario in which multiple candidate samples can be equally good , yet redundant in combination . Furthermore , we investigate a parameter-restricted approach with a single parameter per candidate sample , balancing tractability and exploration ( Kool et al. , 2019 ) . In this case , we adopt a continuous relaxation of top-K sampling to guarantee differentiability ( Plötz & Roth , 2018 ) . Our main contributions are as follows : • DPS : A new regime for task-adaptive subsampling using a novel probabilistic deep learning framework for incorporating a sub-Nyquist sampling scheme into an end-to-end network . • DPS enables joint optimization of a subsampling pattern with a predictive downstream model , without the need for explicit knowledge on a sparsifying basis . • We demonstrate improved performance over strong subsampling baselines in image classification and reconstruction , while sampling both in Fourier and pixel space . 2 RELATED WORK . Recent works have proposed deep-learning-based subsampling methods for fast MRI , historically being one of the most prominent applications of CS . Weiss et al . ( 2019 ) exploit gradient backpropagation to a fixed set of real-valued coordinates , enabled by their subsequent ( limited-support ) interpolation on the discrete k-space grid , and Bahadir et al . ( 2019 ) formulate the sampling problem by learning pixel-based thresholding of i.i.d . samples drawn from a uniform distribution , referred to as LOUPE . Where the former suffers from limited exploratory capabilities ( likely due to its compact support ) , with learned sampling schemes typically not deviating far from their initialization , the latter controls the sample rate only indirectly , through the addition of a sparsity-promoting ` 1 penalty on the sampling mask . Related to our methodology , yet different in purpose , Plötz & Roth ( 2018 ) ; Xie & Ermon ( 2019 ) ; Kool et al . ( 2019 ) leverage an extension to the Gumbel-max trick ( Gumbel , 1954 ) for ( ordered ) subset selection using a categorical distribution with N − 1 free parameters . They rely on top-K sampling , as first proposed by Vieira ( 2014 ) . Based upon the continuous relaxation of the categorical distribution , known in the deep learning community as the Gumbel-softmax trick or the concrete distribution ( Jang et al. , 2017 ; Maddison et al. , 2016 ) , Plötz & Roth ( 2018 ) show that such a relaxation also exists for top-K sampling . In this work , we investigate the use of Gumbel-based relaxations for task-adaptive compressive sampling , and further propose a novel parametrization where we condition the sampling distribution on the output sample index . We differentiate our contribution from deep encoder-decoder methods for data compression ( Baldi & Hornik , 1989 ; Hinton & Zemel , 1993 ; Blier & Ollivier , 2018 ; Habibian et al. , 2019 ) , which do not aim at reducing data rates already at the sensing and digitization stage . Related work by Mousavi et al . ( 2019 ) and Wu et al . ( 2019 ) , focusses on the problem of learning compressive linear encoders/filters , rather than discrete subsampling as addressed here . The authors of Balın et al . ( 2019 ) use a Gumbel-max-based auto-encoder for discrete feature selection , which is however not task-adaptive . Through the lens of contemporary deep learning , subsampling can be interpreted as a form of attention ( Bahdanau et al. , 2014 ; Kim et al. , 2017 ; Parikh et al. , 2016 ; Vaswani et al. , 2017 ) . Rather than attending on intermediate representations , our model “ attends ” directly on the input signal . For subsampling to be effective , sparse weights are essential . In the space of attention , this is known as hard attention ( Xu et al. , 2015 ) , and is typically optimized using the REINFORCE gradient estimator ( Williams , 1992 ) . In contrast to the method of attention as applied in these works , our method aims for a fixed , reduced subsampling rate . 3 METHOD . 3.1 TASK-ADAPTIVE SYSTEM MODEL . We consider the problem of predicting some downstream task s , through a learned subsampling scheme A on a fully-sampled signal x ∈ RN , resulting in measurements y ∈ RM , with M < < N : y = Ax , ( 1 ) where A ∈ { 0 , 1 } M×N is the subsampling measurement matrix . We here concern ourselves specifically with scenarios in which the rows of A are constrained to having cardinality one , i.e . ||am||0 = 1 , ∀m ∈ { 1 , ... , M } . In such cases , A serves as a subset selector , sampling M out of N elements in x . From the resulting low-rate measurements y , we aim at predicting task s through : ŝ = fθ ( y ) , ( 2 ) with fθ ( · ) being a function that is differentiable with respect to its input and parameters θ , e.g . a neural network . Given a downstream task and dataset , we are interested in learning both the optimal processing parameters θ , and the sampling scheme as described in eq . ( 1 ) . To circumvent the non-differential nature of discrete sampling , we will introduce a novel fully probabilistic sampling strategy that allows for gradient-based learning through error backpropagation , on which we detail in the following section . 3.2 DPS : DEEP PROBABILISTIC SUBSAMPLING . Since direct optimization of the elements in A is intractable due to its combinatorial nature , we here instead propose to leverage a tractable generative sampling model that is governed by a learned subsampling distribution , parameterized by Φ : A ( Φ ) ∼ P ( A|Φ ) . ( 3 ) Thus , rather than optimizing A ( Φ ) , we optimize the distribution parameters Φ . To warrant sufficient expressiveness while maintaining tractability , we learn the parameters φm ∈ RN of M independent categorical distributions ( rather than the joint distribution , which scales factorially ) , being the rows of Φ ∈ RM×N . Formally , we define each mth measurement row am ∈ { 0 , 1 } N in A ( Φ ) as a one-hot encoding of an independent categorical random variable rm ∼ Cat ( N , πm ) . We define πm ∈ RN = { πm,1 , . . . , πm , N } , being a vector containing N class probabilities , and parameterize it in terms of its unnormalized logits φm such that : πm , n = exp φm , n∑N i=1 exp φm , i . ( 4 ) We sample from Cat ( N , πm ) by leveraging the Gumbel-max trick ( Gumbel , 1954 ) , a reparameterization of the sampling procedure . It perturbs the unnormalized logits φm , n with i.i.d . Gumbel noise2 samples em , n ∼ Gumbel ( 0 , 1 ) . The argmax of these perturbed logits , results in a sample from the original categorical distribution Cat ( N , πm ) . A realization r̃m is then defined as : r̃m = argmax n { wm−1 , n + φm , n + em , n } , m ∈ { 1 , . . . , M } , ( 5 ) in which wm−1 , n ∈ { −∞ , 0 } for n ∈ { 1 . . . N } mask previous realizations r̃1 , · · · , r̃m−1 , by adding−∞ to the logit of the previously selected category , enforcing sampling without replacement among the M distributions . Introducing the function one hotN ( · ) as the operator that returns a one-hot vector of length N , we finally obtain : am = one hotN { r̃m } = one hotN { argmax n { wm−1 , n + φm , n + em , n } } . ( 6 ) To permit error backpropagation for efficient optimization of Φ , we require ∇φmam to exist ∀m ∈ { 1 , . . . , M } . Since argmax ( · ) is a non-differentiable operator , we adopt the Straight-Through Gumbel Estimator ( Jang et al. , 2017 ; Maddison et al. , 2016 ) as a surrogate for∇φmam : ∇φm am : = ∇φmEem [ softmaxτ ( wm−1 + φm + em ) ] = ∇φmEem [ exp { ( wm−1 + φm + em ) /τ } ∑N i=1 exp { ( wm−1 , i + φm , i + em , i ) /τ } ] , ( 7 ) with ( row operator ) softmaxτ ( · ) as a continuous differentiable approximation of the one-hot encoded argmax ( · ) operation . Appendix A provides the full derivation of∇φmam , and fig . 1 shows a schematic overview of the proposed framework . We refer to sampling using the softmaxτ ( · ) function as soft sampling . Its temperature parameter τ serves as a gradient distributor over multiple entries ( i.e . logits ) in φm . Using a relatively high value enables updating of multiple logits during training , even though a hard sample was taken in the forward pass . In the limit of τ → 0 , soft sampling approaches the one-hot encoded argmax ( · ) operator in eq . ( 6 ) ( Jang et al. , 2017 ; Maddison et al. , 2016 ) . While lowering τ reduces the gradient estimator ’ s bias , it comes at the cost of a higher variance . We find however that using a fixed ( tuned ) value of τ worked well for our experiments , and do not explore methods to reduce the introduced gradient bias ( Grathwohl et al. , 2017 ; Tucker et al. , 2017 ) . 2The Gumbel distribution is typically used to model the maximum of a set of independent samples . Its probability density function is of the form : f ( x ) = 1 β e−z−e −z , with z = x−µ β , where µ and β are the mean and standard deviation , respectively . The standard Gumbel distribution , with µ = 0 and β = 1 is adopted in the Gumbel-max trick . To study the trade-off between parameter efficiency and convergence , we further study a more parameter-restricted version of DPS that shares the weights φm across all M distributions . This allows for a more memory-efficient implementation , by leveraging the Gumbel top-K trick ( Kool et al. , 2019 ) ( rather than Gumbel-max ) , and its corresponding relaxation ( Plötz & Roth , 2018 ) . We refer to this particular model as DPS-topK , and adopt DPS-top1 to indicate the sampling model presented in equations 6 and 7 , where weights are not shared across distributions , and one sample is drawn from each distribution . Algorithm 1 in appendix B describes the DPS algorithm .

Summary: This paper introduces  a novel DPS(Deep Probabilistic Subsampling) framework for the task-adaptive  subsampling case, which attempts to resolve the issue of end-to-end optimization of an optimal subset of signal with jointly learning a sub-Nyquist sampling scheme and a predictive model for downstream tasks. The parameterization is used to simplify the subsampling distribution and ensure an expressive yet tractable distribution.  The new approach contribution is applied to  both reconstruction and classification tasks and demonstrated with a suite of experiments in a toy dataset, MINIST, and COFAR10.