> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/YerevaNN/ChemLactica/llms.txt
> Use this file to discover all available pages before exploring further.

# Sampling Strategies

> Configure sampling parameters to control molecule generation quality and diversity

## Overview

Sampling strategies control how ChemLactica models select tokens during generation, balancing between generating high-quality molecules and exploring diverse chemical space.

## Generation Configuration

Generation parameters are specified in the `generation_config` section of your optimization configuration:

```yaml chemlactica_125m_hparams.yaml:15 theme={null}
generation_config:
  repetition_penalty: 1.0
  max_new_tokens: 100
  do_sample: true
  eos_token_id: 20
```

## Key Parameters

### Temperature

Controls randomness in token selection. Higher values increase diversity.

**temperature** `float` (default: 1.0)

Sampling temperature for generation.

* **Range**: 0.1 to 2.0
* **Lower values** (0.5-0.8): More conservative, higher quality
* **Medium values** (1.0-1.2): Balanced exploration
* **Higher values** (1.3-1.5): More diverse, creative generation

#### Dynamic Temperature

During optimization, temperature increases over time to enhance exploration:

```python optimization.py:169 theme={null}
if num_iter > initial_num_iter:
    config["generation_config"]["temperature"] += \
        config["num_gens_per_iter"] / (oracle.budget - config["num_gens_per_iter"]) * \
        (config["generation_temperature"][1] - config["generation_temperature"][0])
    print(f"Generation temperature: {config['generation_config']['temperature']}")
```

Configuration:

```yaml chemlactica_125m_hparams.yaml:13 theme={null}
generation_temperature: [1.0, 1.5]  # [start, end]
```

> [!TIP]
>
> Start with temperature 1.0 and gradually increase to 1.5 during optimization to balance exploitation early and exploration later.

### Max New Tokens

Maximum number of tokens to generate for each molecule.

**max\_new\_tokens** `int` (default: 100)

Maximum tokens to generate.

* **Typical SMILES**: 50-100 tokens
* **Small molecules**: 50 tokens
* **Drug-like molecules**: 100 tokens
* **Complex molecules**: 150+ tokens

```yaml theme={null}
generation_config:
  max_new_tokens: 100  # Suitable for most drug-like molecules
```

> [!WARNING]
>
> Ensure `max_new_tokens` + prompt length ≤ 2048 (model's context window).

### Repetition Penalty

Penalizes repeated token sequences to encourage diversity.

**repetition\_penalty** `float` (default: 1.0)

Penalty for repeating tokens.

* **1.0**: No penalty (default)
* **1.1-1.2**: Mild penalty for variety
* **1.3+**: Strong penalty (may reduce quality)

```yaml theme={null}
generation_config:
  repetition_penalty: 1.0  # Typically kept at 1.0 for molecules
```

> [!NOTE]
>
> Repetition penalty is usually kept at 1.0 for molecule generation, as molecular structures naturally contain repeated patterns (rings, functional groups).

### Sampling Mode

**do\_sample** `bool` (default: true)

Whether to use sampling or greedy decoding.

* **true**: Sample from probability distribution (diverse outputs)
* **false**: Always select highest probability token (deterministic)

```yaml theme={null}
generation_config:
  do_sample: true  # Required for diverse molecule generation
```

### EOS Token

**eos\_token\_id** `int` (default: 20)

Token ID that signals end of generation.

ChemLactica uses token ID 20 as the EOS token.

```yaml theme={null}
generation_config:
  eos_token_id: 20
```

## Advanced Parameters

Additional parameters available for fine-grained control:

### Beam Search

**num\_beams** `int` (default: 1)

Number of beams for beam search. Set to 1 for sampling.

**num\_beam\_groups** `int` (default: 1)

Number of groups for diverse beam search.

### Sequence Control

**num\_return\_sequences** `int` (default: 1)

Number of sequences to return per prompt.

```python rejection_sampling_configs.py:8 theme={null}
rej_sample_args = {
    "max_new_tokens": 300,
    "temperature": 1.0,
    "repetition_penalty": 1.0,
    "do_sample": True,
    "num_return_sequences": 20,  # Generate 20 molecules per prompt
    "eos_token_id": 20,
}
```

**diversity\_penalty** `float` (default: 0.0)

Penalty for generating similar sequences when `num_return_sequences` > 1.

**length\_penalty** `float` (default: 1.0)

Penalty for sequence length. Values > 1.0 encourage longer sequences.

## Configuration Examples

### Conservative Generation (High Quality)

```yaml theme={null}
generation_config:
  max_new_tokens: 100
  temperature: 0.8
  repetition_penalty: 1.0
  do_sample: true
  eos_token_id: 20
```

Best for:

* Property prediction tasks
* When you need valid, realistic molecules
* Early stages of optimization

### Balanced Generation (Default)

```yaml theme={null}
generation_config:
  max_new_tokens: 100
  temperature: 1.0
  repetition_penalty: 1.0
  do_sample: true
  eos_token_id: 20
```

Best for:

* General molecule generation
* Most optimization tasks
* Standard use cases

### Exploratory Generation (High Diversity)

```yaml theme={null}
generation_config:
  max_new_tokens: 100
  temperature: 1.4
  repetition_penalty: 1.1
  do_sample: true
  eos_token_id: 20
```

Best for:

* Late-stage optimization
* Exploring novel chemical space
* When stuck in local optima

### Rejection Sampling Generation

```python rejection_sampling_configs.py:1 theme={null}
sample_gen_args = {
    "max_new_tokens": 50,
    "temperature": 1.0,
    "repetition_penalty": 1.0,
    "do_sample": True,
    "eos_token_id": 2
}

rej_sample_args = {
    "max_new_tokens": 300,
    "temperature": 1.0,
    "repetition_penalty": 1.0,
    "do_sample": True,
    "num_return_sequences": 20,
    "eos_token_id": 20,
}
```

Best for:

* Optimization with iterative fine-tuning
* Generating many candidates per iteration
* Complex optimization landscapes

## Complete Optimization Configuration

Example configuration from `chemlactica_125m_hparams.yaml`:

```yaml chemlactica_125m_hparams.yaml:1 theme={null}
checkpoint_path: yerevann/chemlactica-125m
tokenizer_path: yerevann/chemlactica-125m
pool_size: 10
validation_perc: 0.2
num_mols: 0
num_similars: 5
num_gens_per_iter: 200
device: cuda:0
sim_range: [0.4, 0.9]
num_processes: 8
generation_batch_size: 200
eos_token: "</s>"
generation_temperature: [1.0, 1.5]

generation_config:
  repetition_penalty: 1.0
  max_new_tokens: 100
  do_sample: true
  eos_token_id: 20

strategy: [rej-sample-v2]

rej_sample_config:
  train_tol_level: 3
  checkpoints_dir: checkpoints
  max_learning_rate: 0.0001
  lr_end: 0
  train_batch_size: 2
  gradient_accumulation_steps: 8
  weight_decay: 0.1
  adam_beta1: 0.9
  adam_beta2: 0.999
  warmup_steps: 10
  global_gradient_norm: 1.0
  dataloader_num_workers: 1
  max_seq_length: 2048
  num_train_epochs: 5
  packing: false
```

## Optimization-Specific Parameters

### Similarity Range

**sim\_range** `list[float]` (default: \[0.4, 0.9])

Range for random similarity values in prompts during generation.

* Encourages exploration within structural constraints
* Lower bound controls minimum similarity
* Upper bound controls maximum similarity

```yaml theme={null}
sim_range: [0.4, 0.9]  # Generate molecules 40-90% similar to references
```

### Generation Batch Size

**generation\_batch\_size** `int` (default: 200)

Number of molecules to generate in parallel per batch.

Limited by GPU memory and model size.

```yaml theme={null}
generation_batch_size: 200  # For 125M model on 16GB GPU
```

### Generations Per Iteration

**num\_gens\_per\_iter** `int` (default: 200)

Total unique molecules to generate per optimization iteration.

```yaml theme={null}
num_gens_per_iter: 200  # Oracle calls per iteration
```

## Usage in Code

### Loading Configuration

```python example_run.py:97 theme={null}
import yaml
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load configuration
config = yaml.safe_load(open("hparams.yaml"))

# Load model
model = AutoModelForCausalLM.from_pretrained(
    config["checkpoint_path"],
    torch_dtype=torch.bfloat16
).to(config["device"])

tokenizer = AutoTokenizer.from_pretrained(
    config["tokenizer_path"],
    padding_side="left"
)

# Generate with config parameters
output = model.generate(
    **data,
    **config["generation_config"]
)
```

### Passing to Optimization

```python example_run.py:108 theme={null}
from chemlactica.mol_opt.optimization import optimize

# Configuration contains all generation parameters
optimize(
    model, tokenizer,
    oracle, config  # Includes generation_config
)
```

## Tuning Guidelines

**1. Start with defaults**

Use `temperature=1.0`, `do_sample=true`, `max_new_tokens=100`

**2. Adjust for task**

* High quality needed: Lower temperature to 0.8
* More diversity needed: Raise temperature to 1.2-1.4

**3. Monitor validity**

Track percentage of valid SMILES generated. If too low, decrease temperature.

**4. Optimize batch size**

Increase `generation_batch_size` until GPU memory is utilized

**5. Dynamic temperature**

Use temperature schedule to start conservative and increase exploration

## Common Issues

> **Low validity rate**
>
> **Problem**: Many invalid SMILES generated
>
> **Solutions**:
>
> * Lower temperature (try 0.8-0.9)
> * Reduce `max_new_tokens`
> * Check prompt format

> **Insufficient diversity**
>
> **Problem**: Generated molecules too similar
>
> **Solutions**:
>
> * Increase temperature (try 1.2-1.4)
> * Widen `sim_range` (e.g., \[0.3, 0.95])
> * Set `repetition_penalty` to 1.1

> **Context length exceeded**
>
> **Problem**: Prompts + generation exceed 2048 tokens
>
> **Solutions**:
>
> * Reduce `max_new_tokens`
> * Limit number of similar molecules in prompt
> * Truncate prompt history

> **OOM errors**
>
> **Problem**: GPU out of memory during generation
>
> **Solutions**:
>
> * Reduce `generation_batch_size`
> * Use smaller model (125M vs 1.3B)
> * Enable gradient checkpointing

## Best Practices

> **Temperature Schedule**
>
> Start conservative (0.8-1.0), increase for exploration (1.2-1.5)

> **Validation**
>
> Always validate generated SMILES with RDKit before scoring

> **Batch Size**
>
> Maximize batch size for GPU efficiency without OOM

> **Token Budget**
>
> Monitor token usage to stay within 2048 context window

## Next Steps

> **[Molecule Generation](/generation/molecule-generation)**
>
> Apply sampling strategies to generate molecules

> **[Optimization](/optimization/overview)**
>
> Use generation in molecular optimization workflows


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.