Embedded Graph Flow

Method

Atoms and atom pairs are one-hot encoded and mapped into continuous node and edge embeddings, x = (x^V, x^E). With probability 1 - p a graph takes the denoising branch: sample t according to flow.t_sampling and ε ~ N(0, I), form the rectified-flow interpolant z_t = t·x + (1-t)·ε, and regress the clean embedding (x-prediction), scored as velocity MSE via v̂ = (x̂ - z_t)/(1-t) against v = x - ε. Otherwise it takes the decoding branch at t = 1: a corrupted x is mapped back to atom and bond logits through the tied unembedding, scored by cross-entropy. One graph transformer serves both, conditioned on t and a binary mode token, and the two branches are packed into a single forward pass. At sampling time the ODE is integrated from t = 0 to t = 1 and the decoding branch turns the endpoint into a discrete graph.

Two configurable choices reduce pressure toward codebook collapse:

• Codebook rows are renormalised to L2 norm √d, aligning the scale of x with the N(0, I) flow source (embedding.normalize).
• The flow target is detached, so the regression loss does not update the codebook; the decoding cross-entropy still shapes its geometry (embedding.detach_target).

Layout

configs/            qm9.yaml, zinc250k.yaml, smoke.yaml
src/embedded_graph_flow/
  config.py         defaults + YAML + key=value overrides
  data/             download, kekulise, dense index tensors, GDSS-compatible split
  models/           embedding (W_A/W_B and the tied W), DiGress-style graph transformer
  flow/             interpolant, corrupt(), the two-branch loss, the ODE sampler
  training/         EMA, AMP, cosine schedule, periodic sampling eval
  metrics/          V.U.N., FCD, NSPDK, scaffold similarity, unified evaluator
  cli/              egf-prepare-data / egf-train / egf-sample / egf-evaluate
tests/              shape, masking, interpolant-identity and metric tests

Install

Requires Python 3.10–3.12 and uv.

uv sync                      # core
uv sync --extra nspdk        # + EDeN, needed for the NSPDK metric
uv sync --extra dev          # + pytest, ruff

Data

uv run egf-prepare-data --dataset all --root data

This downloads the SMILES CSV and validation split, removes hydrogens, kekulises each molecule, and writes padded index tensors with train/test indices. The atom vocabularies and maximum graph sizes match MoFlow, GDSS and GruM (QM9: C N O F, 9 nodes; ZINC250k: C N O F P S Cl Br I, 38 nodes).

To use local files instead:

uv run egf-prepare-data --dataset qm9 \
    --smiles-csv /path/qm9.csv --valid-idx-json /path/valid_idx_qm9.json

Train

uv run egf-train -c configs/qm9.yaml
uv run egf-train -c configs/zinc250k.yaml train.out_dir=runs/zinc-a data.batch_size=96

Any config key can be overridden positionally (train.lr=1e-4, flow.p_decode=0.4, …). Checkpoints, the resolved config and a metrics.jsonl land in train.out_dir. train.resume=<path> continues a run.

A 50-step CPU sanity run:

uv run egf-train -c configs/smoke.yaml

Note on ZINC memory: the DiGress attention tensor is (B, N, N, H, D/H) and N = 38, so it is ~18× larger than on QM9. configs/zinc250k.yaml uses batch_size: 128; raise it if the card allows.

Key settings

key	effect
flow.p_decode	Fraction of graphs routed to the decoding objective.
flow.velocity_weight	When true, weights x-MSE by min(1/(1-t)², flow.weight_max); false uses unweighted x-MSE.
flow.t_sampling	Time distribution for denoising examples: uniform or logit_normal.
corrupt.gaussian_sigma_max, corrupt.swap_prob	Maximum Gaussian noise and label-swap probability for decoder inputs.
embedding.d_node, embedding.d_edge	Node and edge embedding dimensions.
embedding.detach_target	When true, prevents the flow loss from updating the codebook. codebook/*_max_cos logs the maximum off-diagonal cosine similarity.

Sample

uv run egf-sample -k runs/qm9/checkpoints/final.pt --n-samples 10000

Writes the raw dense atom_idx and bond_idx arrays to *_samples_10000.npz, and converted molecules to a .smi file. Since conversion can fail, the SMILES file may contain fewer entries. EMA weights are used by default; --no-ema uses the raw weights.

Evaluate and compare with DiGress / GruM

egf-evaluate scores each sample set with the same metric implementation and reference set:

uv run egf-evaluate --dataset qm9 \
    --samples ours=runs/qm9/checkpoints/final_samples_10000.npz \
    --samples digress=baselines/digress_qm9.npz \
    --samples grum=baselines/grum_qm9.npz \
    --out reports/qm9.txt

Metrics: Valid % (validity without valency correction, the MoFlow / GDSS / GruM convention), Unique %, Novel %, V.U.N. % (the DiGress-style joint fraction over all samples), FCD, NSPDK MMD and Scaf. (cosine similarity of Bemis–Murcko scaffold counts against the test set).

Preparing baseline samples

Use dense .npz input to measure validity without correction on the raw generated graphs. Both baselines construct dense tensors before building RDKit molecules:

• DiGress — in src/analysis/rdkit_functions.py, build_molecule receives atom_types and edge_types per graph. Collect them and np.savez_compressed(path, atom_idx=..., bond_idx=...), with -1 marking padded atoms. Bond codes 0..3 represent no bond, single, double and triple; aromatic code 4 is treated as invalid by the evaluator.
• GruM / GDSS — utils/mol_utils.py:gen_mol holds one-hot x and adj. Convert their argmax outputs before construct_mol: map the final atom class to -1, and adjacency channels 0/1/2/3 to bond codes 1/2/3/0.

A SMILES file also works (--samples grum=path.smi), but then Valid % is the RDKit parse rate because outputs the baseline failed to write are absent.

Tests

uv run pytest

The suite covers the padding/symmetry invariants, the interpolant identity (x̂ - z_t)/(1-t) = x - ε, embedding round-trip losslessness, that padded slots cannot leak into predictions, that the sampler emits well-formed symmetric graphs, and that the validity/uniqueness/novelty/scaffold metrics behave on hand-built cases.

License

This project is available under the PolyForm Noncommercial License 1.0.0. Commercial use requires a separate license from the copyright holder.