PaperlayerView exact revision on arXiv

Source-linked AI summary

CoM$^3$eT: A foundation model for medical image analysis through federated, multidimensional context integration

J. Raphael Schäfer, Kai Geissler, Till Nicke, Chiara Tappermann, Karoline Heber, Eike Petersen, Habib Mergan, Lars Ole Schwen, Nick Weiss, Annika Gerken, Jan Hendrik Moltz, Tom Bisson, Isil Dogan O, Tim-Rasmus Kiehl, Norman Zerbe, Sefer Elezkurtaj, Robin S. Mayer, Nadine Flinner, Peter Wild, Isabel Dahm, Felix Peisen, Heinrich von Busch, Robert Grimm, Sebastian Arndt, Lisa Siegler, Matthias Stefan May, Antje Prasse, Natalia Artysh, Fabian Kiessling, Johannes Lotz

arXiv:2608.16268v1cs.CVcs.LG

TL;DR

Medical foundation models have generally been limited to one specialty and either sparse or dense prediction tasks. CoM³eT unifies pathology and radiology across multidimensional inputs and tasks, outperforming competing models across a broad benchmark while supporting resource-efficient and federated adaptation.

  • Problem

    Existing medical foundation models are typically restricted to a single specialty and either sparse or dense prediction tasks, limiting unified analysis across medical imaging settings.

  • Method

    CoM³eT combines shared image representations with attention-based modules to process multidimensional images and support sparse and dense tasks in one model.

  • Results

    CoM³eT scored 0.442 ± 0.022 versus 0.357 ± 0.014 for a theoretical best competing-model ensemble, achieving the best performance on 6 of 12 tasks.

  • Takeaways & Limitations

    Federated partial fine-tuning achieved performance comparable to centralized training despite cybersecurity, connection, and data-access constraints.

  • Takeaways & Limitations

    Federated deployment still requires manual engineering and supervision because cybersecurity, connection timeouts, data-access policies, and limited standardization remain challenges.

Abstract

from arXiv · show

Medical foundation models improve generalization when training AI models with limited labeled data, but remain confined to a single specialty, such as pathology or radiology, and to either sparse or dense outputs, such as classification or segmentation. Here, we present CoM$^3$eT (Co-representation Multidimensional Multitask Medical Transformer), a medical vision foundation model that unifies pathology and radiology, sparse and dense predictions, and two- and higher-dimensional inputs by modeling multidimensional context with attention. CoM$^3$eT outperformed other medical foundation models in an open competition spanning five tomographic, four whole-specimen, and three two-dimensional datasets, covering sparse and dense prediction tasks as well as report generation. When adapted across diverse clinical applications, training fewer than 2.5% of parameters achieved performance comparable to full fine-tuning, enabling research without access to high-performance GPU clusters. Applied to federated learning across hospitals, this approach achieved performance comparable to pooled-data training over internet connections and with consumer-grade hardware.

Main

CoM³eT is a universal medical vision foundation model for multidimensional images that unifies sparse and dense tasks through shared semantic and spatial co-representations. It was built with multitask pretraining, attention-based components, dimensionality-agnostic positional encoding, and partial fine-tuning for federated learning.

  • Motivation: Existing medical foundation models are commonly restricted to one specialty or modality and cannot jointly support sparse and dense tasks.Examples include pathology- or radiology-focused models and models limited to classification or segmentation.
  • Motivation: Medical vision foundation models must accommodate high-resolution 2D, gigapixel, and volumetric images, but large imaging datasets remain unavailable for many diseases.This data scarcity limits model generalization and motivates pretraining on diverse datasets.
  • Contribution: CoM³eT processes multidimensional images for both classification and segmentation using shared patch and pixel tokens.These co-representations encode semantic and spatial information in one model.
  • Architecture and training: Multitask pretraining scaled to over 60 tasks, while attention-based components shared most parameters across sparse and dense predictions.The model was designed to handle classification and segmentation jointly across multidimensional images.
  • Federated learning: Partial fine-tuning reduced training resources and communication load, enabling federated fine-tuning demonstrated with German university hospitals.The federated setup was conducted in a real-world hospital setting.

Results

CoM³eT was the only model applicable across all benchmarked vision and vision-language tasks and achieved the strongest overall performance. Its multidimensional context mechanisms, partial fine-tuning, and federated training preserved competitive performance across diverse clinical applications.

  • Benchmark performance: CoM³eT applied to all 12 vision and vision-language tasks, achieved the best performance on 6, and exceeded a strongest-per-task ensemble (0.442 ± 0.022 versus 0.357 ± 0.014).The comparison used an independent benchmark spanning pathology and radiology, classification, segmentation, and report generation.
  • Domain performance: CoM³eT achieved the best average performance in both pathology (0.482 ± 0.018) and radiology, leading 4 of 6 pathology tasks.Pathology wins included recurrence prediction, tumor proportion scoring, mitotic figure detection, and tissue type segmentation.
  • Multidimensional context: Adding the pyramid transformer improved uterus segmentation from 75.49% ± 0.5% to 79.4% ± 0.3% 3D Dice, versus 71.0% for 3D U-Net.On vessel segmentation, it also increased CoM³eT from 80.45% to 82.11% 3D Dice, above nnU-Net’s 80.38%.
  • Multidimensional context: For tumor segmentation in 3,936 MRI patients, the pyramid transformer raised CoM³eT from 25.28% to 58.62% 3D Dice, versus 33.13% for nnU-Net.This result tested the unified architecture outside the data-scarce regime.
  • Efficient adaptation: Partial fine-tuning was equivalent to full fine-tuning across tasks (plower < 0.001, pupper = 0.008), with mean performance of 0.839 versus 0.833.It reduced patch-token processing time by more than 99.9%, from 2.3 seconds to 1.5 ms.
  • Federated learning: Federated partial fine-tuning achieved equivalent performance to pooled-data CoM³eT, with C-index 0.743 ± 0.012 versus 0.754 ± 0.012 and AUC 0.820 ± 0.125 versus 0.795 ± 0.170.The equivalence test ruled out a difference of 0.05 or more in either direction (plower = 0.019, pupper = 0.032).

Discussion

CoM³eT addresses fragmented medical vision modeling through unified multidimensional representations and single-stage multimodal pretraining. Its federated partial fine-tuning matched centralized training, although deployment still faces infrastructure and policy challenges.

  • CoM³eT unifies medical image and task types that traditionally required separate model families, spanning pathology, radiology, sparse outputs, dense outputs, and multidimensional inputs.The discussion contrasts this scope with models limited to segmentation, 2D histology, 3D CT, or histology images.
  • A single final pretraining stage combines structured labels and text supervision across available data instead of sequentially adding objectives onto a self-supervised stage.Sequential training may require early stages to preserve all information needed later and can cause catastrophic forgetting.
  • Segmentation objectives train hyperpixel tokens whose pyramid-transformer representations preserve spatial context beyond 2D and 3D images, enabling application to 4D multiparametric MRI.The discussion identifies CoM³eT as the only multipurpose model applicable to multiparametric MRI in the UNICORN competition.
  • Performance comparable to centralized training was achieved with real-world federated partial fine-tuning, despite cybersecurity, connection-timeout, and data-access challenges.Federated learning still requires manual engineering without standardization and trusted training infrastructure.
  • Future applications include conversational assistants and agentic systems, using aligned visual-language representations and reusable auxiliary image-analysis heads.These directions aim to supply specialized visual representations that general models lack.

Methods · Model architecture

CoM³eT combines shared multitask modules with an image transformer for multidimensional patch context and a pyramid transformer for dense multidimensional segmentation. Its design supports sparse and dense tasks across variable-sized 2D inputs and volumetric medical images.

  • Model architecture: The model separates shared modules from task-specific heads across supervision levels, with each task defining an objective, dataset, target, and linear head.Pretraining alternated among tasks so shared modules learned generally applicable features.
  • Model architecture: CoM³eT adds an image transformer for multidimensional patch-token context and a pyramid transformer for dense multidimensional segmentation.Together, these modules model relationships across image patches and propagate contextual information.
  • Model architecture: The vision encoder accepts variable-sized 2D RGB inputs and produces hierarchical feature maps, using Swin Transformer V2 as the selected backbone.Global average pooling and linear projection compact each patch-token representation, particularly for large medical images.
  • Model architecture: The image transformer generates representations for all patch tokens, enabling simultaneous processing of dense tasks such as 3D segmentation and sparse tasks such as classification.It generalizes aggregation methods to multidimensional medical images while modeling interactions between patch tokens.
  • Model architecture: For volumetric images, CoM³eT uses relative symmetric ALiBi positional encoding to represent pairwise slice relationships with variable sequence lengths and short- and long-range dependencies.Positional encoding is required because transformers are permutation-invariant when spatial position is not directly visible.
  • Model architecture: Synthetic MNIST and geometric-shape segmentation tasks tested local context, global context, and patch-token prompting before evaluation on representative medical tasks.The synthetic segmentation dataset contained 10,000 training groups and 1,000 test groups, each containing 3 to 8 images.
  • Model architecture: Real-data evaluations compared transformer-equipped CoM³eT with 2D variants lacking the relevant transformer and with specialized reference methods across sparse and dense tasks.Sparse tasks included whole-slide pathology and MRI tumor classification; dense tasks included MRI uterus and tumor segmentation and CT vessel segmentation.
  • Model architecture: The pyramid transformer contextualizes feature maps by computing patch-token relationships from image-transformer outputs and returning the result through normalization, GELU, layer scaling, and a residual connection.Its operations can be applied wherever feature maps are computed and are intended for shared modules rather than task-specific modules.

Benchmarking CoM³eT for medical imaging: UNICORN

The UNICORN benchmark evaluated frozen medical foundation models across pathology and radiology, spanning sparse and dense, 2D and multidimensional, vision and vision-language tasks. CoM³eT-UNICORN used a memory-constrained hyperpixel configuration and task-specific feature extraction, with performance compared against task-specific and domain-level competitors.

  • Benchmark scope: UNICORN evaluated frozen models across pathology and radiology on sparse and dense, 2D and multidimensional, vision and vision-language medical-imaging tasks.The task set comprised classification, regression, detection, segmentation, and vision-language tasks.
  • Submission configuration: CoM³eT-UNICORN matched CoM³eT-Base except that hyperpixel-token dimensions were reduced to 32 for evaluation-platform RAM constraints.Patch-token dimensions remained 512, and these settings were used across all tasks and domains.
  • Task processing: For 2D tasks, the vision backbone processed the whole image as one patch token; sparse 3D and WSI tasks used up to 10,000 patches and selected the most attended patch for image-level features.This selection produced a single feature vector when the method otherwise generated one feature vector per image patch.
  • Evaluation and comparison: CoM³eT performance was compared with the strongest competing model for each task and with second- and third-strongest models in pathology and radiology aggregates.Confidence intervals in Fig. 3A were averaged across individual tasks using the same procedure as the mean task specification.

Pretraining database

CoM³eT’s pretraining database combined large-scale natural-image data with diverse pathology and radiology datasets spanning classification, detection, segmentation, and vision-language alignment. It used ImageNet21k and COCO alongside multidimensional medical datasets covering tissue, organs, lesions, and other clinical findings.

  • Database composition: The database integrated natural and medical images across classification, object detection, segmentation, and vision-language alignment tasks.Medical data supplied both general pretraining data and domain-specific imaging knowledge, while natural images supported general visual concepts.
  • Natural-image data: ImageNet21k contributed more than 11 million labeled images across roughly 21,000 classes, with all images and labels retained for pretraining.COCO added 118,287 images with 591,753 captions and more than 200,000 annotated images spanning 12 supercategories and 80 object categories.
  • Digital pathology: Digital pathology pretraining covered tissue characterization, tumor detection and grading, mitosis detection, and dense tissue segmentation across colorectal, breast, and prostate pathology.The database included 100,000 colorectal H&E tiles classified into nine tissue classes and 682 gastric image regions from 155 patients for patch-based segmentation.
  • Radiology: Multidimensional radiology pretraining included TotalSegmentator’s five 3D CT segmentation task groups, the Medical Segmentation Decathlon’s 10 3D tasks, and additional anatomical segmentations for prostate, abdominal organs, and spinal structures.TotalSegmentator covered organs, vertebrae, cardiac structures, muscles, and ribs.
  • Additional medical data: Lesion and broader clinical imaging data extended pretraining to CT and MRI lesion detection or segmentation, chest-radiography findings and anatomy, bone-marrow cytology classification, and gastrointestinal polyp segmentation.The lesion data included rectangle annotations, regions of interest, patient-level MRI masks, and private full-body CT lesion segmentation.

Training method

CoM³eT’s training method processes heterogeneous tasks sequentially while discarding task-specific gradients to keep memory requirements constant. Mixed-precision checkpointing and distributed optimizer-state partitioning further support large multidimensional inputs and multi-node training.

  • Task-wise training: Sequential task updates kept memory requirements constant by discarding task-specific gradients after each task, enabling larger batch sizes.Each step processed one task at a time and updated its task head before moving to the next task.
  • Multitask objectives: CoM³eT supported classification, segmentation, and vision-language pretraining with task-specific linear or convolutional heads and SONAR-based caption objectives.Captions were augmented with the Qwen-3 chat model, while SONAR computed sentence embeddings and generated text from embeddings.
  • Memory-efficient training: bfloat16 replaced fp16 for Swin-based training because fp16 produced NaN losses, while mixed precision and gradient checkpointing enabled 2× larger batch sizes.CoM³eT-Large and Base used bfloat16, whereas CoM³eT-Tiny used fp32.
  • Distributed training: ZeRO partitioned optimizer state across processes, while nodes accumulated local task gradients and averaged them with PyTorch DistributedDataParallel using NCCL.This reduced per-device memory relative to standard Adam, whose optimizer states require roughly twice the model size.
  • Memory-efficient training: Gradient checkpointing around the full backbone or individual encoder segments reduced training memory, with segment-level checkpointing more memory efficient but harder to implement.The full-backbone option was simpler to add across architectures.

Partial fine-tuning of CoM³eT

CoM³eT was evaluated with partial fine-tuning by freezing the vision backbone and training the image transformer, pyramid transformer, and decoder modules. Full and partial fine-tuning were compared across representative multidimensional, patch-based, classification, and segmentation tasks under matched training settings.

  • Method: Partial fine-tuning froze the vision backbone while training the image transformer, pyramid transformer, and decoder modules.For patch-level tasks, the image transformer was equivalent to a pretrained MLP because group size was one, classification tasks omitted the decoder, and patch tokens were cached once for reuse.
  • Experimental design: The comparison covered 3D MRI segmentation, gigapixel whole-slide analysis, 2D X-ray classification, and ultrasound and microscopy segmentation.MRI, ultrasound, and microscopy segmentation tasks tested generalization without similar pretraining data.
  • Experimental design: Both strategies used matched training settings with five-fold cross-validation, patient-level sampling, cached patch tokens, and subsampling of 50–100 tokens per patient when necessary.Training used 15 loops of 100 update steps, with eight patients sampled for multidimensional tasks and 64 for patch-based tasks.

Federated fine-tuning of CoM³eT

CoM³eT-FL was developed by federated fine-tuning from a pretrained CoM³eT-Base checkpoint while keeping patient data local. The experiment involved three sites using consumer-grade hardware and constrained hospital networks.

  • Federated fine-tuning: CoM³eT-FL was obtained by federated fine-tuning CoM³eT-Base without biochemical recurrence-specific pretraining data.The image transformer was fine-tuned while the vision backbone remained frozen.
  • Federated fine-tuning: Each site retained patient data locally and shared only updates for trainable parameters through a standard server-based optimization scheme.Federated learning was implemented with Flower for a single downstream application.
  • Real-world deployment: The real-world deployment involved two German university hospitals and a research institute using consumer-grade GPUs and weak internet connections.Communication used gRPC over HTTP/2, with a containerized environment supporting isolated patch-token caching and federated fine-tuning alongside separate read and write storage resources.

Statistics

The study compared full and partial fine-tuning using paired equivalence testing with a predefined 0.05 margin on each task’s main performance metric.

  • Equivalence testing: Equivalence was assessed with two one-sided tests using paired Wilcoxon signed-rank tests and a predefined 0.05 margin for each task’s main metric.Examples of main metrics included AUC for classification and C-index for time-to-event prediction.

Test data

The test suite spans whole-slide pathology, microscopy, radiology, and 2D/3D imaging, covering time-to-event prediction, classification, and segmentation across diverse clinical datasets. It includes tasks designed to test representation quality, multimodal generalization, local 3D continuity, and weak or out-of-domain generalization.

  • Whole-slide pathology: The benchmark includes whole-slide pathology tasks for biochemical-recurrence prediction after prostatectomy and one-year recurrence classification, using multicenter PRAD data and CHIMERA challenge data.Cancer Recurrence uses PRAD data split across five centers, while Surgery Outcome contains 94 patients and 190 images.
  • Segmentation tasks: The suite evaluates segmentation across volumetric MRI, CT angiography, ultrasound, and brightfield microscopy, including uterine structures, aortic vessels, tumors, and precision-cut lung-slice tissue.These tasks cover 3D anatomical continuity, heterogeneous MRI protocols, 2D lesion segmentation, and tissue architecture preserved in ex vivo lung slices.
  • Classification tasks: The test data include 2D and 3D classification tasks for pediatric pneumonia on chest X-rays and benign-versus-malignant breast lesions on dynamic contrast-enhanced MRI.Pneumonia tests weak generalization because pediatric images were absent from pretraining, whereas breast MRI classification spans more than ten institutions and 1,330 studies from 1,299 patients.
  • Generalization scope: The datasets deliberately vary acquisition settings, anatomy, pathology, dimensionality, and label structure to assess generalization beyond the pretraining distribution.Examples include multiple MRI field strengths and vendors, skew toward smaller ultrasound tumors, and pediatric X-rays despite no pediatric pretraining images.

Competing interests

The applicant has a pending patent related to the article’s training algorithm and neural architecture components.

  • Fraunhofer-Gesellschaft zur Förderung der angewandten Forschung eingetragener Verein has filed patent application EP23209015.9, naming J.R.S., T.N., J.L., and F.K. as inventors.

Supplementary information

Supplementary experiments evaluated an early, partially fine-tuned CoM³eT model on external biochemical-recurrence cohorts from the LEOPARD study. CoM³eT achieved the strongest reported single-cohort and weighted multi-cohort results under this federated external evaluation.

  • LEOPARD external evaluation: The early CoM³eT model was pretrained on 42 tasks and partially fine-tuned multitask-wise using LEOPARD and PRAD data with separate task heads.Its pretraining included prostate-cancer tasks such as ISUP grading and whole-slide-image tumor detection, but not biochemical-recurrence-specific data.
  • LEOPARD external evaluation: The external evaluation covered four cohorts of varying size, each contributed by a single hospital, and compared other foundation models with specialized algorithms.The study therefore assessed transfer beyond the training setting across multiple hospital-specific cohorts.
  • LEOPARD external evaluation: 0.771 C-index (95% CI: 0.727 to 0.813) was CoM³eT’s best result on a single private cohort, versus 0.714 (95% CI: 0.669 to 0.761) for CTransPath.CTransPath was the second-best result in this comparison.
  • LEOPARD external evaluation: 0.720 weighted C-index across four cohorts exceeded the next-best results of 0.704 and 0.699.The evaluation included 100 test patients from the training cohort and 723, 427, and 332 from external cohorts.
Loading 2608.16268v1…