Lifelong Histopathology Whole Slide Image Retrieval via Distance Consistency Rehearsal

Continual learning for CBHIR can be defined as training a model on non-stationary data when new data added into the WSI database.

The proposed framework is illustrated in Fig. 1. For an original WSI, we first divided it into several non-overlapping patches with size of $256\times 256$. Then, we utilize pre-trained encoder PLIP [14] to extract patch features, for its excellent performance in understanding image semantic information. Afterwards, a feature cube of a WSI can be represented as $x\in\mathbb{R}^{n_{p}\times d_{f}}$, where $d_{f}$ is the dimension of the feature and $n_{p}$ is the number of patches in the WSI. To alleviate storage pressure, we save feature cubes in the memory bank instead of the original WSIs.

As shown in Fig. 1(I), we define the database as a sequence of datasets $\mathcal{D}=\{\mathcal{D}_{1},\cdots,\mathcal{D}_{T}\}$, where $\mathcal{D}_{t}=\{x_{i},y_{i}\}^{N_{t}}_{i=1}$ is the dataset of task $t$, and $T$ is the total number of tasks. Dataset $\mathcal{D}_{t}$ contains $N_{t}$ labeled WSIs, and $y_{i}$ is the class label of WSI $x_{i}$. In the class-incremental CHBIR scenario, the model needs to return relative WSIs sequence of given query WSIs from both old tasks and new tasks. Universal training process for every task is illustrated in Fig. 1(III). When data of new task coming, described as Fig. 1(a-b), a mini-batch of training data and replay samples from memory bank are input to the WSI encoder to get WSI representations. Then, representations of current task and previous tasks are combined to calculate retrieval relative losses which facilitate retrieval precision of the model. Meanwhile, we execute distance consistency rehearsal on representations of old tasks to maintain the stability of the returned queues, as shown in Fig. 1(IV).

As illustrated in section 1, we design distance consistency rehearsal (DCR) module to maintain stability of result queues for old tasks, which is detailed in Fig. 1(IV) and illustrated in algorithm 1.

In practical scenarios, the stability of the CBHIR system is demonstrated by the consistency of the result queues for old tasks. In terms of feature space, this stability is reflected by maintaining constant distances between instances of old tasks even after learning new tasks. Based on this prior, we propose the distance consistency rehearsal (DCR), which is achieved by minimizing the differential value of distance between representations of current replay samples and distance matrix of corresponding replay samples saved in memory bank. The rehearsal process is elaborated as follows.

First, we construct the target distance matrix of instances saved in memory bank. At the ending of task $t-1$, we execute reservoir sampling algorithm [24] on dataset $\mathcal{D}_{t-1}$ and memory bank $\mathcal{M}_{t-1}$ to get sampled feature cubes $\mathbf{X}_{t-1}\in\mathbb{R}^{n_{t-1}\times n_{p}\times d_{f}}$ and $\mathbf{X}_{t-1}^{m}\in\mathbb{R}^{n_{t-1}^{m}\times n_{p}\times d_{f}}$ for the memory bank $\mathcal{M}_{t}$ of task $t$, where $n_{t-1}$ and $n_{t-1}^{m}$ means sampled number of $\mathcal{D}_{t-1}$ and $\mathcal{M}_{t-1}$ respectively. Then, sampled feature cubes is fed into task $t-1$ encoder $f_{t-1}$ to get feature representations $\mathbf{F}_{t-1}\in\mathbb{R}^{n_{t-1}\times d_{F}}$ and $\mathbf{F}_{t-1}^{m}\in\mathbb{R}^{n_{t-1}^{m}\times d_{F}}$, where $d_{F}$ is the dimension of the whole slide representation. With combination of $\mathbf{F}_{t-1}$ and $\mathbf{F}_{t-1}^{m}$, denoted as $\mathbf{C}_{t-1}\in\mathbb{R}^{n_{t-1}^{d}\times d_{F}}$, we calculate the Euclidean distance between every pair of elements in set $\mathbf{C}_{t-1}$ to obtain the target distance matrix $\mathbf{D}_{t-1}$, which is formulated as

where $\mathbf{D}_{t-1}\in\mathbb{R}^{n_{t-1}^{d}\times n_{t-1}}$, $n_{t-1}^{d}=n_{t-1}+n_{t-1}^{m}$, denotes the total number of instances sampled at the ending of task $t-1$, and $\mathbf{c}_{t-1}^{i},\mathbf{c}_{t-1}^{j}$ represent the $i$-th and $j$-th sampled instance.

Then, distance consistency rehearsal is implemented for current task. In every mini-batch of current dataset $\mathcal{D}_{t}$, we get representations of current dataset and instances sampled from memory bank, denoted as $\mathbf{F}_{t}$ and $\mathbf{F}_{t}^{m}$. For $\mathbf{F}_{t}^{m}$, distance matrix $\mathbf{d}_{t}$ is obtained by the same way in Eq. (1) and illustrated in algorithm 1. To maintain the distance consistency of rehearsal samples, we minimize the Mean Squared Error (MSE) loss between distance matrix of current task and replay samples by the equations

where $n_{t}^{m}$ means total number of feature cubes sampled from memory bank and $\mathbf{d}_{t-1}$ is sub-matrix sampled from $\mathbf{D}_{t-1}$ according to the indexes of $\mathbf{X}_{t}^{m}$.

Above all, we design the distance consistency loss $\mathcal{L}_{DC}$ to maintain the queue stability of a continual CBHIR system. Meanwhile, pair-wise loss and cross-entropy loss have been proven effective in image retrieval task [28, 29]. As a result, the model is optimized in an end-to-end fusion by

where $\mathcal{L}_{P}$ is the pair-wise loss, $\mathcal{L}_{CE}$ is the cross-entropy loss, $\mathcal{L}_{DC}$ is the distance consistency loss, $y_{t}$ and $y_{t}^{pred}$ are the ground-truth label and output logits of input instances, and $\alpha$ is a balancing factor. Through 3, we attempt to not only improve retrieval precision of LWSR but also maintain queue consistency as much as possible.

We first trained a model with all datasets under fully supervision as the upper-bound, which is identified as JointTrain in Table. 2. Subsequently, individual training were sequentially conducted for the four datasets as the lower boundary, which is identified as Finetune. According to Table. 2, model subjected to Finetune exhibits the poorest performance, indicating that catastrophic forgetting occurs as the database size increases. In contrast, our method significantly mitigates catastrophic forgetting, performing nearly as well as the JointTrain approach. Although our method does not achieve as high mAP as JointTrain, it delivers comparable values in R@3 and P@5 metrics with JointTrain, which means the pathologists can efficiently receive diagnostically relevant cases by assessing a few top-returned results. DCR is defined to maintain the distance relationship between the cases from old retrieval tasks, which provides more finer guidance to the model in learning the similarities between WSIs compared to the common pair-wise loss. As highlighted in Table. 2 and illustrated in Fig. 2, DCR markedly improves the precision and consistency of fine-grained retrieval, as demonstrated by enhancements in both the precision and consistency metrics. Moreover, we combined Breakup-Reorganize (BuRo) rehearsal, the main component of ConSlide [13], with our framework, which is presented as LWSR w/ BuRo in Table. 2. BuRo achieves data augmentation by randomly selecting and combining patches from same category WSIs. However, for retrieval tasks, representations are expected to describe the complete semantic information of WSIs. The virtual cases created by BuRo would mislead the retrieval model to describe the real WSIs. It can be obviously observed that BuRo rehearsal method performs poorly in WSI retrieval tasks, with mAP of 0.563, showing a significant performance gap between LWSR w/o DCR and LWSR.

We compare our method with several classic continual learning approaches, involving regularization-based methods LwF [17] and EWC [16] and replay-based methods ER-ACE [2], A-GEM [3] and DER++ [1].

According to Table. 3, regularization based methods fail to effectively alleviate catastrophic forgetting and they only achieve a slight improvement over Finetune. Compared to other replay-based continual learning methods, our method consistently achieve the best performance under various circumstances. The main reason is that these methods are proposed to solve catastrophic forgetting in the domain of classification initially and do not accounted for taking the interaction between the current task and previous tasks into consideration. In comparison with the second-best methods under buffer in size of about 10 WSIs, LWSR achieves an increase of 6.3%/3.0%/3.5% in mAP, R@3 and P@5, and an increase of 4.3%/1.6%/2.6% in site mAP, R@3 and P@5. It has demonstrated the effectiveness of the proposed method for the task of lifelong CBHIR. Furthermore, with the growth of buffer size, our method can achieve further improvements. It is also apparent that the enhancement in performance metrics upon expanding the buffer size from 10 WSIs to 15 WSIs is not as pronounced as the improvement observed from enlarging it from 5 to 10. A 10-WSI buffer balances diversity between current and previous tasks, aiding the model in achieving high retrieval precision after sequential learning. However, a 15-WSI buffer may overly emphasize past tasks at the expense of current ones, leading to decreased focus on the present task.