Stanford CS25: V5 I Multimodal World Models for Drug Discovery, Eshed Margalit of Noetik.ai

Key Concepts

• Multimodal Learning: Integrating information from multiple data modalities (e.g., images, text, audio) to build more robust world models.
• Translation vs. Disambiguation: Two motivations for multimodal learning. Translation aims to capture all information from one modality in another, while disambiguation uses multiple modalities to resolve ambiguities.
• Early vs. Late Fusion: Approaches to multimodal learning. Early fusion combines modalities early in the processing pipeline, while late fusion processes them separately and combines them later.
• Masked Autoencoders: A type of neural network architecture used for learning representations from sparse data by masking out parts of the input and predicting the masked values.
• Spatial Transcriptomics: A technique for measuring gene expression levels while preserving spatial information within a tissue sample.
• Adaptive Layer Normalization: A method for incorporating conditioning information from one modality into another by modulating the parameters of layer normalization.
• Tumor Microenvironment: The cellular environment surrounding a tumor, including immune cells, blood vessels, and other factors that can influence tumor growth and response to therapy.
• Counterfactual Simulations: Using a model to simulate the effects of interventions or perturbations on a system, such as the effect of a drug on a tumor.

Multimodal Model Madness

• Unifying Goal of AI: Building world models that can simulate the future state of the world based on current observations and actions.
• Multimodality in World Models: The world is perceived through multiple modalities (e.g., visual, audio, text), and the best world models can integrate these modalities to make accurate predictions.
• Translation: Capturing all information from one modality in another (e.g., text-to-image systems).
• Disambiguation: Using multiple modalities to resolve ambiguities that cannot be resolved with a single modality (e.g., combining visual and auditory information to determine why people are running out of a building).
• Approaches to Incorporating Multimodal Streams:
  • Learning Joint Embedding Spaces: Using contrastive learning to create a single representation space where different modalities are mapped to similar points (e.g., CLIP, ImageBind).
    • CLIP uses image and text encoders to push corresponding image-text pairs into similar embeddings.
    • ImageBind expands on CLIP by using images as anchors and learning pairwise contrastive relationships between images and other modalities like text, depth maps, and audio.
  • Directly Concatenating Raw Inputs: Stacking raw inputs from different modalities as additional channels (e.g., concatenating a depth image as a fourth channel to an RGB image).
  • Cross Attention: Using attention mechanisms to allow different modalities to interact with each other (e.g., using image tokens to attend to text tokens and vice versa).
    • In cross-attention, one modality generates the queries, while the other provides the keys and values.
  • Concatenating Tokens: Converting inputs from different modalities into tokens and concatenating them into a single sequence (e.g., prepending a class token to a vision transformer, concatenating text and image tokens in DALL-E 1).
  • Adaptive Layer Normalization: Using conditioning information from one modality to modulate the parameters of layer normalization in another modality (e.g., using a linear layer to predict alpha, beta, and gamma parameters for layer norm based on a text label).
    • This approach is parameter-efficient and effective for guiding the processing of the primary sensory stream.

Cancer Biology and Multimodal Learning

• Goal: To build a world model of cancer biology that can simulate the effects of different treatments on a patient's tumor.
• Cancer Immunotherapy: Reactivating the immune system to detect and destroy cancer cells.
• Data Collection:
  • Human lung tumor specimens and other tumors from various organs.
  • Tissue cores (1 mm diameter) are extracted from the tissue blocks.
  • Four data processing pipelines:
    • H&E Images: Relatively cheap and easy to acquire images of tissue morphology.
    • 16-Plex Protein Immunofluorescence Panel: Detecting 16 different proteins using fluorescent antibodies.
    • Spatial Transcriptomics: Detecting thousands of RNA transcripts with spatial localization (XY coordinates).
    • Genetic Sequencing: Whole exome sequencing of the patient's data.
• Spatial Transcriptomics Data:
  • Detecting thousands of genes worth of RNA transcripts with spatial localization.
  • Data is represented as XY coordinates with gene expression levels.
  • Data is very rich, complicated, and expensive to collect.
• Modeling Approach:
  • Masked Autoencoders: Using a transformer backbone to learn representations from sparse gene expression data.
  • Tokenization: Assigning a token to each combination of gene identity and expression level.
  • Partial Masking: Removing a large percentage (e.g., >90%) of the tokens and predicting the masked values.
  • Spatial Context: Incorporating information about the local spatial neighborhood of a cell by forcing the information about the eight nearest neighbors through a bottleneck.
  • Adaptive Layer Normalization: Using the bottlenecked spatial context as an additional input to the backbone through adaptive layer normalization.
• Inference:
  • Taking a real patient sample and picking a location in that sample.
  • Looking up the eight nearest neighbors and feeding their expression into the model.
  • Asking the model to predict the other genes conditioned on the spatial context.
  • Running large-scale inference to simulate what is happening in patient tissue.
• Counterfactual Simulations:
  • Creating synthetic neighborhoods by knocking down a particular target and simulating the effect on the surrounding cells.
  • Using the model to predict the effect of the knockout on the expression of other genes.
• Multimodality as Translation:
  • Using the model to translate from H&E images to spatial transcriptomics data.
  • Passing in the surrounding H&E morphology as an additional input to the model.
  • Training the model to predict gene expression based on the H&E morphology.
  • Using the model to predict the expression of tumor markers from H&E images.
• Protein Image-Based Models:
  • Using masked autoencoders to predict protein images.
  • Feeding in other modalities (H&E, spatial transcriptomics, genetic sequencing) as extra tokens.
  • Simulating the effect of a drug by modifying the spatial transcriptomics data and predicting the effect on the protein image.
• Interpretability Research:
  • Using sparse autoencoders to pull out consistent beans that keep popping out in the data.
  • Building interfaces for biologists to do automatic semantic segmentation with text labels.

Work in Progress

• Moving Toward Raw Data: Training models that can look at the local spatial context in the raw point cloud of RNA transcripts and predict what is going to happen at many points spatially in the sample.
• Breaking Free of Cells: Moving away from the need to think about cells as the atomic unit of simulating the biology.

Conclusion

The presentation highlights the exciting potential of multimodal machine learning, particularly with transformers, in the field of cancer research. By integrating data from various modalities, such as imaging, genomics, and spatial transcriptomics, it's possible to build more robust world models that can simulate the effects of different treatments on a patient's tumor. The work at Noetic demonstrates the power of this approach for drug discovery and personalized medicine. The speaker emphasizes that there are many opportunities for creative work in this area, and the future looks promising for using AI to improve cancer treatment.