DDA: Dimensionality Driven Augmentation Search for Contrastive Learning in Laparoscopic Surgery

In the first stage, for each $\bar{\bm{x}}\in\mathcal{X}$ in a batch of $M$ images from $\mathcal{D}_{u}$, we generate two positive samples ${\bm{x}},{\bm{x}}^{+}$ from the augmentation policy $\mathcal{T}(\bar{\bm{x}})$, and $2(M-1)$ independent negative samples $\{{\bm{x}}_{m}^{-}\}_{m=1}^{2(M-1)}=\{\mathcal{T}({\bm{x}})|{\bm{x}}\neq\bar{\bm{x}}\}$ from the other $M-1$ images. An encoder $f(\cdot)$ maps input images to the representation space $\mathcal{X}\to\mathcal{R}^{d}$ with the representation $\mathbf{z}=f({\bm{x}})$, and projector $g(\cdot)$ obtains embedding $\mathbf{e}=g(\mathbf{z})\in\mathcal{R}^{e}$. A classical optimisation objective for contrastive learning Chen et al. (2020a) is the following,

where $\mathrm{sim}(\cdot)$ is the cosine similarity, and $\eta>0$ is the temperature that controls the smoothness of distance distribution.

In the second stage, after the encoder is trained, an additional classifier or segmentation model $h(\cdot)$ can be attached to the encoder as $h\circ f$, for simplicity, we denote this as $f^{{}^{\prime}}$. The $f^{{}^{\prime}}$ learns to map the input image to the label space $\mathcal{X}\to\mathcal{Y}$ using $\{({\bm{x}},{\bm{y}})\}\in\mathcal{D}_{l}$ following a typical supervised learning setup by minimizing the following objective,

where $\mathcal{L}$ is the objective function for supervised downstream tasks (e.g., cross-entropy).

Each round of a two-stage learning framework mentioned in the previous section is time consuming. This challenge makes it impossible to apply a thorough grid search, especially for a large search space with a diverse number of operations and choices in each augmentation policy. To tackle these challenges, we propose a streamlined search framework, DDA.

Specifically, we setup the augmentation search in a differentiable fashion following existing works Liu et al. (2019); Hataya et al. (2020). Let $\mathbb{O}$ be a set of $K$ image augmentation operations $\mathcal{O}_{k}\in\mathbb{O}:\mathcal{X}\to\mathcal{X}$, and $\mathcal{X}$ is input space. Each operation $\mathcal{O}_{k}(\hskip 4.26773pt\cdot\hskip 4.26773pt;p_{k},\lambda_{k})$ has two parameters: the probability $p_{k}$, and the augmentation magnitude $\lambda_{k}$, for applying the operation. For an input image, ${\bm{x}}\in\mathcal{X}$, the output of applying an augmentation from the set of possible operations $\mathbb{O}$, also known as a sub-policy, depends on weighted sampling. This is defined as, $\bar{\mathcal{O}}({\bm{x}};\bm{p},\bm{\lambda})\in\{\mathcal{O}_{k}({\bm{x}};p_{k},\lambda_{k})|k=1,\dots K\}.$

Let $\tau$ be the set of augmentation sub-policies as $\tau=\{\bar{\mathcal{O}_{1}},\dots\bar{\mathcal{O}_{N_{\tau}}}\}$ contains $N_{\tau}$ consecutive sub-policies. The augmented image is obtained by ${\bm{x}}^{\prime}=\tau({\bm{x}})=(\bar{\mathcal{O}_{N_{\tau}}}\circ\dots\circ\bar{\mathcal{O}_{1}})({\bm{x}};\bm{p_{n}},\bm{\lambda_{n}})$, $n=1,\dots N_{\tau}$. During searching, the output of the $n^{th}$ sub-policy $\bar{\mathcal{O}_{n}}({\bm{x}},\bm{p_{n}},\bm{\lambda_{n}})$ is the weighted sum of all possible operation choices,

The probability $p_{k}$ is convert from learnable real-number parameter, $w_{k}$, using the softmax function as $p_{nk}=\sigma_{\eta}(w_{nk})=\frac{\exp(w_{nk}/\eta)}{\sum_{K}\exp(w_{nk}/\eta)}$, with temperature $\eta>0$. Low temperature will generate the distribution of $p_{k}$ as a one-hot like vector.

The components of the augmentation policy thus become learnable parameters and can be updated by backpropagation. In this case, our augmentation policy is designed as a block of plug-in-and-play operation layers as $\mathcal{T}$ with parameters ${\bm{w}},\bm{\lambda}$. To learn the optimal augmentation policy $\mathcal{T}^{*}$, we further propose a suitable proxy objective in Section 2.3 for the augmentation search, which optimises the contrastive representation on a fixed encoder.

Here we provide an overview of the three-step DDA search framework as illustrated in Figure 1. Firstly, we obtain a contrastive learning encoder $f$ that is trained with the most basic augmentation (e.g. image cropping) with a contrastive projector $g$. Then, we apply our differentiable augmentation search with the proxy objective explained in Section 2.3 to optimise the parameters of the policy only on the fixed $f$. Lastly, the optimal policy $\mathcal{T^{*}}$ will be applied to re-conduct contrastive learning to obtain $g^{*}\circ f^{*}$. We provide a sketch of the pseudocode of the DDA framework in Algorithm A. in Appendix A.

main_search_objective_function can’t be directly optimised due to lacking validation data, and the two-stage training is hard to optimise differentially. We propose a novel objective as a proxy so that the differentiable search only involves contrastive pretraining.

In contrastive learning, the deep representation $\mathbf{z}\in\mathcal{R}^{d}$ extracted from a surgical image needs to be distinguishable from the representations of other images. Our focus lies in understanding the local distribution of deep representations in vicinty of a query image. We define a local neighbourhood-of-interest as a $d$-dimensional sphere with a small radius $r$ centered at the query. Any data within the neighbourhood has a smaller distance than $r$ from the query, and likely has similar visual content as the query. An illustration of the deep representation for a local neighbourhood is shown in Figure 4, in Appendix A.

To evaluate the properties of a query image’s neighborhood, one can estimate the growth rate in the number of nearby data points encountered as the distance from the query increases. This growth rate provides an estimation of the local intrinsic dimensionality (LID) of the query data located in the high-dimensional subspace Houle (2017). Intuitively, the LID assesses the effective number of dimensions (the intrinsic dimension) needed to characterize the local neighbourhood of a query point. If the LID equals 2, the local neighbourhood surrounding the query point behaves like it has two dimensions. LID is thus like a complexity measure and assesses the space-filling properties around a query point. In our scenario, we will assess the LID of each point in the deep representation produced by SSL.

Let $\mathbf{z}$ be the query point, and $\mathbf{s}$ denotes any nearby representation of a different image within its vicinity. The distance between $\mathbf{z}$ and $\mathbf{s}$ is denoted as $r_{s}=l(\mathbf{z},\mathbf{s})$, where $l$ represents a distance measurement. We define the cumulative distribution function (CDF), $F(r)$, of the number of data points concerning the local distance distribution with respect to $\mathbf{z}$ as the probability of the sample distance lying within a threshold $r$, $F(r)\triangleq\mathsf{Pr}[r_{s}\leq r]$.

The local intrinsic dimension (LID) at $\mathbf{z}$ defined as the limit, when the radius $r$ tends to zero as,

In practice, $\operatorname{LID}^{*}_{F}$ needs to be estimated, and it is not an integer. For simplicity, for the rest of the paper, we denote ‘LID’ as the quantity of $\operatorname{LID}^{*}_{F}$.

In our scenario, we will wish to optimise the LID of each data sample, encouraging it to be large. Work by Huang et al. (2024) has shown that it is theoretically desirable to optimise the log transform of the LID in deep learning, rather than the LID. Intuitively, if a query point has a large (log) LID, then its neighbourhood requires more dimensions to characterise. This is a desirable property for SSL, where a key objective is to avoid what is known as dimensional collapse Jing et al. (2022).

This makes LID a suitable proxy objective for the augmentation search. Since 1) it does not require finetuning on downstream tasks, and 2) optimising the LID distribution of samples in the deep space can easily be integrated into a gradient descent framework. Following Huang et al. (2024), to obtain a representation with higher LIDs, we can optimise the Fisher-Rao distance between the LID of the representations $\bm{z}$ and a uniform distance distribution (with $\mathrm{LID}$ of 1) using the following loss function as,

where M is the number of samples in the batch, and $\operatorname{LID}^{*}_{F_{i}}$ can be estimated using any popular estimator for LID in the encoder’s output representation $\bm{z}$.

Different from existing work Huang et al. (2024) that investigated optimisation of model parameters to produce a better contrastive representation, in our work we optimise the augmentation policy with a fixed encoder $f$. We find this simple but important shift in perspective is highly significant. More concretely, we apply differentiable augmentation search to optimise the parameters of the policy with \equationrefeq:abs_fr_distance. as the objective as shown in Algorithm A in the Appendix A. Although it is also possible to apply other proxy functions instead of using LID for DDA, we empirically find that using LID is highly suitable for the augmentation search. To our knowledge, ours is the first work to consider the use of LID as a measure for optimising augmentations.

As shown in Table 1, compared against the original SimCLR augmentation, a large 6% improvement can be observed for our DDA on the linear evaluations. The original SimCLR policy performs similarly to the Base policy (crop and resize) on laparoscopic cholecystectomy datasets. For finetuning on segmentation tasks, augmentations found by our method can outperform the original SimCLR policy by 1-3%. This can be considered as a significant improvement in the context of SSL evaluation He et al. (2020); Bardes et al. (2022). Without the augmentation search, using randomly selected operations with our search space results in a complete collapse of the representations, where the model outputs a constant vector Jing et al. (2022). It has been shown in existing work Jing et al. (2022) that excessive augmentations could cause dimension or complete collapse. It is worth noting that the augmentation policy found by DDA avoids selecting an excessive number of operations by selecting Identical, as shown in Figure 2 and 2. Compared with SelfAugment, our method also demonstrated consistently superior performance. On our private SVHM dataset, the augmentation policy found by SelfAugment caused a complete collpase. For pretraining with Cholec80, DDA outperforms SelfAugment by 13% in the linear evaluations.

In this subsection, we investigate the effect of the different numbers of operations in the policy for the representation quality. We also examine the policy found by DDA. As shown in Figure 2 and 2, the representation quality evaluated by the linear probing shows a consistent improvement over the augmentation used by SimCLR and the Base policy. Similar results for finetuning can be found in Appendix B.4. This consistent improvement indicates that the commonly used augmentations on natural images appear suboptimal for a medical dataset such as laparoscopic cholecystectomy (LC).

To further investigate the difference between the optimal policy for natural images and policies found by DDA, which are suitable for the medical datasets, we plotted the distributions of selected operations in Figure 2 and 2. It can be observed that despite the different pretraining datasets (our private dataset SVHM and Cholec80), the resulting policy is similar. Detailed results for each policy are in Appendix B.3. The policy found by DDA prefers selecting Gaussian Blur, Saturation, Posterize, and Sharpness. One common characteristic regarding these operations is that they do not change the colour profile of the image. In the existing literature, it has been shown that an augmentation resulting in an overlapping view of two different images and semantically similar to each other is beneficial for contrastive learning Cai et al. (2020); Wang et al. (2022); Huang et al. (2023); Joshi and Mirzasoleiman (2023). In the context of LC images, randomly changing the colour profile to other random colours is not ideal for creating such an overlapping view, since the majority of the contents in the images are red. On the other hand, the Gaussian Blur and Posterize could create a blurring effect or decrease the image quality of the image that could easily create an overlapping view. This is because in LC surgery, motion movement of the camera and fog generated during anatomy dissection by diathermy hook can degrade image quality. Similarly, the effect of Sharpness and Saturation could also inherently appear in the dataset. In Figure 3, we provide a visualisation of augmented images using DDA and SimCLR augmentation, where DDA can create semantically similar views while SimCLR augmentation produces different and unrealistic views. Additional augmented images can be found in Appendix C. In summary, the augmentation policy found by DDA is more effective (verified by linear probing and finetuning) and well suited to the LC dataset.