PaperlayerView exact revision on arXiv

Source-linked AI summary

Continual Self-supervised Learning: Towards Universal Multi-modal Medical Data Representation Learning

Yiwen Ye, Yutong Xie, Jianpeng Zhang, Ziyang Chen, Qi Wu, Yong Xia

arXiv:2311.17597v2cs.CV

TL;DR

Medical self-supervised pre-training has largely been modality-specific, while joint multi-modal training can encounter modal data collision and is costly to extend. MedCoSS trains modalities sequentially and uses rehearsal-based retention; the authors report generalization across nine downstream datasets and scalability for adding knowledge.

  • Problem

    Medical SSL has largely focused on individual modalities, while joint multi-modal training can suffer modal data collision and requires costly re-pre-training to incorporate new knowledge.

    Source: §1 · p. 2

  • Method

    MedCoSS pre-trains sequentially by modality and uses a rehearsal buffer with feature distillation and intra-modal mixup to retain prior knowledge.

    Source: §1 · p. 2, §3.3 · p. 4, §3.3.2 · p. 5

  • Results

    MedCoSS consistently outperforms single-modal, joint, and other multi-modal pre-training models across nine downstream datasets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present MedCoSS as a cost-effective, scalable approach for integrating new knowledge from stream-like data sources.

    Source: §1 · p. 2, §5 · p. 9

    MedCoSS may not currently achieve state-of-the-art performance across all tasks.

    Source: §5 · p. 8

Abstract

from arXiv · show

Self-supervised learning is an efficient pre-training method for medical image analysis. However, current research is mostly confined to specific-modality data pre-training, consuming considerable time and resources without achieving universality across different modalities. A straightforward solution is combining all modality data for joint self-supervised pre-training, which poses practical challenges. Firstly, our experiments reveal conflicts in representation learning as the number of modalities increases. Secondly, multi-modal data collected in advance cannot cover all real-world scenarios. In this paper, we reconsider versatile self-supervised learning from the perspective of continual learning and propose MedCoSS, a continuous self-supervised learning approach for multi-modal medical data. Unlike joint self-supervised learning, MedCoSS assigns different modality data to different training stages, forming a multi-stage pre-training process. To balance modal conflicts and prevent catastrophic forgetting, we propose a rehearsal-based continual learning method. We introduce the k-means sampling strategy to retain data from previous modalities and rehearse it when learning new modalities. Instead of executing the pretext task on buffer data, a feature distillation strategy and an intra-modal mixup strategy are applied to these data for knowledge retention. We conduct continuous self-supervised pre-training on a large-scale multi-modal unlabeled dataset, including clinical reports, X-rays, CT scans, MRI scans, and pathological images. Experimental results demonstrate MedCoSS's exceptional generalization ability across nine downstream datasets and its significant scalability in integrating new modality data. Code and pre-trained weight are available at https://github.com/yeerwen/MedCoSS.

1. Introduction

MedCoSS targets the limits of modality-specific and joint medical self-supervised pre-training with a sequential, rehearsal-based approach. The authors report stronger generalization across nine downstream datasets and scalability for adding new knowledge.

1. Introduction

  • Earlier medical SSL work largely focused on one modality or a few paired modalities, while unpaired multi-modal approaches had constrained scope.The authors note that collecting large-scale paired multi-modal data is often impractical.Source: §1 · p. 2
  • Single-modal models falter on other modalities, while joint SSL often underperforms because learning across disparate modalities can cause modal data collision.The authors also describe joint training as costly when integrating new knowledge, requiring pre-training to be repeated.Source: §1 · p. 2
  • MedCoSS consistently outperforms single-modal, joint, and other multi-modal pre-training models across nine downstream datasets.The authors describe it as cost-effective to scale with new knowledge.Source: §1 · p. 2
  • MedCoSS assigns each modality to a sequential training stage and uses rehearsal to mitigate forgetting.It retains a subset of previous-modal data, selected with k-means, and applies feature distillation and intra-modal mixup during later stages.Source: §1 · p. 2

2. Related Work

Prior medical SSL approaches include single-modal, paired, and unpaired multi-modal training, but existing methods face limits in cross-modal generalization and scalability. The paper positions rehearsal-based continual learning as a way to address these challenges.

2.1. Medical self-supervised learning

  • The MedCoSS pipeline trains sequentially by modality and carries a rehearsal buffer forward to support learning at later stages.After each stage, k-means sampling selects current-modality samples for the buffer.Source: Figure 2 · p. 3
  • Single-modal SSL often generalizes poorly across modalities, paired data are resource-intensive to collect, and prior unpaired approaches have limited scope or cross-modal performance.The authors connect the limitations of existing unpaired approaches to modal data collision and difficulty incorporating new data sources.Source: §2.1 · p. 3

2.2. Continual learning for knowledge retention

  • Rehearsal-based continual learning retains previous data or features in a buffer for replay, and DER retains samples with their associated logits.MedCoSS modifies DER with k-means sampling, frozen-encoder logit extraction, and intra-modal mixup.Source: §2.2 · pp. 3–4
  • The rehearsal SSL illustration separates current-modality masked modeling from prior-knowledge retention using a frozen previous model and rehearsal buffer.Feature distillation supports knowledge retention on buffer data.Source: Figure 3 · p. 4

3. Approach

MedCoSS combines a dimension-flexible encoder with modality-by-modality self-supervised pre-training and rehearsal-based retention. Its design uses representative buffer samples, feature distillation, and intra-modal mixup before adapting the encoder to downstream tasks.

3.1. Overview

  • The approach separates unsupervised pre-training from supervised fine-tuning and assigns each modality its own pre-training stage to address collision and forgetting.The five modalities are clinical reports, X-rays, CT, MRI, and pathological images.Source: §3.1 · p. 4

3.2. Universal architecture for multi-modal SSL

  • Three dimension-specific tokenizers convert 1D text and 2D/3D images into sequences for a shared ViT/B encoder.Text modeling masks 15% of words; visual modeling masks 75% of tokens and reconstructs the masked regions.Source: §3.2 · p. 4

3.3. Rehearsal-based continual pre-training

  • At each stage, masked modeling learns from current-modality data, while rehearsal samples are augmented and feature-distilled against a frozen previous model.The rehearsal buffer retains data from prior stages; the learnable and frozen networks’ embeddings are encouraged to match using MSE loss.Source: §3.3 · pp. 4–5

3.3.1 Rehearsal buffer construction

  • The buffer uses k-means over embeddings from the preceding model, selecting samples nearest each cluster center; the cluster count is 1% of subset size.Source: §3.3.1 · p. 5

3.3.2 IMM Augmentation Strategy

  • IMM augments rehearsal data by shuffling duplicated samples and mixing them, using binary element-wise masks for text and per-sample continuous mixing for images.The image mix coefficient is sampled for each item in the batch; text uses a binary mask matching the text batch shape.Source: §3.3.2 · p. 5

3.4. Fine-tuning on downstream tasks

  • For downstream classification and segmentation, the pre-trained encoder is paired with an MLP head or a convolutional-based decoder and segmentation head, respectively.Source: §3.4 · p. 5

4. Experiment

MedCoSS is evaluated across nine downstream datasets spanning five medical modalities, and its results are compared with single-modal, joint, and continual pre-training approaches. The experiments also examine its component strategies, buffer size, knowledge retention, and sensitivity to modality order.

4.1. Datasets

  • The nine downstream datasets cover classification and segmentation across Report, X-ray, CT, MRI, and pathological imaging.The upstream unlabeled data also cover those five modalities.Source: §4.1 · pp. 5–6

4.2. Implementation details

  • Pre-training proceeds through Report, X-ray, CT, MRI, and pathological imaging stages, with performance averaged across three random seeds.Classification uses AUC, accuracy, and F1; segmentation uses Dice and 95% Hausdorff distance.Source: §4.2 · p. 6

4.3. Comparing to state-of-the-art SSL

  • MedCoSS outperforms other continual-learning paradigms across eight datasets in all evaluated metrics and achieves more generalizable performance than joint SSL.Joint SSL often underperforms single-modal models on identical modalities and MedCoSS.Source: §4.3 · p. 7, Table 2 · p. 7

4.4. Ablation studies

  • Feature distillation improves average ACC, AUC, F1, Dice, and HD by 3.19%, 1.08%, 3.39%, 1.47%, and 9.04, respectively.Combining k-means sampling with feature distillation improves generalization across datasets, while adding intra-modal mixup gives the best average classification and segmentation results.Source: §4.4 · p. 7, Table 3 · p. 7

4.5. Impact of buffer size on model performance

  • MedCoSS improves across all five datasets when buffer sampling rises from 1% to 5%, but performance declines at 10% on four datasets; 5% is selected.ER degrades when its sampling ratio increases from 1% to 5%.Source: §4.5 · pp. 7–8, Figure 4 · p. 8

4.6. Knowledge retention in sequential per-training

  • After CT pre-training, sequential SSL accuracy falls from 95.0% to 89.7%, while MedCoSS declines from 95.3% to 94.4%.The comparison is against sequential SSL without knowledge-retention strategies.Source: §4.6 · p. 8, Figure 5 · p. 8

4.7. Pre-training order of modality

  • After changing the modality order, MedCoSS retains high generalization performance, and differences between the last modality and other positions are insubstantial.The modality placed last usually achieves better downstream performance.Source: §4.7 · p. 8, Table 4 · p. 8

5. Conclusion

The paper presents MedCoSS as a sequential approach to multimodal self-supervised learning for unpaired medical data. Its experiments show performance across nine downstream datasets, while the conclusion highlights cost-effective integration of knowledge from real-world stream-like sources.

  • MedCoSS assigns each training stage to a specific modality and uses continual-learning techniques to alleviate forgetting during sequential pre-training.The experiments cover Report, X-ray, CT, MRI, and pathological imaging, with evaluation on nine downstream datasets.Source: §5 · p. 8
  • The authors report MedCoSS's superiority across nine downstream datasets, while noting it may not achieve state-of-the-art performance on every task.Source: §5 · p. 8
  • The conclusion identifies cost-effective integration of knowledge from real-world, stream-like data sources as MedCoSS's distinctive advantage.The authors describe it as adaptable to expanding varieties of clinical multimodal data.Source: §5 · p. 9

A.1. Classification tasks

Classification heads are tailored to each dataset, while segmentation uses dimensionality-matched decoders and heads. The segmentation pipeline encodes images as patch tokens, extracts encoder features, and decodes them into pixel-wise predictions.

A.1. Classification tasks

  • Classification uses two-layer MLP heads for PudMed20k and RICORD, and one-layer heads for ChestXR and NCH.The two-layer heads output 5 and 2 neurons, respectively; the one-layer heads output 3 and 9.Source: §A.1 · p. 12

A.2. Segmentation tasks

  • Segmentation uses tokenizers selected by input dimensionality, an encoder, and randomly initialized decoders and segmentation heads matched to 2D or 3D tasks.The 2D and 3D versions follow [17] [34].Source: §A.2 · p. 12
  • A 2D image is divided into flattened, non-overlapping patch tokens, and encoder outputs from layers 4, 7, 10, and 12 provide segmentation features.The patch count is N = (H × W)/P^2.Source: §A.2 · p. 12
  • The decoder upsamples bottleneck and midstage features, fuses them through residual convolution blocks, and produces pixel-wise predictions at the input resolution.The final prediction layer is a 1 × 1 convolution.Source: §A.2 · p. 12
  • For 3D inputs, the architecture replaces 2D tokenization, convolution, and normalization layers with 3D counterparts.Source: §A.2 · p. 12

B.1. Details of upstream datasets

The upstream data span clinical reports, chest radiographs, CT, MRI, and pathology images, with modality-specific preparation described for the imaging datasets.

B.1. Details of upstream datasets

  • MIMIC-CXR contributes clinical reports and chest radiographs; DeepLesion CT and ADNI MRI scans are cropped into sub-volumes, while TCGA pathology images are tiled into patches.MIMIC-CXR lateral views were excluded; CT scans were resampled before cropping, and TCGA sampling was capped at 100 patches per patient.Source: §B.1 · p. 12
  • Figure 6 illustrates the model architecture used for 2D segmentation tasks.Source: Figure 6 · p. 13
  • Table 5 reports implementation details for nine downstream tasks, including task type, modality, loss, patch size, optimizer, learning rate, batch size, and iterations.Source: Table 5 · p. 13

B.2. Details of downstream datasets

MedCoSS is reported to generalize robustly across tasks and data dimensions compared with prior pre-trained models. For 3D segmentation, combining ViT with ResUNet further improves performance over either backbone alone.

B.2. Details of downstream datasets

  • The study evaluates nine downstream datasets spanning text classification, image classification, and segmentation tasks across modalities.The described datasets include PudMed20k, ChestXR, QaTa, RICORD, LiTS, VS, LA, NCH, and GlaS.Source: §B.2 · p. 13

D. Improving 3D Segmentation

  • Figure 7 illustrates the combination of ResUNet and ViT/B, with learnable ResUNet parameters and a frozen ViT/B.Source: Figure 7 · p. 14

C. Comparing to Pre-trained Models

  • MedCoSS accepts 1D, 2D, and 3D data and consistently achieves robust generalization across tasks, while most comparison models are dimension-limited and often decline on other modalities.Path DINO performs best across metrics on NCH but worst across metrics on ChestXR; Table 6 compares five pre-trained models on four datasets.Source: C. Comparing to Pre-trained Models · p. 14, Table 6 · p. 14

D. Improving 3D Segmentation

  • The LiTS results show that combining ViT and ResUNet improves segmentation performance beyond the pre-trained ResUNet or ViT-based models alone.The ResUNet is obtained from DeSD; the ViT-based models use MedCoSS or single-modal CT pre-training.Source: D. Improving 3D Segmentation · p. 14, Table 7 · p. 14
  • Figure 8 presents segmentation outputs from TFS, Joint SSL*, Joint SSL†, EWC, ER, PackNet, CaSSLe, and MedCoSS alongside ground truths.Organs are red, while tumors and malignant regions are green.Source: Figure 8 · p. 15

E. Visualization of Segmentation Results

Across five datasets, MedCoSS segmentation outputs most closely align with ground truths and avoid over- and under-segmentation. A small liver tumor is uniquely identified by MedCoSS in the illustrated CT example.

E. Visualization of Segmentation Results

Related papers

Loading 2311.17597v2…