System and method for generating feature representations of three-dimensional medical images using self-supervised learning models

Abstract

There is provided a system and method for generating weighting for a machine learning model used for determining feature representations from three-dimensional medical imaging data. The method including: generating augmented versions of the three-dimensional medical imaging data by performing image augmentations; masking patch regions of a portion of the augmented versions; performing pretraining of the machine learning model using a knowledge distillation model including a student network and a teacher network, the augmented versions without masked patch regions are passed through the teacher network to obtain a teacher representation, the augmented versions with masked patch regions are passed through the student network to predict the masked regions in the teacher representation, the student network being trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network.

Description

SYSTEM AND METHOD FOR GENERATING FEATURE REPRESENTATIONS OF THREE-DIMENSIONAL MEDICAL IMAGES USING SELF-SUPERVISED LEARNING MODELSTECHNICAL FIELD

[0001] The following relates, generally, to automated medical imaging interpretation; and more particularly, to a system and method for generating feature representations of three-dimensional medical images using self-supervised learning models.BACKGROUND

[0002] Deep learning (DL) can be used to enhance existing workflows for a variety of clinical tasks involving medical images; including detection, diagnosis, and risk profiling. However, the data-hungry nature of these approaches poses practical challenges for training and generalizability. Creating detailed labels fortraining DL models for three-dimensional (3D) medical imaging modalities is particularly time-consuming and computationally expensive. To address this substantial challenge, self-supervised learning (SSL) approaches have been proposed to reduce reliance on detailed ground truth annotations by leveraging unlabeled datasets. However, most approaches for 3D medical imaging modalities train SSL models on unlabeled datasets that are similar to their downstream applications. Employing SSL-pretrained models on downstream datasets from similar modalities, organs, image characteristics, and / or distributions substantially limit their generalizability and scalability. For example, such single-distribution approaches generally result in much greater training overhead, as separate models would need to be pretrained for each downstream task for optimal performance. This can be further exacerbated by several factors including the rarity of the disease the model is assessing, the ability to acquire high-resolution, multidimensional data for training, and / or the scarcity / cost of certain imaging modalities.SUMMARY

[0003] In an aspect of the present invention, there is provided a method for generating weighting for a machine learning model used for determining feature representations from three-dimensional medical imaging data, the method comprising: receiving three-dimensional medical imaging data; generating augmented versions of the three-dimensional medical imaging data by performing image augmentations on the three-dimensional medical imaging data; masking patchregions of a portion of the augmented versions; performing pretraining of the machine learning model using a knowledge distillation model comprising a student network and a teacher network, the augmented versions without masked patch regions are passed through the teacher network to obtain a teacher representation, the augmented versions with masked patch regions are passed through the student network to predict the masked regions in the teacher representation, the student network being trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network; and outputting weights of the teacher network as a pretrained machine learning model for determining feature representations of previously unseen three-dimensional medical imaging data.

[0004] In a particular case of the method, performing image augmentations on the three-dimensional medical imaging data comprises performing random image augmentations.

[0005] In another case of the method, the augmented versions comprise global crops capturing a whole of the original volume of the imaging data.

[0006] In yet another case of the method, pretraining of the machine learning model comprises jointly optimizes an image-level representation objective on the patched global crops and a patchlevel representation objective on local crops of the augmented versions.

[0007] In yet another case of the method, masking patch regions is performed using a binary mask.

[0008] In yet another case of the method, the teacher network is an exponential moving average of the student network during pretraining.

[0009] In yet another case of the method, the teacher network comprises a vision transformer encoder model.

[0010] In yet another case of the method, generating augmented versions comprises applying three-dimensional data augmentations modeling variability in medical image acquisition.

[0011] In yet another case of the method, the image augmentations comprise, at least, applying a two-dimensional crop in in-plane dimensions and independently sampling a depth dimension.

[0012] In yet another case of the method, the three-dimensional medical imaging data comprises multiple co-registered imaging channels.

[0013] In another aspect of the present invention, there is provided a system for generating weighting for a machine learning model used for determining feature representations from three-dimensional medical imaging data, the system comprising one or more processors and a data storage, the data storage comprising instructions for the one or more processors to execute: an input module to receive three-dimensional medical imaging data; a pretraining module to generate augmented versions of the three-dimensional medical imaging data by performing image augmentations on the three-dimensional medical imaging data, and to mask patch regions of a portion of the augmented versions; a machine learning module to perform pretraining of the machine learning model using a knowledge distillation model comprising a student network and a teacher network, the augmented versions without masked patch regions are passed through the teacher network to obtain a teacher representation, the augmented versions with masked patch regions are passed through the student network to predict the masked regions in the teacher representation, the student network being trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network; and an output module to output weights of the teacher network as a pretrained machine learning model for determining feature representations of previously unseen three-dimensional medical imaging data.

[0014] In a particular case of the system, performing image augmentations on the three-dimensional medical imaging data comprises performing random image augmentations.

[0015] In another case of the system, the augmented versions comprise global crops capturing a whole of the original volume of the imaging data.

[0016] In yet another case of the system, pretraining of the machine learning model comprises jointly optimizes an image-level representation objective on the patched global crops and a patchlevel representation objective on local crops of the augmented versions.

[0017] In yet another case of the system, masking patch regions is performed using a binary mask.

[0018] In yet another case of the system, the teacher network is an exponential moving average of the student network during pretraining.

[0019] In yet another case of the system, the teacher network comprises a vision transformer encoder model.

[0020] In yet another case of the system, generating augmented versions comprises applying three-dimensional data augmentations modeling variability in medical image acquisition.

[0021] In yet another case of the system, the image augmentations comprise, at least, applying a two-dimensional crop in in-plane dimensions and independently sampling a depth dimension.

[0022] In yet another case of the system, the three-dimensional medical imaging data comprises multiple co-registered imaging channels.

[0023] These and other aspects are contemplated and described herein. It will be appreciated that the foregoing summary sets out representative aspects of the system and method to assist skilled readers in understanding the following detailed description.DESCRIPTION OF THE DRAWINGS

[0024] A greater understanding of the embodiments will be had with reference to the Figures, in which:

[0025] FIG. 1 shows a conceptual block diagram of a system for generating feature representations of three-dimensional medical images using self-supervised learning models, according to an embodiment;

[0026] FIG. 2 is a flowchart of a method for generating feature representations of three-dimensional medical images using self-supervised learning models, according to an embodiment;

[0027] FIG. 3 illustrates BraTS segmentation Dice scores and 95th percentile Hausdorff Distance (HD95) in accordance with example experiments;

[0028] FIG. 4 illustrates an example breakdown of large multimodal, multi-organ pretraining dataset of 100,0003D scans with over 10 organs from 35 publicly available and internal studies, in accordance with the example experiments;

[0029] FIG. 5 illustrates examples of an original image, principal component analysis (PCA) on patch-level representations, and multi-head self-attention (MHSA) attention map visualized for three image planes, in accordance with the example experiments;

[0030] FIGS. 6A to 6G illustrate transferring of pretrained models into segmentation architectures, in accordance with the example experiments;

[0031] FIGS. 7A and 7B illustrate two-dimensional and three-dimensional multi-scale deformable attention (MSDA) offset initializations, respectively, in accordance with the example embodiments;

[0032] FIG 8 shows efficient multi-channel interfacing with a ViT-Adapter module, in accordance with the example embodiments;

[0033] FIG. 9 shows BraTS segmentation Dice scores and 95th percentile Hausdorff Distance (HD95) in an evaluation and comparison to other pretrained models, in accordance with the example experiments;

[0034] FIG. 10 shows BTCV segmentation Dice scores and HD95 in the evaluation and comparison to other pretrained models, in accordance with the example experiments;

[0035] FIG. 11 shows ICBM classification AUG and F1 scores in the evaluation and comparison to other pretrained models, in accordance with the example experiments;

[0036] FIG. 12 shows COVID-CT-MD classification AUG and F1 scores in the evaluation and comparison to other pretrained models, in accordance with the example experiments;

[0037] FIG. 13 shows plots comparing per-class Dice / AUC scores for segmentation / classification experiments using the third-largest training dataset size, in accordance with the example experiments;

[0038] FIG. 14 shows normalized and averaged classification confusion matrices using the third-largest training dataset size, in accordance with the example experiments;

[0039] FIG. 15 shows dice scores for linear decoder segmentation experiments, in accordance with the example experiments;

[0040] FIG. 16 shows a visualization of segmentation predictions on BraTS and BTCV datasets, in accordance with the example experiments;

[0041] FIGS. 17A and 17B, FIGS. 18A and 18B, and FIGS. 19A and 19B show class-separated visualizations of segmentation network predictions on BraTS, in accordance with the example experiments;

[0042] FIGS 20A and 20B, FIGS 21 A and 21 B, and FIGS 22A and 22B illustrate class-separated visualizations of segmentation network predictions on the BTCV dataset, in accordance with the example experiments;

[0043] FIGS. 23 to 25 illustrate visualization of classification performance on ICBM age prediction (FIGS. 23 to 25), in accordance with the example experiments;

[0044] FIGS. 26 to 28 illustrate visualization of classification performance on COVID-CT-MD disease prediction, in accordance with the example experiments;

[0045] FIG. 29 illustrates model evaluation on out-of-distribution tasks showing Dice and HD95 scores for left atrium segmentation (LA-SEG) task, in accordance with the example experiments;

[0046] FIG. 30 illustrates LA-SEG visualization when finetuning using 25% of the full labeled dataset, in accordance with the example experiments;

[0047] FIG. 31 illustrates LA-SEG visualization when finetuning using 100% of the full labeled dataset, in accordance with the example experiments;

[0048] FIG. 32 illustrates unsupervised visualizations on random volumes sampled from the LA-SEG, in accordance with the example experiments;

[0049] FIG. 33 illustrates Dice and HD95 scores for the 3D breast ultrasound tumor (TDSC-ABUS) segmentation task, in accordance with the example experiments;

[0050] FIG. 34 illustrates TDSC-ABUS visualization when finetuning using 25% of the full labeled dataset, in accordance with the example experiments;

[0051] FIG. 35 illustrates TDSC-ABUS visualization when finetuning using 100% of the full labeled dataset, in accordance with the example experiments; and

[0052] FIG. 36 illustrates unsupervised visualizations on random volumes sampled from the TDSC-ABUS dataset, in accordance with the example experiments.DETAILED DESCRIPTION

[0053] For simplicity and clarity of illustration, where considered appropriate, reference numerals may be repeated among the Figures to indicate corresponding or analogous elements. In addition,numerous specific details are set forth in order to provide a thorough understanding of the embodiments described herein. However, it will be understood by those of ordinary skill in the art that the embodiments described herein may be practised without these specific details. In other instances, well-known methods, procedures and components have not been described in detail so as not to obscure the embodiments described herein. Also, the description is not to be considered as limiting the scope of the embodiments described herein.

[0054] Various terms used throughout the present description may be read and understood as follows, unless the context indicates otherwise: “or” as used throughout is inclusive, as though written “and / or”; singular articles and pronouns as used throughout include their plural forms, and vice versa; similarly, gendered pronouns include their counterpart pronouns so that pronouns should not be understood as limiting anything described herein to use, implementation, performance, etc. by a single gender. Further definitions for terms may be set out herein; these may apply to prior and subsequent instances of those terms, as will be understood from a reading of the present description.

[0055] Any module, unit, component, server, computer, terminal, engine, or device exemplified herein that executes instructions described herein, may include access computer readable media, such as storage media, computer storage media, or data storage devices (removable and / or nonremovable) such as, for example, magnetic disks, optical disks, or tape. Computer storage media may include volatile and non-volatile, removable and non-removable media implemented in any method or technology for storing information, such as computer readable instructions, data structures, program modules, or other data. Examples include RAM, ROM, EEPROM, flash memory or other non-transitory memory technology, CD-ROM, digital versatile disks (DVD) or other optical storage, magnetic cassettes, magnetic tape, magnetic disk storage or other magnetic storage devices, or any other non-transitory medium that stores the desired data and / or information and can be accessed for execution of an application, module, or both. Any such non-transitory media may be part of the device or accessible thereto or may otherwise provide functionality to the system to enable efficient data storage and retrieval. Processors or controllers set out herein may be implemented as a single processor or a plurality of processors, which may be arrayed or distributed, such as in a cloud-based model, in order to carry out processing functions individually or collectively. Methods, applications, or modules described herein may be implemented using computer readable / executable instructions that may be stored or otherwise held by such non-transitory media and executed by one or more processors to be used in the realization of the embodiments describe herein.

[0056] The present embodiments advantageously provide an approach that can use pretrained weights of general-use to enable deep-learning (DL) in medical imaging applications by greatly boosting the accuracy of DL models in label-scarce applications.

[0057] Some SSL frameworks have been able to scale to larger and highly diverse pretraining datasets, with models able to create image representations that are generalizable to many downstream tasks. While these pipelines can achieve good results for two-dimensional (2D) benchmarks, scaling them to 3D data is computationally prohibitive, requiring large datasets, batch sizes (for example, often ranging from 512-4096), and long train times to learn effectively. One particular way to resolve these constraints can be to cast a 3D SSL task as a 2D task by viewing 3D images slice-by-slice. However, it is generally widely known that keeping a full 3D anatomical context is important when applying DL to medical images.

[0058] To address the need for a scalable 3D pretraining, the present embodiments advantageously provide a memory-efficient framework for 3D medical imaging and provides a general-purpose Vision Transformer (ViT) model. In an example experiment of the present embodiments, this model can be pretrained on an exceptionally large, multimodal, and multiorgan dataset of nearly 100,000 unlabeled 3D medical volumes curated from 35 publicly available and internal data studies, as illustrated in FIG. 3. The example experiments specifically acquired datasets consisting of MRI (N = 70,434) and CT (N = 27,815) volumes, with a small brain PET (N = 566) dataset, as illustrated in FIG. 4. The ViT model backbone was designed to enhance its performance on downstream segmentation tasks by converting an adapter module to 3D inputs (referred to as a “3D ViT-Adapter”). The example experiments show that the present embodiments can be used to create generalizable representations when applied to downstream segmentation and classification tasks in multiple organs and modalities, in most instances significantly improving on other approaches in the art. Thus, while large SSL-pretrained models have been proposed for a single medical imaging modality, the present embodiments advantageously provide a 3D SSL-based medical imaging model that can extract salient features and generalize across multiple modalities simultaneously.

[0059] FIGS. 3 to 5 illustrate an example overview of the present embodiments and large pretraining dataset. FIG. 3 shows that the present embodiments can combine an image-level objective and a patch-level objective. Original volumes can be randomly augmented twice to create global crops, and augmented eight times to yield local crops. The image-level objective is taken by distilling the [CLS] token representations between the student and exponential movingaverage (EMA) teacher networks. The patch-level objective can be determined between patch representations at masked regions in the student network input and corresponding unmasked EMA teacher representations. LCE indicates Cross-Entropy loss, with the final loss consisting of the summed image-level distillation and patch-level reconstruction objectives. FIG. 4 illustrates an example breakdown of large multimodal, multi-organ pretraining dataset of 100,0003D scans with over 10 organs from 35 publicly available and internal studies (the number of volumes per modality per anatomical location / organ; MRI = 70,434 volumes, CT = 27,815, and PET = 566). FIG. 5 illustrates examples of an original image, principal component analysis (PCA) on patchlevel representations, and multi-head self-attention (MHSA) attention map visualized for three image planes. Each row in order: BraTS T1-weighted, T2-weighted, and two patients from Beyond the Cranial Vault (BTCV). PCA visualizations were obtained per image from patch-level representation vectors. The first (in terms of explained variance) PCA component is used to mask image background (white) with a simple threshold, with the next three normalized and mapped to RGB channels. MHSA attention maps are obtained from the [CLS] token of the final layer.

[0060] The present embodiments can be used to greatly reduce the amount of labeled data required to train DL models on medical images. Hence, the present embodiments could be used for, for example:• Initializing models trained on medical images for classification tasks (e.g., diagnosis, subtyping, detection, etc.);• Initializing models trained on medical images for segmentation tasks (e.g., organ segmentation, tumor delineation, etc.);• Using the pretrained DL model as a feature extractor for unsupervised downstream tasks (content-based retrieval, obtaining metrics of image similarity, clustering, biomarker discovery, etc.); and• Using the DL model in conjunction with large amounts of unlabeled data to create 3D foundation models.

[0061] Turning to FIG. 1, shown is a conceptual diagram of a system 100 for generating feature representations of three-dimensional medical images using self-supervised learning models, according to an embodiment. The system 100 can be colloquially referred to as ‘DINOv2’. As shown, the system 100 has a number of physical and logical components, including a processing unit (“PU”) 160, data storage 164, an interface module 168, a network module 176, and a local bus 184 enabling PU 160 to communicate with the other components. PU 160 can include one ormore central processing units, one or more graphical processing unit, microprocessors, dedicated hardware, or other integrated processing circuits. The data storage 164 provides storage of data to the Pll 160 and can include volatile and non-volatile storage for storing data required by the Pll 160, including computer-executable instructions for implementing the methods described herein, as well as any derivative or related data. The interface module 168 enables input to be provided; for example, directly via a user input device, or indirectly, for example, via an external device. The interface module 168 also enables output to be provided; for example, directly via a user display, or indirectly, for example, sent over the network module 176. The network module 176 permits communication with other systems or computing devices; for example, over a local area network or over the Internet.

[0062] In other embodiments, any operating system, programs, or instructions can be executed in hardware, specialized microprocessors, logic arrays, or the like. While FIG. 1 illustrates a system implemented on a single computing device, it is understood that the processing, or any of the functions undertaken by the system 100, can be distributed over multiple computing devices; for example, in a cloud or distributed computing environment.

[0063] In an embodiment, the Pll 160 can be configured to execute a number of conceptual modules 101; for example, an input module 102, a pretraining module 104, a machine learning module 106, and an output module 108. In further cases, functions of the above modules can be combined or executed on other modules. In some cases, functions of the above modules can be executed on remote computing devices, such as centralized servers and cloud computing resources communicating over the network module 176.

[0064] FIG. 2 illustrates a flowchart of a method 200 for generating feature representations of three-dimensional (3D) medical images using self-supervised learning models, according to an embodiment. At block 202, the input module 102 receives original volumes of three-dimensional (3D) medical imaging data from the data storage 164, the interface module 168, and / or the network module 176. Any suitable 3D medical imaging data can be received, for example, 3D X-ray images, magnetic resonance imaging (MRI) images (for example, of the brain, knee, heart, or prostate), computerized tomography (CT) images (for example, of the abdomen, chest, or brain), or positron emission tomography (PET) images. In most cases, the 3D medical imaging data is unlabelled or substantially unlabelled. In some cases, image normalization can be conducted on the medical imaging data; for example, by linearly mapping the 0.05thand 99.95thpercentiles ofintensity to -1 and 1, respectively. In an example, the 3D medical imaging data can include an unlabelled dataset comprising 100,000 medical images or scans.

[0065] At block 204, the pretraining module 104 generates a set of crops by performing image augmentations on the 3D medical imaging data; referred to as augmented versions of the three-dimensional medical imaging data. The image augmentations can be random augmentations and can include, for example, flipping, and blurring. Where the 3D medical imaging data is in color, the random image augmentations can include, for example, image converting to grayscale, solarization, and color jitter. In some cases, medical imaging-related augmentations can be used including, for example, random contrast adjustment, additive noise, gibbs noise, and histogram shift. In some cases, the augmentation of the crops can include cropping two in-plane image dimensions using a two-dimensional random resize cropping technique, and then independently sampling the cross-slice dimension crop size.

[0066] The random image augmentations are particularly advantageous for 3D medical images because such augmentations can be used to mimic potential distribution shifts, noise, artifacts, or other image characteristics that can be found in medical images. Using the random image augmentations will be allow the method 200 to be more generalizable and create better performing feature representations when deployed using images taken from a different medical center or scanner. In particular, the random image augmentation referred to as RandomResizedCrop, described herein, provides a number of such advantages.

[0067] At block 206, the pretraining module 104 masks patch regions for global crops, for example, using a binary mask. In some cases, to simplify 3D masking, the pretraining module 104 can randomly sample masking locations uniformly over all patches instead of blockwise masking.

[0068] At block 208, in some cases, the pretraining module 104 uses an adaptation period using high-resolution inputs of the 3D medical imaging data at the end of pretraining.

[0069] At block 210, the machine learning module 106 performs pretraining using a teacherstudent machine learning model (i.e. , knowledge distillation model). The weights of the teacher machine learning model form the pretrained weights of the model to be used to determine feature representations of inputted 3D medical imaging data (such model referred to herein as “3DINO-ViT”). The augmented versions of the imaging data without masked patch regions are passed through the teacher network to obtain a teacher representation. The augmented versions withmasked patch regions are passed through the student network to predict the masked regions in the teacher representation. Generally, the student network is trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network

[0070] In an example, the teacher machine learning model (i.e. , teacher network) and the student machine learning model (i.e., student network) can be vision transformer encoder networks (ViT) with the same size and configuration, but having different weights. In a further example, the student network can be a simpler model, such as a convolutional neural network (CNN). Generally, both the student network and the teacher network can be randomly initialized. The machine learning module 106 passes global crops through the teacher network and passes the patched regions through the student network. Given the global crop passed through the teacher network, the student network performs predictions of the teacher representation using all other crops. An image-level loss function used by the model can include a cross-entropy function, and because the student network includes patched regions, the loss function adds patch-level loss to image-level loss. In most cases, a pretraining can occur over a predetermined number of iterations (for example, 100,000 iterations).

[0071] In some cases, rather than explicitly training the teacher, the teacher network can be determined using an exponential moving average (EMA) of the student network. In some cases, the teacher network can be adjusted to flatten and project 3D input patches.

[0072] In some cases, instead of using a vision-specific network for a CiT teacher network, a ViT-Adapter can be converted to handle 3D medical imaging inputs. The ViT-Adaptor can be converted to 3D inputs by using 3D convolutions for extracting multi-scale features in 3D. Additionally, Injector and Extractor networks of the ViT can both use Multi-Scale Deformable Attention (MSDA) formulation of sparse attention. The MSDA can be converted to 3D inputs by inputting 3D feature maps, and learning an additional deformable offset for the depth axis.

[0073] At block 212, the machine learning module 106 uses the pre-trained teacher machine learning model to determine feature representations from 3D medical imaging data inputted into the model. In some cases, the feature representations outputted from the pretrained model can be converted by the machine learning module 106 into probability distributions (for example, via a softmax function) and / or sharpened (for example, via a temperature parameter). The feature representations can be any suitable features learned from the inputted 3D medical image(s) and are generalizable to a variety of tasks, including but not limited to: tumor predictions from MRI orultrasound images, age predictions from brain scans, COVID- 19 disease prediction from CT scans of lungs, or the like.

[0074] At block 214, the output module 108 outputs the feature representations of the 3D medical imaging data to the data storage 164, the interface module 168, and / or the network module 176.

[0075] In a particular approach of the system 100, a self-supervised self-distillation can be used that consists of a teacher, g0tand student network, g9sparameterized by 9tand 6Srespectively. From an unlabeled medical image sampled from the dataset, x, two randomly augmented "global crops", xf and x , as well as L randomly augmented "local crops", {x^xf ...,x{} can be generated to create a set of crops C. The global crops are passed through the teacher network, and all crops are passed through the student network. Given a global crop passed through the teacher, x e {xf,x }, the overall task of the student network is to predict the teacher representation using all other crops, x2e C,x2#= x1.

[0076] The feature representation output from the student and teacher can be converted into probability distributions via a softmax function, o-(-), and sharpened via a temperature parameter, T. The student probability distribution can be defined as Ps(x) = with the same formulausing Ptand ztfor the teacher network.

[0077] The following image-level loss function summarizes an objective of the above approach:Where / / (•) is the cross-entropy function. Rather than explicitly training the teacher in this framework, it is obtained as an exponential moving average (EMA) of the student model: 6tE- 6t+ (1 - A)0S. The above approach learns image-level representations by taking features learned from the classification token of a ViT network, which is represented by the [CLS] superscript. Hence, this task forms the "image-level objective".

[0078] The system 100 makes several substantial improvements over the above approach in order to be more scalable and efficient for learning representations from large (medical) images. In a first example, the system 100 can use a masked image modeling (MIM) objective. Such masks patch regions for global crops passed to the student using a binary mask, m e {0, 1}Noverthe N patches comprising the crop. Using the masked crop, x9, the student model is tasked with predicting the teacher representations at the masked regions, leading to the loss function:Where mi is the binary mask value, and P(-)M is the patch token output by the encoding model, both enumerated by patch location, i. By introducing the above objective, patch-level representation quality and robustness to image corruption is greatly improved. This improvement is particularly useful for dense downstream tasks like segmentation and improving representations in the presence of out-of-sample distribution shifts and corruptions / artifacts; which are abundant in medical imaging due to differences in scanners, imaging hardware, acquisition parameters and sequence design. The final loss function can be used to add patchlevel loss to the image-level loss.

[0079] The system 100 provides several improvements on computational and memory efficiency that enable larger batch training. These include improving the efficiency of computing selfattention, allowing nested tensors in self-attention, saving memory in the stochastic depth operation, and taking advantage of fully-sharded data parallel modules. The system 100 also introduces a secondary high-resolution adaptation stage to the pretraining. In segmentation tasks, maintaining image resolution is important for extracting smaller objects and features. However, training self-supervised approaches from scratch, with high-resolution inputs, is highly computationally expensive. Instead, the present inventors found that introducing a short adaptation period using high-resolution inputs at the end of pretraining yields comparable results compared to full training using high-resolution. The adaptation period can include first pretraining using 3DINO on low resolution data, and then taking the resulting weights and pretraining again but for a shorter time period and with higher resolution data. High resolution adaptation generally improves results for downstream tasks with high resolution inputs.

[0080] The above model architecture can be adapted for 3D inputs by adjusting the ViT encoder network to flatten and project 3D input patches. To permit variable-sized inputs, the system 100 can use 3D interpolation of learned position encoding vectors. In some cases, to simplify the 3D masking operation, the system 100 can randomly sample masking locations uniformly over all patches instead of blockwise masking. Proper selection of data augmentation approaches used to create global and local views can be beneficial for generating salient image representations inself-supervised learning (SSL); and thus, particular care should be taken to select data augmentations for 3D adaptation.

[0081] In some cases, since the pretraining dataset can include non-quantitative imaging modalities, image normalization can be performed by linearly mapping, for example, the 0.05thand 99.95thpercentiles of intensity to -1 and 1, respectively. In some cases, random image augmentations can be used for pretraining on red-green-blue (RGB) images, which can include flipping, blurring, converting to grayscale, solarization, and color jitter. Since color-related augmentations generally cannot be implemented for 1-channel inputs, medical imaging-related augmentations can be used that may produce robust representations to domain shift; for example, random contrast adjustment, additive noise, gibbs noise, and histogram shift.

[0082] An augmentation, referred to as ‘RandomResizedCrop’, can be used to randomly crop a portion of an input image and resize it to a specified size, keeping the crop within a range of aspect ratios (i.e. , ensuring that a crop is not too long or wide). Generally, this approach is not easily adapted to 3D, as the images that form the pretraining dataset generally have highly variable slice thicknesses and voxel sizes. For example, the Healthy-Total-Body-CTs dataset contains on the order of -1000 CT slices across a single image. On the other hand, the NYU fastMRI knee dataset contains -30 slices per image. Maintaining a reasonable aspect ratio between the in-plane image dimensions and the out-of-plane (depth) image dimension would be difficult for both datasets simultaneously. Thus, rather than enforcing approximately isotropic spacing and an aspect ratio near one on a randomly cropped and resized volume, the system 100 can crop the two in-plane image dimensions using a 2D RandomResizedCrop operation, and independently sample the cross-slice dimension crop size. While this may mean that the cropped volume resulting from this data augmentation can be stretched or squashed in the out-of-plane axis, using a large variety of pretraining datasets will allow the model to learn to generalize to various volume sizes. This approach additionally preserves the “local-to-global” correspondences, where global views take up a larger portion of the original volume than local views. In an example, the 2D RandomResizedCrop operation can involve taking a random crop out of the original 2D image with a random area and random aspect ratio (the limits of the area and aspect ratio can be preselected or defined by a user), and then resizing the crop to a specific size (which also can be preselected or defined by the user). This augmentation crops out a smaller (e.g., potentially squeezed or stretched) region from the original image.

[0083] ViTs are typically more difficult to train relative to convolutional neural networks (CNNs), especially in a supervised setting with limited training data. This has generally been attributed to the lack of inductive biases in ViTs and their larger number of trainable parameters. In 3D medical imaging, this difficulty has been partially overcome using vision-specific transformer networks, such as the SwinUNETR. However, ViTs have been shown to scale well with dataset sizes. Hence, instead of using a vision-specific network, the system 100 uses a converted ViT-Adapter, a module that injects spatial information into standard ViT networks, to handle 3D medical imaging inputs.

[0084] ViT-Adapters can generally be used for 2D pretrained ViT networks as a way to introduce image-based inductive biases into the network. Using ViT, the system 100 is able to take advantage of large-scale pretraining. For example, using a simple convolutional network called a Spatial Prior Module to extract local multi-scale spatially relevant features from the original input. A Spatial Feature Injector (referred to as ‘Injector’) can be used to introduce extracted multi-scale spatial features into the features obtained from the pretrained ViT. A Multi-Scale Feature Extractor (referred to as ‘Extractor’) can then be used to adapt the multi-scale features based on the pretrained ViT features. By using multi-scale features, the system 100 is able to output a feature pyramid much like typical convolutional encoder networks. Use of the ViT-Adapter greatly improves ViTs for dense segmentation tasks.

[0085] In order to receive 3D inputs, the Spatial Prior Module can be adjusted to use 3D convolutions for extracting multi-scale features in 3D. To avoid resampling errors with the input sizes of the pretrained ViT network (e.g., 112x112x112), instead of using feature maps with, for example, 1 / 8, 1 / 16 and 1 / 32 of the original spatial input size (because 32 does not divide 112 evenly), the system 100 can instead use scales of, for example, 1 / 4, 1 / 8 and 1 / 16. Thus, the .HWD HWD HWD.output of the spatial prior module would be Tspe Jr 43s3i&3> , for height, H, width, W, and depth, D, of the input volume and the Transformer feature size, F.

[0086] The Injector and Extractor networks can both use a Multi-Scale Deformable Attention (MSDA) formulation of sparse attention. Instead of having both query and keys in standard selfattention enumerate all possible spatial locations in an input image, each query in MSDA only attends to a fixed, small number of keys (K=4). Then, the value features are obtained by sampling the feature map at learnable offset locations. MSDA can be converted to 3D inputs by inputting 3D feature maps, and learning an additional deformable offset for the depth axis.

[0087] The above 3D deformable attention operation was adapted by the present inventors to permit 3D inputs and to enable non-integer deformable offsets by performing trilinear interpolation on 3D feature maps. Deformable offsets generally determine which portions on the feature map the model should operate on. As such feature maps have discrete indexes that are integers, linear interpolation on the feature map is necessitated to handle a deformable offset that is not an integer. For example, averaging the feature maps indexed nearest to a particular location of interest weighted on their distance. In order to perform such operation in 3D, a trilinear interpolation is performed (for example, determining an average between eight feature maps that are nearest to the location of interest). In this way, trilinear interpolation allows for non-integer deformable offsets by using the feature map at particular discrete integer indices, and performing averaging that is weighted on distance in order to approximate the feature map at non-integer locations.

[0088] The system 100 implements a substantially amended MSDA when initializing bias parameters for the linear projection in order to predict the deformable offset. For 2D, the work offsets bias for each attention head in selected so that the initial offsets have equal angular separation. For example, with 8 attention heads, the initial offset per head, Oinit-2D, is:Oinit-2D = {(- / -fc)For each key, k e {1,2,.. ,,K , making the angular separation for each offset vector 45 degrees. The key intention of this form of initialization is to ensure more even coverage of the feature map when sampling for value features. With no direct way to extend this into 3D without introducing an intractable number of attention heads, the present inventors determined that it is advantageous to fix the attention heads to 8 (as an example), and then initialize the bias to have initial offsets pointing towards each 3D octant (in the example of 8). Concretely, the initial offset per head in 3D MSDA, Oinit-2D, is:0init-3D = (~k, ~k, —k), (—k, -k, k), (-k, k, -k), (-k, k, k), ( / c, -k, -k), ( / c, -k, k), ( / c, k, -k), ( / c, k, k)}

[0089] Example pretraining experiments were conducted by the present inventors on singlechannel input images. During pretraining, multi-channel unlabelled inputs were split into separate input images. However, some downstream tasks in medical imaging (such as Brain Tumor Segmentation - “BraTS”) benefit from using multiple co-registered modalities to provide complementary information and contrast. The issue of adapting a single-channel pretrained network to multi-channel inputs rarely arises for large 2D natural image datasets, as pretrainingand all downstream tasks tend to remain in RGB color space (3 channels). The present inventors determined that channel mixing can be used to address this challenge having 3D inputs. In most cases, to make full use of the pretrained weights, which are generally specifically tuned on singlechannel inputs, the patch embedding layer of the ViT was not adjusted. Instead, each image channel through the network can be passed individually and a feature vector per channel can be obtained. These features can then be concatenated and passed through a linear layer mapping back to the original transformer feature size, and a Gaussian error linear unit (GELLI) activation function. The resulting patch-level feature vectors can be passed to a decoder network for downstream dense segmentation tasks.

[0090] Practically, in some cases, the operation passing each channel through the network independently can be parallelized by reshaping the input so that channels form part of the minibatch (i.e. , an input of shape ^BXCXHXWXD jSreshaped toforbatchSjze, g) Generally, it can be expected that the spatial information of multi-channel inputs will be relatively similar between co-registered channels. Thus, to maintain tractability and reduce redundancy when training the ViT-Adapter, the multi-channel features output by the frozen pretrained Transformer blocks can be averaged along the channel dimension before being passed to the Injector and Extractor modules. The resulting spatial information output from the Injector can be copied along the channel dimension before being added to the Transformer features. After being passed fully through the ViT, these features can also be concatenated and linearly mapped to the original Transformer feature size. Example experiments illustrated that this approach leads to marked benefits even when single-channel pretrained weights are used on multi-channel segmentation tasks.

[0091] In the example experiments, a multimodal 3D medical imaging pretraining dataset for the system 100 was generated from a variety of publicly available datasets and one internal dataset. The pretraining datasets were filtered for excessively few DICOM slices (>24 slices) to avoid overly pixelated volumes in the cross-slice dimension (lower z-axis resolution). To reduce redundancy and perform a naive form of deduplication, random subsets of a few exceptionally large datasets were taken. Subsets were taken from the FastMRI Knee dataset and the RSNA Intracranial Hemorrhage Detection dataset by randomly sampling half of the dataset. A subset of NLST was taken by randomly sampling 500 patients, and a subset of 4D-Lung was taken by sampling 30-40 volumes per patient. After deduplication and filtering for slice counts, a 3D medical imaging dataset of 98,815 unlabeled volumes was created. The high-resolution adaptationdataset was created by filtering for >48 DICOM slices, which resulted in a higher resolution 53,758 volume subset of the original data.

[0092] The example experiments compared the efficacy of the weights of the model of the present embodiments (referred to as ‘3DINO-ViT’) on downstream tasks against four other pretraining and initialization models. The first comparison randomly initialized the ViT network and trained it end-to-end from scratch (referred to as ‘Random’). As a state-of-the-art pretrained medical imaging backbone, the Sliding Window (referred to as ‘Swin’) ViT was used. To account for differences between model architectures (ViT versus Swin ViT), a randomly initialized Swin ViT was used to evaluate the relative benefits of using pretrained weights (referred to as ‘Swin Random’). Additionally, a more direct comparison was performed the 3DINO pretraining approach by adapting the Swin ViT pretraining framework with a standard ViT (referred to as ‘MONAI-ViT’), and pretrained it on the same unlabeled dataset.

[0093] The Swin Transformer backbone provides an SSL pretraining method for 3D medical images, specifically for CT data. The publicly released pretrained weights associated with such approach was used in the example experiments as one baseline comparison against a pretraining approach (referred to as ‘Swin Transfer’) and took a randomly initialized Swin Transformer (referred to as ‘Swin Random’) to determine the relative benefit of using the pretrained weights. However, since the Swin Transfer network differs from 3DINO-ViT in both model architecture and pretraining dataset, the Swin pretraining implementation was replaced with a standard ViT. Such approach was used to pretrain a standard ViT on the 3DINO-ViT pretraining dataset. This forms a separate comparison that specifically investigates the difference in quality of pretraining algorithms (referred to as ‘MONAI-ViT’). The original Swin ViT pretraining approach uses the inpainting, contrastive coding, and rotational prediction tasks jointly. Similarly, the inpainting task is performed by upsampling the ViT patch tokens using transposed convolutions and comparing the reconstructed output to the original image via an L1 loss. The contrastive and rotational tasks are image-level tasks, hence the output of the ViT [CLS] token was passed to the contrastive coding and rotation prediction linear heads. By doing so, the experiments were able to more explicitly train an image-level representation for classification experiments.

[0094] In the example experiments, the Swin ViT pretraining was minimally adjusted and only additional transforms were added to change intensity scaling to the percentile-based approach used for 3DINO-ViT. A ViT-Large was pretrained on the same pretraining dataset that was used to train 3DINO-ViT. It was pretrained using a patch size of 16x16x16 and an image size of96x96x96. A batch size per GPU of 16 (64 total) was used and pretraining ran for 100,000 iterations over approximately 10 days.

[0095] To evaluate the saliency and generalizability of the pretrained weights, medical image segmentation and classification benchmarks / challenges were used. As a segmentation benchmark, the 2021 Brain Tumor Segmentation (BraTS) challenge for MRI was used, which looks to delineate between Tumor Core, Enhancing Tumor, and Whole Tumor regions in glioblastoma patients. Additionally, the Beyond the Cranial Vault (BTCV) challenge was used for a 14-class CT abdominal organ segmentation task. The example experiments further evaluated the generalizability of the pretraining on unseen (out-of-distribution) organs and 3D modalities with minimal presence in the pretraining dataset. Specifically, performance was evaluated on left atrium MRI segmentation (LA-SEG), an unseen organ, and 3D breast ultrasound tumor segmentation (TDSC-ABUS), a completely unseen modality. For classification tasks, brain age classification was performed on the MRI ICBM dataset and used the COVID-CT-MD lung CT dataset to classify between healthy patients, those with community-acquired pneumonia (CAP), and individuals with Novel Coronavirus (COVID- 19). The example experiments also used different amounts of labeled training data by randomly subsampling a certain percentage of the full labeled dataset.

[0096] The BraTS 2021 training dataset is a widely-used MRI brain segmentation benchmark. This dataset consists of 1251 patients, each with four types routinely required MRI scans: T1-weighted, T2-weighted, T1-weighted with gadolinium contrast, and T2 Fluid-attenuated Inversion Recovery (FLAIR). These scans were skull-stripped and coregistered. For each patient, medical experts manually generated pixel-level segmentation labels that were combined into Whole T umor (WT), T umor Core (TC) and Enhancing T umor (ET) regions. T o form the finetuning dataset for the example experiments, scans in the dataset taken from TCGA-GBM or TCGA-LGG were removed (which are present in the pretraining dataset) to avoid unfair bias in evaluation. This resulted in 1084 patients that were randomly split into train (N = 758), validation (N = 108) and test (N = 218) sets.

[0097] The BTCV dataset is a CT abdominal organ segmentation benchmark. This dataset consists of abdominal CT scans taken from 30 healthy patients with manual labels generated of 13 organs. These were randomly split into train (N = 20), validation (N = 4) and test (N = 6) sets.

[0098] The International Consortium for Brain Mapping (ICBM) dataset is an MRI brain age classification benchmark. This dataset consists of T1 -weighted brain MRI scans of 639 healthy patients with 1339 scans from a variety of ages between 18 and 80. To maintain a reasonably balanced dataset, the example experiments binned data into four bins of width 10 between 20 and 60 years of age, discarding scans that did not fall into this range. The data was split randomly on the patient level to obtain train (N=756), validation (N=151) and test (N=233) scans. Only skullstripping was performed for data preprocessing using the iCVMapp3r pipeline. For the four bins:[20, 30), [30, 40), [40, 50), [50, 60], the train set contained 335, 199, 108, 144 volumes, the validation set contained 63, 25, 39, 24 volumes, and the test set contained 110, 49, 21, 53 volumes respectively.

[0099] The COVID-CT-MD dataset was used in the example experiments as a lung CT classification benchmark between COVID-19, Community Acquired Pneumonia (CAP), and healthy patients. This dataset consisted of 305 patients with one lung CT scan each that was split randomly into train (N=214), validation (N=30) and test (N=61) scans. For the three classes, Healthy, COVID-19, CAP, the train set contained 54, 121, 39 scans, the validation set contained 7, 15, 8, and the test set contained 15, 33, 13 scans respectively.

[0100] The Left Atrium Segmentation Challenge (LA-SEG) dataset was used as a left atrium MRI segmentation benchmark. The heart makes up a very small subset of the pretraining dataset, hence this data is used to evaluate the generalizability of the present embodiments to an out-of-domain organ. This dataset consisted of 154 heart MRI scans from 60 patients and segmented for the left atrial cavity. The challenge dataset was originally split on the patient level-into training (N=100) and test (N=54) sets. The training dataset was randomly split further into subsets used to train (N=80) and validate (N=20) finetuning networks.

[0101] The Tumor Detection, Segmentation and Classification Challenge on Automated 3D Breast Ultrasound (TDSC-ABUS) dataset was used as a 3D breast US lesion segmentation benchmark. Ultrasound images have a very different appearance to MRI and CT images and were not present in the pretraining dataset. Hence, this data was used to evaluate the generalizability of the method to a completely out-of-domain downstream task. The dataset consisted of 100 breast US scans from an unreleased number of patients, with expert segmentations for lesions. The data was split randomly into train (N=70), validation (N=10), and test (N=20) sets.

[0102] For segmentation experiments, such experiments were performed via appending convolutional decoder heads to pretrained encoders and tuning end-to-end. 3DINO-ViT and MONAI-ViT weights were frozen, and a 3D-ViT Adapter module and simple convolutional decoder were tuned. The present embodiments yielded significantly improved segmentation results relative to other techniques on all evaluation metrics in most comparisons (p < 0.05). 3DINO-ViT was able to jointly improve representations for both segmentation tasks in all percentages of labeled data, including when using the full labeled dataset. The pretrained weights significantly improved performance at all percentages of labeled data relative to the Random encoder (BraTS with 10% data: 0.90 (0.88, 0.91) Dice for 3DINO-ViT vs 0.87 (0.85, 0.89) for Random', BTCV with 25%: 0.77 (0.72, 0.81) 3DINO-ViT vs 0.59 (0.53, 0.65) for Random). For both segmentation tasks, 3DINO-ViT trained using less than 50% of all labeled data achieved statistically and visually comparable results to other baselines trained using 100% of labeled data.

[0103] For classification experiments, a linear classifier was trained on top of all pretrained networks without finetuning the pretrained weights. The MONAI-ViT was used as one comparison and the pretrained contrastive head from the Swin T ransfer network was used as another. Despite the tasks’ difficulty and sparser ground truth, 3DINO-ViT of the present embodiments performed universally better than other models (p < 0.05). Averaging over all dataset sizes, 3DINO-ViT obtained an 18.9% higher area under the receiver operating characteristic curve (AUG) on COVID-CT-MD, with a particularly notable increase of 23% AUG on classifying patients with COVID-19 relative to the next best baseline. On ICBM, an average of 5.3% higher AUC was obtained with a 22% AUC improvement for classifying individuals aged [40, 50) years over the next best baseline. These example experiments, conducted with completely frozen pretraining weights, further highlight the saliency of the learned representations for different downstream tasks.

[0104] As a standalone comparison of the effectiveness of patch-level representations for image segmentation, example experiments were performed with a lightweight segmentation decoder. The 3DINO-ViT’s pretrained weights were frozen and a two-layer linear network was finetuned on downstream tasks. The performance was compared against both the Random and MONAI-ViT networks and the example experiments determined that 3DINO-ViT achieved a significant improvement of 61% Dice over the next best baseline on BTCV (p < 0.05), and 15% on BraTS (p < 0.001).

[0105] In the example experiments, on the out-of-distribution tasks, 3DINO-VIT significantly outperformed other tested approaches, with 1.9% improved Dice on left atrium segmentation, and 24% in 3D ultrasound tumor segmentation over the next best baseline, when finetuning with 25% of the labeled dataset. Though this improvement drops to 1.1% and 0.8% respectively, 3DINO-ViT maintains its advantage over other baselines even when tuning with 100% of the labeled dataset. This demonstrates the capability of the present embodiments to create generalizable weights that can be applied to image distributions unseen during pretraining.

[0106] All example experiments used a ViT-Large 29 and a patch size of 16x16x16. Standard pretraining experiments used a batch size per GPU of 128 (512 total), a global crop size of 96x96x96, a local crop size of 48x48x48, and a base learning rate of 0.002. The EMA parameter A is increased from 0.992 to 1.000 in a cosine schedule. Pretraining progresses for 125,000 iterations over approximately nine days. High-resolution adaptation was performed by keeping parameter scheduling the same as pretraining but compressed to progress over 12,500 training iterations instead of the original 125,000 iterations. High-resolution adaptation used a batch size per GPU of 64 (256 total), a global crop size of 112x112x112, a local crop size of 64x64x64, and a base learning rate of 0.001. Adaptation began from the weights learned in the 112,500thpretraining iteration and took approximately two days.

[0107] In the example experiments, example segmentation experiments were performed using weights learned from high-resolution adaptation. The ViT-Adapter module and a U Net-like convolutional decoder were trained on the dense segmentation task. Corresponding to the pretraining input size, these experiments used inputs of size 112x112x112. To evaluate the label efficiency of the pretraining, a random subset of the finetuning train sets were used (with the same random subset taken between experiments).

[0108] The example experiments used a batch size of 8, and were conducted on one A100-SXM4-80GB GPU. The base learning rate was set to 0.0001, and finetuning was conducted for 30,000 iterations (regardless of input dataset size). An AdamW optimizer was used with default and weight decay and a LinearWarmupCosineAnnealing scheduler with 3,000 warmup iterations. For BraTS, a Dice loss function was used, and for BTCV, LA-SEG, and TDSC-ABUS, a Dice-Cross-Entropy was used. For all such experiments, the validation set was used to select the best model epoch from training and reported results on the test set. To create final segmentation logits for testing, the experiments strided a sliding window over the full image with an overlap between images of 0.75.

[0109] In some cases, as in the example experiments, the decoder can take a four-level feature pyramid output from the ViT-Adapter and use UNet-like transposed convolutions for upsampling followed by encoder-decoder connections. The decoder can consist of four layers with feature size 256, 128, 64, 32 before mapping to the number of segmentation classes. The ViT-Adapter can break the encoding layers into four blocks each containing 6 Transformer layers, and use 8 Multi-Scale Dilated Attention (MSDA) heads. The feature size for ViT-adapter operations in the example experiments was 256, or 25% of the full ViT feature size to reduce computational complexity.

[0110] A random encoder network can be converted to perform segmentation by adding a convolutional decoder and adapted to a ViT-Large by taking the output of the 6th, 12th, 18th, and 24th ViT layer. Both the encoder and decoder can be tuned end-to-end. Example segmentation experiments used a Swin Transfer and Swin Random encoders were conducted by attaching a SwinllNETR decoder network. The full SwinllNETR was tuned end-to-end. The input size used for segmentation corresponded to the image size used for pretraining, 96x96x96. The same subsets of the finetuning datasets used for segmentation were taken in these experiments. Experiments tuning the Swin Transfer and Swin Random networks on BTCV used a batch size of 8, with BraTS experiments using a batch size of 4 (largest power of 2 without running into memory errors). All experiments were conducted on one A100-SXM4-80GB GPU. The pretrained weights were originally trained on CT images with intensity normalized to a range of [0, 1], Segmentation experiments using the Swin Transfer network were normalized to this range to better take advantage of pretraining.

[0111] To perform example experiments with a lightweight linear decoder for segmentation, the pretrained network was frozen, and a two-layer linear network was trained. The first linear layer was the multi-channel projection linear layer. The second layer mapped to the number of segmentation classes, taking concatenated patch representations from the final four ViT layers as input. The output of the linear decoder was low-resolution volume of class logits (for example, for an input image of size 96x96x96 and a patch size of 16x16x16, the model would output a 6x6x6 map). This volume was upsampled using trilinear interpolation to the original image size to obtain pixel-wise segmentation logits, which were compared with the ground truth mask. The linear decoder experiments used a base learning rate of 0.001 and a batch size of 16. Other parameters, including optimizer, scheduler, iterations trained, and loss functions remained consistent with other segmentation experiments.

[0112] Example classification experiments were conducted with the three pretrained models: the model of the present embodiments (referred to as ‘3DINO-ViT’), MONAI-ViT, and Swin Transfer. In all cases, linear probing on frozen pretrained weights was performed using a grid search on three key parameters: the learning rate, the number of final ViT layer outputs to concatenate, and whether the averaged patch tokens are concatenated to the [CLS] token. The learning rates were searched in the set {0.0001,0.0002,0.0005,0.001,0.002,0.005,0.01,0.02,0.05,0.1,0.2,0.3,0.5}, number of output layers in {1,4} and averaged patch token concatenation in {True, False}. The best parameters based on validation performance were then used for evaluating on the test set.

[0113] The Swin Transfer network does not train a [CLS] token for forming image-level representations. To probe the image-level representations from these pretrained weights, the output was extracted from the pretrained contrastive coding head of the network. As with segmentation experiments, the input images to the pretrained Swin Transformer network were normalized between [0, 1],

[0114] Classification experiments were conducted for 12,500 iterations with input sizes matching what was used to pretrain the models. All classification experiments used a batch size of 32, and were conducted on one A100-SXM4-80GB GPU. The SGD optimizer was used with a momentum of 0.9 and 0 weight decay, a cosine annealing learning rate scheduler, and Cross-Entropy loss. As with segmentation experiments, the classification experiments extracted a random subset of the finetuning train sets to test reliance on labeled data.

[0115] The example experiments illustrate that the present embodiments provide a host of efficiency improvements, objectives, and regularization approaches that stabilize training and reduce training time by two-times and GPU memory usage by three-times. The pretext formulation combines an image-level objective and a patch-level objective, where original volumes are augmented to generate two global and eight local crops (a total of 10 augmentations for the objectives per scan). The augmentations act on 3D grayscale images, and accordingly, a custom implementation of Random ResizedCrop is used along with other augmentations, such as random contrast adjustment, additive noise, and gibbs noise. The image-specific operations were converted from 2D to 3D, while making minimal changes to the ViT layers. By doing so, the present embodiments take advantage of all high-level and low-level efficiency improvements. The 3DINO-ViT model backbone was advantageously modified to enhance its performance on downstream segmentation tasks by converting an adapter module to 3D inputs {3D ViT-Adapter). This module was employed in 2D images to inject spatial inductive biases into pretrained ViTmodels for dense (pixel-level) tasks. To achieve this, a 3D version of multi-scale deformable attention (MSDA) was used. While large SSL-pretrained models can generally be used for a single 3D medical imaging modality, the present embodiments advantageously provide the first 3D selfsupervised learning (SSL)-based medical imaging model that can extract salient features and generalize across multiple modalities simultaneously.

[0116] The present embodiments provide a computationally efficient 3D SSL framework that performs representation learning on 3D medical images. The example experiments applied the framework to a curated, ultra-large, multi-organ, and multimodal medical imaging dataset to build a general-purpose pretrained ViT model (3DINO-ViT). To enhance pretraining when applied to dense tasks, the 3D ViT-Adapter was developed to evaluate segmentation baselines. Altogether, the example experiments observed that the full pipeline of the present embodiments improved the ViT’s data efficiency and generalizability on both pixel-level (segmentation) and image-level (classification) tasks. Transfer learning with frozen 3DINO-ViT weights improved performance over other approaches on the majority of evaluation metrics at all dataset sizes and for all tasks, including when using the full labeled dataset. The example experiments further demonstrated the system’s 100 capability to generalize to out-of-distribution data from unseen organs and modalities in pretraining.

[0117] On the primary Dice similarity segmentation metric, the present embodiments consistently outperformed over other pretrained models. When comparing the amount of improvement with 3DI NO-ViT initialization over Random initialization relative to Swin T ransfer initialization over Swin Random initialization, it was determined that 3DINO pretraining yielded relatively larger benefits to segmentation performance. The maximum Dice improvement for 3DI NO-ViT over the Random encoder was 13.0% on BraTS and 55.1% on BTCV, whereas the Swin Transfer network versus Swin Random improved 5.1% on BraTS and 1.8% on BTCV, respectively. While MONAI-ViT generally improved over Random initialization, the pretraining of the present embodiments was determined to generate more salient representations for both segmentation and classification tasks due to the quality of the pretext task and the computational efficiency. By greatly simplifying the segmentation setup to train solely a linear network to act on patch-level representations, 3DI NO-ViT features were more directly suitable to segmentation tasks than those output by other ViT initialization approaches.

[0118] The example experiments further determined that 3DI NO-ViT outperformed all other pretrained models in classification tasks for all metrics. Classification with 3DI NO-ViT also yieldedimproved class-wise pseudo-probabilities. For instance, in particular cases on brain age classification on the ICBM dataset where individuals were in between two classification bins, it was observed that the present embodiments were able to successfully classify them and better capture the uncertainty in the prediction. The present embodiments represent the first pretraining approach for 3D medical images that is able to perform both image-level and patch-level tasks, while jointly improving performance on both.

[0119] The example experiments probed the out-of-domain generalizability of the present embodiments using heart MRI and breast US segmentation. The present embodiments significantly improved over other approaches in both of these tasks, demonstrating its ability to learn features that are salient even for unseen distributions. To visually investigate the saliency of representations generated by the present embodiments, the example experiments generated principal component analysis (PCA) and multi-head self-attention (MHSA)-based visualizations of the representations. PCA visualizations demonstrated that common modes of variation for all datasets are between background versus foreground, outlining the surface of the organ, and varying across anatomical axes. Principal components generated on BraTS images found inside the tumor extent were often distinct from other brain tissues.

[0120] The present embodiments were also compared to Segment Anything Models (SAM). Since SAM-like networks are trained using labeled data and 3DINO is a self-supervised pretraining approach, both methods are synergistic. Sam uses weights from SSL pretraining to initialize the image encoder network. In contrast, 3DINO represents a different ViT pretraining approach for 3D inputs that is able to act as an initialization step for 3D SAM or other approaches.

[0121] FIGS. 6A to 6G illustrate transferring pretrained models into segmentation architectures in the example experiments. The illustrated examples include a 3DINO-ViT configuration (FIG.6A), a MONAI-ViT configuration (FIG. 6B), a MIM-ViT configuration (FIG. 6C), a randomly initialized configuration (FIG. 6D), a Swin Transformer random initialization configuration (FIG.6E), a Swin Transformer transfer configuration (FIG. 6F), and a Swin Transformer transfer-FD configuration (FIG. 6G). In the illustrated figures, a blue lock associated with a model indicates that corresponding weights are frozen, whereas the absence of a blue lock indicates that weights are fully tuned.

[0122] FIGS. 7A and 7B schematically illustrate two-dimensional and three-dimensional multiscale deformable attention (MSDA) offset initialization, respectively, in the example experiments.In the illustrated examples, offsets are shown for multiple attention heads, with different colors used for visualization without loss of generality, and a single key location is illustrated.

[0123] FIG. 8 illustrates an example of multi-channel interfacing with a ViT-Adapter module, in the example experiments. The illustrated example shows conversion of separate input channels into batches, followed by channel averaging and channel broadcasting operations. Dimensionality of tensors at each stage is indicated in parentheses, and the example is shown using a multichannel medical imaging input.

[0124] FIGS. 9 to 15 show an evaluation and comparison other pretrained models on the BraTS and BTCV segmentation, and ICBM and COVID-CT-MD classification tasks, in accordance with the example experiments. FIG. 9 shows BraTS segmentation Dice scores and 95th percentile Hausdorff Distance (HD95). FIG. 10 shows BTCV segmentation Dice scores and HD95. FIG. 11 shows ICBM classification AUC and F1 scores. FIG. 12 shows COVID-CT-MD classification AUC and F1 scores. FIGS. 9 to 12 illustrate finetuning results with multiple sizes of labeled dataset, x-axis displays training dataset size in a percentage of the full dataset, with actual number of labeled samples in parentheses. Bar plots compare the third-largest and largest training dataset sizes with error bars in FIGS. 9 and 10 from the 95% bootstrapped confidence interval (Cl) of item-wise metrics and error bars. Bar plots in FIGS. 11 and 12 are obtained from 95% bootstrapped Cl of metrics obtained from five separate experiments with randomly subsampled training sets (100% data comparison is equivalent to adjusting experiment seed). Statistical significance in FIGS. 9 and 10 is computed via a paired nonparametric Wilcoxon test on metrics per item. Significance in FIGS. 11 and 12 are computed via an unpaired Welch’s t-test against metrics per experiment. FIG. 13 shows plots comparing per-class Dice / AUC scores for segmentation / classification experiments using the third-largest training dataset size. FIG. 14 shows normalized and averaged classification confusion matrices using the third-largest training dataset size. FIG. 15 shows dice scores for linear decoder segmentation experiments.

[0125] FIG. 16 shows a visualization of segmentation predictions on BraTS and BTCV datasets, in accordance with the example experiments. Each row corresponds to an original image, ground truth segmentation, and visualized segmentation outputs per pretraining methodology using the third-largest and largest training datasets. The first to fourth rows have BraTS labels being necrotic and non-enhancing tumor core, peritumoral edema, and enhancing tumor. The last two rows have BTCV organs visualized, being right kidney, left kidney, stomach, aorta, inferior vena cava, and pancreas. The numbers above images are Dice segmentation scores obtained on thefull 3D volume. Arrows indicate degraded segmentation outputs in methods relative to ground truth segmentation. The first row shows a contrast-enhanced T1-weighted image from BraTS with segmentation models trained using 10% of the full dataset size, and the second row shows it using 100% of the full dataset size. The third row shows a T2-weighted image from BraTS with segmentation models trained using 10% of the full dataset size, and the fourth row shows it using 100% of full dataset size. The fifth row shows a CT image from BTCV with segmentation models trained using 25% of the full dataset size, and the sixth row shows it using 100% of the full dataset size.

[0126] FIGS. 17A and 17B shows class-separated visualizations of segmentation network predictions on BraTS. Above each image is the Dice score for segmentation per class. Visualized for, in FIG. 17A, models trained using 10% of the labeled finetuning dataset, and, in FIG. 17B, 100% of the labeled finetuning dataset. Whereby TC is T umor Core; WT is Whole T umor; and ET is Enhancing Tumor. FIGS. 18A and 18B, and FIGS. 19A and 19B, show additional class-separated visualizations of segmentation network predictions on BraTS for models trained using 10% and 100% of the labeled finetuning dataset, respectively.

[0127] FIGS. 20A and 20B show class-separated visualizations of segmentation network predictions on the BTCV dataset. Above each image is the Dice score for segmentation per class. FIG. 20A shows models trained using 25% of the labeled finetuning dataset, and FIG. 20B shows models trained using 100% of the labeled finetuning dataset. Visualized organ classes include spleen, kidney, gallbladder, esophagus, stomach, vena cava, portal and splenic veins, pancreas, and adrenal gland. FIGS. 21A and 21 B, and FIGS. 22A and 22B, show additional class-separated visualizations of segmentation network predictions on the BTCV dataset for models trained using 25% and 100% of the labeled finetuning dataset, respectively.

[0128] FIGS. 23 to 28 illustrate visualization of classification performance on ICBM age prediction (FIGS. 23 to 25) and COVID-CT-MD disease prediction (FIGS. 26 to 28), in accordance with the example experiments, for all approaches and all classes at 25%, 50%, 75%, and 100% of finetuning data. FIGS. 23 to 28 are separate cases in the testing dataset. Bar plot depicts the predicted class probabilities output by each model. Ground truth age is depicted above the image. Whereby Pred is Prediction; and Cap is community-acquired pneumonia.

[0129] FIGS. 29 to 36 show model evaluation on out-of-distribution tasks, in accordance with the example experiments. Left atrium segmentation (LA-SEG; unseen organ) and 3D breastultrasound tumor (TDSC-ABUS; unseen modality). FIGS. 29 and 33 show Dice and HD95 scores for the LA-SEG and TDSC-ABUS segmentation tasks, respectively. FIGS. 30, 31, 34 and 35 show an original image, ground truth segmentation, and visualized segmentation per pretraining methodology using third-largest, and largest finetuning dataset subsets, respectively. Arrows indicate degraded model outputs relative to ground truth segmentations. The numbers above images are Dice segmentation scores obtained on the full 3D volume. FIG. 30 shows LA-SEG visualization when finetuning using 25% of the full labeled dataset. FIG. 31 shows LA-SEG visualization when finetuning using 100% of the full labeled dataset. FIG. 34 shows TDSC-ABUS visualization when finetuning using 25% of the full labeled dataset. FIG. 35 shows TDSC-ABUS visualization when finetuning using 100% of the full labeled dataset. FIGS. 32 and 36 shows unsupervised visualizations on random volumes sampled from the LA-SEG and TDSC-ABUS datasets, respectively.

[0130] The present embodiments address key prior limitations of SSL pipelines for 3D medical imaging in terms of generalizability and computational complexity, leading to advantageous results in downstream tasks and intuitive unsupervised representations. By leveraging 3DINO-ViT, the amount of labeled data needed for diverse downstream medical imaging tasks can be significantly reduced without requiring expensive model retraining on in-domain unlabeled datasets, enabling generalizable and data-efficient models. The present embodiments are highly beneficial when finetuned across a wide variety of challenging applications and tasks in medical imaging, especially in environments with limited access to detailed annotations and resources.

[0131] The computational complexity and data processing requirements of the described embodiments significantly exceed what could be performed effectively through human mental processes, underscoring the technical nature and necessity of computer implementation. The invention provides a solution to self-supervised learning in medical imaging, comprising advanced machine learning models and techniques that process vast amounts of data with high efficiency and accuracy. In some cases, the described embodiments leverage advanced and sophisticated mathematical models and techniques that surpass human cognitive capabilities in speed, accuracy, and scalability, in order to achieve the technical advancements. In some cases, embodiments of the present disclosure can be implemented on specialized hardware and software platforms designed to handle the extensive computational load, thereby offering significant improvements over conventional human-performed tasks. In this way, the interaction between these components / steps and the technical outcomes achieved by the presentembodiments, are impractical for human cognition to replicate due to the scale and complexity of the operations involved.

[0132] Although the foregoing has been described with reference to certain specific embodiments, various modifications thereto will be apparent to those skilled in the art without departing from the spirit and scope of the invention as outlined in the appended claims.

Claims

Claims1. A method for generating weighting for a machine learning model used for determining feature representations from three-dimensional medical imaging data, the method comprising:receiving three-dimensional medical imaging data;generating augmented versions of the three-dimensional medical imaging data by performing image augmentations on the three-dimensional medical imaging data;masking patch regions of a portion of the augmented versions;performing pretraining of the machine learning model using a knowledge distillation model comprising a student network and a teacher network, the augmented versions without masked patch regions are passed through the teacher network to obtain a teacher representation, the augmented versions with masked patch regions are passed through the student network to predict the masked regions in the teacher representation, the student network being trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network; andoutputting weights of the teacher network as a pretrained machine learning model for determining feature representations of previously unseen three- dimensional medical imaging data.

2. The method of claim 1, wherein performing image augmentations on the three- dimensional medical imaging data comprises performing random image augmentations.

3. The method of claim 1, wherein the augmented versions comprise global crops capturing a whole of the original volume of the imaging data.

4. The method of claim 3, wherein pretraining of the machine learning model comprises jointly optimizes an image-level representation objective on the patched global crops and a patch-level representation objective on local crops of the augmented versions.

5. The method of claim 1, wherein masking patch regions is performed using a binary mask.

6. The method of claim 1, wherein the teacher network is an exponential moving average of the student network during pretraining.

7. The method of claim 1 , wherein the teacher network comprises a vision transformer encoder model.

8. The method of claim 1, wherein generating augmented versions comprises applying three-dimensional data augmentations modeling variability in medical image acquisition.

9. The method of claim 1, wherein the image augmentations comprise, at least, applying a two-dimensional crop in in-plane dimensions and independently sampling a depth dimension.

10. The method of claim 1, wherein the three-dimensional medical imaging data comprises multiple co-registered imaging channels.

11. A system for generating weighting for a machine learning model used for determining feature representations from three-dimensional medical imaging data, the system comprising one or more processors and a data storage, the data storage comprising instructions for the one or more processors to execute:an input module to receive three-dimensional medical imaging data; a pretraining module to generate augmented versions of the three-dimensional medical imaging data by performing image augmentations on the three- dimensional medical imaging data, and to mask patch regions of a portion of the augmented versions;a machine learning module to perform pretraining of the machine learning model using a knowledge distillation model comprising a student network and a teacher network, the augmented versions without masked patch regions are passed through the teacher network to obtain a teacher representation, the augmented versions with masked patch regions are passed through the student network to predict the masked regions in the teacher representation, the student network being trained over multiple pretraining iterations to arrive at a loss function used for weighting the teacher network; andan output module to output weights of the teacher network as a pretrained machine learning model for determining feature representations of previously unseen three-dimensional medical imaging data.

12. The system of claim 11, wherein performing image augmentations on the three- dimensional medical imaging data comprises performing random image augmentations.

13. The system of claim 11, wherein the augmented versions comprise global crops capturing a whole of the original volume of the imaging data.

14. The system of claim 13, wherein pretraining of the machine learning model comprises jointly optimizes an image-level representation objective on the patched global crops and a patch-level representation objective on local crops of the augmented versions.

15. The system of claim 11, wherein masking patch regions is performed using a binary mask.

16. The system of claim 11, wherein the teacher network is an exponential moving average of the student network during pretraining.

17. The system of claim 11, wherein the teacher network comprises a vision transformer encoder model.

18. The system of claim 11, wherein generating augmented versions comprises applying three-dimensional data augmentations modeling variability in medical image acquisition.

19. The system of claim 11, wherein the image augmentations comprise, at least, applying a two-dimensional crop in in-plane dimensions and independently sampling a depth dimension.

20. The system of claim 11, wherein the three-dimensional medical imaging data comprises multiple co-registered imaging channels.

Images