Skip to content

predict_structure leaves ETKDG randomSeed unset, so embed_smiles is not reproducible #2

Description

@leifsieben

Thanks for releasing the code and checkpoint — being able to run the encoder end to end made this easy to track down.

Summary: predict_structure never sets randomSeed on the ETKDGv3 params, so conformer generation draws from the global RNG. Because embed_smiles forks workers that inherit one RNG state, the conformer a molecule receives depends on which worker happens to pull its chunk — which is timing-dependent. Two identical calls to embed_smiles therefore return different embeddings.

Where

https://github.com/blazejba/Monroe/blob/57238ed/monroe/model/featurizer.py#L378-L385

def predict_structure(mol: Chem.Mol, n_confs: int = 1) -> Chem.Conformer:
    params = getattr(rdDistGeom, "ETKDGv3")()
    params.enforceChirality = True
    params.useRandomCoords = True
    params.numThreads = 1
    params.maxIterations = 500
    ...

enforceChirality, useRandomCoords, numThreads and maxIterations are set; randomSeed is not. RDKit's default is -1 (seed from the global RNG) — confirmed on rdkit 2026.03.6:

>>> from rdkit.Chem import rdDistGeom
>>> rdDistGeom.ETKDGv3().randomSeed
-1

What we measured

Commit 57238ed, bundled checkpoint/, CPU, macOS, rdkit 2026.03.6:

test result
40 molecules, n_workers=8, run twice identical
600 molecules, n_workers=8, run twice differ on 479 / 600 rows, max abs diff 0.397
600 molecules, n_workers=1 vs n_workers=8 differ, max abs diff 0.152
600 molecules after setting params.randomSeed identical across 1, 4 and 8 workers (max abs diff 0.0)

Per-molecule cosine similarity between two runs of the same molecule is 0.9926–1.0, so the effect is small per molecule but touches almost every row.

The 40-molecule case passing is what makes this easy to miss: with few chunks each worker takes at most one, so the assignment is stable. It only shows up once there are more chunks than workers and the hand-out order starts depending on timing.

Reproducer

import os
if not hasattr(os, "sched_getaffinity"):            # macOS
    os.sched_getaffinity = lambda pid: set(range(os.cpu_count() or 1))
import numpy as np, torch
_r = torch.load
torch.load = lambda *a, **k: _r(*a, **{**k, "map_location": "cpu"})

from monroe.model.ckpt import load_ckpt
from monroe.eval.embed import embed_smiles

enc = load_ckpt("checkpoint").to("cpu").eval()
S = [...]                                            # any 600 distinct SMILES
a = embed_smiles(S, enc, device="cpu", n_workers=8)
b = embed_smiles(S, enc, device="cpu", n_workers=8)
A = np.stack([a[s] for s in S if s in a])
B = np.stack([b[s] for s in S if s in b])
print(np.array_equal(A, B), np.abs(A - B).max())     # False 0.39...

Suggested fix

params.randomSeed = 0xC0FFEE   # or any fixed value, ideally configurable

That restored bitwise agreement across repeat runs and across worker counts in our testing. Exposing it as an argument on embed_smiles / featurize_smiles would let callers vary it deliberately when they want conformer variance as a source of ensembling.

Why it may matter for the paper's numbers

Table 1 and the ablations report mean ± std across 3 seeds. As far as we can tell that seed does not reach conformer generation, so this variance sits inside each reported run rather than being averaged over it, and it moves with machine load and worker count rather than with the seed. It would also mean the published results are not bitwise reproducible from the released code, which seems worth knowing given how carefully the rest of the evaluation is specified.

For context on why we cared: we are adding Monroe as an arm to an external benchmark where each embedding cache is keyed by a content hash of its definition, so a cache that cannot be regenerated breaks the guarantee that hash is supposed to provide. Happy to share more detail or test a patch.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions