Improving the efficiency of conformal predictors via test-time augmentation

Thanks for taking the time to share constructive, thoughtful feedback. We are glad to hear you believe the results are strong and the problem we focus on is important. Our responses to your suggestions are below:

Analytically determining augmentation weights: We’ve included additional experiments comparing TTA-Learned to TTA-Acc-Weighted (each augmentation is weighted by its accuracy) and TTA-Err-Weighted (each augmentation is weighted inversely to its classification error). An augmentation which achieves 70% accuracy compared to other augmentations (which each achieve, say, 90% accuracy) should not receive a non-zero weight, which it would under both TTA-Err-Weighted and TTA-Acc-Weighted. Exact results are included below for the expanded augmentation policy; Table 10 contains results for both the expanded and simple augmentation policy.

On parametrizing the augmentations: Thanks for the suggestion; this could certainly improve performance, but would also require more labeled data. Our results are meant to illustrate the surprising effectiveness of test-time augmentation given two relatively simple and general policies; what you suggest would be a good extension of our work.

Directly optimizing efficiency objective: Good point; we ran two experiments to study this. In the first, we compare against the use of focal loss, since this loss is known to produce better calibrated classifiers, which correlates directly with smaller prediction set sizes. In the second, we compare against a conformal training loss, which directly optimizes for smaller sets. We found that neither produced significantly smaller prediction set sizes than using the cross-entropy loss. We’ve also included these results in the supplement of the updated manuscript (Table 3, reproduced below).

Presentation: Thanks for these edits, we’ve made changes to address them in the manuscript.

On improvements to top-k accuracy: Thanks for the suggestion; we have included a comparison of the optimal k required to achieve coverage using the original probabilities, TTA-Average transformed probabilities, and TTA-Learned probabilities in Figure 10. TTA-Learned produces significantly lower values for the optimal k on the two datasets we see the largest improvements: ImageNet & iNaturalist.

Thanks for your positive and constructive review. We agree that the proposed method is both simple and powerful, with strong empirical evidence that it produces consistent improvements over a variety of datasets. We also agree that the paper would be improved by addressing the issues you raise.

The authors should clarify why their idea is different from replacing the underlying model with an ensemble method.

Thanks for this suggestion. In short, our method is different because it does not require multiple models, or the computational costs associated with training multiple models. This is what makes test-time augmentation an attractive approach in image classification; in our work, we show that it is a wise choice in conformal prediction.

The difference would be clear if the aggregation weights were trained by optimizing the CP efficiency directly. But the learning strategy is "minimizing the cross entropy loss with respect to the true labels on the calibration set".

Your point is valid. We have added experiments that learn the TTA policy by minimizing two alternate losses: the focal loss and a conformal training loss, both of which have more direct ties to small set sizes. We found that neither loss significantly improved upon using cross-entropy loss (Table 3, reproduced below for convenience).

The aggregation function is trained by "minimizing the cross entropy loss with respect to the true labels on the calibration set". Does this preserve the marginal validity of the prediction sets?

Excellent point; thank you for bringing it up. We have replicated all experiments by learning the TTA policy on a distinct set of examples from those used to identify the conformal threshold and find no significant changes to results. We've updated the methods section to reflect this data split, and include an additional experiment showing that TTA's results are robust to any particular choice of data split.

Have you compared with any adaptive CP approaches like [1]?

We compare to [1] in Table 6 (APS) and work that builds on it in Table 1 (RAPS). We have made this more clear in the updated manuscript.

Thanks for providing a thoughtful review. We believe this paper to be an important contribution to the literature on making conformal prediction practically useful, for the reason you describe: the method proposed is simple and effective.

On the novelty w.r.t. conformal prediction: We agree that it is well known that better classifiers lead to better conformal predictors. We did not intend to convey the impression that this was a contribution of the paper. Instead, our goal is to demonstrate the surprising effectiveness of a simple and efficient method to improve the underlying model. Additionally, while there is literature on ensembles in conformal prediction, none present the possibility of a test-time ensemble, which is computationally cheaper than alternate approaches.

On the novelty w.r.t learned test-time augmentation policies: We agree that the value of learned test-time augmentation policies has been established; however, a majority of the literature emphasizes the value of TTA to Top-1 classification accuracy. In contrast, our work shows that the benefits of test-time augmentation are not limited to the highest predicted probability; improvements in top-k accuracy produce significantly more efficient conformal predictors.

On using calibration set once: This is a valid complaint. We’ve replicated all experiments by learning the TTA policy on a subset of the calibration set distinct from the subset used to identify the threshold. There are no significant changes to the results. We include an additional experiment in which we vary the percentage of data used to learn the test-time augmentation policy and find that results are largely not sensitive to the % of data used to train the TTA policy (Figure 8).

Thanks for spending the time to write a constructive and detailed review; we are glad to hear that you think the approach is interesting, that its model-agnosticity is a notable strength, and that you believe the experimentation to be thorough.

On the choice of augmentation policy: We drew our augmentation policies from the test-time augmentation literature; one is inspired by the classic augmentation policy used to boost image classification accuracy (flips and crops), and the other is drawn from prior work on learning train-time augmentation policies [1]. We have made these connections more clear in the updated manuscript. Our results show that relatively simple and general test-time augmentation policies can produce large improvements to conformal predictors. That isn’t to say that other augmentations might not prove a fruitful direction for future work.

On computational efficiency: good point. Inference scales linearly with the number of augmentations. One can mitigate this cost through parallelism, or by not generating augmentations with an assigned weight of 0. This is covered in the revised manuscript.

On real-world scenarios: We believe that small set size is especially critical in healthcare, where the human bandwidth to consider a large number of options is limited. We are working to incorporate an experiment on a dermatology classification dataset.

On potential limitations: Our proposed method has at least three limitations. The first is, as you describe, computational inefficiency. The second is the dependence on using labeled data to train the TTA policy, which reduces the data available to estimate the threshold. There is a tradeoff between the two, but our empirical results suggest that there is a large region of data splits for which TTA produces improvements. The third is the restricted search space of transformations; one can imagine stronger results with a more flexible space of transforms, but this would also likely require more labeled data.

[1] Cubuk, Ekin D., et al. "Autoaugment: Learning augmentation strategies from data." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.

Thanks for providing a detailed, constructive review. We’re glad you appreciate the simplicity of the approach and the detailed experimentation. Test-time augmentation’s simplicity (both in implementation and intuition) is one of the key strengths we wanted to demonstrate in the context of conformal prediction.

On fixing the experiments to use 3 separate sets of examples: Excellent point; thank you for bringing it up. We have replicated all experiments by using three sets (the original training set, 20% of original calibration set for TTA, 80% of the original calibration set for identifying the threshold). We found that splitting the data has no significant impact on the reported results. We have modified the methods section to include this data split as a requirement, and have described the trade-off this incurs. We have also included additional experiments that measure the tradeoff between data used to learn the test-time augmentation policy and data used to identify the threshold (Figure 8); in short, the benefits of TTA are robust to any particular split, which suggests that the benefits TTA confers outweigh the cost to identifying the correct threshold.

On using a different loss to learn the TTA policy: Yes, it could be that there are more appropriate losses to use to train the test-time augmentation policy. To study this, we ran additional experiments training the TTA policy using the focal loss and using the conformal training loss proposed by [1]. We found no significant improvements to the prediction set size when compared to the cross-entropy loss (Table 3 in the revised manuscript, reproduced below).

On producing a rigorous proof that learning the weights could maintain validity: we will think further on this. It is an exciting suggestion that we will consider beyond the rebuttal period.

Thanks for taking the time to provide a detailed review. We’re glad you found the experiments extensive, the presentation well-organized, and the explanations intuitive. We also appreciate your raising questions that should be addressed.

On using distinct examples to train the TTA policy: Good point; if the learned TTA policy is overfit to the calibration set, conformal scores computed on the calibration set will differ from conformal scores computed on the test set. This difference is practically important in the context of small calibration sets, which is an important distinction to make. We reran all experiments by training the TTA policy on 20% of the calibration set and identifying the threshold based on the remaining 80% of examples and found largely no meaningful change in results. We have updated the methods section to reflect this change.

On why TTA-Learned does not outperform TTA-Avg when combined with APS: Although TTA-Learned improves Top-K accuracy compared to TTA-Avg (Figure 10), this occurs at K too low to affect the performance of APS. As earlier work has described [1], APS produces large prediction sets because of noisy estimates of small probabilities, which are then included in the prediction sets to meet the conformal threshold. Both TTA-Learned and TTA-Avg smooth the probabilities: they reduce the number of low-probability classes by aggregating predictions over perturbations of the image. The benefit that both TTA-Learned and TTA-Avg add to APS is thus similar to how RAPS penalizes classes with low probabilities. We have added this discussion to the revised manuscript.

[1] Angelopoulos, Anastasios Nikolas, et al. "Uncertainty Sets for Image Classifiers using Conformal Prediction." International Conference on Learning Representations. 2020.