Conformal confidence sets for biomedical image segmentation

We are very pleased that the reviewer enjoyed reading the paper and are grateful for their comments and questions which we address below. We have taken the reviewer’s advice on board and, in order to improve the quality of the manuscript, have included applications to two new datasets involving segmentation in the context of brain imaging and dentistry. Our results on these datasets show the robustness and wide applicability of our approach. See the relevant Section 4, 5, A.7 and A.8 of the updated paper for the results and application examples. Regarding the need for quantitative metrics we have now included dice, precision and recall metrics, in Section A.9, for the 3 different segmentation models used in the paper. These metrics correlate with the performance of the distance transformed scores but not necessarily with other score transformations. Moreover we would like to clarify that evaluation of the inclusion specified in equations 1 and 2 is done in the validations in Section 3.3 (and for the new datasets in Sections A.7.4 and A.8.4). These validations subsample the data with replacement (each time dividing into a calibration and a test set) and check whether the inclusions 1 and 2 hold in order to establish that the methods have the right coverage rate. They show that for each of the datasets considered the confidence sets provide coverage at the nominal rate for interesting coverage levels. In the first version of the paper we compared to the bounding box approach of [1] as this is the main other approach we are aware of which controls the same error rate. Other methods used in conformal image segmentation typically consider weaker error rates as these are easier to satisfy whilst being less meaningful. However score transformations such as the distance transformation can be very helpful when using these other methods for the same reasons they are helpful in our context. We shall prepare and include an illustration of the resulting benefits of doing so, for other methods such as conformal risk control [2], for the final version of the paper. We shall also show quantitatively the degree to which sets derived using risk control (with the expected proportion of false non-discoveries) provide (severe) undercoverage when considering inclusion coverage rates. [1] Andéol, Léo, et al. "Confident Object Detection via Conformal Prediction and Conformal Risk Control: an Application to Railway Signaling." Conformal and Probabilistic Prediction with Applications. PMLR, 2023. [2] Angelopoulos, Anastasios N., et al. "Conformal risk control." ICLR, 2024. We agree that there are other score transformations which can be considered. In particular as the reviewer remarks smoothing the score contributions via a smoothing kernel is a good idea. We illustrate this in the new applications to brain imaging and dental records, see Sections 4, 5, A.7 and A.8. Here we compare the results of smoothing the scores using a Gaussian kernel with varying levels of applied smoothness. In the brain imaging application we see that this leads to a big improvement over the use of the original scores (which perform quite poorly). However in this setting the improvement is not as great as using the distance transformed scores. Instead for the dental application, smoothing is very helpful and in fact provides the largest inner confidence sets, which we then use in practice. For this application it also helps to provide tight outer sets. These can in fact be tighter than those provided by the distance transformation however tend to have extra blobs which do not correspond to teeth which is why we settled on the distance transformation for the final calibration. Instead for the polyps application we found that smoothing did not significantly improve the quality of the inner and outer sets on the learning dataset, likely because the score contributions from the model are already smooth (see e.g. the surface plot of the scores in Figure 2). We will add the results of applying smoothing in the polyps application to the final version of the manuscript. We have added labels to the rows/columns of the figures displaying the confidence sets throughout the main text and the appendix and thank the reviewer for this suggestion as it greatly helps to improve the clarity. Moreover we would like to apologize for the spelling error of polyps which we have now corrected in the updated draft and appreciate that this was spotted. We have also replaced "... the set a side [num] images ...", with the “… [num] images which we set aside” or another appropriate variant.

Regarding the reviewers remaining questions. What we referred to as the original scores are indeed the scores which result from using the identity transformation. In order to improve the clarity of this in the paper we now refer to these scores as the logit scores or the untransformed logit scores throughout the paper instead of as the original scores. Furthermore the dashed lines in Figure 4 provide 95% uncertainty bands for the coverage, we have now clarified this in the caption of Figure 4. Instead the grey dashed line in Figure 5 indicates the value 1 at all levels, this is included for comparison because the best possible value of the inner and outer ratio in the respective plots is 1. We thank the reviewer once more for their helpful comments and look forward to hearing their thoughts on our response and discussing any follow up questions that they may have.

We are grateful for the comments and thoughts of the reviewer and for the opportunity to clarify our contributions. 1- The distance transformation is indeed a sensible choice of score transformation. However as far as we are aware other papers have not considered it in the context of conformal inference for image segmentation. Given how necessary this transformation turns out to be in some applications (see e.g. the new brain imaging example in which the untransformed scores provide extremely uninformative bounds), to us this is an important gap to fill in the literature. We also regard the theory which we derive surrounding inner and outer sets, including the newly added results, Theorems 2.8 and A.4, as a key contribution. 2 - We would like to clarify further that we in fact do compare to the results of other existing methods. In particular the bounding box approach of [1] is compared to on the learning dataset and the testing datasets, for the polyps application, and shown to perform less well than the use of the distance transformation. This is shown visually in Figure 2 and Figures A8-12. We also compared to the precision of this approach in Figure 5 and included it in our validations in Figure 4. We explain the relationship with [1] in Section 2.5 of the manuscript. [1] Andéol, Léo, et al. "Confident Object Detection via Conformal Prediction and Conformal Risk Control: an Application to Railway Signaling." Conformal and Probabilistic Prediction with Applications. PMLR, 2023. Our existing results in fact also compare to the result of applying [2], the paper mentioned by the reviewer. This is because for our problem setting the approach of [2] is equivalent to empirical risk control [3] with the binary loss function which we showed can be used to derive valid inner and outer sets in Section A.2. We have clarified this in Remark 2.4. Applying the method of [2], without modification, in our context would result in the blue outer set obtained from the identity score transformation which is typically very wide and not useful. This is exemplified in the brain imaging application, see Figure A20, in which the blue outer set (which would be the result of applying the algorithm in [2]) obtained from using the untransformed scores is extremely uninformative. Indeed [2] observed very poor performance with the binary loss function, noting that the resulting “prediction set will be theoretically valid but not very informative”. The use of the score transformations and the distance transformation in particular is thus crucial in improving the width of the confidence sets. As far as we are aware our paper is the first (other than the bounding box approach of [1] which we compare to) to provide informative conformal confidence sets which are guaranteed to fully contain the segmented outcome (rather than controlling another weaker error rate), [2] Mossina et al. Conformal Semantic Image Segmentation: Post-hoc Quantification of Predictive Uncertainty, CVPR Workshops, 2024, [3] Angelopoulos, Anastasios N., et al. "Conformal risk control." ICLR, 2024. 3 - We have now included two additional applications involving brain imaging and dentistry. These show that the performance of the model indeed generalizes across datasets. It also helps to emphasize the need for score transformations. The distance transformation does particularly well on the brain imaging dataset. We have performed validations on these datasets, see Sections A.7.4 and A.8.4, which show that the model correctly controls the coverage rate in these settings. 4- Regarding the reviewer’s question about segmentation of multiple structures. This is indeed an interesting question. The segmentation problem for each one of these multiple structures is itself a binary segmentation problem. As such corresponding results for multiple structures follow as a corollary to our results. Joint coverage over the structures can then be obtained by jointly sampling the maximum of the scores over the different classes. We shall formalize this and add an application for the final version of the paper. 5- In order to clarify what the algorithm does during test time we have included a formal algorithm describing the steps taken by the model. See Algorithm 1 in Appendix A.5, now referenced in Section 3.2. Inner and outer thresholds are in fact computed separately based on the inner and outer scores respectively during calibration. When applying the distance transformation the distance is computed relative to the predicted mask obtained by thresholding the logit scores at 0 not to the inner set. Then at test time transformed inner and outer scores are calculated and compared to the calculated threshold. We hope that the provided algorithm helps to make the steps taken clearer.

6- Uncertainty quantification for our method is indeed quantified by the margin between the inner an outer confidence sets. We do not rely on a posterior distribution, instead using the calibration set to calculate the inner and outer thresholds. As such our method does not make assumptions on the distribution of the data in order to provide valid uncertainty. In particular the width of the confidence bands directly depends on the quality of the neural network. I.e. as the predicted segmented mask approaches the ground truth mask in Hausdorff distance both inner and outer sets will converge to the ground truth mask. In order to formalize this we have added Theorems 2.8 and A.4 which show that if the Hausdorff distance between predicted and ground truth masks on the calibration dataset is bounded then confidence sets for new observations are precise. Importantly this result does not hold for the original untransformed scores which can give very wide and uninformative confidence sets even when the neural network provides very good predictions. This very well illustrated in the brain imaging application, see Figure A20 in Appendix A.7. In the case that experts disagree on the true segmented mask we would for now recommend using a consensus mask which is a function of the masks produced by each expert. In that case the method would provide confidence bands relative to this consensus mask. The method is only as good as the quality of the expert calculated masks and relies strongly on a good quality ground truth. We do not directly incorporate the uncertainty in the ground truth masks in our approach but it would be very interesting to do so, as we now observe in the discussion. 7- The size margin between the inner and outer boundaries (for the confidence sets obtained from using the distance transformed scores) depends on the application setting and quality of the image segmentation algorithm, as shown in Theorems 2.8 and A.4 and discussed in the response to (6) above. The width of the uncertainty bands helps to visually capture the uncertainty of the model and in our view allows practitioners to better understand the limitations of these models. It is indeed the case that not all segmentations within the margin will be equally plausible. Because we are not working with a posterior distribution it is not possible to obtain samples from the model. Instead obtaining a set of plausible segmentations within these bounds would in our view require additional biological information to be taken advantage of. We have added a comment to the discussion on this point as an interesting direction for future research. ------------------------------------------------------------------------------------------------------------ We would direct the reviewer to our responses to 3 and 4 (and the response to all reviewers) with regards to their first question, to 6 for the second question and thirdly to 5 and 7 and below for the response to their third question. Regarding how the method can/should be used in practice. This depends on the application setting. For polyps segmentation the method could be used to rule over regions of the image where the polyps could lie. We can be sure, up to the guarantee provided by the model that there are no polyps outside of the blue set meaning that practitioners could deprioritize looking for polyps within those regions. Instead for instance in the brain imaging application it is important to detect locations which lie within/outside the brain for follow up analyses. Within the inner set we can be sure to find areas inside the brain which could help with alignment further down the pipeline and the detection of activation (e.g. when using fMRI). Instead the outer set can be used to mask out areas where we can be sure that there is no brain, and thus no activation. Having precise confidence bounds on this is important because otherwise we risk missing areas of the brain. We thank the reviewer once more for their helpful comments and look forward to hearing their thoughts on our response and discussing any follow up questions that they may have.

We would like to thank the reviewer once again for their feedback. We admit that we had been hoping that the reviewer would consider increasing their score slightly in light of the changes which we have made. The addition of the new applications and theory – prompted by the helpful comments of the reviewer (and other reviewers) – have in our view greatly strengthened the paper. With regards to the reviewer continued concerns. We shall address these in reverse order. With regards to the reviewer’s second point. It is our strongly held view that deep learning models are widely used without proper uncertainty quantification. Indeed as discussed in [4], “neural networks do not deliver certainty estimates or suffer from over- or under-confidence”. These models are powerful but are essentially black boxes the outputs of which are difficult to properly understand. We would strongly argue that whenever it is useful to use a deep learning based segmentation model it is important to have proper confidence bands on the output in order to better understand the limitations of the model. As discussed in [5], “…medical AI, especially in its modern data-rich deep learning guise, needs to develop a principled and formal uncertainty quantification (UQ)”. Deep learning models are widely used without much thought or consideration of uncertainty and in our view the field is in great need of a change. We regard our work as a step along this path in the right direction. [4] Gawlikowski, Jakob, et al. "A survey of uncertainty in deep neural networks." Artificial Intelligence Review 56.Suppl 1 (2023): 1513-1589. [5] Begoli, Edmon, Tanmoy Bhattacharya, and Dimitri Kusnezov. "The need for uncertainty quantification in machine-assisted medical decision making." Nature Machine Intelligence 1.1 (2019): 20-23. The reason for which uncertainty quantification is necessary in deep learning is the same as the reason for which statisticians advise against providing point estimates on effect sizes without confidence bands. In our setting the predicted output of the neural network is the point estimate of the segmented outcome and our confidence bands provide the necessary uncertainty. As discussed in [6], “Point estimates alone can be misleading, as they do not quantify the variability or reliability of the estimate.” In particular , “confidence intervals (and by extension, bands) offer a way to convey the precision of estimates, reminding us that data are noisy and estimates are not exact”, and they “help us see how useful a model might be by explicitly recognising its limitations” as discussed in [7] and [8] respectively. [6] Casella, George, and Roger Berger. Statistical inference. CRC Press, 2024. [7] Freedman, David A. Statistical models: theory and practice. Cambridge University Press, 2009. [8] Box, George EP, and Norman R. Draper. Empirical model-building and response surfaces. John Wiley & Sons, 1987. We would be happy to add the details discussed in the above paragraphs (and more) to the paper if the reviewer feels that this would better help to explain and motivate the importance of our work. To the reviewer’s first point. In our view the distance transformation is not the only contribution of the paper, the others being the theory developed (including the new results i.e. Theorems 2.8 and A.4) and the emphasis on the choosing the optimal score transformation based on a learning dataset, (and justifying this theoretically) which as far as we are aware has not been suggested in the literature on conformal inference for images. This is important no matter what error rate is being controlled. However the distance transformation is an important part of the paper. It has the distinct advantage of allowing conformal sets to work very well where other existing approaches such as [2] fail, sometimes extremely badly (c.f. Figure A20). Fixing other existing methods in this manner is, in our view, an important contribution. [2] Mossina et al. Conformal Semantic Image Segmentation: Post-hoc Quantification of Predictive Uncertainty, CVPR Workshops, 2024. We look forward to hearing the reviewer's thoughts on our response and thank them once more for the time taken to read our work.

We are very grateful for the reviewer’s comments and remarks. Our response is below. Regarding the reviewers first concern. We in fact view the adaptability of the choice of score function to the dataset/model as a strength not a weakness of the method because different datasets/models have different features that mean the optimal score transformation may vary. This approach (of learning transformations on an independent dataset) has been previously used and theoretically justified in [1] in the context of conformal inference for time series data, in which the optimal copula was chosen based on a learning dataset (as we mention in Section 2.4). In the new datasets provided the optimal score transformations are different than for the polyps dataset and indeed certain choices (such as the original scores) can perform very badly (e.g. in the brain imaging application, see Figures A18 and A20 in the updated manuscript). This helps to illustrate the need to optimize the score functions. [1] Sun, Sophia, and Rose Yu. "Copula conformal prediction for multi-step time series forecasting." ICLR, 2024. We indeed regard one of the main contributions to be the theory developed. We would like to clarify that learning the score functions on an independent learning is theoretically valid. Crucially, as for the similar approach taken in [1] in the time series setting, the independence of the learning data set from the calibration and testing datasets guarantees the vaildity of the optimally chosen score function. Regarding the reviewers second concern that we only consider a single dataset we thank the reviewer for this comment. In order to address this we have added two additional datasets involving brain imaging and teeth segmentation. These new datasets help to illustrate the robustness and usefulness of our method. Regarding the reviewers questions. 1- The dataset is indeed public but is independent of the data used to train the original polyp segmentation model. 2- In the original dataset there were a few images from the same video frames however we removed these for the purposes of our analysis. This is important our model assumes exchangeability which would be violated if there was dependence between the images. We now clarify in Section 3 that the images used come from different patients. 3- Regarding the use of the word tumor, we apologize for this oversight and have mostly removed the word throughout the paper, except in one setting in which we are not referring to polyps in particular. We thank the reviewer for pointing this out. 4- The reviewer is right to note that the fact that the data is from different centres may influence the annotations. We rely strongly on a good quality ground truth and the model is only as good as the ground truth available. Where possible we would recommend taking a consensus rating by combining the annotations of multiple annotators and then using this consensus in combination with our model for best results. 5 – We have now included a table in Appendix A.9 illustrating the performance of the different segmentation models used, measured in terms of dice, precision and recall scores. This table helps to show how improvements in these metrics correspond to improvements in the performance of our method. In particular for the best performing model based on these metrics (the HDBET model designed for brain extraction), has relatively tight confidence sets. This is a relationship which we have now formalized in Theorems 2.8 and A.4 results which give guarantees on the size of the resulting confidence sets related to the performance of the model. Other choices of score function do not correlate with these metrics, indeed the original untransformed scores perform notably badly for the brain imaging application despite the high performance on the metrics (see e.g. Figure A20). As such, even for well performing segmentation models, appropriate score transformations are required in order to obtain tight confidence bounds. Theorems 2.8 and A.4 show that the model handles perfect segmentations very well as they have a Hausdorff distance of 0 from the true mask. Instead complete misses will typically increase the size of the confidence bands, appropriately so as they indicate a failure of the model in that instance.

6- We agree that using the learning dataset the training data may not provide the optimal score transformations. This is a not a problem for validity as the training data is assumed to be independent of the calibration and test datasets and so the results of Sections 2.2 and 2.3 still ensure that we can make inclusion statements with confidence. However doing so may impact the choice of score functions. As we now take greater care to emphasise in Section 2.4, we do not recommend using the training data as part of the learning dataset if there is a large amount of data available. However in cases where there is limited data (such as the teeth segmentation problem which we now consider), learning the score function on the training data may still be helpful. In particular the training data may still contain information which allow us to distinguish between different score functions. We saw in the application to teeth segmentation in which the score functions on the learning dataset (which was made up of the training data) had a similar performance on the test dataset. Doing so means that we are not required to give up any of the calibration data used, or make the decision to train the model using fewer images. This is a trade off that must be decided upon carefully by the researcher and where possible we recommend that researchers use a learning dataset which is independent of the training data (as we do with the polyps and brain imaging data settings which we consider). Where possible we thus recommend splitting the data into independent training, learning, calibration and test datasets. 7- Our method is indeed a lightweight addition to any existing black box image segmentation model and is relatively easy to apply to additional datasets. The new datasets and applications which we have added to the paper help to illustrate this, showing that the model is generally applicable, informative and valid in these settings. In particular, for each of the models considered, we perform validations in which we resample with replacement from the data in order to check the coverage rate of the method, see Sections 3.3, A.7.4 and A.8.4 of the updated manuscript. We would like to clarify that the guarantees are in fact that 100% of true mask is included 90% of the time rather than that 90% of the mask is included 90% of the time. This guarantee allows full coverage and means that the resulting confidence sets are more meaningful. We shall include validations across additional settings for the camera ready version of the paper. We thank the reviewer once more for their helpful comments and look forward to hearing their thoughts on our response and discussing any follow up questions that they may have.

Dear Reviewer Bd8E, We were wondering whether you had had a chance to have a look at the new changes which we have made to the paper in our response. In our view the new version of the paper has been greatly improved thanks to your review. Do you feel that we have addressed your comments? Thanks in advance!

We are very grateful for the reviewer’s comments and remarks. We agree that providing uncertainty quantification for black box neural network models is an interesting problem. Regarding additional datasets we have taken the reviewer’s advice on board and have now included extensive analysis of two further datasets. The first is a brain imaging dataset. And, the second, following the reviewer’s suggestion, involves teeth segmentation. As now shown in the main text (Sections 4 and 5) our method works very well in these scenarios providing meaningful confidence sets which have robust confidence guarantees. This demonstrates that our method extends robustly to other settings and models. Further analysis is shown in Sections A.7 and A.8. We also now included Dice, Precision and Recall metrics (evaluated on the corresponding validation dataset) for each of the 3 segmentation algorithms considered. (I.e. PraNet, HDBET and the UNET based GAN model we used for teeth segmentation). See the relevant table Section A.9 in the updated draft for full details. The results are very helpful in understanding how to performance of the models affects the performance of the confidence sets. In particular improvements in these metrics correspond to improvements in the quality of the confidence sets based on the distance transformed scores. The HDBET model has the highest dice score and has the tightest confidence sets as a result. Note that other score transformations such as the identity (which yields the original logit scores) do not have this monotonicity property. Indeed Figure A20 shows that the untransformed logit scores can be very uninformative, but the degree to which this is true depends on the application. In order to formalize the relationship between the distance transformed scores and the quality of the model we have provided new results (Theorems 2.8 and A.4) which further motivate the use of the distance transformation. Comparison of between the metrics now shown in Section A.9 and the performance of the confidence sets helps to illustrate this result. Regarding the comparison to other baseline models, we shall include measures of the performance (e.g. relative to SANet and UACAnet and others) in the final version of the manuscript. Regarding the technical terms (FWER/FDR) we have fully written out their acronyms for clarity where they are introduced and have included a new section of the Appendix (Section A.10) in which these are formally defined and where we discuss the relationship between them and different measures of coverage in the segmentation setting. We would like to apologize for the spelling error of polyps which we have now corrected in the updated draft and we thank the reviewer for pointing this out. We thank the reviewer once more for their helpful comments and look forward to hearing their thoughts on our response and discussing any follow up questions that they may have.

Dear Reviewer Jpsi, We were wondering whether you had had a chance to have a look at the new changes which we have made to the paper in our response. In our view the new version of the paper has been greatly improved thanks to your review. Do you feel that we have addressed your comments? Thanks in advance!