Machine learning systems for processing multi-modal patient data

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for classifying a patient. In one aspect, a method comprises: receiving multi-modal data characterizing a patient, wherein the multi-modal data comprises a respective feature representation for each of a plurality of modalities; processing the multi-modal data characterizing the patient using an encoder neural network to generate an embedding of the multi-modal data characterizing the patient; determining a respective classification score for each patient category in a set of patient categories based on the embedding of the multi-modal data characterizing the patient; and classifying the patient as being included in a corresponding patient category from the set of patient categories based on the classification scores.

1 . (canceled) 
     
     
         2 . A method comprising:
 jointly training an encoder neural network having a set of encoder parameters and a decoder neural network having a set of decoder parameters, comprising, at each of a plurality of training iterations:
 obtaining a batch of training examples, wherein each training example corresponds to a respective subject and includes multi-modal data that characterizes the subject; 
 processing the multi-modal data from each training example using the encoder neural network, in accordance with current values of the encoder parameters, to generate a respective embedding of the multi-modal data from each training example in a latent space; 
 processing the embedding of the multi-modal data from each training example using the decoder neural network, in accordance with current values of the decoder parameters, to generate a respective reconstruction of the multi-modal data from each training example; 
 clustering a set of embeddings into a plurality of clusters of embeddings, wherein each cluster of embeddings includes a plurality of embeddings from the set of embeddings, and wherein the set of embeddings includes the respective embedding of the multi-modal data from each training example; 
 determining a clustering loss based on the clustering of the set of embeddings into the plurality of clusters of embeddings; and 
 updating the current values of the set of encoder parameters and the current values of the set of decoder parameters using gradients of an objective function that depends on: (i) a respective error in the reconstruction of the multi-modal data from each training example, and (ii) the clustering loss. 
   
     
     
         3 . The method of  claim 2 , wherein each embedding in the set of embeddings is associated with a cluster label that identifies a cluster that includes the embedding, and wherein determining the clustering loss based on the clustering of the set of embeddings into the plurality of clusters of embeddings comprises:
 designating a proper subset of the set of embeddings as being training embeddings;   training a classification machine learning model, comprising, for each training embedding, training the classification machine learning model to process the training embedding to predict the cluster label of the training embedding; and   after training the classification machine learning model, determining the clustering loss using the classification machine learning model.   
     
     
         4 . The method of  claim 3 , wherein determining the clustering loss using the classification machine learning model comprises:
 designating a proper subset of the set of embeddings as validation embeddings;   evaluating a classification accuracy of the classification machine learning model on a task of processing each validation embedding to predict the cluster label of the validation embedding; and   determining the clustering loss based on the classification accuracy of the classification machine learning model.   
     
     
         5 . The method of  claim 4 , wherein the set of validation embeddings excludes any training embeddings. 
     
     
         6 . The method of  claim 4 , wherein updating the current values of the set of encoder parameters using gradients of the objective function that depends on the clustering loss encourages an increase in the classification accuracy of the classification machine learning model. 
     
     
         7 . The method of  claim 3 , wherein the classification machine learning model comprises a neural network model. 
     
     
         8 . The method of  claim 2 , wherein each embedding in the set of embeddings is associated with: (i) a cluster label that identifies a cluster that includes the embedding, and (ii) a set of confounding features;
 wherein determining the clustering loss based on the clustering of the set of embeddings into the plurality of clusters of embeddings comprises:
 designating a proper subset of the set of embeddings as being training embeddings; 
 training a classification machine learning model, comprising, for each training embedding, training the classification machine learning model to process the set of confounding features corresponding to the training embedding to predict the cluster label of the training embedding; and 
 after training the classification machine learning model, determining the clustering loss using the classification machine learning model. 
   
     
     
         9 . The method of  claim 8 , wherein determining the clustering loss using the classification machine learning model comprises:
 designating a proper subset of the set of embeddings as validation embeddings;   evaluating a classification accuracy of the classification machine learning model on a task of processing the set of confounding features corresponding to each validation embedding to predict the cluster label of the validation embedding; and   determining the clustering loss based on the classification accuracy of the classification machine learning model.   
     
     
         10 . The method of  claim 8 , wherein the set of confounding features are designated as being substantially irrelevant to a medical condition. 
     
     
         11 . The method of  claim 8 , wherein the set of confounding features are designated as being substantially irrelevant to a treatment for a medical condition. 
     
     
         12 . The method of  claim 8 , wherein for each embedding, the set of confounding features are not included in multi-modal data processed by the encoder neural network to generate the embedding. 
     
     
         13 . The method of  claim 8 , wherein for each embedding, the corresponding set of confounding features comprise: features of a sensor that captured sensor data included in the multi-modal data processed by the encoder neural network to generate the embedding, or features of an acquisition protocol used to acquire a portion of the multi-modal data processed by the encoder neural network to generate the embedding, or both. 
     
     
         14 . The method of  claim 9 , wherein updating the current values of the set of encoder parameters using gradients of the objective function that depends on the clustering loss encourages a decrease in the classification accuracy of the classification machine learning model. 
     
     
         15 . The method of  claim 9 , wherein updating the current values of the set of encoder parameters using gradients of the objective function that depends on the clustering loss encourages confounding features corresponding to embeddings with different cluster labels to become more entangled in a confounding feature space. 
     
     
         16 . The method of  claim 2 , wherein clustering the set of embeddings into the plurality of clusters of embeddings comprises applying a k-means clustering operation to the set of embeddings. 
     
     
         17 . The method of  claim 2 , further comprising outputting the encoder neural network and the decoder neural network after the joint training of the encoder neural network and the decoder neural network. 
     
     
         18 . The method of  claim 2 , wherein for each training example, the multi-modal data included in the training example comprises a respective feature representation for each of a plurality of modalities. 
     
     
         19 . The method of  claim 18 , wherein for each training example, the plurality of modalities include a functional magnetic resonance imaging (fMRI) modality, wherein the feature representation for the fMRI modality is derived from a series of fMRI images that each correspond to a respective time point in a sequence of time points and characterize blood flow in a brain of the corresponding subject at the time point. 
     
     
         20 . A system comprising:
 one or more computers; and   one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations comprising:   jointly training an encoder neural network having a set of encoder parameters and a decoder neural network having a set of decoder parameters, comprising, at each of a plurality of training iterations:
 obtaining a batch of training examples, wherein each training example corresponds to a respective subject and includes multi-modal data that characterizes the subject; 
 processing the multi-modal data from each training example using the encoder neural network, in accordance with current values of the encoder parameters, to generate a respective embedding of the multi-modal data from each training example in a latent space; 
 processing the embedding of the multi-modal data from each training example using the decoder neural network, in accordance with current values of the decoder parameters, to generate a respective reconstruction of the multi-modal data from each training example; 
 clustering a set of embeddings into a plurality of clusters of embeddings, wherein each cluster of embeddings includes a plurality of embeddings from the set of embeddings, and wherein the set of embeddings includes the respective embedding of the multi-modal data from each training example; 
 determining a clustering loss based on the clustering of the set of embeddings into the plurality of clusters of embeddings; and 
 updating the current values of the set of encoder parameters and the current values of the set of decoder parameters using gradients of an objective function that depends on: (i) a respective error in the reconstruction of the multi-modal data from each training example, and (ii) the clustering loss. 
   
     
     
         21 . One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:
 jointly training an encoder neural network having a set of encoder parameters and a decoder neural network having a set of decoder parameters, comprising, at each of a plurality of training iterations:
 obtaining a batch of training examples, wherein each training example corresponds to a respective subject and includes multi-modal data that characterizes the subject; 
 processing the multi-modal data from each training example using the encoder neural network, in accordance with current values of the encoder parameters, to generate a respective embedding of the multi-modal data from each training example in a latent space; 
 processing the embedding of the multi-modal data from each training example using the decoder neural network, in accordance with current values of the decoder parameters, to generate a respective reconstruction of the multi-modal data from each training example; 
 clustering a set of embeddings into a plurality of clusters of embeddings, wherein each cluster of embeddings includes a plurality of embeddings from the set of embeddings, and wherein the set of embeddings includes the respective embedding of the multi-modal data from each training example; 
 determining a clustering loss based on the clustering of the set of embeddings into the plurality of clusters of embeddings; and 
 updating the current values of the set of encoder parameters and the current values of the set of decoder parameters using gradients of an objective function that depends on: (i) a respective error in the reconstruction of the multi-modal data from each training example, and (ii) the clustering loss.