CoolFace
Modelpublic

ddidacus/RS-mol-llama-1b

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes24downloads
Model Card

*Mol-MoE*: Training Preference-Guided Routers for Molecule Generation

Diego Calanzone (1, 2), Pierluca D'Oro (2), Pierre-Luc Bacon (1, 2) <br> (1) Universite de Montreal, (2) Mila Quebec AI Institute <br> arXiv: https://arxiv.org/abs/2502.05633

Abstract: Recent advances in language models have enabled framing molecule generation as sequence modeling. However, existing approaches often rely on single-objective reinforcement learning, limiting their applicability to real-world drug design, where multiple competing properties must be optimized. Traditional multi-objective reinforcement learning (MORL) methods require costly retraining for each new objective combination, making rapid exploration of trade-offs impractical. To overcome these limitations, we introduce Mol-MoE, a mixture-of-experts (MoE) architecture that enables efficient test-time steering of molecule generation without retraining. Central to our approach is a preference-based router training objective that incentivizes the router to combine experts in a way that aligns with user-specified trade-offs. This provides improved flexibility in exploring the chemical property space at test time, facilitating rapid trade-off exploration. Benchmarking against state-of-the-art methods, we show that Mol-MoE achieves superior sample quality and steerability.

How to use this model

This LM is fine-tuned to generate molecules in the SMILES format wrt. desired properties. For unconditioned SMILES generation, use the BOS token <s>. <br> For conditioned generation, please refer to the paper and the official codebase to derive different conditioned models. <br> This model is the merging result of 5 fine-tuned versions (JNK3, DRD2, GSK3B, CYP2D6, CYP2D19) with equal interpolation weight: w_i = 0.2.

An example of the generation pipeline:

from transformers import AutoTokenizer, AutoModelForCausalLM
import re

# Setup
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained("ddidacus/RS-mol-llama-1b")
model = AutoModelForCausalLM.from_pretrained("ddidacus/RS-mol-llama-1b")
generation_kwargs = {
    "max_new_tokens": 128,
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 0.9,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "temperature": 1.0
}

# Inference
query = "<s>"
toks = tokenizer([query], return_tensors="pt")["input_ids"].to(device)
output = model.generate(toks, **generation_kwargs)
output = tokenizer.batch_decode(output)

# Parsing
filter = r'<s>(.*?)</s>'
molecule = re.findall(filter, output[0], re.DOTALL)

Model Description

This model is a fine-tuned version of LLaMa 3.2 1B through two stages:

  1. 1.Fine-tuning on ~3.5M molecules extracted from: ZINC 250K, MOSES, CHEMBL
  2. 2.RLHF-tuning using RLOO on 5 distinct reward functions from PyTDC [1]
  • —Developed by: Diego Calanzone (diego.calanzone@mila.quebec)
  • —Model type: Decoder-only Transformer
  • —Finetuned from model [optional]: LLaMA 3.2 1B

Read the paper for further details.

Sources

[1] https://tdcommons.ai/singlepredtasks/overview

<!--

Model Sources [optional]

  • —Repository: [More Information Needed]
  • —Paper [optional]: [More Information Needed]
  • —Demo [optional]: [More Information Needed]

Uses

Direct Use

[More Information Needed]

Downstream Use [optional]

[More Information Needed]

Out-of-Scope Use

[More Information Needed]

Bias, Risks, and Limitations

[More Information Needed]

Recommendations

Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.

How to Get Started with the Model

Use the code below to get started with the model.

[More Information Needed]

Training Details

Training Data

[More Information Needed]

Training Procedure

Preprocessing [optional]

[More Information Needed]

Training Hyperparameters
  • —Training regime: [More Information Needed]
Speeds, Sizes, Times [optional]

[More Information Needed]

Evaluation

Testing Data, Factors & Metrics

Testing Data

[More Information Needed]

Factors

[More Information Needed]

Metrics

[More Information Needed]

Results

[More Information Needed]

Summary

Model Examination [optional]

[More Information Needed]

Environmental Impact

Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).

  • —Hardware Type: [More Information Needed]
  • —Hours used: [More Information Needed]
  • —Cloud Provider: [More Information Needed]
  • —Compute Region: [More Information Needed]
  • —Carbon Emitted: [More Information Needed]

Technical Specifications [optional]

Model Architecture and Objective

[More Information Needed]

Compute Infrastructure

[More Information Needed]

Hardware

[More Information Needed]

Software

[More Information Needed]

Citation [optional]

BibTeX:

[More Information Needed]

APA:

[More Information Needed]

Glossary [optional]

[More Information Needed]

More Information [optional]

[More Information Needed]

Model Card Authors [optional]

[More Information Needed]

Model Card Contact

[More Information Needed] -->