TASK-ADAPTIVE COMPRESSED SENSING
Year: 2020

Abstract:
The field of deep learning is commonly concerned with optimizing predictive models using large pre-acquired datasets of densely sampled datapoints or signals. In this work, we demonstrate that the deep learning paradigm can be extended to incorporate a subsampling scheme that is jointly optimized under a desired sampling rate. We present Deep Probabilistic Subsampling (DPS), a widely applicable framework for task-adaptive compressed sensing that enables end-to-end optimization of an optimal subset of signal samples with a subsequent model that performs a required task. We demonstrate strong performance on reconstruction and classification tasks of a toy dataset, MNIST, and CIFAR10 under stringent subsampling rates in both the pixel and the spatial frequency domain. Thanks to the data-driven nature of the framework, DPS is directly applicable to all realworld domains that benefit from sample rate reduction. The code used for this paper is made publicly available1.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD

heading: 3.1 TASK-ADAPTIVE SYSTEM MODEL
text: 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.

heading: 3.2 DPS: DEEP PROBABILISTIC SUBSAMPLING
text: 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.

heading: 4 EXPERIMENTS
text: We test the applicability of the proposed task-adaptive DPS framework for three datasets and two distinct tasks: image classification and image reconstruction. We explore subsampling in pixelcoordinate space as well as in k-space. The latter is relevant for scenarios in which data is acquired in the frequency domain, such as (but not limited to) magnetic resonance imaging (MRI).

heading: 4.1 MNIST CLASSIFICATION
text: Experiment setup Classification performance was tested on the MNIST database (LeCun et al., 1998), comprising 70,000 28 × 28 grayscale images of handwritten digits 0 to 9. We split the dataset into 50,000 training images, 5,000 validation, and 5,000 test images. We train DPS-top1 and DPS-topK to take partial measurements in either the image or Fourier domain, and process them through the task model to yield a classification outcome. Results are compared to those obtained using uniformly distributed pixel/Fourier samples, a sampled disk/low pass filter, and the data-driven LOUPE method (Bahadir et al., 2019).
Besides, to show the benefit of task-adaptive sampling, i.e. jointly training sampling with the task, we add a case in which we disjointly learn the classification task from the sampling pattern. As such, we train DPS-topK for MNIST reconstruction, after which the learned sampling pattern is frozen, and the classifier is trained.
Task model After sampling M elements, all N zero-masked samples (or 2N in the case of complex Fourier samples) are passed through a series of 5 fully-connected layers, having N , 256, 128, 128 and 10 output nodes, respectively. The activations for all but the last layer were leaky ReLUs, and 20% dropout was applied after the first three layers. The 10 outputs were normalized by a softmax function to yield the respective classification probabilities. Table 1 in appendix C shows a detailed overview of the complete architecture of the task model. Zero-filling and connecting all possible samples, rather than only connecting the M selected samples, facilitated faster co-adaptation of the network to different sampling patterns during training.
Training details We train the network to maximize the log-likelihood of the observations D = {(xi, si) | i ∈ 0, . . . , L} through minimization of the categorical cross-entropy between the predictions and the labels, denoted by Ls. We moreover promote training towards one-hot distributions Cat(N,πm) by penalizing high entropy, adopting:
Le = − M∑ m=1 N∑ n=1 πm,n log πm,n, (8)
with πm,n defined as in eq. (4). The total optimization problem is thus:
Φ̂, θ̂ = argmin Φ,θ
{ E(x,s)∼pDLs + µLe } , (9)
where pD is the data generating distribution. Penalty multiplier µ was set to linearly increase 1e−5 per epoch, starting from 0.0. Increasing this entropy penalty enforces one-hot distributions, and therefore less variations in the realizations, i.e. sampling patterns, when training evolves. The temperature parameter τ in eq. (7) was set to 2.0, and the sampling distribution parameters Φ were initialized randomly, following a zero-mean Gaussian distribution with standard deviation 0.25. Equation 9 was optimized using stochastic gradient descent on batches of 32 examples, approximating the expectation by a mean across the train dataset. To that end, we used the ADAM solver
(β1 = 0.9, β2 = 0.999, and = 1e−7 ) (Kingma & Ba, 2014), and we trained until the validation loss plateaued. We adopted different learning rates for the sampling parameters Φ and the parameters of the task model θ, being 2e− 3 and 2e− 4, respectively. Results The results presented in fig. 2a show that image-domain sampling using DPS significantly outperforms the fixed sampling baselines (uniform and disk). The data-driven LOUPE method is outperformed in case of a strong subsampling factor. The resulting patterns qualitatively demonstrate how, for this task, a sensible selection of pixels that are most informative was made (slightly slanted, and capturing discriminative areas). Notably, partial Fourier measurements (fig. 2b) allowed for a much greater reduction of the number samples, with DPS sampling outperforming uniform sampling and LOUPE, and showing similar performance as fixed disk sampling. Interestingly, the DC and very-low frequency components were consistently not selected.
Figure 3 shows the learned sampling pattern in pixel space (96.8% pixels removed), when training DPS-topK for the reconstruction task. The adopted reconstruction network architecture is the same as the one used for CIFAR10 reconstruction (see section 4.3). The subsequently trained classifier resulted in a classification error of 9.3%, compared to 6.6% for task-adaptive learning with DPS-topK. This case clearly illustrates the benefit of task-adaptive learning of a sampling pattern.

heading: 4.2 ‘LINES AND CIRCLES’ IMAGE RECONSTRUCTION
text: Experiment setup To evaluate reconstruction of structured images from highly undersampled partial Fourier (k-space) measurements (keeping 3.1% of the coefficients), we generated synthetic toy data comprising images that each contain up to 5 horizontal lines and randomly-sized circles. Lines and circles were placed at random positions and their pixel intensity was drawn from a uniform distribution between 1 and 10. Examples were generated in an on-line fashion during training. A
pre-generated hold-out test set of 1000 randomly generated examples was used for all cases. Two illustrative test examples, along with their Fourier-domain representations, are given in Figure 4(a,b). We compare the results to those obtained using three fixed partial Fourier measurement baselines, following uniform, random, and low pass subsampling patterns, respectively, and the data-driven LOUPE method.
Task model Image reconstruction from partial Fourier measurements was performed by following the methodology in Zhu et al. (2018). We use a deep neural network consisting of two subsequent fully-connected layers with tanh activations that map the 2N (zero-filled) Fourier coefficients (stacked real and imaginary values) to a vector of length N . This vector was subsequently reshaped into a √ N × √ N image, and processed by 3 convolutional layers. Table 2 in appendix C provides the layer parameters of this task network.
Training details The optimization problem is the same as in eq. (9), however with the loss Ls now defined as the negative log-likelihood:
Ls = E(x,s)∼pD ‖fθ(A(Φ)x)− s‖ 2 2. (10)
The learning rates for Φ and θ were 1e− 3 and 1e− 4, respectively, and µ and τ were respectively set to 2e − 4 and 5.0. The ADAM optimizer (with settings as provided in section 4.1) was used to train in mini-batches of 128 examples. Training was stopped when the validation loss plateaued.
Results An overview of the results is given in fig. 4. As expected, uniform subsampling leads to strong spatial aliasing that can not be recovered by the task model due to violation of the Nyquist criterion. In turn, random subsampling introduces an incoherent aliasing pattern, that can only partly be recovered. Although not suffering from aliasing, low pass sampling deteriorates resolution, of which the effect is particularly evident for the broadband/sharp horizontal lines. In contrast, data-driven sampling methods (LOUPE, DPS-top1, and DPS-topK) prevent aliasing by learning an appropriate sampling pattern. DPS sampling specifically, yields high-resolution accurate reconstructions with an improved PSNR evaluated on the entire test set compared to the baselines. Note how the learned
sampling patterns (top row) have evolved from a random initialization, similar to that of (c), to the ultimate task-adaptive scheme (e-g).

heading: 4.3 CIFAR10 IMAGE RECONSTRUCTION
text: Experiment setup The CIFAR10 database (Krizhevsky et al., 2009) contains 60,000 images of 32 × 32 pixels in 10 different classes. We converted all images to grayscale, and subsequently split them into 50,000 training images, 5,000 validation and 5,000 test images. We again learn partial Fourier sampling and image reconstruction, keeping only 12.5% of the Fourier coefficients, and compare ourselves to the three fixed partial Fourier measurement baselines and LOUPE, as described in section 4.2.
Task model The challenging reconstruction task for CIFAR10 motivates the adoption of a structured model to enable strong reconstruction. We draw inspiration from iterative proximal-gradient schemes (Parikh et al., 2014), which are dedicated to solving the ill-posed linear measurement problem in eq. (1). To that end, we unfold K = 5 such iterations, learning an adequate image-domain proximal mapping P(k)θ and stepsize α(k) at each fold:
ŝ(k+1) = P(k)θ ( ŝ(k) − α(k)FHAT(Φ) ( A(Φ)Fŝ (k) −A(Φ)x )) , (11)
where F ∈ CN×N is a discrete Fourier transform (DFT) matrix, and (·)H denotes the Hermitian (conjugate transpose). In the above formulation, at each fold a step is taken towards the samplingconsistent subspace that adequately represents the physical measurement of s by A(Φ)F. The trained proximal operator P(k)θ , a 7-layer convolutional network, then projects this onto the manifold of visually plausible images (Mardani et al., 2018), removing noise, aliasing, or blurring artifacts. Table 3 in appendix C provides provides details on the layers in the task model.
Training details Optimization settings were similar to those in section 4.2, leveraging a meansquared-error (negative log-likelihood) reconstruction cost and a distribution entropy penalty onπm. To promote visually plausible reconstructions, we added an adversarial (Ledig et al., 2017) cost by
adopting a discriminator network Dψ(·) that aims to discriminate between images reconstructed from partial Fourier measurements (ŝ) and actual images s. The discriminator comprised 3 convolutional layers with leaky ReLU activations, followed by global average pooling, 40% dropout, and a logistic binary classification model (parameters are provided in table 3, appendix C). The samplingand task model parameters were then trained to both minimize the negative log-likelihood and discriminator loss LDψ , in addition to the entropy penalty Le:
LDψ = Es,ŝ [log(Dψ(s)) + log(1−Dψ(ŝ))] , (12)
Φ̂, θ̂, ψ̂ = argmin Φ,θ argmax ψ
{ E(x,s)∼pD [ ‖fθ(A(Φ)x)− s‖ 2 2 ] + µLe + λLDψ } , (13)
where λweighs adherence to the data-driven MSE loss and the discriminator loss. It was empirically set to 0.004. The learning rates for {Φ, ψ} and θ were 1e − 3 and 2e − 4, respectively, and µ was empirically set to 1e − 6, balancing sharpness and high frequency hallucinations. Temperature parameter τ was set at a constant value of 2.0. Training was performed using the ADAM optimizer (with settings as given in section 4.1) in batches of 8 images, until the validation loss plateaued.
Results Figure 5 shows how task-adaptive sampling significantly outperforms all baselines. Both the uniform and random subsampling schemes suffer from severe aliasing which the task model is not able to adequately restore. While low pass sampling does not lead to aliasing, the absence of real high-frequency information causes the task model to ‘hallucinate’ such (super-resolution) content. Interestingly, LOUPE (e) learned a rather compact, i.e. around DC frequencies, sampling scheme, resulting in blurry images, but DPS learned a more wide-spread pattern, also capturing higher frequencies. The result of this is particularly notable for the example displayed in the bottom row of fig. 5, being less structured (or predictable) than that in the row above it. DPS thus enabled high-quality, high-resolution reconstructions that go beyond those obtained with the other methods, reaching a PSNR of 27.1 dB and 28.3 dB, and a structural similarity index (SSIM) of 0.88 and 0.90 across the 5000 test examples, for topK and top1 respectively. For the most competitive methods, i.e. fixed sampling low pass and data-driven LOUPE, these statistics were 25.3 dB and 0.82, and 25.3 dB and 0.86, respectively.

heading: 5 CONCLUSIONS
text: We have introduced Deep Probabilistic Subsampling (DPS), a framework that enables jointly learning a data- and task-driven sampling pattern, with a subsequent task-performing model. The framework is generic and can be combined with any network architecture that performs the required task using the subsampled set of signal elements. Empirically we find the method to perform strongly on toy datasets and canonical deep learning problems. Further work can explore the effectiveness of DPS in real-world problems such as MRI scanning. Even though LOUPE (Bahadir et al., 2019) learned similar subsampling patterns in all experiments, compared to DPS, its performance on the downstreak task showed to be lower. We argue that LOUPE probably suffers from the fact that samples in the training phase are relaxed. As such, the task network is optimized on a relaxed subsampled pattern, whereas a hard thresholded pattern is used for inference on the test set to make fair comparisons to DPS.
In comparing the parameter-restricted DPS-topK versus the more expressive DPS-top1, we find the former to perform better in our experiments. Studying the sampling distributions after training, we find that DPS-top1 exhibits more uncertainty. The insets in fig. 2 show this by the higher variance present among the 1000 realizations for DPS-top1 compared to DPS-topK. We tentatively attribute this due to difficulties with convergence. Further work will explore if the DPS-top1’s more expressive parametrization can improve over DPS-topK in more complex problems that might benefit from a broader hypothesis space during training, as well as under more carefully tuned softmax temperature annealing schemes.
Our data-driven method does not explicitly need knowledge regarding the, often unknown, sparsifying basis, whereas conventional CS algorithms do require this. Like all data-driven optimized methods, a learned sampling scheme is at risk of overfitting to a small training dataset. Although we
did not observe this issue in our experiments, careful regularization might be required to ensure that this effect is minimal in such high-risk tasks. The fully-differentiable DPS framework allows for flexibility, and interesting extensions can be explored in future work. Rather than learning a fixed set of parameters for the subsampling distribution, a neural network can be used to predict the parameters instead, conditioned on contextual data or the samples acquired so far. Finally, our method currently requires the desired sampling rate to be predetermined as a hyperparameter. Future work can explore if this rate can be jointly optimized to incorporate optimization of the subsampling rate.

heading: ACKNOWLEDGMENTS
text: This research was supported in part by Philips Research. It is also part of a research program Rubicon ENW 2018-3 with project number 019.183.EN.014, which is financed by the Dutch Research Council (NWO). We thank Wouter Kool and the anonymous reviewers for their valuable feedback.

heading: A GRADIENT OF GUMBEL-SOFTMAX SAMPLING
text: For any mth pair of rows (am,φm), any nth element am,n in am can be differentiated towards all elements in φm through:
∇φmam,n
= ∇φmEem [ softmaxτ (wm−1 + φm + em) ∣∣∣ n ] = Eem [ ∇φmsoftmaxτ (wm−1 + φm + em) ∣∣∣ n
] = Eem [ ∇φm
exp{(wm−1 + φm + em)/τ}∑N i=1 exp{(wm−1,i + φm,i + em,i)/τ} ∣∣∣ n ] (14)
Gumbel noise vector em can be reparametrized as a function of a uniform noise vector m ∼ U(0, 1) i.i.d., through: em = − log(− log( m)). (15) This allows rewriting eq. (14) into:
∇φmam,n = E m [ ∇φm
exp{(wm−1 + φm − log(− log( m)))/τ}∑N i=1 exp{(wm−1,i + φm,i − log(− log( m,i)))/τ} ∣∣∣ n ]
= ∫ ∞ −∞ . . . ∫ ∞ −∞ P [ m = [k1, · · · , kN ] ] ∇φm exp{(wm−1 + φm − log(− log(k)))/τ}∑N i=1 exp{(wm−1,i + φm,i − log(− log(ki)))/τ} ∣∣∣ n dkN · · · dk1
= ∫ 1 0 . . . ∫ 1 0 P [ m,1 = k1 ] P [ m,2 = k2 ] · · ·P [ m,N = kN ] ·
∇φm exp{(wm−1 + φm − log(− log(k)))/τ}∑N
i=1 exp{(wm−1,i + φm,i − log(− log(ki)))/τ}
∣∣∣ n dkN · · · dk1
= ∫ 1 0 . . . ∫ 1 0 1 · ∇φm exp{(wm−1 + φm − log(− log(k)))/τ}∑N
i=1 exp{(wm−1,i + φm,i − log(− log(ki)))/τ}
∣∣∣ n dkN · · · dk1.
(16)

heading: B ALGORITHM DESCRIPTION
text: Algorithm 1 Deep Probabilistic Subsampling (DPS) Require: Training dataset D, Number of iterations niter, temperature parameter τ , initialized train-
able parameters Φ and θ. Ensure: Trained logits matrix Φ and task network parameters θ.
for i = 1 to niter do - Draw mini-batches (xi,si): a random subset of D - ComputeAΦ = [a1; . . . ;aM ] using: if DPS-top1 then
- Initialize mask: w0 = 0 for m = 1 to M do
- Draw i.i.d. Gumbel noise samples em ∈ RN - Sample from the distribution : r̃m = argmax
n
{ wm−1,n + φm,n + em,n } - Create one-hot vector: am = one hotN (r̃m) - Take current mask: wm = wm−1 - Update mask: wm,r̃m = −∞
end for else if DPS-topK then
- Draw i.i.d. Gumbel noise samples e1 ∈ RN - Sample M samples from the distribution and create an M-hot vector: M-hot = TopK
n
{ φ1,n + e1,n } - Expand the M-hot vector in M one-hot vectors am, with m ∈ {1, . . . ,M} - Create mask per sample: wm = −∞ · am,∀m ∈ {1, . . . ,M} - Create the cumulative masks wm = ∑m i=1 wi, ∀m ∈ {1, . . . ,M}.
end if - Subsample the signal: yi = AΦxi - Compute the task: ŝi = fθ(yi) - Compute loss using : Li = ‖si − ŝi‖22 + Lpen - ∀m ∈ {1, . . . ,M}, redefine ∇φmam := ∇φmEem [ softmaxτ (wm−1 + φm + em) ] , where φm = φ1 for DPS-topK - Use Adam optimizer to update Φ and θ
end for

heading: C HYPERPARAMETERS OF ADOPTED TASK MODELS
text: This appendix contains tables that provide details on the adopted task model architectures for the different experiments. We use the abbreviation FC, which stands for ‘fully-connected layer’.
D EXAMPLE TRAINING GRAPH FOR TRAINING OF ‘LINES AND CIRCLES’ RECONSTRUCTION USING DPS-TOPK

Review summaries:
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.