Paper: Improving the efficiency of conformal predictors via test-time augmentation (ICLR-2024)

Reviewer comment (Evaluation):
Lack of explanation: the idea is pretty simple but quite effective by looking at results, but the paper does not fully explain why this kind of method works so well. Simple test time augmentation does not bring so much improvement to me. I suggest providing top-k accuracy and adding more explanations on why it benefits conformal prediction so much accordingly.

Author response:
**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.

Reviewer comment (Presentation):
Presentation: a) typo: $k_reg$ in the related work part; b) in section 3, indexing from 0 to N should induce N+1 samples; c) a large part in Section 4 is preliminary of conformal prediction, which should be introduced in Section 3 or a separate preliminary part.

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

Reviewer comment (Theory):
Can we analytically write out the optimal aggregation weights based on the empirical utility of each augmentation? (quite feasible to me) For example, augmentation with higher accuracy should have a larger aggregation weight. If that is the case, there is no need for optimization of those weights.

Author response:
**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.

| Dataset   |   alpha | Method                | Average Prediction Set Size   |
|:----------|--------:|:----------------------|:------------------------------|
| ImageNet  |    0.01 | RAPS+TTA-Acc-Weighted | 37.115 (4.112)                |
| ImageNet  |    0.01 | RAPS+TTA-Err-Weighted | 36.012 (3.501)                |
| ImageNet  |    0.01 | RAPS+TTA-Learned      | 31.248 (2.177)                |
| ImageNet  |    0.05 | RAPS+TTA-Acc-Weighted | 5.258 (0.171)                 |
| ImageNet  |    0.05 | RAPS+TTA-Err-Weighted | 5.352 (0.366)                 |
| ImageNet  |    0.05 | RAPS+TTA-Learned      | 4.889 (0.168)                 |
| ImageNet  |    0.1  | RAPS+TTA-Acc-Weighted | 2.443 (0.068)                 |
| ImageNet  |    0.1  | RAPS+TTA-Err-Weighted | 2.416 (0.076)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned      | 2.312 (0.054)                 |

Reviewer comment (Experiments):
Why don't you parameterize the augmentations and also optimize the weights of augmentations?

Author response:
**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.

Reviewer comment (Experiments):
Do you think about directly optimizing the efficiency objective (i.e., set size) as conformal training papers?

Author response:
**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).

| Dataset   |   alpha | Method                     | Average Prediction Set Size   |
|:----------|--------:|:---------------------------|:------------------------------|
| CUB-Birds |    0.01 | RAPS+TTA-Learned+CE        | 14.045 (1.323)                |
| CUB-Birds |    0.01 | RAPS+TTA-Learned+Conformal | 13.776 (2.198)                |
| CUB-Birds |    0.01 | RAPS+TTA-Learned+Focal     | 13.416 (1.991)                |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+CE        | 3.571 (0.576)                 |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+Conformal | 3.302 (0.312)                 |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+Focal     | 3.194 (0.202)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+CE        | 1.893 (0.187)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+Conformal | 1.865 (0.163)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+Focal     | 1.791 (0.102)                 |
| ImageNet  |    0.01 | RAPS+TTA-Learned+CE        | 31.248 (2.177)                |
| ImageNet  |    0.01 | RAPS+TTA-Learned+Conformal | 32.257 (3.608)                |
| ImageNet  |    0.01 | RAPS+TTA-Learned+Focal     | 32.612 (3.799)                |
| ImageNet  |    0.05 | RAPS+TTA-Learned+CE        | 4.889 (0.168)                 |
| ImageNet  |    0.05 | RAPS+TTA-Learned+Conformal | 4.867 (0.122)                 |
| ImageNet  |    0.05 | RAPS+TTA-Learned+Focal     | 4.906 (0.195)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+CE        | 2.312 (0.054)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+Conformal | 2.308 (0.068)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+Focal     | 2.363 (0.085)                 |

Reviewer comment (Novelty):
It is well-known that a better classifier (as a scoring function) provides smaller prediction sets in size (e.g., Table 1 of Angelopoulos et al., 2022 — only considering papers with deep learning experiments); without these examples, it is clear that if we have a perfect classifier, the expected prediction set size is the smallest value (which is one). Given these, this paper revisited that a better classifier provides a more efficient prediction set, which is not new.

Author response:
**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.

W1 -> R2: > The paper simply shows that the more-certain score function has smaller prediction sets, which is a trivial consequence of conformal prediction.

We would say we go beyond this and can clarify what we mean. TTA-Learned is distinct from other methods which produce more-certain score functions (e.g. training a fancy new architecture or collecting additional data) because it requires only the ingredients available to split conformal predictors (i.e. a labeled calibration set, a pre-trained model, and a score function). Despite its simplicity, we show that TTA-Learned results in significant reductions in set size across datasets, models, and coverage guarantees. This was previously unknown, and is of practical importance to those who wish to apply conformal prediction. 

Our results also speak to the tradeoff between using labeled data to improve the conformal score (as with TTA-Learned) and using labeled data to identify the conformal threshold. We show that you can achieve much better efficiency by using a portion of the labeled data to improve the conformal score, instead of reserving all examples to identify the conformal threshold. We see our work as a step towards thinking about the question of how to best use labeled data in conformal prediction, which we think will be key to real-world adoption. We appreciate the discussion, which has convinced us to make this point more clear in the manuscript.

Reviewer comment (Novelty):
TTA with learnable parameters is at least firstly introduced in Shanmugam et al., 2021, which can be seen as making a better classifier from a base classifier by augmentations with learned weights over augmentations. Therefore, the provided way of using TTA in this paper is not new.

Author response:
**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.

W2 -> R2: > Finally, for "our goal is to demonstrate the surprising effectiveness of a simple and efficient method to improve the underlying model", it is already demonstrated by the TTA paper.

The TTA paper you reference demonstrates the effectiveness of learned TTA policies for improving Top-1 image classification accuracy of trained models. For example, they show that learned test-time augmentation policies introduce an improvement of ~1% in classification accuracy for ResNet-50+ImageNet. For the same combination, we show a 16% reduction in prediction set size at a coverage guarantee of 99%. While the works share a focus on learned test-time augmentation policies, our results do not follow directly from theirs.

The insight in our work is that TTA changes much more than the highest predicted probability; indeed, it significantly improves the optimal k at which Top-k sets achieve coverage (Figure 10 in the revised manuscript). In short, our findings on the relationship between TTA and conformal predictors are distinct from prior lines of work and their focus on the maximum predicted probability.

Reviewer comment (Experiments):
This paper uses a calibration twice — for learning weight parameters for augmentation in (5) and for choosing a threshold for conformal prediction — which constitutes “double-dipping” and should be avoided. The reviewer believes the results would not change much but recommends using a calibration only once for choosing the threshold for conformal prediction.

Author response:
**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).

Reviewer comment (Novelty):
It is not easy to accept the paper’s claim that it found a novel way to improve the efficiency of conformal prediction via TTA, since it is well-known that a better classifier provides efficient conformal predictors and the provided TTA approach is not new. The reviewer requests highlighting the novel points of the paper.

Author response:
**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.

W4 -> R2: > The paper simply shows that the more-certain score function has smaller prediction sets, which is a trivial consequence of conformal prediction.

We would say we go beyond this and can clarify what we mean. TTA-Learned is distinct from other methods which produce more-certain score functions (e.g. training a fancy new architecture or collecting additional data) because it requires only the ingredients available to split conformal predictors (i.e. a labeled calibration set, a pre-trained model, and a score function). Despite its simplicity, we show that TTA-Learned results in significant reductions in set size across datasets, models, and coverage guarantees. This was previously unknown, and is of practical importance to those who wish to apply conformal prediction. 

Our results also speak to the tradeoff between using labeled data to improve the conformal score (as with TTA-Learned) and using labeled data to identify the conformal threshold. We show that you can achieve much better efficiency by using a portion of the labeled data to improve the conformal score, instead of reserving all examples to identify the conformal threshold. We see our work as a step towards thinking about the question of how to best use labeled data in conformal prediction, which we think will be key to real-world adoption. We appreciate the discussion, which has convinced us to make this point more clear in the manuscript.

> Finally, for "our goal is to demonstrate the surprising effectiveness of a simple and efficient method to improve the underlying model", it is already demonstrated by the TTA paper.

The TTA paper you reference demonstrates the effectiveness of learned TTA policies for improving Top-1 image classification accuracy of trained models. For example, they show that learned test-time augmentation policies introduce an improvement of ~1% in classification accuracy for ResNet-50+ImageNet. For the same combination, we show a 16% reduction in prediction set size at a coverage guarantee of 99%. While the works share a focus on learned test-time augmentation policies, our results do not follow directly from theirs.

The insight in our work is that TTA changes much more than the highest predicted probability; indeed, it significantly improves the optimal k at which Top-k sets achieve coverage (Figure 10 in the revised manuscript). In short, our findings on the relationship between TTA and conformal predictors are distinct from prior lines of work and their focus on the maximum predicted probability.

Reviewer comment (Novelty):
The authors should clarify why their idea is different from replacing the underlying model with an ensemble method.

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Experiments):
Have you compared with any adaptive CP approaches like Romano, Yaniv, Matteo Sesia, and Emmanuel Candes. "Classification with valid and adaptive coverage." Advances in Neural Information Processing Systems 33 (2020): 3581-3591?

Author response:
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.

Reviewer comment (Theory):
The biggest weakness of the approach is that the conformal guarantee is broken, or at least it has not been formally proved to hold for the learned setting. It is not clear whether learning the augmentation weights using the calibration set yields valid coverage. Since the resulting classifier $\hat{g}$ uses calibration data to learn the weights $\Theta$, the exchangeability between the calibration set and the test set is broken because information from the calibration set "leaks" into the weights $\Theta$. If one swaps a calibration point with a test point, the learned weights will be different. This is equivalent to why one cannot use the same dataset to both train and calibrate the base classifier $f$ and why Split CP or the full conformal approach must be used. To maintain validity under the split framework, three datasets are necessary: one for training, one for learning the augmentation policy, and one for the final calibration. However, this leads to a trade-off of smaller sets, the same limitation standard Split CP suffers from.

Author response:
**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.

Reviewer comment (Experiments):
The approach trains the augmentation weights to maximize accuracy by minimizing cross-entropy loss against the true labels, but this is not necessarily aligned with the actual goal of conformal prediction. The smallest prediction sets are obtained when predicted probabilities for all classes match the ground-truth oracle probabilities, as shown in prior work (e.g., [1] or APS). Cross-entropy loss tends to produce over-confidence for the true class and does not encourage proper calibration for the rest of the probabilities. Alternative losses, such as the one proposed by Stutz et al. (2022) mentioned in the related work or [1], could lead to further improvements. Given the limited technical contributions, such further experimental analysis is warranted.

Author response:
**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).


| Dataset   |   alpha | Method                     | Average Prediction Set Size   |
|:----------|--------:|:---------------------------|:------------------------------|
| CUB-Birds |    0.01 | RAPS+TTA-Learned+CE        | 14.045 (1.323)                |
| CUB-Birds |    0.01 | RAPS+TTA-Learned+Conformal | 13.776 (2.198)                |
| CUB-Birds |    0.01 | RAPS+TTA-Learned+Focal     | 13.416 (1.991)                |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+CE        | 3.571 (0.576)                 |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+Conformal | 3.302 (0.312)                 |
| CUB-Birds |    0.05 | RAPS+TTA-Learned+Focal     | 3.194 (0.202)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+CE        | 1.893 (0.187)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+Conformal | 1.865 (0.163)                 |
| CUB-Birds |    0.1  | RAPS+TTA-Learned+Focal     | 1.791 (0.102)                 |
| ImageNet  |    0.01 | RAPS+TTA-Learned+CE        | 31.248 (2.177)                |
| ImageNet  |    0.01 | RAPS+TTA-Learned+Conformal | 32.257 (3.608)                |
| ImageNet  |    0.01 | RAPS+TTA-Learned+Focal     | 32.612 (3.799)                |
| ImageNet  |    0.05 | RAPS+TTA-Learned+CE        | 4.889 (0.168)                 |
| ImageNet  |    0.05 | RAPS+TTA-Learned+Conformal | 4.867 (0.122)                 |
| ImageNet  |    0.05 | RAPS+TTA-Learned+Focal     | 4.906 (0.195)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+CE        | 2.312 (0.054)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+Conformal | 2.308 (0.068)                 |
| ImageNet  |    0.1  | RAPS+TTA-Learned+Focal     | 2.363 (0.085)                 |

Reviewer comment (Theory):
My biggest question about this work is on the assumption of data exchangeability in TTA-Learned. As the weights of augmentations in TTA-Learned are obtained by minimizing the cross entropy loss on the *calibration set*, there is a statistical dependence between the data in the calibration set and the learned weights of augmentations that are used to compute the conformal score. This means that calibration data and unseen examples are not exchangeable with respect to the computation of conformal score. So I do not see how the assumption of exchangeability is preserved. If this point could be clarified, I would be willing to reconsider my score.

Author response:
**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.

Reviewer comment (Evaluation):
As pointed out by the authors, it is understandable that TTA-Learned works better than TTA-Avg for the expanded augmentation policy as it allows the adjustment of the weights associated to the augmentations not included in the training. Meanwhile the results using APS reported in Table~4 show a close match between TTA-Learned and TTA-Avg. Could the authors provide some intuition behind that?

Author response:
**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.