Self-supervised learning for artificial intelligence-based systems for medical imaging analysis
Systems and methods for training an artificial intelligence-based system using self-supervised learning are provided. For each respective training medical image of a set of unannotated training medical images, the following steps are performed. A first augmented image is generated by applying a first augmentation operation to the respective training medical image. A second augmented image is generated by applying a second augmentation operation to the respective training medical image. A first representation vector is created from the first augmented image using an encoder network. A second representation vector is created from the second augmented image using the encoder network. The first representation vector is mapped to first cluster codes. The second representation vector is mapped to second cluster codes. The encoder network is optimized using the first and second representation vectors and the first and second cluster codes.
1. A computer-implemented method comprising:
for each respective training medical image of a set of unannotated training medical images:
generating a first augmented image by applying a first augmentation operation to the respective training medical image,
generating a second augmented image by applying a second augmentation operation to the respective training medical image,
creating a first representation vector from the first augmented image using an encoder network,
creating a second representation vector from the second augmented image using the encoder network,
mapping the first representation vector to first cluster codes, and
mapping the second representation vector to second cluster codes; and
optimizing the encoder network by 1) calculating a first loss based on the first representation vectors and the second cluster codes and 2) calculating a second loss based on the second representation vectors and the first cluster codes.
2. The computer-implemented method of claim 1 , wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
optimizing the encoder network to find cluster codes that maximize a similarity between representation vectors and cluster prototypes.
3. The computer-implemented method of claim 1 , wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
optimizing the encoder network further by combining the first loss and the second loss.
4. The computer-implemented method of claim 1 , wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
optimizing a loss function by clustering based on a modality of each of the training medical images.
5. The computer-implemented method of claim 1 , wherein one of the first augmentation operation or the second augmentation operation comprises an energy-based augmentation performed by:
decomposing an image into a plurality of energy bands;
computing an energy value for each of the plurality of energy bands; and
transforming the image based on the energy values.
6. The computer-implemented method of claim 1 , wherein the set of unannotated training medical images comprises a plurality of subsets of training medical images each associated with different sites, and wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
continuously optimizing the encoder network with the plurality of subsets of training medical images using federated learning and continual learning.
7. The computer-implemented method of claim 1 , further comprising
receiving annotated training data comprising annotated training medical images and corresponding annotations; and
mapping the annotated training medical images to predicted annotations using a machine learning based network based on features of the encoder network,
wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises optimizing the encoder network to minimize a distance of the predicted annotations to the corresponding annotations.
8. The computer-implemented method of claim 1 , further comprising:
receiving an input medical image;
creating a representation vector from the input medical image using the optimized encoder network;
performing a medical imaging analysis task for the input medical image based on the representation vector; and
outputting results of the medical imaging analysis task.
9. The computer-implemented method of claim 8 , wherein performing a medical imaging analysis task for the input medical image based on the representation vector comprises:
decoding the representation vector using a decoder network to perform the medical imaging analysis task.
10. An apparatus comprising:
for each respective training medical image of a set of unannotated training medical images:
means for generating a first augmented image by applying a first augmentation operation to the respective training medical image,
means for generating a second augmented image by applying a second augmentation operation to the respective training medical image,
means for creating a first representation vector from the first augmented image using an encoder network,
means for creating a second representation vector from the second augmented image using the encoder network,
means for mapping the first representation vector to first cluster codes, and
means for mapping the second representation vector to second cluster codes; and
means for optimizing the encoder network by 1) calculating a first loss based on the first representation vectors and the second cluster codes and 2) calculating a second loss based on the second representation vectors and the first cluster codes according to a contrastive loss function using the first and second representation vectors and the first and second cluster codes.
11. The apparatus of claim 10 , wherein the means for optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
means for optimizing the encoder network to find cluster codes that maximize a similarity between representation vectors and cluster prototypes.
12. The apparatus of claim 10 , wherein the means for optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
means for optimizing the encoder network further by combining the first loss and the second loss.
13. The apparatus of claim 10 , wherein the means for optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
means for optimizing a loss function by clustering based on a modality of each of the training medical images.
14. The apparatus of claim 10 , wherein one of the first augmentation operation or the second augmentation operation comprises an energy-based augmentation performed by:
means for decomposing an image into a plurality of energy bands;
means for computing an energy value for each of the plurality of energy bands; and
means for transforming the image based on the energy values.
15. A non-transitory computer readable medium storing computer program instructions, the computer program instructions when executed by a processor cause the processor to perform operations comprising:
for each respective training medical image of a set of unannotated training medical images:
generating a first augmented image by applying a first augmentation operation to the respective training medical image,
generating a second augmented image by applying a second augmentation operation to the respective training medical image,
creating a first representation vector from the first augmented image using an encoder network,
creating a second representation vector from the second augmented image using the encoder network,
mapping the first representation vector to first cluster codes, and
mapping the second representation vector to second cluster codes; and
optimizing the encoder network by 1) calculating a first loss based on the first representation vectors and the second cluster codes and 2) calculating a second loss based on the second representation vectors and the first cluster codes.
16. The non-transitory computer readable medium of claim 15 , wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
optimizing the encoder network to find cluster codes that maximize a similarity between representation vectors and cluster prototypes.
17. The non-transitory computer readable medium of claim 15 , wherein the set of unannotated training medical images comprises a plurality of subsets of training medical images each associated with different sites, and wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises:
continuously optimizing the encoder network with the plurality of subsets of training medical images using federated learning and continual learning.
18. The non-transitory computer readable medium of claim 15 , the operations further comprising
receiving annotated training data comprising annotated training medical images and corresponding annotations; and
mapping the annotated training medical images to predicted annotations using a machine learning based network based on features of the encoder network,
wherein optimizing the encoder network by 1) calculating the first loss based on the first representation vectors and the second cluster codes and 2) calculating the second loss based on the second representation vectors and the first cluster codes comprises optimizing the encoder network to minimize a distance of the predicted annotations to the corresponding annotations.
19. The non-transitory computer readable medium of claim 15 , the operations further comprising:
receiving an input medical image;
creating a representation vector from the input medical image using the optimized encoder network;
performing a medical imaging analysis task for the input medical image based on the representation vector; and
outputting results of the medical imaging analysis task.
20. The non-transitory computer readable medium of claim 19 , wherein performing a medical imaging analysis task for the input medical image based on the representation vector comprises:
decoding the representation vector using a decoder network to perform the medical imaging analysis task.