D EEP P ROBABILISTIC SUBSAMPLING FOR T ASK - ADAPTIVE C OMPRESSED S ENSING
A BSTRACT
The ﬁeld 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 classiﬁcation 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 beneﬁt from sample rate reduction.
The code used for this paper is made publicly available 1 .
I NTRODUCTION
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, efﬁcient 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 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 (Pl ¨ otz & 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 classiﬁcation and reconstruction, while sampling both in Fourier and pixel space.
R ELATED 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 ﬁxed 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 (cid:96) 1 penalty on the sampling mask.
Related to our methodology, yet different in purpose, Pl¨otz & 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 ﬁrst 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), Pl¨otz & 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.
and Wu et al.
(2019), focusses on the problem of learning compressive linear encoders/ﬁlters, rather than discrete subsampling as addressed here.
The authors of Balın et al.
𝒆𝒆 𝑚𝑚 ∼ 𝜏𝜏 →𝑨𝑨 �𝒔𝒔 𝒚𝒚 𝒔𝒔 𝒂𝒂 𝑚𝑚
Forward pass Backpropagation ∀𝑚𝑚 ∈ {1,…, 𝑀𝑀 } 𝜏𝜏 >𝒙𝒙
Generative sampling model ℒ ( 𝑠𝑠 , ̂𝑠𝑠 ) 𝝓𝝓 𝑚𝑚 𝑒𝑒 𝑧𝑧 𝑖𝑖 / 𝜏𝜏 ∑ 𝑗𝑗 𝑒𝑒 𝑧𝑧 𝑗𝑗 / 𝜏𝜏 𝑓𝑓 𝜃𝜃 𝒚𝒚 …
Gumbel(0,1)
Sampling
Task model +
Desired task [Figure] digit b. c.
𝑨𝑨 𝑨𝑨 𝑓𝑓 𝜃𝜃 𝑓𝑓 𝜃𝜃
Task Sampling
Task Sampling Input
Input
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 ﬁxed, reduced subsampling rate.
M ETHOD 3. T ASK - 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 ∈ R N , resulting in measurements y ∈ R M , with M << N :
[Equation]
where A ∈ { 0 , 1 } M × N is the subsampling measurement matrix.
We here concern ourselves speciﬁcally with scenarios in which the rows of A are constrained to having cardinality one, i.e.
|| a m ||
= 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:
[Equation]
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.
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. DPS: D EEP P ROBABILISTIC S UBSAMPLING
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 Φ :
[Equation]
Thus, rather than optimizing A ( Φ ) , we optimize the distribution parameters Φ .
To warrant sufﬁcient expressiveness while maintaining tractability, we learn the parameters φ m ∈ R N of M independent categorical distributions (rather than the joint distribution, which scales factorially), being the rows of Φ ∈ R M × N .
Formally, we deﬁne each m th measurement row a m ∈ { 0 , 1 } N in A ( Φ ) as a one-hot encoding of an independent categorical random variable r m ∼ Cat ( N, π m ) .
We deﬁne π m ∈ R N = { π m, 1 , . . .
, π m,N } , being a vector containing N class probabilities, and parameterize it in terms of its unnormalized logits φ m such that:
[Equation]
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 noisesamples e m,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 deﬁned as:
[Equation]
in which w m − 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 hot N ( · ) as the operator that returns a one-hot vector of length N , we ﬁnally obtain:
[Equation]
To permit error backpropagation for efﬁcient optimization of Φ , we require ∇ φ m a m 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 ∇ φ m a m :
[Equation]
with (row operator) softmax τ ( · ) as a continuous differentiable approximation of the one-hot encoded argmax( · ) operation.
Appendix A shows the full derivation of ∇ φ m a m , and ﬁg. 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 ﬁnd however that using a ﬁxed (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).
To study the trade-off between parameter efﬁciency 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-efﬁcient implementation, by leveraging the Gumbel top-K trick (Kool et al., 2019) (rather than Gumbel-max), and its corresponding relaxation (Pl¨otz & 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.
E XPERIMENTS
We test the applicability of the proposed task-adaptive DPS framework for three datasets and two distinct tasks: image classiﬁcation 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).
4. MNIST CLASSIFICATION Experiment setup
Classiﬁcation 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 classiﬁcation outcome.
Results are compared to those obtained using uniformly distributed pixel/Fourier samples, a sampled disk/low pass ﬁlter, and the data-driven
LOUPE method (Bahadir et al., 2019).
Besides, to show the beneﬁt of task-adaptive sampling, i.e. jointly training sampling with the task, we add a case in which we disjointly learn the classiﬁcation task from the sampling pattern.
As such, we train DPS-topK for MNIST reconstruction, after which the learned sampling pattern is frozen, and the classiﬁer is trained.
Task model
After sampling M elements, all N zero-masked samples (or 2 N 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 ﬁrst three layers.
The 10 outputs were normalized by a softmax function to yield the respective classiﬁcation probabilities.
Table 1 in appendix C shows a detailed overview of the complete architecture of the task model.
Zero-ﬁlling 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 = { ( x i , s i ) | i ∈ 0 , . . .
, L } through minimization of the categorical cross-entropy between the predictions and the labels, denoted by L s .
We moreover promote training towards one-hot distributions Cat ( N, π m ) by penalizing high entropy, adopting:
[Equation]
with π m,n deﬁned as in eq.
The total optimization problem is thus:
[Equation] [Equation]
where p D is the data generating distribution.
Penalty multiplier µ was set to linearly increase 1 e −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 [Figure] D [Figure]
- t o p K
[Figure]
x a m p l e i npu [Figure] s [Figure]
LOUPEUniformDisk
D P S - t o p
( β 1 = 0 .
9 , β 2 = 0 .
999 , and (cid:15) =
1 e − 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 2 e − 3 and 2 e − 4 , respectively.
Results
The results presented in ﬁg. 2a show that image-domain sampling using DPS signiﬁcantly outperforms the ﬁxed 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 (ﬁg. 2b) allowed for a much greater reduction of the number samples, with DPS sampling outperforming uniform sampling and LOUPE, and showing similar performance as ﬁxed 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 classiﬁer resulted in a classiﬁcation error of 9.3%, compared to 6.6% for task-adaptive learning with DPS-topK. This case clearly illustrates the beneﬁt of task-adaptive learning of a sampling pattern.
4. ‘L INES AND C IRCLES ’ IMAGE RECONSTRUCTION Experiment setup
To evaluate reconstruction of structured images from highly undersampled partial Fourier (k-space) measurements (keeping 3 . 1% of the coefﬁcients), 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
PSNR : 35.5 dB (e) PSNR : 28.7 dB (d) PSNR : 27.4 dB (c)
[Figure] PSNR : 16.7dB (b) Target [Figure]
transform
[Figure]
[Figure]
40dB
Uniform
Random
Low pass
Proposed
LOUPE
DPS-top
DPS-topK
PSNR: 37.6 dB (f)
PSNR: 39.6 dB (g) 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 ﬁxed 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.
We use a deep neural network consisting of two subsequent fully-connected layers with tanh activations that map the 2 N (zero-ﬁlled) Fourier coefﬁcients (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 L s now deﬁned as the negative log-likelihood:
[Equation]
The learning rates for Φ and θ were 1 e − 3 and 1 e − 4 , respectively, and µ and τ were respectively set to 2 e − 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 ﬁg. 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 speciﬁcally, yields high-resolution accurate reconstructions with an improved PSNR evaluated on the entire test set compared to the baselines.
Note how the learned [Figure] sampling patterns (top row) have evolved from a random initialization, similar to that of (c), to the ultimate task-adaptive scheme (e-g).
4. CIFAR10 IMAGE RECONSTRUCTION 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 coefﬁcients, and compare ourselves to the three ﬁxed 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.
To that end, we unfold K = 5 such iterations, learning an adequate image-domain proximal mapping P ( k ) θ and stepsize α ( k ) at each fold:
[Equation]
where F ∈ C N × 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 classiﬁcation 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 L D ψ , in addition to the entropy penalty L e :
[Equation]
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 1 e − 3 and 2 e − 4 , respectively, and µ was empirically set to 1 e − 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 signiﬁcantly 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 ﬁg. 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.
across the 5000 test examples, for topK and top1 respectively.
For the most competitive methods, i.e. ﬁxed sampling low pass and data-driven LOUPE, these statistics were 25.3 dB and 0.82, and 25.3 dB and 0.86, respectively.
C ONCLUSIONS
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 ﬁnd 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 ﬁnd the former to perform better in our experiments.
Studying the sampling distributions after training, we ﬁnd that DPS-top1 exhibits more uncertainty.
The insets in ﬁg. 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 difﬁculties 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 beneﬁt 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 overﬁtting 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 ﬂexibility, and interesting extensions can be explored in future work.
Rather than learning a ﬁxed 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
A CKNOWLEDGMENTS
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 ﬁnanced by the Dutch Research
Council (NWO).
We thank Wouter Kool and the anonymous reviewers for their valuable feedback
A G RADIENT OF G UMBEL -S OFTMAX SAMPLING
For any m th pair of rows ( a m , φ m ) , any n th element a m,n in a m can be differentiated towards all elements in φ m through:
[Equation]
Gumbel noise vector e m can be reparametrized as a function of a uniform noise vector (cid:15) m ∼ U (0 , 1) i.i
[Equation]
d., through:
[Equation]
[Equation]
This allows rewriting eq
[Equation] (14) into:
[Equation]
B A LGORITHM DESCRIPTION
Require:
Training dataset D , Number of iterations n iter , temperature parameter τ , initialized trainable parameters Φ and θ .
Ensure:
Trained logits matrix and task network parameters to then
[Equation] where φ m = φ 1 for DPS-topK - Use Adam optimizer to update Φ and θ end for
C H YPERPARAMETERS OF ADOPTED TASK MODELS
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’.
[Table]
[Table]
[Table]
GlobalAveragePooling2D
[Table]
D E XAMPLE TRAINING GRAPH FOR TRAINING OF ‘ LINES AND CIRCLES ’ RECONSTRUCTION USING DPS- TOP K
[Figure]