Screener: Learning Conditional Distribution of Dense Self-supervised Representations for Unsupervised Pathology Segmentation in 3D Medical Images

Dear Reviewer 5vEP, thank you for dedicating your time to reviewing our submission and for offering thoughtful and generally positive feedback. We are encouraged that you find our approach innovative and consider our results to be promising. We share your opinion that some additional experiments and benchmarks could be added to demonstrate the practical benefits of our method. We comment on these and other concerns below.

Dice score comparison with supervised SOTA

The key problem is that the existing CT datasets provide annotations only for specific classes of findings. For example, LIDC includes lung cancer masks, MIDRC provides pneumonia masks, KiTS offers kidney tumor masks, and LiTS contains liver tumor masks. However, other pathologies present in these datasets remain unlabeled.

This limitation introduces two key issues:

1. Evaluation challenges: Standard segmentation metrics like the Dice score or object detection metrics cannot be reliably used. True positive predictions for unannotated pathologies are incorrectly counted as "false positives" when compared to the incomplete ground truth annotations. For instance, in Figure 2, our model correctly identifies pneumothorax in the second image, but because the ground truth mask does not include this finding, metrics like the Dice score are significantly penalized.
2. Limited supervised models: Supervised models can only be trained for a single annotated pathology (e.g., lung cancer). Comparing our self-supervised model, which is designed to segment all pathologies, with such supervised baselines would be inherently unfair, as they address different tasks.

To address the first issue, we draw inspiration from the MVTecAD benchmark, which uses pixel-level AUROC and AUPRO metrics to evaluate unsupervised visual segmentation. These metrics are more robust to a small number of "false positive" errors, as they rely on pixel-level recall (or finding-level recall for AUPRO) and pixel-level specificity.

For the second limitation, instead of directly comparing our model to supervised baselines, we propose an alternative experiment. We could distill our composite model into a standard U-Net architecture and use it as a pre-trained checkpoint. This pre-trained U-Net could then be fine-tuned for specific tasks such as lung cancer segmentation and compared against supervised baselines trained either from scratch or using other pre-trained initializations. While we were unable to include these experiments in the current submission due to time constraints, we plan to add them to the GitHub repository. Thank you for inspiring us with this idea!

Domain Gap Influence

We trained Screener on three datasets: one lung cancer screening dataset (NLST) and two abdominal datasets (AbdomenAtlas and AMOS). We then evaluated its performance on four separate datasets: two chest CT datasets (LIDC and MIDRC-RICORD-1a) and two abdominal datasets (KiTS and LiTS). While there are noticeable domain gaps between the training and testing datasets, as well as among the testing datasets themselves, the results in Table 2 demonstrate strong generalization. Screener achieves pixel-level AUROCs of approximately 0.9–0.95 on each of four testing datasets, indicating its robustness in handling diverse chest and abdominal CT data.

Overlapping Test and Training Data

We argue that our training and test datasets are entirely distinct, with no overlap. In particular, in the paper about the AbdomenAtlas dataset, authors contrast it to the pre-existing KiTS and LiTS datasets. Therefore, we assume that there is no intersection between AbdomenAtlas and either the KiTS or LiTS datasets.

Architecture and Sampling Details

All the implementation details, e.g. neural network architectures, preprocessing, hyperparameters will be available in the github repository.

Benchmarking on brain MRI datasets

We acknowledge that benchmarking our model on brain MRI datasets is desirable. Unfortunately, we cannot include it in the current version, due to the time limitations.

Dear Reviewer px7H, thank you for reviewing our submission and for offering thoughtful feedback. We greatly value your insights and have addressed your concerns below.

Clarifying novelty & improving method presentation

You correctly pointed out that our novel contributions were not clearly distinguished from prior work. Additionally, our citation of SimCLR in the context of dense SSL was misleading, and terms “local patterns” and “global context” were not well explained.

To address this, we have made the following revisions:

1. Added a Background & notation section (Section 2) to outline the existing density-based UVAS framework, involving supervised descriptor model, density model, and conditioning on vanïla sin-cos positional encodings.
2. In the Background & notation section we also discuss works on self-supervised learning (e.g., SimCLR and VICReg), and dense self-supervised learning (e.g., VADER, DenseCL, VICRegL), which are closely related to our method. We also replace misleading citations in lines 81-82 with correct citations of dense SSL methods.
3. In Method section (Section 3), we emphasize our two key contributions: self-supervised descriptor model and condition model. Suitability of self-supervised features within the density-based UVAS framework is a non-trivial finding which contributes to the field. The concept of a condition model which learns data-driven variables for conditioning is also novel within the context of UVAS literature.
4. Added Figure 2, illustrating the main components of our method: descriptor model, condition model and density model.
5. In Section 3.1 we provide intuition and details of the pre-training procedure for our self-supervised descriptor model. It produces dense feature maps that capture all the mutual information between different augmented crops preserving “local” image content. For example, the lung cancer mass visible in both crops in the upper part of Figure 2. Therefore, descriptor model’s feature maps contain information about the presence of pathologies.
6. In Section 3.2 we describe our condition model. Similarly to the descriptor model, it produces dense feature maps capturing all the mutual information between different augmented masked crops. We call this information “global”, because it can be inferred from masked image views, which miss the “local” image content. For example, note that the lung cancer mass is not visible in the first masked crop in the middle part of Figure 2. Thus, information about the presence of pathologies is excluded from the condition model’s feature maps.
7. In Section 3.3 we describe the details of our density model, which can be viewed as the predictive model, predicting the descriptor model’s feature maps from the condition model’s feature maps.

Metrics

Your concerns regarding our evaluation metrics are completely valid and understandable. We acknowledge that pixel-level AUROC is not the most representative metric to assess the segmentation quality. Ideally, we would prefer to evaluate our model using segmentation metrics such as Dice score or Hausdorff distance, object detection metrics like finding-level Average Precision, and classification metrics such as patient-level AUROC.

However, it is challenging to fairly estimate these metrics due to the limitations of the available test datasets. These datasets provide annotations only for specific classes of findings — for instance, LIDC includes only lung cancer masks, MIDRC contains only pneumonia masks, KiTS provides kidney tumor masks, and LiTS offers liver tumor masks — while other pathologies present in these datasets remain unlabeled. As a result, our model’s true positive predictions for unannotated pathologies are inadvertently counted as “false positives” when compared with the available partial annotations. For example, in Figure 2, our model correctly predicts pneumothorax in the second image, but the ground truth mask does not include this finding, which can significantly reduce metrics such as the Dice score.

To address this limitation, we draw inspiration from the MVTecAD benchmark, which uses pixel-level AUROC and AUPRO metrics to assess the quality of unsupervised visual segmentation. These metrics are less sensitive to a small number of “false positive” errors because they depend on the pixel-level recall (or finding-level recall in the case of AUPRO) and pixel-level specificity. We calculate recall with respect to only the annotated pathologies. We estimate specificity using a random subsample of voxels that do not belong to the annotated pathologies. This yields quite a tight lower bound since most of these voxels are indeed normal.

Improving Figure 1

Following your suggestion, we updated Figure 1 to highlight that our model correctly identifies pneumothorax in the second image. The ground truth mask misses this pathology, as the image originates from the LIDC dataset, which provides only lung cancer annotations.

Dear Reviewer jnQD, thank you for taking the time to review our submission and for providing valuable feedback. We appreciate your insights and have addressed your concerns below.

Clarity and Illustration

We acknowledge that our initial submission may not have fully conveyed the high-level concept of our method. To enhance clarity, we have added Figure 2 illustrating the overall workflow of our approach, as you recommended. The high-level concept is further elaborated in Section 2.1, where we outline how our method builds upon the existing density-based framework for unsupervised visual anomaly segmentation (UVAS).

Novelty of our method is two-fold:

1. Self-Supervised Descriptor model: Existing density-based UVAS methods employ supervised feature extractors, which we replace with a self-supervised one. We demonstrate that density estimation in a self-supervised feature space is an effective approach to anomaly segmentation — a non-trivial finding that contributes to the field.
2. Self-Supervised Condition Model: Existing density-based UVAS methods are either unconditional or use very simple conditioning on standard sin-cos positional encodings. To estimate density of visual features they need to employ highly expressive Normalizing Flows. We introduce a trainable self-supervised condition model that learns masking-invariant feature maps. Conditioning on these data-driven conditions allows us to achieve SotA results using a simple Gaussian density model.

Detailed descriptions of the descriptor, condition, and density models are provided in Sections 2.2, 2.3, and 2.4, respectively. We believe these sections offer comprehensive explanations of our method's components. If there are specific aspects that remain unclear, we would greatly appreciate further guidance on how we can clarify them.

Experiments demonstrating effectiveness of the components

To demonstrate the effectiveness of the proposed components, we have conducted extensive experiments for component ablation study:

• Descriptor Model Ablation: Table 4 presents an ablation study of the descriptor model, including a comparison with MSFlow — a state-of-the-art density-based UVAS method that employs a supervised descriptor model.
• Condition and Density Models Ablation: Table 3 provides an ablation study of the condition and density models.

These experiments highlight the contributions of each component to the overall performance of our method. We hope these results address your concerns regarding the experimental validation of our approach. If there are additional experiments or comparisons you believe would strengthen our work, we welcome your suggestions.

Dataset Statistics

We would like to clarify that we do not introduce a new dataset in our work. Instead, we train our model on a collection of publicly available datasets, specifically NLST, AMOS, and AbdomenAtlas. We have updated Table 1 to include the numbers of images with non-zero pathology masks.

Regarding pathology distribution, calculating statistics is challenging due to the following reasons:

• The training datasets do not contain annotations of pathologies — neither labels nor masks (that's why self-supervised learning is employed).
• Each test dataset provides annotations for only a specific class of pathologies (e.g., LIDC contains only lung cancer masks, MIDRC contains only pneumonia masks, KiTS contains only kidney tumor masks, and LiTS contains only liver tumor masks).

While we acknowledge that large, uncurated CT datasets inherently include cases with various pathologies, detailed statistics on pathology distribution are not feasible without comprehensive annotations. For example, in Figure 1, the second image from the left displays a pneumothorax (the completely black region in the lung), which is not annotated in the ground truth mask because the LIDC dataset provides annotations only for lung cancer nodules.

Thank you again for your thoughtful review. We believe that these revisions have significantly improved the clarity and completeness of our manuscript. If there are specific areas where you feel further improvements can be made, we would be grateful for your additional feedback.

Dear Reviewer jnQD,

As described in Sections 2.2 and 2.3, we train both descriptor and condition model to produce feature maps (maps of pixel-level embeddings) using dense contrastive (or VICReg) objective. These objectives enforce embeddings in positive pairs to be similar and embeddings in negative pairs to be apart from each other.

We define a positive pair as any pair of pixel-level embeddings obtained from different augmented crops but having the same absolute position w.r.t. the seed image (we depict such embeddings by the same color in Figure 2). Thus, pooling together embeddings in positive pairs ensures that feature maps are equivariant to image crops and invariant to color variations.

We call a pair of pixel-level embeddings negative if they have different absolute position w.r.t. to the seed image or if they are obtained from different seed images. Pushing apart embeddings in negative pairs helps to avoid representation collapse (constant non-informative feature maps), thus making feature maps informative.

The only difference between our descriptor and condition model is the augmentations that we use for positive pairs generation.

When generating positive pairs for the condition model, we use random masking of image blocks in addition to crops and color jitter. Thus, condition model learns masking-invariant informative feature maps. This means that these feature maps are predictable from all masked image views. However, presence of anomalies is not predictable from all masked image views (for example lung mass is not visible in the first masked augmented crop in Figure 2). Therefore, masking-invariant feature maps do not contain information about presence of anomalies.

On the contrary, when training the descriptor model, we ensure that all the augmentations preserve local image semantics. For example, lung mass is visible in both unmasked augmented crops in Figure 2. Thus, descriptor model's feature maps are allowed to distinguish between lung mass and normal lung.

As a result, you can see in Figure 2 that feature maps produced by descriptor model are more fine-grained, while condition model's feature maps are much smoother.

As we mention in Section 2.1, density model can be viewed as a predictive model that tries to predict descriptor model's feature maps from condition model's feature maps. High anomaly scores can be interpreted as high prediction errors. Our empirical results show that

1. Descriptor model indeed produces different features in normal and abnormal image regions (otherwise, density model would not assign different anomaly scores to them).
2. Condition model does not contain information about presence of pathologies (otherwise, density model would yield low prediction errors in pathological regions)
3. Our self-supervised condition model contains more information than baseline condition models (standard positional encodings, APE), which simplifies conditional distributions of descriptor model's pixel-level embeddings given condition model's pixel-level embeddings, allowing to use simply gaussian density model.

Thank you for your response. I appreciate that Figure 2 has been added to illustrate the pipeline of this work. However, some of my concerns remain only partially addressed:

As shown in Figure 2, the training process for both the depictor model and the conditional model appears to follow standard procedures, incorporating two augmentation methods respectively. How do the authors ensure that the depictor model extracts informative feature maps that are invariant to image crops and color variations? Additionally, the feature maps on the right side of Figure 2 for both the depictor model and the conditional model seem identical, which may require further clarification or revision.

Nonetheless, I have slightly raised my score from 3 to 5.