Official repository of MELD (Masked Element-wise Learnable Diffusion) introduced in Learning Flexible Forward Trajectories for Masked Molecular Diffusion.
MELD is a masked diffusion model tailored for molecular generation that resolves the state-clashing problem through learnable forward trajectories, showing robust performance in both conditional and unconditional generation.
Install required packages:
pip install -r requirements_mcmg.txt
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
pip install torch-sparse==0.6.18 torch-scatter==2.1.2 -f https://data.pyg.org/whl/torch-2.1.0+cu121.htmlIf you face errors when installing mini-moses, try this command:
GIT_LFS_SKIP_SMUDGE=1 pip install git+https://github.com/igor-krawczuk/mini-moses@1eda77fcc13045a93024835d403ff61634ac4b69Due to storage limit, we provide preprocessed datasets and pretrained oracles via Google Drive.
Please download the data folder and place it in the main MELD directory.
To train MELD on unconditional molecular generation (e.g., ZINC or QM9):
# task_name: 'ZINC', 'QM9'
python src/main.py --config-name=config.yaml \
model.schedule_degree=2 \
model.diffusion_noise_schedule='meld-schedule' \
model.ensure_connected=True \
dataset.task_name={target} \
general.wandb='enabled' \
general.wandb_key={wandb_key} \
general.exp_name={exp_name} \
general.save_model=True \
train.batch_size={bs} \
train.n_epochs={epochs} \
train.lr={lr} \
model.use_provided_drop_ids=True \
model.diffusion_steps={steps} \
model.unconditioned=True \
train.use_ema=True
To train MELD on multi-conditional molecular generation (e.g., polymer dataset O2-N2-CO2):
# task_name: 'O2-N2-CO2' (polymer dataset)
python src/main.py --config-name=config.yaml \
model.schedule_degree=2 \
model.diffusion_noise_schedule='meld-property-schedule' \
model.ensure_connected=True \
dataset.task_name={target} \
general.wandb='enabled' \
general.wandb_key={wandb_key} \
general.exp_name={exp_name} \
general.save_model=True \
train.batch_size={bs} \
train.n_epochs={epochs} \
train.lr={lr} \
model.use_provided_drop_ids=True \
model.diffusion_steps={steps}
🚀 Visit our project page for more examples and visualizations.
If you use our code or model, please cite our work:
@misc{seo2025learningflexibleforwardtrajectories,
title={Learning Flexible Forward Trajectories for Masked Molecular Diffusion},
author={Hyunjin Seo and Taewon Kim and Sihyun Yu and SungSoo Ahn},
year={2025},
eprint={2505.16790},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2505.16790}
}