Optimization Algorithm

The optimize() function implements a genetic-like optimization algorithm that iteratively generates and refines molecules based on oracle feedback.

The optimize() Function

Located in chemlactica/mol_opt/optimization.py, this is the main entry point:

from chemlactica.mol_opt.optimization import optimize

optimize(
    model,              # Pre-trained ChemLactica model
    tokenizer,          # Corresponding tokenizer
    oracle,             # Your custom oracle
    config,             # Hyperparameter configuration
    additional_properties={},  # Optional additional properties
    validate_smiles=lambda x: True  # Optional SMILES validation
)

Parameters

AutoModelForCausalLM

required

The ChemLactica model (125M, 1.3B, or Chemma-2B)

AutoTokenizer

required

Tokenizer corresponding to the model

Oracle

required

Your custom oracle implementing the oracle interface

dict

required

Hyperparameter configuration dictionary (from YAML)

dict

default:"{}"

Optional additional properties to include in prompts (e.g., molecular weight, logP)

callable

default:"lambda x: True"

Optional function to validate generated SMILES before scoring

Algorithm Workflow

The optimization process follows these steps:

Initialize Pool

Create an empty pool to store the top-performing molecules

pool = Pool(config["pool_size"], validation_perc=config["validation_perc"])

The pool maintains:

Top pool_size molecules sorted by score
Training/validation split for fine-tuning
Diversity filtering to avoid duplicates

Generation Loop

Generate num_gens_per_iter new molecules per iteration:a) Create prompts from pool molecules:

prompts = [
    optim_entry.to_prompt(
        is_generation=True,
        include_oracle_score=prev_train_iter != 0,
        config=config,
        max_score=max_score
    )
    for optim_entry in optim_entries
]

Each prompt includes:

Similar molecules with Tanimoto similarities
Desired oracle score (random between max_score and max_possible_oracle_score)
Additional properties if specified

b) Generate molecules:

output = model.generate(**data, **config["generation_config"])
output_texts = tokenizer.batch_decode(output)

c) Parse and validate:

molecule = create_molecule_entry(output_text, validate_smiles)

Extract SMILES between [START_SMILES] and [END_SMILES] tags

Oracle Evaluation

Score unique molecules with the oracle:

if oracle.takes_entry:
    oracle_scores = oracle([optim_entry.last_entry for smiles in unique_smiles])
else:
    oracle_scores = oracle(unique_smiles)

Track best score and log results:

if score > max_score:
    max_score = score
    new_best_molecule_generated = True

Update Pool

Add new molecules to the pool:

pool.add(list(iter_unique_optim_entries.values()))

The pool:

Sorts molecules by score (descending)
Removes duplicates and highly similar molecules
Keeps only top pool_size molecules
Maintains train/validation split

Adaptive Fine-tuning (Optional)

If using rej-sample-v2 strategy and no improvement for train_tol_level iterations:a) Prepare datasets:

train_entries, validation_entries = pool.get_train_valid_entries()

train_dataset = Dataset.from_dict({
    "sample": [
        entry.to_prompt(is_generation=False, include_oracle_score=True, ...)
        for entry in train_entries
    ]
})

b) Fine-tune model:

trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
    ...
)
trainer.train()
model.load_state_dict(best_model_state_dict)

The model learns to generate molecules that match the oracle’s preferences.

Adjust Temperature

Increase sampling temperature over time for exploration:

config["generation_config"]["temperature"] += 
    (temp_end - temp_start) / num_iterations

Starts at generation_temperature[0] and increases to generation_temperature[1]

Check Stopping Condition

Stop when oracle budget is exhausted:

if oracle.finish:
    break

Key Components

Pool Management

The Pool class (mol_opt/utils.py:88) maintains top molecules:

class Pool:
    def __init__(self, size, validation_perc: float):
        self.size = size  # Maximum pool size
        self.optim_entries: List[OptimEntry] = []
        self.num_validation_entries = int(size * validation_perc + 1)
    
    def add(self, entries: List, diversity_score=1.0):
        # Sort by score
        self.optim_entries.extend(entries)
        self.optim_entries.sort(key=lambda x: x.last_entry, reverse=True)
        
        # Remove duplicates and similar molecules
        new_optim_entries = []
        for entry in self.optim_entries:
            insert = True
            for e in new_optim_entries:
                # Skip if duplicate or too similar
                if (entry.last_entry == e.last_entry or
                    tanimoto_dist_func(entry.last_entry.fingerprint, 
                                     e.last_entry.fingerprint) > diversity_score):
                    insert = False
                    break
            if insert:
                new_optim_entries.append(entry)
        
        # Keep top molecules only
        self.optim_entries = new_optim_entries[:min(len(new_optim_entries), self.size)]

Prompt Construction

The OptimEntry.to_prompt() method creates prompts for generation: During generation (is_generation=True):

</s>[SIMILAR]CCO 0.65[/SIMILAR][SIMILAR]CC(C)O 0.58[/SIMILAR][PROPERTY]oracle_score 85.23[/PROPERTY][START_SMILES]

During training (is_generation=False):

</s>[SIMILAR]CCO 0.65[/SIMILAR][SIMILAR]CC(C)O 0.58[/SIMILAR][PROPERTY]oracle_score 85.23[/PROPERTY][START_SMILES]CC(O)CO[END_SMILES]</s>

Molecule Entry

The MoleculeEntry class represents a molecule:

class MoleculeEntry:
    def __init__(self, smiles, score=0, **kwargs):
        self.smiles = canonicalize(smiles)  # Canonical SMILES
        self.score = score                   # Oracle score
        self.mol = Chem.MolFromSmiles(smiles)  # RDKit mol
        self.fingerprint = get_morgan_fingerprint(self.mol)  # For similarity
        self.similar_mol_entries = []        # Similar mols in prompt
        self.add_props = kwargs              # Additional properties

Tolerance Mechanism

The algorithm tracks iterations without improvement:

tol_level = 0  # Tolerance counter

if new_best_molecule_generated:
    tol_level = 0  # Reset if we find a better molecule
else:
    tol_level += 1  # Increment if no improvement

# Trigger fine-tuning after train_tol_level iterations without improvement
if tol_level >= config["rej_sample_config"]["train_tol_level"]:
    # Fine-tune model...
    tol_level = 0  # Reset after fine-tuning

This adaptive mechanism ensures:

Continued exploration when making progress
Model adaptation when stuck
Efficient use of oracle budget

Generation Temperature Schedule

Temperature increases linearly to encourage exploration:

# Initial temperature
config["generation_config"]["temperature"] = config["generation_temperature"][0]

# Each iteration
config["generation_config"]["temperature"] += 
    config["num_gens_per_iter"] / (oracle.budget - config["num_gens_per_iter"]) * 
    (config["generation_temperature"][1] - config["generation_temperature"][0])

Why increase temperature?

Early: Low temperature (1.0) focuses on likely, high-quality molecules
Late: High temperature (1.5) explores diverse, creative solutions
This schedule balances exploitation and exploration

Diversity Filtering

The pool removes molecules that are too similar:

for entry in self.optim_entries:
    insert = True
    for existing in new_optim_entries:
        similarity = tanimoto_dist_func(
            entry.last_entry.fingerprint,
            existing.last_entry.fingerprint
        )
        if similarity > diversity_score:  # Default: 1.0 (only exact duplicates)
            insert = False
            break

By default, only exact duplicates are removed. You can increase diversity_score in the pool to filter more aggressively.

Algorithm Pseudocode

Here’s the complete algorithm in pseudocode:

INPUT: model, tokenizer, oracle, config
OUTPUT: Optimized molecules in oracle.mol_buffer

1. Initialize:
   pool ← empty Pool(pool_size)
   max_score ← 0
   tol_level ← 0
   temperature ← config.generation_temperature[0]

2. While not oracle.finish:
   
   a. Generate molecules:
      unique_molecules ← {}
      while len(unique_molecules) < num_gens_per_iter:
          prompts ← create_prompts_from_pool(pool)
          outputs ← model.generate(prompts, temperature)
          molecules ← parse_and_validate(outputs)
          unique_molecules.add(new molecules not in oracle.mol_buffer)
   
   b. Evaluate with oracle:
      scores ← oracle(unique_molecules)
      if max(scores) > max_score:
          max_score ← max(scores)
          tol_level ← 0
      else:
          tol_level += 1
   
   c. Update pool:
      pool.add(molecules with scores)
      pool ← top pool_size molecules (sorted, filtered for diversity)
   
   d. Fine-tune if needed (rej-sample-v2 only):
      if tol_level >= train_tol_level:
          train_data, val_data ← pool.get_train_valid_split()
          model ← fine_tune(model, train_data, val_data)
          tol_level ← 0
   
   e. Increase temperature:
      temperature ← temperature + delta

3. Return oracle.mol_buffer

Comparison with Other Approaches

vs. Genetic Algorithms

Similarities:

Maintains a pool of top candidates
Iterative generation and selection
Fitness-based ranking

Differences:

Uses LLM generation instead of crossover/mutation
Prompts guide generation with similarity constraints
Optional model fine-tuning for adaptation

vs. Reinforcement Learning

Similarities:

Oracle acts as reward function
Model learns to maximize rewards
Exploration-exploitation tradeoff

Differences:

No RL training loop (uses supervised fine-tuning)
Simpler implementation
Faster convergence on many benchmarks

vs. Bayesian Optimization

Similarities:

Efficient use of oracle budget
Balances exploration and exploitation

Differences:

No surrogate model
LLM directly generates candidates
Better handles discrete molecular space

Performance Characteristics

Memory Usage

Pool: ~10-100 molecules (negligible)
Model: 125M (~0.5GB), 1.3B (~5GB), 2B (~8GB)
Fine-tuning: +2x model size during training

Oracle Calls

Typical budget: 1K - 10K calls
Batch size: 200 molecules/iteration
PMO benchmark: ~5K calls to SOTA

Runtime

Generation: ~1-2 sec/batch (200 molecules)
Fine-tuning: ~30 sec - 2 min per round
Total: Minutes to hours depending on oracle complexity

Success Rate

QED optimization: 99% success (10K calls)
PMO benchmark: 17.5 avg score
Docking: 3-4x fewer calls than SOTA

Next Steps

Configure Hyperparameters

Learn how to tune pool size, temperature, and fine-tuning settings

Complete Examples

See full working examples with different oracles

Getting Started

Core Concepts

Training

Molecular Optimization

Generation

Guides

Optimization Algorithm

The optimize() Function

Parameters

Algorithm Workflow

Key Components

Pool Management

Prompt Construction

Molecule Entry

Tolerance Mechanism

Generation Temperature Schedule

Diversity Filtering

Algorithm Pseudocode

Comparison with Other Approaches

Performance Characteristics

Memory Usage

Oracle Calls

Runtime

Success Rate

Next Steps

Configure Hyperparameters

Complete Examples

​The optimize() Function

​Parameters

​Algorithm Workflow

​Key Components

​Pool Management

​Prompt Construction

​Molecule Entry

​Tolerance Mechanism

​Generation Temperature Schedule

​Diversity Filtering

​Algorithm Pseudocode

​Comparison with Other Approaches

​Performance Characteristics

Memory Usage

Oracle Calls

Runtime

Success Rate

​Next Steps

Configure Hyperparameters

Complete Examples

The optimize() Function

Parameters

Algorithm Workflow

Key Components

Pool Management

Prompt Construction

Molecule Entry

Tolerance Mechanism

Generation Temperature Schedule

Diversity Filtering

Algorithm Pseudocode

Comparison with Other Approaches

Performance Characteristics

Next Steps