Generating Joint Transcriptomic and Morphological Responses to Drug Perturbations via Rectified Flow

This paper introduces PertFlow, a unified computational framework that simultaneously predicts drug-induced transcriptomic profiles and generates cellular morphology images by leveraging rectified flow to capture complex molecular-phenotypic dependencies across multiple cell lines and compounds.

The Gist

Imagine you are a chef trying to predict exactly how a new spice will change a soup.

In the world of drug discovery, scientists face a similar challenge. They want to know: If we add this specific drug to a cell, how will the cell's internal chemistry change, and how will the cell itself look different?

For a long time, scientists had to guess these two things separately. They had one tool to predict the chemical changes (the "recipe") and another to guess the visual changes (the "appearance"). But cells are complex; the chemistry causes the appearance. If you ignore one, you miss the full story.

Enter PertFlow, a new AI system created by researchers at Purdue University. Think of PertFlow as a super-smart "Time-Traveling Chef" that can predict both the recipe and the final dish simultaneously.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Split Screen" Issue

Imagine you are watching a movie on a split screen. On the left, you see the script (the genes talking to each other). On the right, you see the actors' faces (the cell's shape).

• Old AI models were like watching only the script or only the faces. They couldn't see how a line of dialogue (a gene turning on) made the actor cry (the cell changing shape).
• The Goal: We need a model that watches the whole movie at once, understanding that the script drives the acting.

2. The Solution: The "Universal Translator"

PertFlow is built to look at a "Control" cell (a healthy cell doing nothing) and a "Drug" (the spice). It then predicts what that cell will look like and what its genes will say after the drug is added.

It uses three main tricks to do this:

• The Shared Language (Cross-Modal Attention):
  Imagine the cell's genes are speaking English and the cell's shape is speaking French. They are in different rooms and can't talk. PertFlow builds a universal translator in the middle. It forces the "Gene Speaker" and the "Shape Speaker" to sit at the same table and learn each other's language. This way, when the drug comes in, the AI understands that a change in the "Gene language" must result in a specific change in the "Shape language."

• The "Straight-Line" Shortcut (Rectified Flow):
  Usually, AI models that generate images (like creating a new picture from scratch) are like a drunk person walking home. They take a winding, zig-zag path, trying random steps until they get it right. This is slow and can get messy.
  PertFlow uses a technique called Rectified Flow. Imagine instead of a drunk walk, the AI draws a straight, high-speed highway from the "Healthy Cell" to the "Sick Cell." It knows the most direct route, making the prediction faster and much more accurate.

• The Knowledge Library (PrimeKG):
  The AI doesn't just guess; it reads a massive encyclopedia of biology. It knows that "Drug A usually breaks microtubules" (like snapping the scaffolding of a building). By feeding this real-world knowledge into the model, it ensures the predictions make biological sense, not just mathematical sense.

3. The Results: A Crystal Ball for Cells

The researchers tested PertFlow on thousands of cells and 40 different drugs. Here is what happened:

• The "Recipe" Prediction: It predicted the gene changes with 78% accuracy. That's like guessing the exact ingredients in a soup just by smelling the broth.
• The "Picture" Generation: It generated images of the cells after treatment that were so realistic, board-certified pathologists (human experts) gave them a score of 7.1 to 7.9 out of 10. They couldn't easily tell the AI-generated cells from real ones!
• The "Aha!" Moments: The AI correctly figured out that certain drugs would cause cells to stop dividing (like a construction site shutting down) or that others would damage DNA. It even recreated the specific "look" of a cell dying, matching real-world observations perfectly.

Why Does This Matter?

Think of drug discovery as trying to find a key that fits a lock.

• Before: Scientists had to try thousands of keys, break the lock, and then check if the door opened. It was slow, expensive, and destructive.
• With PertFlow: Scientists can now simulate the lock picking in a virtual world. They can ask, "What happens if we try this key?" and the AI instantly shows them the chemical reaction and the visual result.

This allows scientists to:

1. Screen drugs faster: Test thousands of ideas in a computer before ever touching a petri dish.
2. Understand mechanisms: See why a drug works (or fails) by watching the "movie" of the cell's reaction.
3. Personalize medicine: Imagine predicting how your specific cells will react to a drug before you even take it.

The Bottom Line

PertFlow is a bridge between the invisible world of genes and the visible world of cell shapes. By using a "straight-line" AI approach and forcing these two worlds to talk to each other, it gives us a powerful new crystal ball to see how drugs will change our bodies, potentially speeding up the discovery of life-saving medicines.

Technical Summary

1. Problem Formulation

The central challenge addressed by this work is the lack of computational methods capable of jointly predicting both transcriptomic (gene expression) and morphological (cellular image) responses to chemical drug perturbations.

• Current Limitations: Existing models typically treat these modalities in isolation. Transcriptomic models (e.g., scGen, CPA, GEARS) predict gene expression but ignore morphological changes. Imaging models (e.g., PhenDiff) predict cellular morphology but lack transcriptomic context. Cross-modal frameworks exist but focus on measurement or translation rather than generative synthesis of both modalities simultaneously from a control state.
• The Gap: There is no unified framework that takes a control cellular state (RNA-seq + Image) and drug metadata (compound, dose, time, cell line) to simultaneously generate the resulting treatment gene expression profile and the corresponding cellular morphology image.
• Key Challenges:
  1. Aligning data across fundamentally different representational spaces (discrete high-dimensional vectors vs. continuous 2D/3D images).
  2. Capturing complex conditioning involving compound structure, dosage, cell type, and time.
  3. Ensuring biological realism and cross-modal consistency (i.e., the generated image must match the generated gene expression).

2. Methodology: PertFlow Architecture

The authors introduce PertFlow, a unified generative framework based on Rectified Flow and Cross-Modal Attention.

A. Data and Preprocessing

• Dataset: Utilized the Ginkgo Data Platform (GDP), comprising 17,242 paired samples from 3 cell lines and 40 compounds.
• Inputs: Control RNA-seq ($x_{ctrl}^{rna}$), Control Cell Painting images ($x_{ctrl}^{img}$), and Drug Conditioning ($c$: compound ID, cell line, concentration, timepoint).
• Drug Encoding: Compounds are represented via Morgan/RDKit fingerprints (structural) and physicochemical descriptors (2D/3D properties).
• Preprocessing: RNA-seq is normalized and log-transformed; images are scaled and interpolated to uniform dimensions.

B. Model Architecture

PertFlow employs a shared representation learning paradigm with three main components:

1. Modality Encoders:
   • RNA Encoder: Uses multi-layer self-attention to capture gene-gene interactions.
   • Image Encoder: A ResNet-style CNN with global pooling to extract hierarchical visual features.
   • Knowledge Graph Integration: Both modalities are enhanced by integrating PrimeKG (a biological knowledge graph) via Graph Neural Networks (GNNs) to incorporate protein-protein interactions and drug-target relationships.
2. Cross-Modal Attention Mechanism:
   • Encoded tokens from RNA and Image modalities undergo self-attention followed by bidirectional cross-attention. This allows RNA tokens to attend to image features and vice versa, aligning the two modalities in a shared latent space.
   • Drug conditioning is fused into this space via a dedicated encoder.
3. Generation Heads:
   • Transcriptomic Head: A fully connected regression head that directly predicts the treatment RNA-seq profile from the shared embedding ($h_{shared}$).
   • Morphological Head: Adapts Rectified Flow dynamics. Instead of standard diffusion, it learns a straight-line transport path between a noise distribution ($z_0$) and the target treatment image ($x_{treat}^{img}$). A UNet predicts the velocity field ($v = x_{treat}^{img} - z_0$) conditioned on the noisy state, timestep, and $h_{shared}$. An adaptive ODE solver (DOPRI5) is used for inference.

C. Loss Functions

The model is trained end-to-end with a composite loss function:

• Transcriptomic Loss ($L_{rna}$): A weighted sum of Mean Squared Error (MSE) for gene-level accuracy and Pearson Correlation loss to preserve relative expression patterns.
• Image Loss ($L_{img}$): Rectified flow objective minimizing the velocity prediction error.
• Triplet Contrastive Loss ($L_{triplet}$): Enforces consistency by penalizing misaligned RNA-image pairs and rewarding aligned pairs, ensuring the generated image corresponds to the generated transcriptome.

3. Key Contributions

1. First Joint Generative Framework: PertFlow is the first method to simultaneously synthesize treatment gene expression and cellular morphology from control states and drug metadata.
2. Rectified Flow for Biology: Successfully adapts Rectified Flow (known for computational efficiency and straight-line transport) to the complex task of generating biological images conditioned on multi-modal data.
3. Cross-Modal Consistency: Introduces a triplet contrastive loss and bidirectional attention to ensure that the generated transcriptomic and morphological outputs are biologically coherent with each other.
4. Knowledge Graph Integration: Effectively integrates external biological knowledge (PrimeKG) to improve generalization and biological plausibility.

4. Experimental Results

The model was evaluated on the GDPx3 dataset (3 cell lines, 40 drugs).

A. Transcriptomic Performance

• Metrics: Achieved a Pearson correlation of 0.780 ± 0.264 and Spearman correlation of 0.792.
• Comparison: Outperformed single-modality baselines (MLP, VAE-PRNet) and ablation studies. The joint model maintained competitive RNA prediction accuracy while adding image generation capabilities.

B. Morphological Performance

• Metrics: Achieved a Fréchet Inception Distance (FID) of 24.06, significantly outperforming diffusion baselines (e.g., PertDiff variants had FIDs >50) and the PhenDiff baseline (FID 62.50).
• Visual Quality: Generated images showed high Structural Similarity (SSIM: 0.205) and Peak Signal-to-Noise Ratio (PSNR: 11.66).
• Expert Evaluation: Board-certified pathologists rated the generated images with median similarity scores of 7.11–7.89/10, outperforming the PhenDiff baseline (6.11–6.33).

C. Biological Validity

• Mechanism Recovery: The model successfully recovered known drug mechanisms, including microtubule disruption (Cevipabulin, Podophyllotoxin), DNA damage (Camptothecin), and MAPK pathway inhibition (Dabrafenib).
• Gene Enrichment: Gradient-based feature importance analysis confirmed activation of expected pathways (e.g., p53/apoptosis for DNA damage, EMT for stress responses) consistent with the cell line and drug type.
• Embedding Analysis: UMAP and t-SNE visualizations demonstrated that cross-modal attention successfully aligned RNA and image features into a unified latent space where biologically similar samples clustered together regardless of modality.

5. Significance and Impact

• Virtual Drug Screening: PertFlow enables the in silico prediction of both molecular mechanisms and phenotypic outcomes for new compounds, potentially accelerating drug discovery by reducing reliance on wet-lab experiments.
• Mechanistic Insight: By generating paired data, the model helps researchers understand the causal link between gene regulation and cellular morphology under specific perturbations.
• Foundation for Multi-Modal AI: This work establishes a new benchmark for joint generative modeling in biology, demonstrating that rectified flow and cross-modal attention can effectively bridge the gap between discrete genomic data and continuous imaging data.
• Future Directions: The authors note that while promising, generalization to truly novel compounds and unseen cell lines remains a challenge, suggesting future integration of more granular chemogenomic databases and modeling of tissue-level interactions.

In summary, PertFlow represents a significant advancement in computational biology, offering a unified, biologically consistent framework for predicting the dual molecular and phenotypic effects of drug perturbations.