Paper: Temporally Equivariant Contrastive Learning for Disease Progression (ICLR-2024)

Reviewer comment (Theory):
Regarding equation 4, the regularization term: Would utilizing a normal L2 loss on the MLP, adjusted by a coefficient τ, be feasible? The rationale behind adding a constant 1 in this equation is unclear.

Author response:
The problem with L2 loss is that it does not prevent the collapse solution which manifests as 0 norm of the displacement map. That is true that it will prevent large displacement maps, but our main concern is to prevent the collapse (manifests as 0 norm displacement map) and L2 promotes the collapse solution by encouraging 0 norm displacement map. We added a derivation of the regularization term from the pairwise ranking loss in Appendix 4. Also we would like to highlight that without the adding constant 1, the loss term is 0 when the displacement norm is 0, which is exactly the opposite of its intention. And it would promote very large displacement map norms, because $\lim_{x \to \infty} \log\left(e^{-\frac{x}{\tau}}\right)$ is obviously -inf, whereas $\lim_{x \to \infty} \log\left(1+e^{-\frac{-x}{\tau}}\right)$   is 0.

Reviewer comment (Reproducibility):
The structure of ψ in equation 2 needs elaboration. Given its crucial role in predicting future representations, it seems unlikely to be just an MLP.

Author response:
$h_ψ$ is the temporal representation prediction module, which include a displacement map prediction 2 layer MLP $p_ψ$. Although a more elaborate prediction network would improve the results, such as a very recent work [Garrido, Q., et al. (2023) Self-supervised learning of Split Invariant Equivariant representations.], we followed the MLP architecture described in Equimod and AugSelf for predicting the representation from the Siamese Network in order to highlight our reformulation of the equivariance task. Thank you for highlighting the shortcomings of the architecture, we are planning to build a more specific network for the equivariance prediction as a future step.

Reviewer comment (Experiments):
Would utilizing negative time in training, aiming to predict previous images given a current one, contribute to enlarging the dataset?

Author response:
With our current time interval selection (30-270 days), the HARBOR dataset yields 180 pairs per patient from 24 visits. Additionally, in terms of clinical application, it is more interesting to predict only the future, such as for planning future visits. Given additional computational overhead, and the extra complexity from trying to reverse a non-linear transformation, we believe such dataset enlargement would not be beneficial

Reviewer comment (Theory):
Could you clarify the definition of 'time' in this context? In scenarios where a healthy individual undergoes multiple scans, the images might appear identical, whereas for a patient, variations are expected. This could potentially lead to the collapse of the time equivariance module with limited patient data.

Author response:
Your concern is on the point and correct. In this study we aim to address the degenerative disease datasets which the disease generally progresses with varying degrees, and rarely stagnates. As you highlighted if there are healthy patients in the dataset, calculating displacement map would lead to trivial solutions. On the other hand, unsupervised split of HARBOR trial dataset includes late-stage wet-AMD patients. In this trial, those patients are treated with an intravitreal injection called anti-VEGF. It reduces  As a future work, we are planning to address reversed and null changes in time.

Reviewer comment (Experiments):
Is there a scope for the model to handle both 3D scans and 2D images as inputs? This might enable the model to learn a richer set of representations.

Author response:
Our internal experiments with 3D OCT scans achieve higher than what is reported with 2D slices (B-Scans). But OCT scans are highly dense with 128 B-Scans, which increase the computational substantially. Especially, EquiMod has 2 contrastive tasks, which makes it almost unfeasible to train. We could have fit only batch size of 8 (very low even for non-contrastive loss functions), and it took ~11 days to complete a single pretraining on a single Nvidia A100 80gb. Thus, we reverted back to using 2D retinal central B-Scan (called fovea), which is the most relevant position for Age-related Macular Degeneration [Andrew C Lin et al. Assessing the clinical utility of expanded macular octs using machine learning].

Reviewer comment (Writing):
Could you please explain the term "CIn testing" mentioned in section 6? It appears to be a novel concept without a provided definition.

Author response:
Thanks for highlighting the phrase, it was a typo. It is changed to CIp as it is described in Section 4.3.2

Reviewer comment (Novelty):
While the TC method introduces a promising approach for handling irreversible diseases with discretely acquired scans, its utility seems confined to this specific scenario. The real world of medical conditions is vast, and many diseases do not follow a strictly irreversible path. The narrow scope potentially limits the method's broader applicability across the diverse landscape of medical imaging and conditions.

Author response:
In our setup, we targeted degenerative diseases that affect the retina. In AMD, the patient monitoring begins with a diagnosis of the intermediate stage. And the disease only progresses from that point. Although the application seems limited, AMD is the leading cause of blindness, affecting dozens of millions in the world. There is not much clinical insight on the biomarkers that cause the stage conversion. This makes good predictive models necessary for timely delivery of the available drugs and reducing the monitoring load of the hospitals, saving billions in health care cost.

Reviewer comment (Theory):
The method heavily relies on the assumption that degenerative diseases follow a monotonic non-decreasing function over time. This is a strong assumption that might not hold across all conditions or datasets. Diseases can have varying trajectories, with periods of stability, rapid progression, or even temporary reversal. Basing a method on this assumption might lead to inaccuracies in real-world applications.

Author response:
Unsupervised part of the Harbor dataset includes late-stage patients that are treated with anti-VEGF injection. The injection reverses some symptoms and reduces the fluid pockets. But it does not prevent the disease progression. We specifically did not remove those cases in order to have a more realistic set-up. In terms of optimization, they are similar to noisy samples (also there are other artifacts from erroneous scanning as well), and the models should be robust to it. The displacement map is conditioned on both the time difference and the current image, which is necessary to capture the individual trajectories and progression speeds. To make a comparison, the input image is not necessary to define to rotate the image. The 3x3 rotation matrix is enough. But in order to rotate an object within a 2D image, the rotation degree and the input image have to be provided.

Reviewer comment (Experiments):
Medical images are notorious for their variability and potential anomalies, especially given the variation between patients, imaging equipment, and acquisition protocols. The paper does not delve deeply into how the TC method would handle potential anomalies, outliers, or inconsistencies in the longitudinal medical images.

Author response:
The Harbor dataset have scans after treatment creating temporal inconsistencies. We aim to extract the available temporal information from the unsupervised temporal dataset, hence the noise and the anomalies were expected. Similarly, we chose the SIGF as the second dataset, because the image acquisition intervals are highly irregular.

Reviewer comment (Evaluation):
While the paper alludes to reduced complexity in comparison to other models, a more explicit discussion or breakdown of computational costs, resource requirements, and scalability considerations is conspicuously missing. For real-world application, especially in clinical settings, understanding computational overhead is crucial.

Author response:
Compared to EquiMod, another method for predicting the transformed representations, TC has substantially fewer parameters due to the fact that TC predicts the next time step in the representation space rather than a high dimensional projection space. Additionally EquiMod relies on additional contrastive task to prevent trivial equivariant solution, contributing to its higher traning time, unlike TC. The models relying on transformation parameter recovery, ESSL and AugSelf, have fewer parameter compared to TC and EquiMod, since their predictor outputs a scalar rather than a representation vector, but they have worse execution time because they both require additional transformed images to predict the transformation parameter. Consideting most of the medical datasets consist of volumes the efficiency of TC in terms of trainable parameters and time, makes it more scalable to higher dimensional datasets.
We added this discussion and the cost table in Appendix A.4.

Reviewer comment (Theory):
The introduction of a regularization loss term, crucial for maintaining time parameter sensitivity, is presented without an in-depth rationale or derivation. A clearer explanation of this term's derivation, its impact on the model's performance, and sensitivity analyses would provide more confidence in the method's robustness.

Author response:
We added the full derivation of the term in Appendix A.5.

Reviewer comment (Experiments):
There are only 3 baselines tested for the experiments and for these the performance improvement is minor. Error bars should be added to identify whether these performance improvements are significant and how robust to different random seeds is the model compared to the other baselines.

Author response:
We extended our tables with mean and std of the metrics. Additional for Harbor dataset we added PRAUC as an additional metric since the dataset is heavily imbalanced (1:20 for the positive class)

Reviewer comment (Experiments):
It is unclear how the hyperparameters are selected. Specifically, for VICReg, hyperparameters are selected to bring the magnitude of the loss terms in the same range. What does this mean? What happens when you use the hyperparameters from the original VICReg implementation?

Author response:
In pretext tasks such as the rotation prediction, it is possible to monitor the pretext task accuracy performance and optimize the hyperparameters accordingly. But the contrastive methods do not rely on a pretext task. In order to avoid biasing for the supervised downstream task, we only monitored the values of the loss terms during pretraining and ensured that they converges in a similar rate.

Reviewer comment (Theory):
The authors note that the displacement term needs to be larger than 0 in order to prevent collapse. However, in many medical settings it is possible that there is no change. However, it looks like this model requires a small amount of change. Since the severity of the disease is monotonically non-decreasing, it seems like it is also the case here. If that is true, this method seems very limited in its application.

Author response:
It is corrected that the model requires a change, the displacement map still can take very small values. In our setup, we targeted degenerative diseases that affect the retina. In AMD, the patient monitoring begins with a diagnosis of the intermediate stage. And the disease only progresses from that point. Although the application seems limited, AMD is the leading cause of blindness, affecting millions in the world. There is not much clinical insight on the biomarkers that cause the stage conversion. This makes good predictive models necessary for timely delivery of the available drugs and reducing the monitoring load of the hospitals

Reviewer comment (Experiments):
How were hyperparameters selected? Specifically, could you provide some intuition for why you selected the hyperparameters you chose?

Author response:
In pretext tasks such as the rotation prediction, it is possible to monitor the pretext task accuracy performance and optimize the hyperparameters accordingly. But the contrastive methods do not rely on a pretext task. In order to avoid biasing for the supervised downstream task, we only monitored the values of the loss terms during pretraining and ensured that they converges in a similar rate.

Reviewer comment (Writing):
In the first paragraph, the text says that "However, sensitivity to some of these transformations may be crucial for specific downstream tasks, such as color information for flower classification or rotation for traffic sign detection." This is in reference to, if I understand correctly, how data augmentation is used standardly in contrastive learning (e.g., SimCLR, supervised contrastive learning). However, when using such contrastive learning approaches, it is standard to make sure that the random perturbations/transformations used for data augmentation only consist of changes that the model shouldn't care about whereas the ones that we actually want the latent embedding representation to learn should not be used in data augmentation. Am I missing something here?

Author response:
Since contrastive learning is an unsupervised step, it is hard to determine a common set of image perturbations that are not relevant for all possible downstream tasks. To give an example, in  "What should not be contrastive in contrastive learning."  Xiao, Tete, et al. showed that invariance to color perturbations is harmful when classifying flowers, while it is useful in coarse-grained animal classification task. Similarly, we can argue that rotation invariance is harmful for the traffic sign classification since some signs differ only in orientation, but their classes are different. Both tasks can be defined on the ImageNet dataset. But when using contrastive pretraining, there should be a bias for the target downstream task to prevent unintended invariances.  Similar motivation can be found in ESSL, AugSelf, and EquiMod. Since our focus is on progression of the degenerative diseases, we hypothesized that time sensitive representations are useful in predictive downstream tasks.

Reviewer comment (Novelty):
How does the proposed method TC relate to steerable equivariant representation learning (Bhardwaj et al 2023)?

Author response:
Thanks for highlighting another related work. It tries to find a balance between invariance and equivariance within the same representation space. We believe that this approach is hard to optimize, since learning complete invariance is the trial solution for the equivariance. EquiMod solves this problem by enforcing invariance and equivariance in 2 distinct projection spaces after the representations. But it requires calculation of 2 different costly contrastive losses. We argued that the second projection space is unnecessary, such that if the representation space is equivariant, this space can be projected to an invariant space trivially on which we calculate the contrastive task for the invariance. Also we proposed the displacement map regularization term instead of a second contrastive task (similar to EquiMod).

Reviewer comment (Experiments):
Please update the experimental setup so that error bars could be reported (such as running experimental repeats with different random seeds and reporting mean/std dev of different achieved evaluation scores).

Author response:
We extended our tables with mean and std of the metrics. Additional for Harbor dataset we added PRAUC as an additional metric since the dataset is heavily imbalanced (1:20 for the positive class)

Reviewer comment (Evaluation):
Overall I would like to see a much more detailed discussion of precisely why the authors think that the proposed method at times outperforms existing 3 equivariant contrastive methods.

Author response:
Firstly, we believe our method enforces equivariance directly by predicting the future representation from a previous one, similar to EquiMod. The transformation  parameter recovery methods, ESSL and SelfAugs, do not guarantee that kind of relation in the representational space. Secondly, we formulated directly on the representation space rather than on a projection space as it is in the EquiMod, thus our representations are directly optimized to be equivariant.

Reviewer comment (Novelty):
There is a large body of literature on medical image registration. I think some discussion of this literature would be helpful --- do state-of-the-art methods from this literature simply not work well for the problem setup considered here?

Author response:
In deep learning, medical image registration can be performed either by learning a linear affine transformation or generating the registered images using generative models. Our main goal is to learn time equivariant representations such that they are useful in the downstream task without generating any image and relying on any labels. We showed that the time equivariance property is useful for the disease progression task, on top of invariance to random perturbations.