Structure-enhanced Multimodal Medical Concept Quantization for Medication Recommendation (SCQMed)

This repository implements a Structure-enhanced Multimodal Medical Concept Quantization for Medication Recommendation (SCQMed) for jointly learning discrete representations from two aligned embedding modalities (e.g., knowledge graph embeddings and text embeddings) in healthcare data. The model quantizes each modality independently using a shared residual vector quantization (RQ) structure and aligns them via contrastive learning. In this codebase, Modality 1 = text and Modality 2 = KG.

---

📁 Project Structure

.
├── mm/                    # Model package directory
│   ├── datasets_mm.py     # Custom dataset loader for paired embeddings
│   ├── layers.py          # MLP layers, activation functions, K-means, Sinkhorn
│   ├── rq_mm.py           # Residual Vector Quantizer with uniqueness loss
│   ├── vq_mm.py           # Single-layer VectorQuantizer with constrained K-means init & Sinkhorn
│   ├── rqvae_mm.py        # Main SE-RQVAE model definition
│   ├── trainer_new_mm.py  # Training loop, validation, checkpointing
│   └── utils.py           # Helper utilities (logging, dir creation, etc.)
├── main_new_MM.py         # Entry point: argument parsing & training setup
├── scripts/               # Pre-configured training scripts
└── data/                  # Expected data directory layout

---

🚀 Quick Start: Training

1. Run Training

Use one of the provided shell scripts from the scripts/ folder:

# For MIMIC-III
bash scripts/train_diag_rqvae.sh
bash scripts/train_proc_rqvae.sh
bash scripts/train_med_rqvae.sh

# For MIMIC-IV
bash scripts/train_diag_rqvae_m4.sh
bash scripts/train_proc_rqvae_m4.sh
bash scripts/train_med_rqvae_m4.sh

🔧 Key Arguments:

• --num_emb_list: Codebook sizes per RQ layer (e.g., [256,256,256,256])
• --e_dim: Embedding dimension for each codebook vector
• --layers: Hidden dimensions of encoder/decoder MLPs
• --recon: Weight for reconstruction loss
• --align: Weight for InfoNCE alignment loss between modalities

---

⬇️ Downstream Task: Medication Recommendation

This project includes a downstream medication recommendation framework that leverages the quantized multi-modal embeddings (or semantic IDs) generated by the SE-RQVAE. The core model is a Gated Recurrent Unit (GRU) enhanced with DDI (Drug-Drug Interaction) awareness and modality alignment.

📁 Downstream Directory Structure

The downstream/ folder contains the following key files:

downstream/
├── main.py                # Entry point for the downstream task
├── config.py              # Configuration for data paths, model hyperparameters, etc.
├── utils.py               # Utilities for data loading, metrics, and preprocessing
└── trainer.py             # Trainer class for downstream model

💡 Note: The downstream code expects preprocessed MIMIC-III or MIMIC-IV datasets in a specific format (see Data Preparation below).

---

📥 Data Preparation

Your downstream dataset directory (e.g., ../downstream_dataset/mimic-iv_data/) should contain the following files:

• records_final.pkl: A list of patient records. Each record is a list of admissions, and each admission is a tuple (diagnoses, procedures, medications) where each element is a list of integer codes.
• voc_final.pkl: A dictionary containing three Vocab objects: {"diag_voc": ..., "pro_voc": ..., "med_voc": ...} for mapping between codes and indices.
• ddi_A_final.pkl: A Drug-Drug Interaction adjacency matrix (NumPy array or sparse matrix).

You can place your SE-RQVAE output files (e.g., diag_fused_quantized_best.pt) in a subfolder like mmrq_embedding/ within this directory.

---

▶️ Running the Downstream Task

1. Navigate to the downstream/ directory.

   cd downstream

2. Ensure your config.py is correctly set, especially:

   • The ROOT and ontology_ROOT paths.
   • The paths to your SE-RQVAE outputs (HIDVAE_DIAG_EMB, etc.).
   • The correct TASK ('MIII' or 'MIV').

3. Run the main script.

   python main.py

The script will:

• Load the patient records, vocabulary, and DDI matrix.
• Load your pre-trained SE-RQVAE embeddings or SIDs.
• Initialize and train the downstream model.
• Evaluate the model on the test set, reporting metrics like Jaccard, PRAUC, DDI Rate, and Avg. # of Meds.
• Save model checkpoints and training history in the ./log/ directory.