Stable-Drift: A Patient-Aware Latent Drift Replay Method for Stabilizing Representations in Continual Learning

AI summary

Overview

• Research area: Continual learning / lifelong learning for medical imaging (computer vision), specifically replay-buffer construction under cross-hospital domain shift for COVID-19 chest CT classification.
• Technical level: Intermediate. The core idea is intuitive (track how much a sample's internal features change, replay the most-changed ones), but the paper assumes familiarity with continual learning terminology (catastrophic forgetting, replay buffers, backward/forward transfer), CNN and Vision Transformer backbones, and distance metrics.
• One-sentence scope: The paper proposes a replay-buffer selection strategy that stores the per-patient CT slices whose internal feature representations shift most when a source-trained model is naively fine-tuned on a new hospital's data, and shows this reduces forgetting relative to naive fine-tuning and random replay.

What This Paper Is About

When a deep learning model is trained on data from one hospital and then fine-tuned on data from another, it can abruptly lose the diagnostic knowledge it had before, a failure called catastrophic forgetting. This paper's goal is to choose which old samples to keep in a small replay memory so that this forgetting is minimized, using a signal called "latent drift" that measures how much each sample's internal representation changes during naive domain adaptation. The authors test this on a cross-hospital COVID-19 CT classification task with both a CNN (ResNet50) and a Vision Transformer (Swin Transformer).

Key Contributions

1. Latent drift as a replay signal. The paper introduces representation drift, the change in a sample's feature vectors between a source-domain model and a naively fine-tuned target-domain model, as a practical and interpretable criterion for identifying samples at high risk of being forgotten.
2. Patient-aware, multi-layer buffer construction. Rather than selecting top-scoring slices globally, the method aggregates drift per patient, ranks patients by average drift, and stores the top 30 highest-drift slices from each patient until the buffer is full. Drift is averaged over the final two backbone layers (L and L−1) using cosine distance.
3. A cross-hospital benchmark result. The framework is evaluated on COVID-19 CT data from two hospitals (H1 as source, H2 as target), and the authors report that it outperforms naive fine-tuning and random replay baselines under both ResNet50 and Swin Transformer backbones.
4. Ablations isolating design choices. The paper separately tests patient-aware versus global/center slice selection, multi-layer versus single-layer drift, cosine versus Euclidean versus Mahalanobis distance, and a drift-plus-entropy hybrid score.

Main Findings

• Proposed method gives the best stability-plasticity trade-off (per-patient accuracy). On Swin Transformer: 92.45% on H1 and 93.75% on H2. On ResNet50: 88.13% on H1 and 89.29% on H2.
• Naive fine-tuning demonstrates severe catastrophic forgetting. Per-patient accuracy on H1 drops to 51.80% (Swin) and 71.58% (ResNet50), while H2 accuracy is 95.31% (Swin) and 69.87% (ResNet50).
• Random replay helps but less than the proposed method. Per-patient: 92.21% H1 / 90.70% H2 for Swin, and 87.29% H1 / 79.01% H2 for ResNet50. The paper reports the proposed method improves over random replay by +0.24 percentage points (pp) on H1 and +3.05 pp on H2 for Swin, and by +0.84 pp on H1 and +10.28 pp on H2 for ResNet50.
• Source-only model is highly stable but poorly adapted. Per-patient: 94.24% H1 / 88.17% H2 (Swin) and 93.17% H1 / 88.66% H2 (ResNet50).
• Patient-aware selection beats global and center slice selection. For Swin, per-patient H1 accuracy is 82.01% for global slice selection and 92.09% for center slice selection, versus 92.45% for the proposed patient-aware method. The paper argues global selection risks redundant buffers by over-sampling a few "difficult" patients, and that center slices are informative but do not achieve the highest overall performance.
• Multi-layer drift beats single-layer drift. For Swin, per-patient H1 retention improves from 91.37% (single-layer) to 92.45% (multi-layer). Single-layer per-patient results are 93.51% H2 / 91.37% H1 (Swin) and 81.03% H2 / 90.29% H1 (ResNet50).
• Alternative distance metrics and hybrid scores are competitive but not superior overall. Euclidean distance achieves the highest H2 per-patient accuracy for Swin (94.42%), which the authors say shows the drift concept is robust to the distance metric. The latent drift + entropy hybrid (α = 0.7, β = 0.3) performs well for ResNet50 at 91.01% H2 per-patient, but for Swin the authors report the cosine-based pure-drift method gives the best overall balance.
• Backward transfer is best for the proposed method. Per-patient BWT: −0.424 (Swin) and −0.216 (ResNet50) for naive fine-tuning; −0.020 and −0.059 for random replay; −0.018 and −0.050 for the proposed method. Forward transfer is identical across methods because all share the same H1 initialization: 0.029 (Swin) and 0.768 (ResNet50).
• Swin Transformer outperforms ResNet50 in this scenario. The paper attributes this to hierarchical self-attention being better suited to modeling long-range contextual features that define an imaging domain.
• High retention efficiency with a small buffer. The buffer holds 30,000 samples, approximately 10% of the H1 training set, class-balanced with 15,000 COVID and 15,000 non-COVID samples.

Methodology in Plain English

The method runs in three stages.

Stage 1 — Measure instability. Train a model on Hospital 1 data until convergence (the "source model"). Make a copy and fine-tune it directly on Hospital 2 data (the "naively fine-tuned model"). For every Hospital 1 training slice, compute how much its feature vector changed between the two models. That change is the Multi-Layer Latent Drift (MLD): the average cosine distance between the two models' features at the final two backbone layers. Cosine distance is chosen over Euclidean distance because it captures a change in the direction (semantic meaning) of the feature vector rather than its magnitude, and averaging over two layers makes the score more stable than using one layer alone. A high MLD score marks a slice as representationally unstable, hence at high risk of being forgotten.

Stage 2 — Build a patient-aware buffer. Group all slices by patient ID and compute each patient's average MLD. Rank patients from highest to lowest average drift. Walking down that ranked list, take the top 30 slices (by individual MLD) from each patient and add their raw images and labels to the buffer, stopping once the buffer reaches its target size K (30,000). This prevents the buffer from being dominated by a handful of patients and keeps it diverse across clinical cases.

Stage 3 — Replay-based continual learning. Start again from the converged source model, then train on Hospital 2 data with mini-batches that mix new H2 samples and replayed H1 buffer samples (50% probability of drawing from either source), using a combined loss.

Implementation details: ImageNet-pretrained ResNet50 and Swin Transformer backbones; AdamW optimizer with learning rate 5e-5 for Swin and 1e-5 for ResNet50; batch size 32; 15 epochs on H1 and 10 on H2; Weighted Random Sampler plus Focal Loss to handle class imbalance (H1 is roughly 60% COVID, H2 roughly 10% COVID); random horizontal flips and rotations for augmentation; splits made at the patient level to prevent leakage. Evaluation uses per-slice and per-patient (majority vote) accuracy, BWT, FWT, and Latent Representation Shift (LRS), the average MLD on the H1 test set between the source model and the final continual learning model. Experiments ran on a server with 8 × NVIDIA Tesla V100 GPUs.

Why This Matters

Impact on research. The paper reframes replay-buffer construction as a representation-stability problem: instead of sampling randomly or by gradient/uncertainty heuristics, it uses the observed drift in a model's own internal features as the selection criterion, and it shows that aggregating that signal at the patient level (a structure unique to clinical data) matters. It also provides a head-to-head comparison of a CNN and a Vision Transformer under the same cross-hospital continual learning protocol. The ablation results are notable for what they complicate: Euclidean distance outperformed cosine on Swin H2 accuracy (94.42%), and drift-plus-entropy improved ResNet50 H2 accuracy (91.01%), so the paper's headline configuration is not uniformly best on every single metric.

Real-world applications:

• Cross-hospital deployment of diagnostic AI, where a model trained at one institution must be updated for a new site without discarding validated diagnostic knowledge.
• Privacy- and storage-constrained clinical settings, where only a small fraction of historical data (here, approximately 10% of the source training set) can be retained on-device or transferred.
• Triage and quality-control tooling that flags which past cases are most at risk of being degraded by a model update, since the latent drift signal can be read as a diagnostic indicator on its own.
• Chest CT classification pipelines for COVID-19 and comparable imaging tasks built on multi-slice patient scans, where per-patient aggregation matches how radiologists actually read studies.

Industry relevance. Any vendor or health system running models that are periodically retrained on incoming data faces the stability-plasticity trade-off directly. The paper's framing, a compact, class-balanced, patient-aware replay buffer that preserves nearly the same H1 performance as models with full access to historical data, speaks to practical constraints around storage, data transfer, and privacy, and to the regulatory need for demonstrable retention of prior diagnostic performance.

Future Directions

1. Multi-domain sequences. Extend the method from two hospitals to longer sequences of several different hospitals, where forgetting compounds and drift signals may interact across tasks.
2. 3D volumetric CT. Apply drift-based selection to volumetric data rather than individual slices; the authors expect this may reveal different data structures and forgetting dynamics.
3. Formal uncertainty quantification. Integrate the drift signal with principled uncertainty estimation, building on the preliminary drift-plus-entropy hybrid the paper tests with α = 0.7 and β = 0.3.
4. Federated learning. Use the drift signal in a federated setting so that replay-based retention can happen across institutions without centralizing patient data.

Target Audience

Researchers and practitioners in continual learning and medical image analysis, particularly those working on cross-institution domain shift, catastrophic forgetting, or replay-buffer design. It is also relevant to clinical AI engineers and health-system technical leads evaluating how to keep deployed diagnostic models current without losing validated performance, and to readers interested in comparing Vision Transformer versus CNN backbones under domain shift. Readers without a background in continual learning will need to pick up terms such as backward transfer, forward transfer, and replay buffer from the related-work section.

Authors’ abstract

When deep learning models are sequentially trained on new data, they tend to abruptly lose performance on previously learned tasks, a critical failure known as catastrophic forgetting. This challenge severely limits the deployment of AI in medical imaging, where models must continually adapt to data from new hospitals without compromising established diagnostic knowledge. To address this, we introduce a latent drift-guided replay method that identifies and replays samples with high representational instability. Specifically, our method quantifies this instability via latent drift, the change in a sample internal feature representation after naive domain adaptation. To ensure diversity and clinical relevance, we aggregate drift at the patient level, our memory buffer stores the per patient slices exhibiting the greatest multi-layer representation shift. Evaluated on a cross-hospital COVID-19 CT classification task using state-of-the-art CNN and Vision Transformer backbones, our method substantially reduces forgetting compared to naive fine-tuning and random replay. This work highlights latent drift as a practical and interpretable replay signal for advancing robust continual learning in real world medical settings.