Localization with Sampling-Argmax

Abstract: Soft-argmax operation is commonly adopted in detection-based methods to localize the target position in a differentiable manner. However, training the neural network with soft-argmax makes the shape of the probability map unconstrained. Consequently, the model lacks pixel-wise supervision through the map during training, leading to performance degradation. In this work, we propose sampling-argmax, a differentiable training method that imposes implicit constraints to the shape of the probability map by minimizing the expectation of the localization error. To approximate the expectation, we introduce a continuous formulation of the output distribution and develop a differentiable sampling process. The expectation can be approximated by calculating the average error of all samples drawn from the output distribution. We show that sampling-argmax can seamlessly replace the conventional soft-argmax operation on various localization tasks. Comprehensive experiments demonstrate the effectiveness and flexibility of the proposed method. Code is available at https://github.com/Jeff-sjtu/sampling-argmax


Reviews and author responses:

The paper proposes a sampling-based soft argmax formulation that is differentiable. The method focuses on the fact that the conventional softmax algorithm focuses on having the expectation of a heatmap to be in a desired location, which may not be what one desires. Instead, the paper proposes to consider the expected error to the desired location, should one sample according to the heatmap. To implement this idea, the paper proposes a two-stage sampling strategy that utilized the Gumbel softmax and Kernel Density Estimation. The method is shown to outperform a naive softmax strategy in various tasks.  The strategy proposed in the paper seems effective and can be used as a drop-in replacement over existing soft argmax implementations. The motivation behind the paper is also interesting albeit simple. It does point out an overlooked problem with conventional soft argmax formulations. However, the paper in its current form has critical issues that need to be resolved before it could be considered for publication.

1. Why sampling?

Although the authors motivate that the variance of the formulation in Eq 5 would be large, this is never demonstrated. REINFORCE in some cases does not suffer from the variance issue and can be trained without too much trouble, for example as shown by [a]. Thus, to have the sampling-based scheme justified, either theoretical or empirical analysis should be provided. The paper provided neither.

[a] Michal et al., "DISK: Learning local features with policy gradient", NeurIPS 20

2. Isn't 3.1 Kernel Density Estimation (KDE)?

The formulation shown in 3.1 seems to be KDE, also known as Parzen windowing. This formulation is simply a sum of kernel functions, which all are friendly to backpropagation. Also, the formulations shown here are cases when these kernel bandwidths are set to behave as interpolation kernels. As such, I also find it hard to get convinced that approximating this with sampling is a better idea than simply using the analytical gradient.

In fact, the temperature used in Eq. 11 can be incorporated into the analytical solution as kernel bandwidths, which would effectively behave the same, but less randomly and as if one took infinite samples. As shown by the results in Table 2, more samples seem to lead to a better outcome, thus the analytical solutions seem like they would be better.

3. Experiments required to answer 1 and 2

Thus, simply put, the following baselines are missing to demonstrate the effectiveness of the method.
- Method trained with Eq. 5
- Method trained by directly using the KDE formulations instead of sampling

4. Is it truly better than regularization?

Another natural question for the method is whether the method is truly better than regularization. For example, `Soft w/ V.R` seems to perform quite well in most of the experiments, almost as good as the proposed method (except Table 6). The differences are small enough to require multiple runs. While the paper argues in the checklist that multiple runs would be costly, at this point where the performance difference is around 1% difference relative, it is hard to see if the proposed method brings significant advantage without it. Moreover, considering that the regularization parameters can affect the performance quite a lot, the gain from the proposed method is not convincing.

An addition to this is when considering results in Table 1. In [b], a heatmap-based supervision method is employed, which would be somewhat similar to `Soft w/ D.R.`, in the sense it is enforcing a certain distribution. In [b], the method reports 73.7 mAP, which is significantly higher than the 69.8 mAP of the proposed method. 

[b] Li et al., "Simple Pose: Rethinking and Improving a Bottom-up Approach for Multi-Person Pose Estimation", AAAI 20

5. Presentation

The quality of the presentation needs improvement. The paper does not read well in some cases, with quite a few grammatical errors. For example line 1 "the detection-based method"; line 7 it is unclear at this point what error the paper is talking about; line 16 "from the input is"; line 61 categorical weights is unclear at this point since mixture distributions could be anything; line 62 the same goes for the term "sub-distribution"; line 62 it is unclear what category means at this point; line 103 "we find the model is slow to converge". This is not an exhaustive list.

==== Post Rebuttal Update ====

After the rebuttal and the discussions, I am now convinced that the method greatly outperforms the closed-form solution without sampling (I would love to try this for myself as I find this gap still very surprising!). However, I am still concerned about the discussion on KDE, and the lack of significance tests, which make me hesitant to cross the acceptance threshold. 

 While the paper lays out the limitations of using KDE and mention normalizing flows, but I am not sure how this could be done. For example, this would require that the method is aware of the distributions that can occur, while the aim of the method is to be used as a generic soft argmax algorithm. To me, it seems like a contradiction here.

While I do not believe it is a major issue, in the Broader Impact section, the paper discusses reducing the computational burden, but the proposed method itself would require more computation than a typical soft argmax. 


 Thank you for your update.

1. We check our implementation of training with Equation 4 many times and we are sure that the experiments are conducted properly. Besides, we test both L1 and L2 loss in Equation 4. The results are 30.9mAP and 29.8 mAP, respectively. Moreover, we also followed some methods to reduce the variance. For example, subtracting a constant value in Equation 5:
$$
 = \mathbb{E}_y [(d(y_t, y)-b) \nabla_\theta \log \pi_y]
$$
$$
 = \mathbb{E}_y [d(y_t, y) \nabla_\theta \log \pi_y] - b\mathbb{E}_y [\nabla_\theta \log \pi_y] = \mathbb{E}_y [d(y_t, y) \nabla_\theta \log \pi_y] - b\sum_i \nabla_\theta \pi_i
$$
$$
= \mathbb{E}_y [d(y_t, y) \nabla_\theta \log \pi_y] - b\nabla_\theta 1 = \mathbb{E}_y [d(y_t, y) \nabla_\theta \log \pi_y].
$$
Subtracting a constant will make the gradient more stable and reduce the variance. To obtain a minimal variance, the optimal `b` is:
$$
b = \frac{\mathbb{E}_y[d(y_t, y) || \nabla_\theta \log \pi_y ||^2]}{\mathbb{E}_y[|| \nabla_\theta \log \pi_y ||^2 ]}.
$$
We found the loss can reduce faster but the final results won't get much improvement (31mAP). Therefore, we think it is difficult to directly use the discrete distribution (Equation 4 and 5) and there are still many works that need to be done to facilitate its training process. In our work, we provide a more reliable and effective method without complicated tuning for detection-based localization methods.

2. I understand your point about using combinations of linear basis functions. However, when we implement the loss function, the first thing to do is to calculate the expectation of the distribution. It is equivalent to calculate the expectation of each basis function and add them together with weights. And the expectation of the basis function $f_i(y)$ is just $y_i$. Therefore, when writing codes, using combinations of linear basis without sampling is exactly the same as the implementation of conventional soft-argmax.

3. For the problem "Is it truly better than regularization", our method has three advantages as mentioned in the response to Reviewer mj3d:

    1. Easy to use and without tedious tuning. Regularization needs tedious tuning of parameters, which cost time and computational resources.
    2. Consistent improvement across all different tasks. The performances of regularization are inconsistent in different tasks. When facing a new problem or a new dataset, we don't know whether we can use regularization until we spend lots of time tuning the loss function.
    3. Sampling-argmax achieve better performance than regularization in all tasks and datasets without any computational overhead (when using only one sample). 

Therefore, we think sampling-argmax is better than regularization in all dimensions.

 Thank you for the rebuttal.

I find it strange that training with Equation 4 only gives such bad results. Although I do not have much experience with the Reinforce esitmator, this could also mean that the experiment was simply not performed properly.

And I realise that using combinations of linear basis functions without sampling still corresponds to the normal soft-argmax. My point was more that I wanted to see if this modelling choice results in an improvement on the soft-argmax baseline as well.

I also think that Reviewer mj3d raised some valid points, and they were not answered that well (ie "Is it truly better than regulariation" point). As a result, I am keeping my rating as is.



 We sincerely thank you for your helpful feedback and insightful comments. Below we address your comments and questions:
***

**Q1**: Insights relating to the theoretical justification remains unclear.

**A1**: In this paper, we study the problem of using probability maps (heatmaps) for localization. Based on the loss function, existing works can be divided into three categories: i) pixel-wise supervision, ii) coordinate supervision (soft-argmax), iii) and maximum likelihood. In human pose estimation, these three types of methods are well studied. Directly using maximum likelihood obtains inferior performance to pixel-wise supervision and coordinate supervision. I think the reason is that in the high-resolution heatmap (e.g. 64\*64=4096 in 2D, 64\*64\*64=262144 in 3D), the supervision signal of maximum likelihood is too weak. Maximum likelihood only maximizes the probability of one pixel and treat all other pixels as negative samples. Besides, all negative samples are regarded as the same, ignoring the fact that some of them are closer to the ground truth and some are farther away. Therefore, pixel-wise supervision and coordinate supervision are more commonly used.

Compared to the other two types of methods, soft-argmax can predict coordinates in a differentiable manner, which is beneficial for the end-to-end learning of many downstream tasks. Besides, soft-argmax is robust to the change of input resolution and can reduce the quantization error, which is necessary for many real-world applications where computation resources are limited. Despite these advantages, weaknesses of soft-argmax are rarely studied. In this paper, to address the weaknesses of soft-argmax, we propose sampling-argmax to provide well-calibrated probability maps and obtain higher localization accuracy.

**Q2**: Grammatical errors.

**A2**: Thanks for your valuable comments. After proofreading, we have addressed the typos, grammatical errors and unclear descriptions in our paper.


 We sincerely thank you for your helpful feedback and insightful comments. Below we address your comments and questions:
***

**Q1**: Use linear combination of basis functions without sampling.

**A1**: Thanks for your suggestion. If we use a linear combination of basis functions and the standard soft-argmax as the loss function, the results will be the same as the conventional soft-argmax. The loss function can be written as:

$$
\mathcal{L} = d(y_t, \mathbb{E}[y]) = d(y_t, \int y \cdot p(y)dy) = d(y_t, \sum_{i=1}^n \int y \cdot w_i \cdot f_i(y)dy) = d(y_t, \sum_{i=1}^n w_i \mathbb{E}_{y\sim f_i(y)}[y]).
$$

Since the basis function $f_i(y)$ is centred at the position $y_i$, we have

$$
\mathbb{E}_{y \sim f_i(y)}[y] = y_i.
$$

Therefore, $\mathcal{L} = d(y_t, \sum_{i=1}^n w_iy_i)$,
which is the same as the conventional soft-argmax. We will state it more clearly in the final version.

**Q2**: Use Eq. 4 for training without the basis functions.

**A2**: Thanks for your suggestion. We train the model with Eq. 4 and find it only obtain 30.9 mAP on COCO Keypoint, while conventional soft-argmax obtains 64.5 mAP and our method obtains 69.8 mAP. We will add this experiment in our final version to make our contributions clearer.

**Q3**: The computational cost of using more samples.

**A3**: Thanks for your suggestion. We will provide the computational cost of using different numbers of samples in the final version.

**Q4**: Informative captions for tables.

**A4**: Thanks for your suggestion. We will improve the readability of our experiments in the final version.

**Q5**: Why COCO observes large improvements but not on Human3.6M?

**A5**: COCO Keypoint contains in-the-wild images and multi-person scenarios, which is more challenging. On the other hand, the Human3.6M dataset is a single-person dataset captured in the indoor laboratory environment, which is less challenging and the conventional soft-argmax can predict accurate results.

Besides, in addition to a more accurate localization performance, our method can predict well-calibrated probability maps and provide more reliable confidence scores. On Human3.6M (Table 4), the MPJPE metric only reflects the localization performance and ignore confidence scores. Thus the improvement is not significant. However, in many real-world applications and downstream tasks, a reliable confidences score is very important and necessary. On COCO Keypoint, the mAP metric takes the confidence scores into account to evaluate multi-person pose estimation. Thus, combining the accurate localization and the reliable confidence score, the improvement is more significant. More discussion will be added in the final version.

**Q6**: Why the baseline performs very badly in Table 5?

**A6**: The size of the MSHC dataset for retina segmentation is limited (medical data is not easy to collect). Therefore, a smaller capacity backbone model is used to prevent overfitting. We think that due to limited data and model capacity, the conventional soft-argmax is unconstrained and obtains poor performance. We carefully checked our implementation and NO error was found. Detailed discussion will be added in the final version.

**Q7**: Why does Soft-argmax with distributional regularization do so badly in Table 7?

**A7**: The variants of soft-argmax (using regularization terms) need laborious tuning of parameters during training, such as the weight of the regularization term and the variance of the target distribution. The performance of these variants is not consistent across different tasks even after laborious tuning. For example, `Soft w/ V.R.` performs well in 3D pose estimation, retina segmentation and unsupervised object keypoint estimation, but it obtains poor performance in 2D pose estimation and supervised object keypoint estimation. I think Gaussian distribution is not suitable for unsupervised object keypoint estimation. Therefore, forcing the heatmap to be a Gaussian (`Soft w/ D.R.`) will damage the model performance.

**Q8**: Report Brier Score and Expected Calibration Error.

**A8**: Brier score and expected calibration error are used to measure the accuracy of probabilistic predictions in the classification task. The predicted confidence is used to indicate the probability of accuracy. For example, $c=0.7$ indicates that the classification result has a 70% chance of being correct. However, for the regression task, it is impossible to tell whether a result is correct. We only measure the distance between the prediction and the ground truth. A good regressor can predict a distance that closed to zero, but the distance will not be exactly zero. Although we can set a threshold and regard the distance that is smaller than the threshold as correct, setting a reasonable threshold is difficult. Besides, if the network predicts confidence score $c=0.7$, it is inappropriate to say the distance between the prediction and ground truth has a 70% chance of being smaller than the threshold. Because the network does not know what the threshold is. Choosing different thresholds will significantly affect the results of brier score and expected calibration error. Therefore, the brier score and expected calibration error can not be used directly in the regression task. What we expect is that the higher the confidence, the smaller the distance. In our experiment, we measure the correlation between confidence and distance.

**Q9**: Writing issues.

**A9**: Thanks. After proofreading, we have addressed the typos and unclear descriptions in our paper.

 We sincerely thank you for your helpful feedback and insightful comments. Below we address your comments and questions:
***

**Q1**: Notation in Eq. 3.

**A1**: In Eq. 3, we want to present the idea of optimizing "the expectation of the error". This idea does not require a specific form of the output distribution. The distribution could be either discrete or continuous. Then in Eq. 4, we present how to calculate the expectation of the discrete distribution. Thanks for your valuable comment, we will update the notation and descriptions to avoid confusion.

**Q2**: Why the SF estimator has high variance?

**A2**: Thanks for your suggestion. When training with SF estimator (i.e. Eq. 4 in the paper), the model only obtains 30.9 mAP on COCO Keypoint with the same network and training settings, while conventional soft-argmax obtains 64.5 mAP and our method obtains 69.8 mAP. We will add this experiment in our final version to show SF estimator is not easy to train.

**Q3**: Figure 1 can be replaced with one sentence.

**A3**: Thanks for your suggestion. We will update our final version.

**Q4**: Explain how the network predictions are normalized.

**A4**: Thanks for your suggestion. We normalize the output with the softmax operation. We will elaborate it clearly in the final version.

**A5**: Show SOTA results.

**Q5**: Thanks for your suggestion. We implement our method with HRNet on COCO Keypoint, which obtains 75.4 mAP. On the Human3.6M dataset, we implement the SOTA method (HemletsPose) and find our method brings 1.2 mm improvements, which shows that our method is also effective for newer approaches. We will provide SOTA results in the final version.

**Q6**: Soft-argmax is not well cited.

**A6**: Thanks for your suggestion. We will complete the reference in the final version.


 We sincerely thank you for your helpful feedback and insightful comments. Below we address your comments and questions:
***

**Q1**: Why sampling?

**A1**: Thanks for your suggestion. The model can be trained with REINFORCE but it obtains poor performance. We conduct an experiment of training the model with Eq. 5. The model only obtains 30.9 mAP on COCO Keypoint with the same network and training settings, while conventional soft-argmax obtains 64.5 mAP and our method obtains 69.8 mAP. This result shows that using discrete distribution is hard to train. We will add this experiment in our final version for demonstration.

**Q2**: Isn't 3.1 Kernel Density Estimation (KDE)?

**A2**: Thanks for your suggestion. But our method is different from the standard KDE. The standard formulation of KDE is $p(y) = \frac{1}{nh}\sum_{i=1}^nK(\frac{y - y_i}{h})$. Here, $(y_1, y_2, \cdots, y_n)$ is independent and identically distributed samples draw from the underlying distribution $p(y)$. KDE is used to estimate the density function based on the finite data sample. Instead, in our formulation (Eq. 6-9), $y_i$ denotes the grid position on the probability map (i.e. heat map). We treat the probability map as the discrete version of the density function and want to recover its continuous form, which is different from the motivation of KDE. Besides, KDE simply adds all kernel functions together to estimate the underlying density, while our formulation represents the density as the weighted sum of the basis functions. Nevertheless, we can consider the weight $w_i$ as the frequency of the sample and transform our formulation to KDE.

Moreover, the kernel bandwidth $h$ in KDE is used to control the scale of the kernel function ($\frac{y-y_i}{h}$), which plays the same role as the deviation $\sigma$ in Eq. 9 ($\frac{y-y_i}{\sigma}$), not the temperature $\tau$ in Eq. 11.

**Q3**: Using the analytical formulation instead of sampling.

**A3**: Under the analytical formulation, the expectation of the error can be written as $\mathcal{L}=\mathbb{E}_y [d(y_t, y)]=\int p(y) |y_t - y| dy$. If we use the uniform basis, the loss function is:

$$
\mathcal{L}=\sum_{i=1}^{n}\int_{y_i - c/2}^{y_i + c/2} w_i |y_t - y| dy.
$$
The integral result is the same as Eq. 4, except the case that $y_i - c/2 < y_t < y_i + c/2$. We conduct experiments to evaluate the performance of the analytical formulation. We find the model obtains 31.2 mAP, which is similar to directly training with Eq. 4.

If we use the Gaussian basis, the loss function can be formulated as $\mathcal{L} = \frac{1}{\sigma \sqrt{2\pi}} \sum_{i=1}^{n} \int w_i \cdot \exp({-\frac{1}{2}(\frac{y-y_i}{\sigma})^2}) \cdot |y_t - y| dy$. Eq. 4 can be considered as the special case that $\sigma$ is approaching $0$. This integral is complicated but still tractable. However, it is difficult to implement with batch operations, which will increase the training time and the performance is still poor (30.2 mAP on COCO).

Therefore, the proposed sampling-argmax has advantages of convenience and good performance. When training with only $1$ sample ($N_s = 1$), sampling-argmax has the same computation cost as conventional soft-argmax and achieve superior performance.

**Q4**: Is it truly better than regularization?

**A4**: Although the variants of soft-argmax (using regularization terms) can also improve conventional soft-argmax in some cases, they need laborious tuning of parameters, such as the weight of the regularization term and the variance of the target distribution. The best parameters for different tasks are different. In our experiment, we tune the loss weight ranging from `0.1` to `10` and the variance ranging from `1` to `5` for each task. 

Besides, the best parameters for variance regularization and distribution regularization are also different, which increases the effort needed for the process of parameters tunning. After laborious tuning, the performance of these variants is still not consistent across different tasks and is inferior to the performance of our method, while our method is out-of-the-box and free from parameters tuning. For example, `Soft w/ V.R.` performs well in 3D pose estimation, retina segmentation and unsupervised object keypoint estimation, but it obtains poor performance in 2D pose estimation and supervised object keypoint estimation. The tuning process cost lots of time and computational resources. On the other hand, training sampling-argmax with only $1$ sample does not introduce and computational overhead and brings consistent improvement. Therefore, we think our method is effective and general to different cases. We will add more descriptions about tuning the variants of soft-argmax in our final version.

**Q5**: Compared with heatmap-based supervision.

**A5**: The heatmap-based method can provide pixel-wise supervision during training, while the supervision signal for the regression-based method is only the coordinates. Although `Soft w/ D.R.` can restrict the shape of the density to be Gaussian distribution, its performance is still inferior to the heatmap-based method. In [b], they use a stronger backbone and larger input size to obtain 73.7 mAP. We train sampling-argmax with a stronger backbone (HRNet) and our method can obtain 75.4 mAP. Besides, the heatmap-based method suffers from quantization error, which means that they need a large input size to maintain high accuracy. Our regression-based method can predict continuous outputs and obtain better performance in low-resolution input, which is important in real-world applications.

[b] Li et al., "Simple Pose: Rethinking and Improving a Bottom-up Approach for Multi-Person Pose Estimation", AAAI 20

**Q6**: Presentation.

**A6**: Thanks for your valuable comments. After proofreading, we have addressed the typos, grammatical errors and unclear descriptions in our paper.

**Q7**: How to use normalizing flows?

**A7**: Here we present an intuitive solution. During training, we need to draw samples from the pre-defined basis density function. Then we can use normalizing flows to transform the samples, which is equivalent to directly drawing from the learn basis function. In this case, the normalizing flows can learn the optimal density function that minimizes the prediction error.

**Q8**: Computation burden.

**A8**: When training with only $1$ sample ($N_s = 1$), sampling-argmax has the same computation cost as conventional soft-argmax and achieve superior performance (Table 2, 6 and 7). Ablation experiments on other datasets will be provided in the final version.

The authors propose a novel Sampling-Argmax operator, that they demonstrate empirically out-performs the traditional Soft-Argmax operator on a number of localization tasks. Firstly, the authors propose a modified loss which encourages probability maps to have high likelihood near the ground truth location (i.e. the loss adds sensitivity to pixel location), while not overly constraining the predicted distribution. Next, the authors acknowledge that this loss has high gradient variance, and so propose a sampling method to stabilize training and reduce grad variance.  A few minor points first:
For some reason I found the notation in (3) confusing, while the written out discrete distribution in equation (4) was immediately obvious. In particular, the fact that the network predictions \pi_{y_i} are not present in the equation makes it rather difficult to parse.
Line 107: To make the paper more self contained, I would have appreciated an explanation for why the SF estimator has high variance rather than a citation.
Figure 1 seems superfluous. I struggled to figure out what the insightful point was until I realized that all this is trying to convey is that 2 distributions can have the same mean. This could have been done with one sentence.
\pi_{y_i} is a normalized distribution. It’s probably obvious to the authors how the network predictions are normalized, but in the interest of being more self contained, I would state this exactly (i.e. there is a exp()/sum(exp()) in there).
It’s not necessarily an issue, but the quantitative results are effectively all ablations. For example, on COCO Keypoint, the authors do not show the performance of other methods (e.g. SoTA large-scale detection). I think these numbers are important to calibrate the reader and decide if the improvements are statistically significant compared to other methods.

However, the paper is well written. The problem is well scoped and the solution seems reasonable and empirically justified. While the technical contribution feels a little light (IMO) - e.g. there’s a few tricks being borrowed from pretty standard papers at this point (e.g. gumbel softmax), the method as presented does seem novel (from a quick literature review) and I can’t think of any major reason NOT to accept this paper. Furthermore, soft-argmax has a pretty loyal following in the community (ironically, I’m not one of them), especially those who work in compute-constrained environments, and I think this paper would be of interest to that community.

One other minor point is that I think soft-argmax is not well cited. There are two concurrent works that first presented the idea: Goroshin et al., Learning to Linearize Under Uncertainty and Finn et al., Deep Spatial Autoencoders for Visuomotor Learning. These should be added. No obvious limitations that I can think of and that the authors haven't already addressed (e.g. that their mixture distribution might not represent the true underlying distribution).

This method presents a loss function that can be used on a wide range of localization tasks in different domains and with different input modalities. In that sense, it’s impact could be large as it benefits multiple topics.

This paper proposes a different loss function for training localisation models in computer vision. Currently, many keypoint-localisaiton tasks use the soft-argmax to predict the location of a keypoint from heatmaps / probability maps predicted by the network. The authors note that the soft-argmax only encourages the mean of the predicted heatmap to match the ground truth ("error of the expectation"). The authors then propose a sampling-argmax which minimises the "expectation of the error" by computing the probability-weighted sum of errors at all locations. As this method has high variance and trains poorly, the authors introduce additional structure into the predicted probability maps by representing them as a linear combination of basis functions, and show how this can be sampled from in a differentiable manner.

The authors present strong results on multiple keypoint localisation tasks where they outperform the soft argmax, and other regularisation methods for heatmaps [21]. 
  The authors have two main contributions in their loss: representing the predicted probabilty map as a linear combination of basis functions, and then sampling from this to compute the expectation of the proposed loss (Eq 4). However, the effects of these two contributions are not really disentangled from each other.

It is possible to use a linear combination of basis functions for the probability map, and the standard soft-argmax as the loss function (ie no stochasticity in the loss). This ablation is not performed at all, and I believe it would strengthen the paper considerably. Without it, it is not clear to me if the "sampling" part of the contribution is necessary at all.

And it is also possible to use the sampling-based loss without the basis functions. The authors mention on Page 3 that "we find the model is slow to converge by training with Equation 4", but do not include actual experimental results on it, which I believe would be useful to the reader.

Table 2 shows that using more samples in the loss improves final accuracy, in line with what I would expect. However, the computational cost of using more samples is not detailed at all, and is an important detail.

I have some further comments about the experiments:

- In general, the tables do not have informative captions. And it is not immediately clear what "Soft", "Soft w/ V.R " etc are as these abbreviations are not described anywhere.
- A strength of the paper is that the authors have experimented on multiple different keypoint localisation tasks. However, the resutls are not all necessarily consistent across them, and there is no real explanation for this
  - The sampling-argmax gives large improvements on COCO keypoints, but not on Human 3.6M. Why is this?
  - In Table 5, the soft-argmax baseline performs very badly. Why is this? Is this an error with the implementation rather than the method?
  - Why does Soft-argmax with distributional regularization do so badly in Table 7, unlike all the other experiments?
- The Brier Score and Expected Calibration Error are more standard metrics for comparing the calibration of a network. I think it would be better to also report these in Table 3, unless the authors have reasons for not doing so.


There are a few writing issues in the paper too:

- Line 1. "Soft-argmax operation is commonly adopted ..." -> "The soft-argmax operation ..."
- Line 77. "Preliminary" --> "Preliminaries"


My main concern about the paper is that the effects of the main contributions (basis functions and sampling to compute the expectation) are not disentangled from each other. I am willing to increase my rating if this is addressed in the rebuttal.
 Yes. Limitations are included in the main paper, and societal impact in the supplementary.

The paper proposes a method for learning better calibrated probability distributions for detection (i.e., localisation and classification) tasks. Rather than directly apply the maximum likelihood principle to learn the parameters of the model that maximizes the probability of the ground-truth detection, the approach introduces a differentiable sampling operation with mixture kernels to better control the "shape" of the learned probability distribution. This replaces the argmax delta function approximation and is compared against other regularisation mechanisms such as variance regularisation (L2) and distribution regularisation (Jensen-Shannon divergence). While I have some concerns about the theoretical motivation for the proposed approach (see below), the empirical results are compelling for the tasks considered in the experiments (human pose estimation, retina segmentation and facial landmark localisation), noting that these are all structured prediction tasks in some sense.  There are several grammatical errors but overall the paper is straightforward to understand. My main concern is the theoretical positioning of the work and the discussion around the term "shape". The maximum likelihood principle for learning model parameters is well-established and theoretically sound, and any deviation from it needs to be justified clearly and carefully, e.g., insufficient data and the need for model priors (maximum a posteriori estimation), poor model calibration or approximation due to improve tractability. While the paper introduces a nice differentiable approach that improves on some experimental benchmarks, insights relating to the theoretical justification for this remains unclear.

In particular, for a typical single object detection task the probability distribution appears to be unimodal and highly peaked---the object appears at a single (well-defined) location within the image. Here standard regularisation techniques appear to handle the small ambiguity in the neighbourhood of the mode quite well. For a multimodal scenario the situation is different and a model that learns to match a single mode may not perform well. The tasks explored experimentally in this paper are all structured prediction tasks in that they predict on multiple output variables simultaneously, where exploring multiple modes may be what is leading to the improved empirical results.

In short, the paper argues for constraining the "shape" of the learned probability distribution in deep learning architectures and provides a differentiable sampling approach for doing so. However, it does not make clear under what scenarios this is the right thing to do compared with standard maximum likelihood. A discussion in the rebuttal and indication of how this will be addressed in the paper will help me to refine my recommendation for the paper. I cannot see any negative social impact of this work. Minor technical limitations are mentioned but the bigger question raised above is not addressed.

Meta-review: This paper received mixed ratings, with one reviewer recommending acceptance and three rejection. The paper was thoroughly discussed by the reviewers but the authors' feedback did not convince the negative reviewers to recommend acceptance. In particular, the reviewers' main concerns include the motivation behind the proposed sampling strategy and its theoretical justification, the lack of comparison with some baselines (although the authors provided some of these comparisons in their feedback, the reviewers remained unconvinced), and the significance of the improvements obtained in some experiments. Based on this, we therefore believe that this paper is not ready for acceptance to NeurIPS but encourage the authors to revise it based on the feedback and resubmit to a future venue.

Decision: accept