Add GEOM-Drugs explicit-H epoch 50 with verified inference assets

geom/explicit_h/README.md

@@ -0,0 +1,12 @@
+# GEOM-Drugs with explicit hydrogens — epoch 50
+
+Continued from the best checkpoint after 48 completed epochs. The original Lightning checkpoint includes current weights, EMA weights and optimizer state. The inference loader constructs only the selected network (EMA by default).
+
+```python
+from morph.inference import load_pretrained
+
+model = load_pretrained("geom-explicit-h", epoch=50, device="cuda")
+trajectories, validity, vocab, coordinate_std = model.generate(n_mols=16)
+```
+
+The training subset contains at most 72 atoms including hydrogen. Coordinates returned by generation are in Angstrom. Config, vocabularies and training distributions are included; no dataset is needed for inference.

geom/explicit_h/atom_tokens.txt

@@ -0,0 +1,10 @@
+Br
+C
+Cl
+F
+H
+I
+N
+O
+P
+S

geom/explicit_h/charge_tokens.txt

@@ -0,0 +1,5 @@
+-2
+-1
+0
+1
+2

geom/explicit_h/config.yaml

@@ -0,0 +1,345 @@
+representation: molecule
+seed: 42
+heads:
+  heads:
+    _target_: morph.model.layers.heads.MultiHeadModule
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    heads_configs:
+      graph_aggregation: mean
+      node_heads:
+        atom_type_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.atom_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        charge_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.charge_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        ins_rate_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_rate: ${model.head_initialization.ins_rate}
+        do_del_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_del_probability}
+        do_sub_a_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_sub_a_probability}
+      edge_heads:
+        edge_type_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: ${len:${data.vocab.edge_tokens}}
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+        do_sub_e_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+          initial_probability: ${model.head_initialization.do_sub_e_probability}
+  ins_edge_head:
+    _target_: morph.model.layers.heads.InsertionEdgeHead
+    _recursive_: false
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_edge_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    n_edge_types: ${len:${data.vocab.edge_tokens}}
+    n_atom_types: ${len:${data.vocab.atom_tokens}}
+    n_charge_types: ${len:${data.vocab.charge_tokens}}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+insertion_head:
+  mode: unimodal
+  dimension: 3
+  num_components: 1
+  model:
+    _target_: morph.model.layers.heads.UnimodalInsertionHead
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_adapter_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    N_a: ${len:${data.vocab.atom_tokens}}
+    N_c: ${len:${data.vocab.charge_tokens}}
+model:
+  module:
+    optimizer_config:
+      implementation: torch
+      optimizer:
+        execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+        lr: 0.01
+        momentum: 0.95
+        weight_decay: 0.0
+        nesterov: true
+        ns_steps: 5
+        eps: 1.0e-07
+        adjust_lr_fn: original
+      aux_optimizer:
+        lr: 0.0003
+        betas:
+        - 0.9
+        - 0.95
+        eps: 1.0e-10
+        weight_decay: 0.0
+      gradient_clip_val: 1.0
+      gradient_clip_algorithm: norm
+      scheduler:
+        _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+        warmup_steps: 2000
+        total_steps: 200000
+        min_lr_fraction: 0.05
+        last_epoch: -1
+      interval: step
+      monitor: null
+    loss_weights:
+      do_sub_a: 0.2
+      sub_a_class: 0.03
+      do_sub_e: 0.25
+      sub_e_class: 0.5
+      do_del: 0.4
+      ins_rate: 0.15
+      ins_state: 0.01
+      ins_e: 0.4
+      x: 2.0
+      c: 0.5
+      ins_e_ii: 0.03
+    _target_: morph.model.lightning.base.LightningModuleRates
+    insertion_head_mode: ${insertion_head.mode}
+    ins_noise_scale: ${data.ins_noise_scale}
+    ema_decay: 0.999
+    ema_warmup_steps: 5000
+    use_ema_for_eval: true
+    time_dist: ${if:${eq:${model.sample_time_strategy}, log}, ${model.beta_sched},
+      ${model.linear_sched}}
+    integrator: ${model.integrator}
+    vocab: ${data.vocab}
+    use_learnable_loss_weights: false
+    use_time_weights: false
+    model:
+      _target_: morph.model.backbones.dit.DiTBackboneWithHeads
+      _recursive_: false
+      atom_type_embedding_args: ${embeddings.atom_type_embedding}
+      edge_type_embedding_args: ${embeddings.edge_type_embedding}
+      time_embedding_args: ${embeddings.time_embedding}
+      node_count_embedding_args: ${embeddings.node_count_embedding}
+      incident_bond_embedding: false
+      backbone_model_args: ${model.backbone}
+      heads_args: ${heads.heads}
+      insertion_head_args: ${insertion_head.model}
+      ins_edge_head_args: ${heads.ins_edge_head}
+  rotate: true
+  n_augmentations: 1
+  integrator: ${integrator.integrator}
+  beta_sched:
+    _target_: torch.distributions.Beta
+    concentration1: 2.0
+    concentration0: 1.0
+  linear_sched:
+    _target_: torch.distributions.Uniform
+    low: 0.0
+    high: 1.0
+  sample_time_strategy: linear
+  head_node_dim: ${model.backbone.out_node_nf}
+  head_edge_dim: ${model.backbone.out_edge_nf}
+  head_node_adapter_dim: 256
+  head_edge_adapter_dim: 256
+  head_insertion_adapter_dim: 256
+  head_insertion_edge_dim: 256
+  head_independent_adapters: true
+  head_initialization:
+    output_std: 0.001
+    coordinate_output_std: 0.0
+    do_del_probability: 0.5
+    do_sub_a_probability: 0.5
+    do_sub_e_probability: 0.5
+    ins_rate: 0.6931471805599453
+  token_weighting:
+    type_loss_token_weights: uniform
+    weight_alpha: 0.4
+  lora:
+    enabled: false
+    rank: 8
+    alpha: 8.0
+    dropout: 0.0
+    targets:
+    - qkv
+    - proj
+    - mlp.gate_value
+    - mlp.out_proj
+    freeze_base: true
+  backbone:
+    _target_: morph.model.backbones.dit.DiTBackbone
+    _recursive_: false
+    d_input: ${embeddings.atom_type_embedding.out_dim}
+    d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+    d_model: 512
+    out_node_nf: 256
+    out_edge_nf: 256
+    num_layers: 8
+    nhead: 16
+    mlp_ratio: 4.0
+    ffn_hidden_dim: null
+    ffn_multiple_of: 64
+    qk_norm: true
+    coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+    center_coordinate_output: true
+    in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+    edge_pair_dim: 256
+    pair_bias_hidden_dim: 32
+    num_atom_types: ${len:${data.vocab.atom_tokens}}
+    lora: ${model.lora}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+schedule:
+  _target_: morph.flow_matching.schedules.SmoothstepSchedule
+  shift: 0.8
+trainer:
+  trainer:
+    deterministic: warn
+    devices: 1
+    accelerator: auto
+    strategy: auto
+    precision: 32-true
+    max_epochs: -1
+    accumulate_grad_batches: 1
+    num_sanity_val_steps: 0
+    gradient_clip_val: null
+    gradient_clip_algorithm: norm
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+    limit_val_batches: null
+    limit_predict_batches: 4
+    max_time: null
+  do_train: false
+  do_validate: false
+  do_predict: true
+  checkpoint_path: null
+  final_checkpoint_path: null
+callbacks: {}
+logging:
+  enabled: false
+  project: morph-${data.dataset_name}
+  name: geom-explicit-h
+  save_dir: .
+  offline: false
+integrator:
+  integrator:
+    _target_: morph.flow_matching.integration.RateIntegrator
+    vocab: ${data.vocab}
+    distributions: None
+    insertion_head_mode: ${insertion_head.mode}
+    insertion_dimension: ${insertion_head.dimension}
+    insertion_num_components: ${insertion_head.num_components}
+    time_strategy: log
+    time_beta_alpha: 2.0
+    time_beta_beta: 1.0
+    num_integration_steps: 100
+    ins_noise_scale: ${data.ins_noise_scale}
+    renoise_insertions: true
+    exact_schedule_integration: true
+    del_schedule: ${data.interpolator.del_schedule}
+    ins_schedule: ${data.interpolator.ins_schedule}
+    sub_schedule: ${data.interpolator.sub_schedule}
+    sub_e_schedule: ${data.interpolator.sub_e_schedule}
+data:
+  root_path: .
+  preprocess_workers: 0
+  preprocess_chunk_size: 128
+  statistics_workers: ${data.preprocess_workers}
+  hydrogen_mode: explicit
+  cache_name: ${data.dataset_name}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}
+  source_dir: null
+  max_atoms: 72
+  allowed_atom_symbols:
+  - H
+  - C
+  - 'N'
+  - O
+  - S
+  - F
+  - Cl
+  - Br
+  - P
+  - I
+  n_atoms_strategy: uniform
+  uniform_n_atoms_padding: 5
+  optimal_transport: equivariant
+  ins_noise_scale: 0.65
+  move_noise_scale: 0.5
+  allow_charged: true
+  preprocessing:
+    _target_: morph.dataset.core.preprocessing.Preprocessing
+    _recursive_: false
+    root: .
+    num_workers: 0
+    train_dataset: null
+    atom_tokens_path: ${data.preprocessing.root}/atom_tokens.txt
+    edge_tokens_path: ${data.preprocessing.root}/edge_tokens.txt
+    charge_tokens_path: ${data.preprocessing.root}/charge_tokens.txt
+    distributions_path: ${data.preprocessing.root}/distributions.pt
+  vocab: null
+  interpolator:
+    _target_: morph.flow_matching.interpolation.Interpolator
+    distributions: null
+    optimal_transport: ${data.optimal_transport}
+    ins_noise_scale: ${data.ins_noise_scale}
+    move_noise_scale: ${data.move_noise_scale}
+    atom_type_noise_prob: 0.0
+    pre_align_iters: 2
+    c_move: 1.0
+    c_sub: 0.0
+    c_ins: 6.125
+    c_del: 6.125
+    del_schedule: ${schedule}
+    ins_schedule: ${schedule}
+    sub_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+    sub_e_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+  datamodule: null
+  dataset_name: geom
+  dataset_target: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+embeddings:
+  atom_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.atom_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  edge_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.edge_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  time_embedding:
+    _target_: morph.model.layers.embedding.TimeEmbedding
+    embedding_dim: 128
+    out_dim: 128
+  node_count_embedding:
+    _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+    embedding_dim: 64
+    out_dim: 128
+    min_period: 128.0
+    max_period: 1024.0
+    max_count: 96.0

geom/explicit_h/distributions.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40a31704ab58d0349daff35083d43dcf4771ef368fd9c8a8503c94fe844016c5
+size 63581

geom/explicit_h/edge_tokens.txt

@@ -0,0 +1,4 @@
+<NO_BOND>
+1
+2
+3

geom/explicit_h/epoch_050.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:016c4a5bbb4f5ef36d1d62dff772ecf8d384e8ab3e2c94ebeb4c4357183f690a
+size 509181385

geom/explicit_h/manifest.json

@@ -0,0 +1,602 @@
+{
+  "README.md": {
+    "bytes": 701,
+    "sha256": "9947149d5b311da2f381b57f135811eb1eade2dba84cdf452ae69f5e11ba1016"
+  },
+  "atom_tokens.txt": {
+    "bytes": 22,
+    "sha256": "c9b712563ad2509593aa48fddbb5145aa42a70c4791255b8eea2984f549cae6a"
+  },
+  "charge_tokens.txt": {
+    "bytes": 12,
+    "sha256": "e2533aa090751f1006a74d165aa2d4a7ac23dbf73a65eb624c2a81216cc7dff5"
+  },
+  "config.yaml": {
+    "bytes": 10648,
+    "sha256": "e448c580096659af3b7e7e718e395565aa397b7b6666bc148576c7a455940a66"
+  },
+  "distributions.json": {
+    "bytes": 220286,
+    "sha256": "17bcf85c09486e5070aa8b3be4f4411fa01e6b3cf1b39c5488a5f46149f52e0f"
+  },
+  "distributions.pt": {
+    "bytes": 63581,
+    "sha256": "40a31704ab58d0349daff35083d43dcf4771ef368fd9c8a8503c94fe844016c5"
+  },
+  "edge_tokens.txt": {
+    "bytes": 16,
+    "sha256": "42b635ef938d2ec5a601044e4b01f3f84ed121f1bfaac984d86f180a8ccba3fd"
+  },
+  "epoch_050.ckpt": {
+    "bytes": 509181385,
+    "sha256": "016c4a5bbb4f5ef36d1d62dff772ecf8d384e8ab3e2c94ebeb4c4357183f690a"
+  },
+  "metadata.json": {
+    "bytes": 1094,
+    "sha256": "b23b8944369765d12b2cbae0f5d48614225e44cb59deff8bb0e49783318ddb31"
+  },
+  "runtime/configs/callbacks/callbacks.yaml": {
+    "bytes": 699,
+    "sha256": "dce0adebf5956dce687d64a074238f9c12866373bfe21df6658d967d6bfd1f90"
+  },
+  "runtime/configs/data/chembl3d/chembl3d.yaml": {
+    "bytes": 278,
+    "sha256": "a96974d90353a850c4a1d0bac27de5f1fc8724b81cc5bbc122aab6f76f3051c5"
+  },
+  "runtime/configs/data/chembl3d/implicit_h.yaml": {
+    "bytes": 81,
+    "sha256": "08da196d0288cf1fe4dba8dd407c49d9bee003607504beb59f79e5d6b9eaef0c"
+  },
+  "runtime/configs/data/chembl3d/max100.yaml": {
+    "bytes": 246,
+    "sha256": "052dba7b3f3efce990019518823476bd6bf1381bb5f606822c9fc7a485e6d0bf"
+  },
+  "runtime/configs/data/chembl3d/max100_implicit_h.yaml": {
+    "bytes": 79,
+    "sha256": "152325f11c7d1500469d4ee381100ae12247e6a2e0b90e5ef7501b405144bb2d"
+  },
+  "runtime/configs/data/geom/geom.yaml": {
+    "bytes": 428,
+    "sha256": "736c04e86c029bed65ce3ed022f937dd418ddc70d4b9fc064b38d979bf29ee6d"
+  },
+  "runtime/configs/data/geom/implicit_h.yaml": {
+    "bytes": 73,
+    "sha256": "c6b3ff5496af997de0c28e5aa7265903d6d0c8b59aab001ab21d6a64baeb772a"
+  },
+  "runtime/configs/data/geom/scaffold_decoration.yaml": {
+    "bytes": 1107,
+    "sha256": "d317728a98961e98a1f03bcc3b0f8c513169ca23b68ea07ba807b99066e7ca80"
+  },
+  "runtime/configs/data/molecular.yaml": {
+    "bytes": 5005,
+    "sha256": "f79ade5ddbf7d08f2eecbb2aa9c3e6b182b8a8e7e48960aff5e5edda5601bdc9"
+  },
+  "runtime/configs/data/qm9/qm9.yaml": {
+    "bytes": 3648,
+    "sha256": "3beea93868a1e889b049dcc2ccbc3b7489fa4672e964138654bc512ea01ab9f2"
+  },
+  "runtime/configs/data/qm9/scaffold_decoration.yaml": {
+    "bytes": 890,
+    "sha256": "40b36f2bc25bf4052cdd676a26d9e528cb4f7b13deddbb1660f454f97b159c40"
+  },
+  "runtime/configs/default.yaml": {
+    "bytes": 1294,
+    "sha256": "bcf9744acbf53afa4305ffd6e9cf95e6f99571750488eaa8cad85896a51aca66"
+  },
+  "runtime/configs/defaults/chembl3d/chembl3d.yaml": {
+    "bytes": 90,
+    "sha256": "f09553bc670e6c87ac8c1f8bbd0a8df73f1e83b02a3dcf9d6d06b915ac3bce93"
+  },
+  "runtime/configs/defaults/chembl3d/implicit_h.yaml": {
+    "bytes": 102,
+    "sha256": "a0ee1440984eb4c9d3c96a5d1c18df9f59b8e42c499f8a340ccd8deb8c454cd8"
+  },
+  "runtime/configs/defaults/chembl3d/implicit_h_2x4090.yaml": {
+    "bytes": 331,
+    "sha256": "518571405333a40e1602ee4c77fd153b56fd53e77aaa2aaf530796142f58d4b8"
+  },
+  "runtime/configs/defaults/chembl3d/train_2x4090.yaml": {
+    "bytes": 826,
+    "sha256": "0e6b01b264ab1ae70ae7f144ecd0a03d242e24f02feb087e3d5e2b7bc24790da"
+  },
+  "runtime/configs/defaults/core.yaml": {
+    "bytes": 431,
+    "sha256": "0dcfb68e2fece1dc277f647d3b3769af882cf8175e0c3b479bc50029e2538b1c"
+  },
+  "runtime/configs/defaults/geom/geom.yaml": {
+    "bytes": 425,
+    "sha256": "f447a1bdd39e211b5cf43c4cce56e708da6472d58f79e87885621ec4899c1c8e"
+  },
+  "runtime/configs/defaults/geom/implicit_h.yaml": {
+    "bytes": 98,
+    "sha256": "9d3a188e6bd2a2b01609ba0d652184bb8a2f7e27fd89d4a3031568051cae605e"
+  },
+  "runtime/configs/defaults/geom/scaffold_decoration.yaml": {
+    "bytes": 828,
+    "sha256": "bcce546bce2b0b7bfaf7c0663e05423fd03c952bf6e52cbbf31976723661f0d1"
+  },
+  "runtime/configs/defaults/qm9/qm9.yaml": {
+    "bytes": 132,
+    "sha256": "f83bd8e10802df07c2b1933dd40597807e3ee8a8cc39349ab14eb6f722f8475e"
+  },
+  "runtime/configs/defaults/qm9/scaffold_decoration.yaml": {
+    "bytes": 348,
+    "sha256": "3bb8b6a757ea4921c31df2c0869949baef703df980b964d9636f52b64ac39267"
+  },
+  "runtime/configs/embeddings/geom.yaml": {
+    "bytes": 803,
+    "sha256": "06793f4d446d899197e3db28d064592af9049612ee78cb0718a63e3d6fdc3151"
+  },
+  "runtime/configs/embeddings/qm9.yaml": {
+    "bytes": 849,
+    "sha256": "de57af7f1ff9b44dbb952f852c0177789196105b77f396298ff8e0218c0626b7"
+  },
+  "runtime/configs/heads/head.yaml": {
+    "bytes": 2475,
+    "sha256": "238e2346041e0aedccd16a6ed4066e163004c0e45a2cbc6fd1f6d17e16eb44ed"
+  },
+  "runtime/configs/insertion_head/gmm.yaml": {
+    "bytes": 488,
+    "sha256": "1e26914eed9b4971e0f6ff0b61bbdb34103233d346719a55ba6273dd3d785c05"
+  },
+  "runtime/configs/insertion_head/unimodal.yaml": {
+    "bytes": 495,
+    "sha256": "547a6391c3b341e3e369f8f4c7aab9752c67dfbbb533388d25956e9011b305e8"
+  },
+  "runtime/configs/integrator/integrator.yaml": {
+    "bytes": 1580,
+    "sha256": "8f0d9aafb0b2727cd58154e1fd97df693c8b7d4f5a3fb5efcb9188553f54660f"
+  },
+  "runtime/configs/logging/wandb_logger.yaml": {
+    "bytes": 139,
+    "sha256": "a9a2787e048d4eb70f2ff3c49d4355f1a5a382047c8619203b1f9131d829928d"
+  },
+  "runtime/configs/loss_weights/calibrated_geom.yaml": {
+    "bytes": 644,
+    "sha256": "d4763c84aa9e1eb77fb2b456384332a4b3a84b072d560918b2357222a43f4ca1"
+  },
+  "runtime/configs/loss_weights/calibrated_geom_frozen.yaml": {
+    "bytes": 1141,
+    "sha256": "a08acb61d67a852051dbded08f8bb67f21101888fd0a22a96a522781d28c570d"
+  },
+  "runtime/configs/loss_weights/tuned_weights.yaml": {
+    "bytes": 420,
+    "sha256": "2bb5884465666d6fc638f787fd79257e4126c69a5cc55b55615fc45e02e65894"
+  },
+  "runtime/configs/loss_weights/uniform.yaml": {
+    "bytes": 336,
+    "sha256": "3c14d6443048f69d436fe03792703a33c67f368ecbb349de2fd145da4c96eabc"
+  },
+  "runtime/configs/model/backbones/dit.yaml": {
+    "bytes": 4039,
+    "sha256": "65b0f08be8fb307b8b01e7b460fd083fe60d48148fbece378d7e5a23dd4a9fb0"
+  },
+  "runtime/configs/model/finetune/lora_finetune.yaml": {
+    "bytes": 328,
+    "sha256": "6334a5112499c0f3bc1fe5de224e2b6b3265dc524c8d6b9048113868fb864558"
+  },
+  "runtime/configs/model/finetune/scaffold_finetune.yaml": {
+    "bytes": 5659,
+    "sha256": "1aad0d33e9b5ae95d9f50b5201a29bb8677fe42ea2bcf59e0e61d053ef1dfddb"
+  },
+  "runtime/configs/optimizer/muon_aux_adam.yaml": {
+    "bytes": 1212,
+    "sha256": "3b0a6189f3ba81f85f7d26c6f6ad7f64dcb25098a14e411211d7a62983a7180c"
+  },
+  "runtime/configs/resume_geom.yaml": {
+    "bytes": 13525,
+    "sha256": "abdca1624682973bc89829625c2c01abdaa3f318daa51a08da189967c8253195"
+  },
+  "runtime/configs/schedule/beta.yaml": {
+    "bytes": 78,
+    "sha256": "75c485edfd70016b255282adcf5de25dfb95c9e32bb6587017fb8935ffe64eaa"
+  },
+  "runtime/configs/schedule/cubic.yaml": {
+    "bytes": 54,
+    "sha256": "a998a2d3390ab5daf4df079852ae8daa907687cf059f55c500409c9493949c0c"
+  },
+  "runtime/configs/schedule/fast_power.yaml": {
+    "bytes": 126,
+    "sha256": "22f36106bd43ca253ffea412e00539f2b07efbc49109b97c1232df6a827a2baf"
+  },
+  "runtime/configs/schedule/linear.yaml": {
+    "bytes": 55,
+    "sha256": "f7881af26080ab5569e67b39f6eae2733c033b9623020e6271959f09fac45fa0"
+  },
+  "runtime/configs/schedule/smoothstep.yaml": {
+    "bytes": 70,
+    "sha256": "f6e8651953d93f0d18c997d40c199b3d80ddd352a2041cc51e29adc9a7f22778"
+  },
+  "runtime/configs/trainer/trainer.yaml": {
+    "bytes": 1402,
+    "sha256": "6f78e83ad46529bee7a2878b20727ff403420a515c1629bd73a608a2fb654f5f"
+  },
+  "runtime/pyproject.toml": {
+    "bytes": 4464,
+    "sha256": "d4854128f81144905cef68d37222d0d06fb372e2bd148e49177e0ff082a7b01f"
+  },
+  "runtime/requirements-tested.txt": {
+    "bytes": 177,
+    "sha256": "fac030b53775fd8ed8f89b17050ea74180f99dfab270b1d75d4ca3c0573ede76"
+  },
+  "runtime/src/morph/__init__.py": {
+    "bytes": 79,
+    "sha256": "bb402ea5c6f79e569205cc2bf4168d06430753ce8499ca4ffb8614743626fd59"
+  },
+  "runtime/src/morph/__main__.py": {
+    "bytes": 149,
+    "sha256": "95a7983479559ca2b15c7afa3881b3142ef4636583501184469b972df7646d4f"
+  },
+  "runtime/src/morph/_vendor/__init__.py": {
+    "bytes": 72,
+    "sha256": "8480cd966a03870a7142e1763094bf00e221d2c3c2b1af1d4d458fb8b7f37a20"
+  },
+  "runtime/src/morph/_vendor/geom_drugs_preprocessing.py": {
+    "bytes": 11328,
+    "sha256": "33a9c21d29df269a5f342a0b0e8bced0fc620559f0cb3a7a4cfabd5c369ceccb"
+  },
+  "runtime/src/morph/_vendor/muon.py": {
+    "bytes": 16405,
+    "sha256": "c61af8c2c4a635cc682b84e279851cb54fedc6fb935d8eb11aafbeb1c5e81fc6"
+  },
+  "runtime/src/morph/cli.py": {
+    "bytes": 9723,
+    "sha256": "27eee3fd6d6410747f12b19b670a74f50a51b842a50b1eff4e75c46d0d08a986"
+  },
+  "runtime/src/morph/config.py": {
+    "bytes": 704,
+    "sha256": "bd3562c030611f87cc514308f155ad93740dd59927732f71b2e15009012d78ff"
+  },
+  "runtime/src/morph/dataset/__init__.py": {
+    "bytes": 620,
+    "sha256": "44ae493354fed0fec4b410f46e0c1fb0351b5fbe5f404ff50f77181e60a39cee"
+  },
+  "runtime/src/morph/dataset/conditional/__init__.py": {
+    "bytes": 238,
+    "sha256": "e46c747069acbf3d82407ccfcf96f2e33608c010a429f9abdb064bb71795d134"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/__init__.py": {
+    "bytes": 167,
+    "sha256": "6ef414b31cceb7303c174d720e0d0666efd6f6678c97b3cf86d33c799d2153dd"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/cutters.py": {
+    "bytes": 12282,
+    "sha256": "b29e0d7b1758461cda868bcba56d58b21ce76a45a9365bed73f7a2735103616f"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/datamodule.py": {
+    "bytes": 1402,
+    "sha256": "ecd4a5819d3654f8f76e3e0d059a60bfc9d0ac2ec343514d45517e7094c7c282"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/lmdb.py": {
+    "bytes": 4243,
+    "sha256": "54ea075ebb4a8ed481f93f57e9a5e25504918b6595bde106f2489e21e7e73b5b"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/wrappers.py": {
+    "bytes": 7720,
+    "sha256": "a44d61452b676dc1dfd8d6bd44c09b22d463c8129b23cf98ec3903d10a693f2c"
+  },
+  "runtime/src/morph/dataset/core/__init__.py": {
+    "bytes": 376,
+    "sha256": "7bab9c675f0850efc50e14a9898accfde5f8d1c1e2420414ac8c0986822d27e0"
+  },
+  "runtime/src/morph/dataset/core/atomic_data.py": {
+    "bytes": 708,
+    "sha256": "7d59ec16ed5b3cf82b57b49dd4d7799862709353a65f4d473d4571ba91b86fa6"
+  },
+  "runtime/src/morph/dataset/core/batch.py": {
+    "bytes": 23924,
+    "sha256": "7b2b2afd749b621b1ca890d192e4dc9378b61ead9ef34b90f04dc449e6bad369"
+  },
+  "runtime/src/morph/dataset/core/containers.py": {
+    "bytes": 19637,
+    "sha256": "adae656c3fbbe5cfe32f378a419c7e1076f653163ffe7d2c1d1d0abc71a99cd7"
+  },
+  "runtime/src/morph/dataset/core/conversion.py": {
+    "bytes": 2224,
+    "sha256": "0842565c503048a86a517c202dd0cc7c07e9296a841601b178ec38c43f903a6a"
+  },
+  "runtime/src/morph/dataset/core/flow_matching_wrapper.py": {
+    "bytes": 7568,
+    "sha256": "7672ce6dc416bbccae1471b8e4fb66f44adf7eaf8f0f2fa151f1268e9513a17d"
+  },
+  "runtime/src/morph/dataset/core/insertion_data.py": {
+    "bytes": 14393,
+    "sha256": "7e99e4d59b29ea537708b5f7887ed3e61b4b31ebf80d55fe5a0c6823ff8a439c"
+  },
+  "runtime/src/morph/dataset/core/insertion_targets.py": {
+    "bytes": 6614,
+    "sha256": "2e62e89125310572735d916e97c8a544cfb3eb76fb167ecd134b5f8b0424df8c"
+  },
+  "runtime/src/morph/dataset/core/lightning_datamodule.py": {
+    "bytes": 7057,
+    "sha256": "4d525cbcde4a64d7cac41837e0379c40097cc38fbbc2f54ff30ff7e3ab83de4c"
+  },
+  "runtime/src/morph/dataset/core/lmdb.py": {
+    "bytes": 10111,
+    "sha256": "748b7a10f8fb20f74c267de393cc3e1933ea86debde99d8fdaaef94e0a9afd29"
+  },
+  "runtime/src/morph/dataset/core/preprocessing.py": {
+    "bytes": 20145,
+    "sha256": "da9cffd6beb6da63ee749dd8b60d4d0509bcd937778014739b04ffc599dde066"
+  },
+  "runtime/src/morph/dataset/core/properties.py": {
+    "bytes": 1952,
+    "sha256": "4dc63574f5fde284245078adf50fb2d8a9ac16443b8276154c03b7703cc2e937"
+  },
+  "runtime/src/morph/dataset/core/representation.py": {
+    "bytes": 6591,
+    "sha256": "0ecac55e1faf9748826ec8f4c7f1a274263214ea3f68cc69de7047b55dd6a65f"
+  },
+  "runtime/src/morph/dataset/core/tokenized_decode.py": {
+    "bytes": 6559,
+    "sha256": "483afb3d468113da33cc8699998e6439fcde2264ee7ca479724206360584a2b8"
+  },
+  "runtime/src/morph/dataset/core/vocab.py": {
+    "bytes": 2236,
+    "sha256": "f8618488f7bcdc8a58d8b9f3443b764eac60985a33ead1bcc2d17d3bdc213329"
+  },
+  "runtime/src/morph/dataset/unconditional/__init__.py": {
+    "bytes": 403,
+    "sha256": "d7cd3b85745ad1a4d884c0a11a64b3ab9cedcc64e30f604dbe83f5beb14fa97b"
+  },
+  "runtime/src/morph/dataset/unconditional/chembl3d.py": {
+    "bytes": 9854,
+    "sha256": "d5cfd1ee4ddd3085b95b9a678d761d4f0beef99ac5ecfb21d3dfb3cae59653eb"
+  },
+  "runtime/src/morph/dataset/unconditional/geom.py": {
+    "bytes": 4071,
+    "sha256": "12bc792c2a61b3e6c29ae265b9997ba610765e268408e873918b5a65225945be"
+  },
+  "runtime/src/morph/dataset/unconditional/qm9.py": {
+    "bytes": 9520,
+    "sha256": "06b6086cb1234b283cd3a45daad2e06d3bdb5d45810b873f860c82969ccbfe53"
+  },
+  "runtime/src/morph/flow_matching/__init__.py": {
+    "bytes": 78,
+    "sha256": "674f4630aa8421d7e75c913741c6efb1a163fbc38ce909c6cdb011f9ba9ef48e"
+  },
+  "runtime/src/morph/flow_matching/assignment.py": {
+    "bytes": 16050,
+    "sha256": "b6ad66538f5bd7f14e0eee2a0ffd25a4110d1504e057a0a9eb63a354d9f6aadb"
+  },
+  "runtime/src/morph/flow_matching/integration.py": {
+    "bytes": 33356,
+    "sha256": "f64b7b9ed2a79c11932b72e498e57eef924b5c85a05aa6893ddfe11e2d4595e9"
+  },
+  "runtime/src/morph/flow_matching/interpolation.py": {
+    "bytes": 22169,
+    "sha256": "6ef2e03626fd4b3194c68df6e2410ab6e630f4891f97a8328b4eb0b3e18c7dcd"
+  },
+  "runtime/src/morph/flow_matching/prior.py": {
+    "bytes": 2389,
+    "sha256": "0f9b77274238b23be7a2ad88ee8e6dcbb785817fcfa3a35aa863b489d2a6e05d"
+  },
+  "runtime/src/morph/flow_matching/sampling_step.py": {
+    "bytes": 2806,
+    "sha256": "c768682a47da1b7852838e20b39317528717fc0485cbc35030b3ca71d7d845bb"
+  },
+  "runtime/src/morph/flow_matching/schedules.py": {
+    "bytes": 6776,
+    "sha256": "6fd8396a05052ae7831babc75ca40a7e748fb9edffe82bb2559533afa949b2a9"
+  },
+  "runtime/src/morph/inference/__init__.py": {
+    "bytes": 732,
+    "sha256": "63047340b4c7ed784e969b7d54035ab737282ee2acc6eeb26ba6cb3d52b4f573"
+  },
+  "runtime/src/morph/inference/bundles.py": {
+    "bytes": 6395,
+    "sha256": "f6e2d46e6e428452a58208e78c83da01d6c786f10f7062643a093fb6750dafe6"
+  },
+  "runtime/src/morph/inference/catalog.py": {
+    "bytes": 6215,
+    "sha256": "53a27b86d39641fd4964da82f4b1875598a0cdb1533f981074c6b7a14b4174b7"
+  },
+  "runtime/src/morph/inference/generate.py": {
+    "bytes": 11889,
+    "sha256": "c3181fbf6187cf22b476431ed7c0fdea4c74ea8e97acf270f486318e43646d19"
+  },
+  "runtime/src/morph/inference/hub.py": {
+    "bytes": 9739,
+    "sha256": "fa1f43937d66f0ab4b7b9c11ed749aef8d23444b2df3257c36ff72ecd3b821ca"
+  },
+  "runtime/src/morph/inference/module.py": {
+    "bytes": 3820,
+    "sha256": "415e0f9c624c191855874b8dd88d5e7aa18941632bc403a9c0b685bbf74f1d60"
+  },
+  "runtime/src/morph/inference/registry.py": {
+    "bytes": 4314,
+    "sha256": "8ca3e4246e00707ed0d646e6e760acebd5e3d3443e5e8d5f71f3b54653993fa9"
+  },
+  "runtime/src/morph/model/README.md": {
+    "bytes": 10319,
+    "sha256": "9c06b5b90300f7449048640417116aa8869dc1a75f106cbdd9344a89f6e446d3"
+  },
+  "runtime/src/morph/model/__init__.py": {
+    "bytes": 80,
+    "sha256": "1c7c3875387d5ba7f9166a804ea5a7ef7bd469728c1ad5c5003150d44e33be51"
+  },
+  "runtime/src/morph/model/backbones/__init__.py": {
+    "bytes": 64,
+    "sha256": "a1a1419590c11aa6ed429f8b6b293470e30fe767f5571d2517b0883efb071e0c"
+  },
+  "runtime/src/morph/model/backbones/conditional_dit.py": {
+    "bytes": 25451,
+    "sha256": "411fe1d0fd17c24e03e153d57f299aa0436bc5749e2064d5574681878429795d"
+  },
+  "runtime/src/morph/model/backbones/dit.py": {
+    "bytes": 31262,
+    "sha256": "e2798baaf31bc8004a17c2432c56af9c4b12651d259cc0deaf5f18c83c6ddd33"
+  },
+  "runtime/src/morph/model/build.py": {
+    "bytes": 3442,
+    "sha256": "e69f4231815f09c094d787f11edc49e222cbdc970edc6e416b6b94a9b88d4b21"
+  },
+  "runtime/src/morph/model/finetune/__init__.py": {
+    "bytes": 282,
+    "sha256": "8746b1e29e5a3b3b1805389257270b90f5d30995783f2cf41b9fb533c485fc0c"
+  },
+  "runtime/src/morph/model/finetune/base/__init__.py": {
+    "bytes": 803,
+    "sha256": "6e0ded0b3a5c2a76ccb85b05cf63bcc54e145bd1be07d933a737d6151433d856"
+  },
+  "runtime/src/morph/model/finetune/base/conditional_backbone.py": {
+    "bytes": 2853,
+    "sha256": "1353e96443eacf3ce88473a30d02ab4f1c8b6a9c90763ea61c58d353b806b378"
+  },
+  "runtime/src/morph/model/finetune/base/context_encoder.py": {
+    "bytes": 15124,
+    "sha256": "8d5dfc91ec73120161ed0a51dcc7edac442004ad4d5213ba34d9d11cfd1eee15"
+  },
+  "runtime/src/morph/model/finetune/base/finetune_base.py": {
+    "bytes": 9526,
+    "sha256": "2efa00fd8eed1096a9cf8799dee4abfd8b27de1c0600e22a906dddaa089c4ca0"
+  },
+  "runtime/src/morph/model/finetune/scaffold/__init__.py": {
+    "bytes": 366,
+    "sha256": "6345e71c8b03ee5eb0dcd105c0532acf73a1795bf6360b2847be600bcecd40ba"
+  },
+  "runtime/src/morph/model/finetune/scaffold/scaffold.py": {
+    "bytes": 1984,
+    "sha256": "6de7a6083b4303ae422f03f15b5b40b1233d48bfc2dd1e0aa83a4eaa8adf40dd"
+  },
+  "runtime/src/morph/model/finetune/scaffold/scaffold_context.py": {
+    "bytes": 6086,
+    "sha256": "d9446254770e3c3ea2c634a0636ab82a718514a7cb69760d756e9235cfbe2814"
+  },
+  "runtime/src/morph/model/layers/__init__.py": {
+    "bytes": 189,
+    "sha256": "79671fbdea358c330a45b7fd3d842119b32b1d548804b3d441bbfb7dbd703b69"
+  },
+  "runtime/src/morph/model/layers/attention.py": {
+    "bytes": 21038,
+    "sha256": "7efc416ac8fec0236909661ad0693a0aef613ff97a964fb6d8252cb49c60b5c6"
+  },
+  "runtime/src/morph/model/layers/embedding.py": {
+    "bytes": 14356,
+    "sha256": "efe924190da5904e05495f421dcc313134cebe0dee9bf665060e1e5781cd95ef"
+  },
+  "runtime/src/morph/model/layers/gmm.py": {
+    "bytes": 6753,
+    "sha256": "628e90c73bdbb3e70104214f80bbf8a94896ddbe22a9a3fc40cce34b20d2565f"
+  },
+  "runtime/src/morph/model/layers/heads.py": {
+    "bytes": 30913,
+    "sha256": "7113c29a6be81c93ca980dc19d0847d2015ac36702683e4f1adedd549aed46ae"
+  },
+  "runtime/src/morph/model/layers/kernels.py": {
+    "bytes": 12483,
+    "sha256": "be0c52b69c7612946244755e2865de0491440a41163cf828168eb750a8f929cf"
+  },
+  "runtime/src/morph/model/layers/lora.py": {
+    "bytes": 5947,
+    "sha256": "7194750605e7b5df9a02e4c4a39cf29c982f117c40198c99c73211c85ecdfe4b"
+  },
+  "runtime/src/morph/model/lightning/__init__.py": {
+    "bytes": 339,
+    "sha256": "c53943a4eb1e45568cebfb538655cc5eb7ead46284bbcd515800102c89177e52"
+  },
+  "runtime/src/morph/model/lightning/_optimizer.py": {
+    "bytes": 19578,
+    "sha256": "f77519db8d3458c05aec1543cb4cd487995cf398abcc71b62ea1835cb132165e"
+  },
+  "runtime/src/morph/model/lightning/_sampling.py": {
+    "bytes": 15344,
+    "sha256": "2694ba5cd0b315147285ba268ed5ad246d6091d2426fcd7b21f78d0eb0a8fca1"
+  },
+  "runtime/src/morph/model/lightning/_setup.py": {
+    "bytes": 6537,
+    "sha256": "835f510094110c4dbc590a88d9b94d41474ae87b442785a949328feda5b9afce"
+  },
+  "runtime/src/morph/model/lightning/_training.py": {
+    "bytes": 22479,
+    "sha256": "0451c3c4f1ad46cc5380615a3fbf8452d1fb4a0d27d69c885a411a3612c0997c"
+  },
+  "runtime/src/morph/model/lightning/_validation.py": {
+    "bytes": 11509,
+    "sha256": "03890457e504eed96f23dd8633a6af4317a9ed61e80b3e4ec935aaf9e63038d7"
+  },
+  "runtime/src/morph/model/lightning/base.py": {
+    "bytes": 9294,
+    "sha256": "ce9159e7b387ece0a3ce010b164f675b81980d2f617b5b05d00f96695ee33e28"
+  },
+  "runtime/src/morph/model/losses/__init__.py": {
+    "bytes": 190,
+    "sha256": "bb9de5950496ed184a32e11834d8e6f3b7d406773509ffa6c406a86bb5249fae"
+  },
+  "runtime/src/morph/model/losses/functional.py": {
+    "bytes": 11765,
+    "sha256": "42ae52f604a0cb8b7fe08955cfd273945ca068727994f20b02c80cdee3a00e1b"
+  },
+  "runtime/src/morph/utils/__init__.py": {
+    "bytes": 78,
+    "sha256": "eb76ca86f046931eaea693a696317e8b3ce0b8ad8695a92e3744fd1892f2b9b8"
+  },
+  "runtime/src/morph/utils/checkpoints.py": {
+    "bytes": 2301,
+    "sha256": "750302eb039e3947a5d4c84e4e7d04eb46e974cc0031ef1d197a3573e20eae4b"
+  },
+  "runtime/src/morph/utils/edge_alignment.py": {
+    "bytes": 4896,
+    "sha256": "0edcde93cfa3a2d825e783ea5cf64715aa248e59ffb8c6a23d4490d74976a71e"
+  },
+  "runtime/src/morph/utils/graph_math.py": {
+    "bytes": 6755,
+    "sha256": "2bae326c3538d757aac66dcfc4a0b5d023b17bd987d814d3dab9a57f9ce60552"
+  },
+  "runtime/src/morph/utils/loss_accumulation.py": {
+    "bytes": 15115,
+    "sha256": "0f3eca42145e715290576dd86f4bcfb1553714ec4cabeda9246cfaf1e4f4323c"
+  },
+  "runtime/src/morph/utils/lr_schedulers.py": {
+    "bytes": 5554,
+    "sha256": "b954c12df413ca6676263ef1af960e2a4ff52692acfd4223e72a6e02122a3638"
+  },
+  "runtime/src/morph/utils/metrics/__init__.py": {
+    "bytes": 6730,
+    "sha256": "abec143d21e36343fff54db0ca8042111b9e70da426567b13ec36d6cb05696b6"
+  },
+  "runtime/src/morph/utils/metrics/chemistry.py": {
+    "bytes": 34152,
+    "sha256": "61b598b99fad41571f9334b28f94fbeeeb0282d1b80423b1f2bc8f14180884e8"
+  },
+  "runtime/src/morph/utils/metrics/distribution.py": {
+    "bytes": 3348,
+    "sha256": "0a6c25c4e65843308059a3effc44e1072837e747926074c1d7c47ba7ee72cb46"
+  },
+  "runtime/src/morph/utils/metrics/plotting.py": {
+    "bytes": 4270,
+    "sha256": "686d491d7fbc83cf35f772649a869dc1a271966d2a4fa6162d6e2334db735723"
+  },
+  "runtime/src/morph/utils/metrics/posecheck_strain.py": {
+    "bytes": 6695,
+    "sha256": "60a4751cce0e23d87d8602135de0c9850182fe97f20632cb9095445bdc6b0365"
+  },
+  "runtime/src/morph/utils/metrics/tensor.py": {
+    "bytes": 13503,
+    "sha256": "3fcb867cde5d47f81725acc3b862fc56af5ecce23dee3a11ebe6d05c1772c6fe"
+  },
+  "runtime/src/morph/utils/rdkit_utils.py": {
+    "bytes": 9018,
+    "sha256": "0fddc055c6de60893272fe8c50c92de29c9e72b7819827caac753e4cd0c6f1e5"
+  },
+  "runtime/src/morph/utils/run_config.py": {
+    "bytes": 8025,
+    "sha256": "0c289120887a051379523af528ebc9310f356f3b92efb898ddaf773dbf698aa2"
+  },
+  "runtime/src/morph/utils/vocab_utils.py": {
+    "bytes": 6343,
+    "sha256": "7b5506971b5092865b86311173b24d7fb70a881f123c57fd8696c412a25d618b"
+  },
+  "runtime/uv.lock": {
+    "bytes": 611263,
+    "sha256": "6fd6609500ea661b3b5c9f17d0997c5709295ca2cbbe16c97fc8259178fa931e"
+  },
+  "training_config.yaml": {
+    "bytes": 13428,
+    "sha256": "4fd11b415e9b4a599344431218db7f9960ba443dc29787410a4566ce5dacb434"
+  },
+  "validation.json": {
+    "bytes": 294,
+    "sha256": "5acbbb78be426e03abd58c3dc062b30b000e91e815c329dfa7f1422c9866b740"
+  }
+}

geom/explicit_h/metadata.json

@@ -0,0 +1,48 @@
+{
+  "dataset": "geom",
+  "variant": "explicit_h",
+  "includes_explicit_hydrogens": true,
+  "source_max_total_atoms_including_h": 72,
+  "training_run": "geom-explicit-e48-resume-16520626",
+  "resumed_from_completed_epoch": 48,
+  "source_snapshot": "geom_resume_e48_20261008",
+  "checkpoint_format": "Original Lightning checkpoint, current + EMA + optimizer state",
+  "epoch_numbering": "Completed epochs (one-based); original checkpoint epoch is zero-based",
+  "checkpoints": [
+    {
+      "filename": "epoch_050.ckpt",
+      "original_filename": "epoch=49-step=441050.ckpt",
+      "completed_epochs": 50,
+      "checkpoint_epoch_index": 49,
+      "global_step": 441050,
+      "bytes": 509181385,
+      "sha256": "016c4a5bbb4f5ef36d1d62dff772ecf8d384e8ab3e2c94ebeb4c4357183f690a"
+    }
+  ],
+  "atom_tokens": [
+    "Br",
+    "C",
+    "Cl",
+    "F",
+    "H",
+    "I",
+    "N",
+    "O",
+    "P",
+    "S"
+  ],
+  "edge_tokens": [
+    "<NO_BOND>",
+    "1",
+    "2",
+    "3"
+  ],
+  "charge_tokens": [
+    "-2",
+    "-1",
+    "0",
+    "1",
+    "2"
+  ],
+  "coordinate_std_angstrom": 2.4406259059906006
+}

geom/explicit_h/runtime/configs/callbacks/callbacks.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+
+## Put your callback args here
+callbacks:
+  monitor_metric: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},val/validity,val/batch/pair_dist_l1}
+  monitor_metric_mode: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},max,min}
+
+  #   model_checkpoints         -> best-by-metric (top-1, monitored)
+  #   every_n_epochs_checkpoint -> latest (no monitor; overwrites every epoch)
+  model_checkpoints:
+    save_top_k: 1
+    save_last: False
+    verbose: False
+
+  every_n_epochs_checkpoint:
+    every_n_epochs: 1
+    save_top_k: 1
+    save_last: False
+    verbose: False
+    save_on_train_epoch_end: True # Force per-train-epoch firing

geom/explicit_h/runtime/configs/data/chembl3d/chembl3d.yaml

@@ -0,0 +1,8 @@
+# LoQI ChEMBL3D source format, without dataset-specific size/element filters.
+defaults:
+  - /data/molecular@_here_
+  - _self_
+
+dataset_name: chembl3d
+dataset_target: morph.dataset.unconditional.chembl3d.FlowMatchingChEMBL3DDataset
+source_dir: ${oc.env:CHEMBL3D_SOURCE_DIR,null}

geom/explicit_h/runtime/configs/data/chembl3d/implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/chembl3d/chembl3d@_here_
+  - _self_
+
+hydrogen_mode: implicit

geom/explicit_h/runtime/configs/data/chembl3d/max100.yaml

@@ -0,0 +1,7 @@
+# Inclusive total-atom limit, applied before hydrogen removal.
+defaults:
+  - /data/chembl3d/chembl3d@_here_
+  - _self_
+
+max_atoms: 100
+cache_name: ${data.dataset_name}_max${data.max_atoms}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}

geom/explicit_h/runtime/configs/data/chembl3d/max100_implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/chembl3d/max100@_here_
+  - _self_
+
+hydrogen_mode: implicit

geom/explicit_h/runtime/configs/data/geom/geom.yaml

@@ -0,0 +1,11 @@
+# GEOM-Drugs source format and default molecule subset.
+defaults:
+  - /data/molecular@_here_
+  - _self_
+
+dataset_name: geom
+dataset_target: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+source_dir: ${oc.env:GEOM_SOURCE_DIR,null} # Optional prefiltered pickle directory.
+# Preserve the established GEOM subset. Limits apply before hydrogen removal.
+max_atoms: 72
+allowed_atom_symbols: [H, C, N, O, S, F, Cl, Br, P, I]

geom/explicit_h/runtime/configs/data/geom/implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/geom/geom@_here_
+  - _self_
+
+hydrogen_mode: implicit

geom/explicit_h/runtime/configs/data/geom/scaffold_decoration.yaml

@@ -0,0 +1,24 @@
+# GEOM-Drugs data config for the **scaffold decoration** task.
+#
+# Mirrors configs/data/qm9/scaffold_decoration.yaml but on GEOM-Drugs: source
+# is the selected scaffold of each target; target is the full molecule. The model
+# learns to grow it by inserting decoration atoms around the fixed scaffold.
+#
+# Inherits from configs/data/geom/geom.yaml and overrides:
+#   - datamodule._target_ -> ScaffoldLightningDataModule
+#
+# The scaffold cache is built lazily into ``scaffold_groups_dir`` on first run
+# (one file per split and scaffold kind) and reused afterwards.
+defaults:
+  - /data/geom/geom@_here_
+  - _self_
+
+scaffold_groups_dir: ${data.preprocessing.root}/processed/scaffold_groups
+scaffold_kind: rgroup # murcko, rgroup (BRICS), heavy (generate H), linker
+# Evaluate each distinct molecule/scaffold pair; training always draws randomly.
+enumerate_slicings: false
+datamodule:
+  _target_: morph.dataset.conditional.scaffold.datamodule.ScaffoldLightningDataModule
+  scaffold_groups_dir: ${data.scaffold_groups_dir}
+  scaffold_kind: ${data.scaffold_kind}
+  enumerate_slicings: ${data.enumerate_slicings}

geom/explicit_h/runtime/configs/data/molecular.yaml

@@ -0,0 +1,136 @@
+# Shared drug-like molecular training and preprocessing configuration.
+
+root_path: ${oc.env:PROJECT_ROOT}/data
+preprocess_workers: 10
+preprocess_chunk_size: 128
+statistics_workers: ${data.preprocess_workers}
+hydrogen_mode: explicit
+# The mode is part of the cache path; changing it discovers fresh vocab/statistics.
+cache_name: ${data.dataset_name}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}
+source_dir: null
+max_atoms: null
+allowed_atom_symbols: null
+
+n_atoms_strategy: "uniform" # "approx", "empirical", "learned"
+# Uniform source-count support padding; explicit so inference and training can
+# be matched or ablated without code changes.
+uniform_n_atoms_padding: 5
+optimal_transport: "equivariant"
+# Physical Angstroms (converted to normalized coordinates internally). These
+# reproduce the former normalized 0.25 / 0.2 scales at GEOM's training-set std.
+# Physical Angstroms. Insertion bridge noise vanishes as (1-t); move noise is
+# the persistent training-time jitter floor for matched atoms and insertions
+# that are already present.
+ins_noise_scale: 0.65
+move_noise_scale: 0.5
+
+# Drug-like datasets contain charged species; allow them in validity metrics.
+allow_charged: true
+
+
+# Preprocessing configuration - computes tokens from data and distributions from training dataset only
+preprocessing:
+  _target_: morph.dataset.core.preprocessing.Preprocessing
+  _recursive_: false
+  root: ${data.root_path}/${data.cache_name}
+  num_workers: ${data.statistics_workers}
+  train_dataset: ${data.datamodule.datasets.train}
+  atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+  edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+  charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+  distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+  # Tokens and distributions will be discovered from data and saved to files
+  # On subsequent runs, they will be loaded from the files
+
+vocab: null
+
+interpolator:
+  _target_: morph.flow_matching.interpolation.Interpolator
+  distributions: null
+  optimal_transport: ${data.optimal_transport}
+  ins_noise_scale: ${data.ins_noise_scale}
+  move_noise_scale: ${data.move_noise_scale}
+  # Training-only categorical corruption; pairs are scaled per molecule to
+  # match the expected number of changed atoms. Clean targets stay unchanged.
+  atom_type_noise_prob: 0.0
+  # OT matching (cost always squared-Euclidean; move cost in Angstrom). c_move
+  # scales the squared move distance (A^2); c_ins/c_del/c_sub are fixed A^2 costs,
+  # directly comparable. At 6.125 each, delete+insert costs 3.5^2, so atoms
+  # farther than 3.5 A apart are left unmatched by default.
+  pre_align_iters: 2   # equivariant-OT inner loop (Hungarian+Kabsch) iters
+  c_move: 1.0
+  c_sub: 0.0
+  c_ins: 6.125
+  c_del: 6.125
+  del_schedule: ${schedule}
+  ins_schedule: ${schedule}
+  sub_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+  sub_e_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+
+datamodule:
+  _target_: morph.dataset.core.lightning_datamodule.LightningDataModule
+  # tokens and distributions will be passed from preprocessing in run.py
+
+  vocab: ${data.vocab}
+  distributions: null
+  representation: ${representation}
+
+  datasets:
+    train:
+      _target_: ${data.dataset_target}
+      root: ${data.preprocessing.root}
+      preprocess_workers: ${data.preprocess_workers}
+      preprocess_chunk_size: ${data.preprocess_chunk_size}
+      hydrogen_mode: ${data.hydrogen_mode}
+      source_dir: ${data.source_dir}
+      max_atoms: ${data.max_atoms}
+      allowed_atom_symbols: ${data.allowed_atom_symbols}
+      vocab: ${data.vocab}
+      distributions: null
+
+    val:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+
+    test:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+
+  interpolator: ${data.interpolator}
+
+  rotate: ${model.rotate}
+  n_augmentations: ${model.n_augmentations}
+
+  num_workers:
+    train: 10
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 320
+    val: 128
+    test: 64
+
+  n_atoms_strategy: ${data.n_atoms_strategy}
+  uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+  time_dist: ${model.module.time_dist}

geom/explicit_h/runtime/configs/data/qm9/qm9.yaml

@@ -0,0 +1,106 @@
+## QM9 data and preprocessing configuration.
+
+dataset_name: qm9
+root_path: ${oc.env:PROJECT_ROOT}/data
+preprocess_workers: 0
+
+n_atoms_strategy: "uniform" # "approx", "empirical", "learned"
+# Expand both ends of the empirical count support for the uniform flow source.
+# QM9's minimum is already the hard lower bound of 3, so positive padding is
+# necessarily asymmetric; set to 0 for an unpadded 3--29 diagnostic prior.
+uniform_n_atoms_padding: 5
+optimal_transport: "equivariant"
+# Physical Angstroms (converted to normalized coordinates internally). Insertion
+# bridge noise vanishes as (1-t); move noise is the persistent training-time
+# jitter floor for matched atoms and insertions that are already present.
+ins_noise_scale: 0.35
+move_noise_scale: 0.3
+
+# QM9 contains only neutral molecules; charged species are treated as invalid.
+allow_charged: false
+
+
+# Preprocessing configuration - computes tokens from data and distributions from training dataset only
+preprocessing:
+  _target_: morph.dataset.core.preprocessing.Preprocessing
+  _recursive_: false
+  root: ${data.root_path}/qm9
+  train_dataset: ${data.datamodule.datasets.train}
+  # Canonical (representation-agnostic) cache. One preprocessing per dataset;
+  # representation-specific projection happens at runtime in the wrapper.
+  atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+  edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+  charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+  distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+
+vocab: None
+
+interpolator:
+  _target_: morph.flow_matching.interpolation.Interpolator
+  distributions: None
+  optimal_transport: ${data.optimal_transport}
+  ins_noise_scale: ${data.ins_noise_scale}
+  move_noise_scale: ${data.move_noise_scale}
+  # OT matching (cost always squared-Euclidean; move cost in Angstrom). c_move
+  # scales the squared move distance (A^2); c_ins/c_del/c_sub are fixed A^2 costs,
+  # directly comparable. At 6.125 each, delete+insert costs 3.5^2, so atoms
+  # farther than 3.5 A apart are left unmatched by default.
+  pre_align_iters: 2   # equivariant-OT inner loop (Hungarian+Kabsch) iters
+  c_move: 1.0
+  c_sub: 0.0
+  c_ins: 6.125
+  c_del: 6.125
+  del_schedule: ${schedule}
+  ins_schedule: ${schedule}
+  sub_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+  sub_e_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+
+datamodule:
+  _target_: morph.dataset.core.lightning_datamodule.LightningDataModule
+  # tokens and distributions will be passed from preprocessing in run.py
+
+  vocab: ${data.vocab}
+  distributions: None
+  representation: ${representation}
+
+  datasets:
+    train:
+      _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+      root: ${data.root_path}/qm9
+      vocab: ${data.vocab}
+      distributions: None
+
+    val:
+      - _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+        root: ${data.root_path}/qm9
+        vocab: ${data.vocab}
+        distributions: None
+
+    test:
+      - _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+        root: ${data.root_path}/qm9
+        vocab: ${data.vocab}
+        distributions: None
+
+  interpolator: ${data.interpolator}
+
+  rotate: ${model.rotate}
+  n_augmentations: ${model.n_augmentations}
+
+  num_workers:
+    train: 10
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 1024
+    val: 128
+    test: 32
+
+  n_atoms_strategy: ${data.n_atoms_strategy}
+  uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+  time_dist: ${model.module.time_dist}

geom/explicit_h/runtime/configs/data/qm9/scaffold_decoration.yaml

@@ -0,0 +1,22 @@
+# QM9 data config for the **scaffold decoration** task.
+#
+# Source is the selected scaffold of each target; target is the full molecule.
+# The model learns to grow the molecule by inserting decoration atoms
+# around the fixed scaffold.
+#
+# Inherits from configs/data/qm9/qm9.yaml and overrides:
+#   - datamodule._target_ → ScaffoldLightningDataModule
+#
+defaults:
+  - /data/qm9/qm9@_here_
+  - _self_
+
+scaffold_groups_dir: ${data.root_path}/qm9/processed/scaffold_groups
+scaffold_kind: rgroup # murcko, rgroup (BRICS), heavy (generate H), linker
+# Evaluate each distinct molecule/scaffold pair; training always draws randomly.
+enumerate_slicings: false
+datamodule:
+  _target_: morph.dataset.conditional.scaffold.datamodule.ScaffoldLightningDataModule
+  scaffold_groups_dir: ${data.scaffold_groups_dir}
+  scaffold_kind: ${data.scaffold_kind}
+  enumerate_slicings: ${data.enumerate_slicings}

geom/explicit_h/runtime/configs/default.yaml

@@ -0,0 +1,25 @@
+defaults:
+  - _self_
+  # Single-knob presets bundle
+  # `data` + `embeddings` (with dataset-tuned ExtrapolatableCount params)
+  # plus the shared `core` group selections. Override individual groups as
+  # usual, e.g. `defaults=geom/geom model=backbones/dit embeddings=qm9`.
+  - defaults: qm9/qm9
+
+# What the model trains on. Datasets declare their CAPABILITIES; run.py validates
+# the requested mode against them before preprocessing.
+#   pointcloud         — x, a only (charges and edges are dummy singletons)
+#   charged_pointcloud — x, a, c   (edges are dummy singletons)
+#   geometric_graph    — x, a, edge_index, e (charges are dummy neutral)
+#   molecule           — full chemistry (current behavior)
+representation: molecule
+seed: 42
+
+# Run directory: <root>/<dataset>/<cfg_method>/<date>/<wandb_id>. Root is
+# MORPH_RUN_ROOT (scratch, set in .env) so every run's outputs/checkpoints/wandb
+# land off-repo; falls back to <cwd>/outputs for ad-hoc runs with no env set.
+# wandb id is MORPH_RUN_ID (same as WANDB_RUN_ID). Set in the shell before
+# torchrun, or set automatically for single-process runs in run.py before Hydra composes.
+hydra:
+  run:
+    dir: ${oc.env:MORPH_RUN_ROOT,${hydra:runtime.cwd}/outputs}/${hydra:runtime.choices.data}/${now:%Y-%m-%d}/${oc.env:MORPH_RUN_ID}

geom/explicit_h/runtime/configs/defaults/chembl3d/chembl3d.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - core
+  - /data: chembl3d/chembl3d
+  - /embeddings: geom

geom/explicit_h/runtime/configs/defaults/chembl3d/implicit_h.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: chembl3d/implicit_h
+  - /embeddings: geom

geom/explicit_h/runtime/configs/defaults/chembl3d/implicit_h_2x4090.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+defaults:
+  - /defaults/chembl3d/train_2x4090
+  - override /data: chembl3d/max100_implicit_h
+  - _self_
+
+# Train from scratch: the atom vocabulary/statistics differ from the H model.
+# Preserve global batch 640, eval every epoch, both checkpoint callbacks every 5.
+logging:
+  project: morph-chembl3d-implicit-h

geom/explicit_h/runtime/configs/defaults/chembl3d/train_2x4090.yaml

@@ -0,0 +1,40 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: chembl3d/max100
+  - /embeddings: geom
+  - _self_
+
+# Profiled on max100 ChEMBL: 2 GPUs * 320 * 1 = 640 effective global batch.
+data:
+  datamodule:
+    train_prefetch_factor: 1
+    batch_size:
+      train: 320
+      val: 32
+      test: 32
+    num_workers:
+      train: 20
+      val: 2
+      test: 2
+
+trainer:
+  trainer:
+    accelerator: gpu
+    devices: 2
+    accumulate_grad_batches: 1
+    max_epochs: -1
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+  # Validation runs during fit; leave wall-clock headroom for final checkpointing.
+  do_validate: false
+  do_predict: false
+  final_checkpoint_path: final.ckpt
+
+callbacks:
+  model_checkpoints:
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link

geom/explicit_h/runtime/configs/defaults/core.yaml

@@ -0,0 +1,13 @@
+# @package _global_
+# Shared default groups that don't change between QM9 / GEOM runs.
+# Dataset-specific presets (defaults/qm9/qm9.yaml, defaults/geom/geom.yaml) inherit
+# from this and only add `data` + `embeddings`.
+defaults:
+  - /heads: head
+  - /insertion_head: unimodal
+  - /model: backbones/dit
+  - /schedule: smoothstep
+  - /trainer: trainer
+  - /callbacks: callbacks
+  - /logging: wandb_logger
+  - /integrator: integrator

geom/explicit_h/runtime/configs/defaults/geom/geom.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+# GEOM preset: shared core + GEOM data and embeddings.
+defaults:
+  - core
+  - /data: geom/geom
+  - /embeddings: geom
+  - _self_
+
+# Evaluate every epoch; keep every fifth completed epoch for releases/resuming.
+trainer:
+  trainer:
+    check_val_every_n_epoch: 1
+
+callbacks:
+  model_checkpoints:
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link

geom/explicit_h/runtime/configs/defaults/geom/implicit_h.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: geom/implicit_h
+  - /embeddings: geom

geom/explicit_h/runtime/configs/defaults/geom/scaffold_decoration.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+# GEOM-Drugs **scaffold decoration** fine-tune preset.
+#
+# Mirrors defaults/qm9/scaffold_decoration.yaml on GEOM-Drugs. Bare-scaffold
+# source -> full-molecule target: the model learns to grow a molecule by
+# inserting decoration atoms around the fixed scaffold.
+#
+# Differences from the QM9 scaffold preset:
+#   - data       -> geom/scaffold_decoration (GEOM base datasets)
+#   - embeddings -> geom
+#
+# Reuses the scaffold_finetune model config (loads the pretrained unconditional
+# DiT into the clean lane + adds the node/edge conditioning and the gated
+# cross-attn blocks).
+#
+# Pair with a GEOM unconditional DiT checkpoint via
+# ``model.module.pretrained_ckpt=...``.
+defaults:
+  - core
+  - /data: geom/scaffold_decoration
+  - /embeddings: geom
+  - override /model: finetune/scaffold_finetune
+  - _self_

geom/explicit_h/runtime/configs/defaults/qm9/qm9.yaml

@@ -0,0 +1,6 @@
+# @package _global_
+# QM9 preset: shared core + QM9 data and embeddings.
+defaults:
+  - core
+  - /data: qm9/qm9
+  - /embeddings: qm9

geom/explicit_h/runtime/configs/defaults/qm9/scaffold_decoration.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+# QM9 **scaffold decoration** fine-tune preset.
+#
+# Bare-scaffold source → full-molecule target. The model learns to grow
+# a molecule by inserting decoration atoms around the fixed scaffold.
+defaults:
+  - core
+  - /data: qm9/scaffold_decoration
+  - /embeddings: qm9
+  - override /model: finetune/scaffold_finetune
+  - _self_

geom/explicit_h/runtime/configs/embeddings/geom.yaml

@@ -0,0 +1,28 @@
+# Shared embedding stack tuned for GEOM-Drugs.
+# - node_count_embedding uses ExtrapolatableCountEmbedding with periods/max_count
+#   scaled up for the much larger atom counts in GEOM-Drugs.
+
+atom_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.atom_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+edge_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.edge_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+time_embedding:
+  _target_: morph.model.layers.embedding.TimeEmbedding
+  embedding_dim: 128
+  out_dim: 128
+
+node_count_embedding:
+  _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+  embedding_dim: 64
+  out_dim: 128
+  min_period: 128.0
+  max_period: 1024.0
+  max_count: 96.0

geom/explicit_h/runtime/configs/embeddings/qm9.yaml

@@ -0,0 +1,29 @@
+# Shared embedding stack tuned for QM9.
+# - node_count_embedding uses ExtrapolatableCountEmbedding so the model can
+#   generalise slightly outside the training atom-count support.
+# - Defaults sized for QM9 (training counts in [3, 29]).
+
+atom_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.atom_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+edge_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.edge_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+time_embedding:
+  _target_: morph.model.layers.embedding.TimeEmbedding
+  embedding_dim: 128
+  out_dim: 128
+
+node_count_embedding:
+  _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+  embedding_dim: 64
+  out_dim: 128
+  min_period: 64.0
+  max_period: 512.0
+  max_count: 32.0

geom/explicit_h/runtime/configs/heads/head.yaml

@@ -0,0 +1,75 @@
+# Standard molecular DiT readouts. The insertion-content head is an independent
+# Hydra choice (``insertion_head=unimodal`` by default, or ``=gmm``).
+
+heads:
+  _target_: morph.model.layers.heads.MultiHeadModule
+  output_init_std: ${model.head_initialization.output_std}
+  independent_adapters: ${oc.select:model.head_independent_adapters,false}
+
+  heads_configs:
+    graph_aggregation: mean
+
+    node_heads:
+      atom_type_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: ${len:${data.vocab.atom_tokens}}
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: state
+
+      charge_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: ${len:${data.vocab.charge_tokens}}
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: state
+
+      ins_rate_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_rate: ${model.head_initialization.ins_rate}
+
+      do_del_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_probability: ${model.head_initialization.do_del_probability}
+
+      do_sub_a_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_probability: ${model.head_initialization.do_sub_a_probability}
+
+    edge_heads:
+      edge_type_head:
+        input_dim: ${model.head_edge_dim}
+        output_dim: ${len:${data.vocab.edge_tokens}}
+        adapter_dim: ${model.head_edge_adapter_dim}
+        group: state
+
+      do_sub_e_head:
+        input_dim: ${model.head_edge_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_edge_adapter_dim}
+        group: state
+        initial_probability: ${model.head_initialization.do_sub_e_probability}
+
+ins_edge_head:
+  _target_: morph.model.layers.heads.InsertionEdgeHead
+  _recursive_: false
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_edge_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  n_edge_types: ${len:${data.vocab.edge_tokens}}
+  n_atom_types: ${len:${data.vocab.atom_tokens}}
+  n_charge_types: ${len:${data.vocab.charge_tokens}}
+  rbf_embedding_args:
+    _target_: morph.model.layers.embedding.RBFEmbedding
+    num_rbf: 128
+    rbf_dmax: 3.5
+    out_dim: 64
+    dropout: 0.0
+    trainable: false

geom/explicit_h/runtime/configs/insertion_head/gmm.yaml

@@ -0,0 +1,16 @@
+# Optional mixture-density insertion-content head.
+mode: gmm
+dimension: 3
+num_components: 10
+
+model:
+  _target_: morph.model.layers.heads.GMMInsertionHead
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_adapter_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  num_components: ${insertion_head.num_components}
+  sigma_min: 0.05
+  sigma_max: 2.0
+  sigma_init: 0.5
+  N_a: ${len:${data.vocab.atom_tokens}}
+  N_c: ${len:${data.vocab.charge_tokens}}

geom/explicit_h/runtime/configs/insertion_head/unimodal.yaml

@@ -0,0 +1,13 @@
+# Default insertion-content head: one nearest pending atom per spawning node.
+mode: unimodal
+dimension: 3
+num_components: 1
+
+model:
+  _target_: morph.model.layers.heads.UnimodalInsertionHead
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_adapter_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  independent_adapters: ${oc.select:model.head_independent_adapters,false}
+  N_a: ${len:${data.vocab.atom_tokens}}
+  N_c: ${len:${data.vocab.charge_tokens}}

geom/explicit_h/runtime/configs/integrator/integrator.yaml

@@ -0,0 +1,30 @@
+integrator:
+  _target_: morph.flow_matching.integration.RateIntegrator
+  vocab: ${data.vocab}
+  distributions: None
+  insertion_head_mode: ${insertion_head.mode}
+  insertion_dimension: ${insertion_head.dimension}
+  insertion_num_components: ${insertion_head.num_components}
+  # strategy used for sampling time during **integration**
+  # NOTE can be different from the time strategy used for the model (see below)
+  time_strategy: log # "linear", "log", or "beta"
+  # shape params for time_strategy: beta (Beta inverse-CDF spacing).
+  # defaults match the model's beta_sched (concentration1=2.0, concentration0=1.0)
+  time_beta_alpha: 2.0
+  time_beta_beta: 1.0
+  num_integration_steps: 100
+  ins_noise_scale: ${data.ins_noise_scale}
+  # When true, freshly-inserted nodes are pushed to the t+dt marginal (kappa-mix
+  # on atom/bond types, (1-t)*ins_noise on x) to match training. Charge remains
+  # a clean prediction target rather than an interpolated state. When false,
+  # insertions enter at the model's clean prediction.
+  renoise_insertions: true
+  # Integrate the known categorical schedule hazard over every time interval
+  # instead of using its left-endpoint Euler approximation.
+  # With either method, the final interval always completes pending discrete
+  # transitions using the learned edit probabilities (no extra model call).
+  exact_schedule_integration: true
+  del_schedule: ${data.interpolator.del_schedule}
+  ins_schedule: ${data.interpolator.ins_schedule}
+  sub_schedule: ${data.interpolator.sub_schedule}
+  sub_e_schedule: ${data.interpolator.sub_e_schedule}

geom/explicit_h/runtime/configs/logging/wandb_logger.yaml

@@ -0,0 +1,6 @@
+## Put your model args here
+enabled: true
+project: morph-${data.dataset_name}
+name: ${oc.env:MORPH_RUN_ID,base}
+save_dir: .
+offline: false

geom/explicit_h/runtime/configs/loss_weights/calibrated_geom.yaml

@@ -0,0 +1,18 @@
+# @package model.module.loss_weights
+# Balance-only (geometric-mean=1) calibration weights, scaled down by a global
+# factor of 0.2 (1/5). At the original geometric-mean=1 scale the realized
+# full-model grad-norm was ~5, far above gradient_clip_val=1.0, so clipping
+# fired on every step. The 0.2 factor brings the realized norm to ~1.0 so the
+# clip becomes a spike safety-valve again. Re-running calibration now applies
+# this scaling automatically (see trainer.target_total_grad_norm).
+c: 0.164
+do_del: 0.20
+do_sub_a: 0.30
+do_sub_e: 0.24
+ins_e: 0.21
+ins_e_ii: 0.22
+ins_state: 0.14
+ins_rate: 0.20
+sub_a_class: 0.15
+sub_e_class: 0.19
+x: 0.22

geom/explicit_h/runtime/configs/loss_weights/calibrated_geom_frozen.yaml

@@ -0,0 +1,25 @@
+# @package model.module.loss_weights
+# Balance-only (geometric-mean = 1) loss weights — ratios only. The learning
+# rate owns step size and gradient_clip_val stays a pure spike catch, so these
+# are NOT scaled to a target grad-norm (unlike the older calibrated_geom.yaml,
+# which baked in a ~0.2 global scale to hit the clip target).
+#
+# Source: frozen init-state calibration (trainer.calibrate_mode=loss,
+# calibrate_freeze_model=true), geom / DiT / uncond,
+# run 2026-06-07 (outputs/geom/uncond/cal_then_lr/2026-06-07_12-29-25/calibrate).
+# Because the model is held at init, every probe batch is an unbiased snapshot of
+# the initial per-term losses — so the three binary do_* losses come out equal
+# (all BCE at p=0.5 = ln2), unlike the previous 2500-step *training* probe
+# (calibrated_geom.yaml) which gave them different weights purely from differing
+# training rates. See notebooks/lr_calibration_analysis.ipynb for the comparison.
+c: 0.745146
+do_del: 1.730107
+do_sub_a: 1.730107
+do_sub_e: 1.730106
+ins_e: 0.867232
+ins_e_ii: 0.867232
+ins_state: 0.519379
+ins_rate: 1.128104
+sub_a_class: 0.522109
+sub_e_class: 0.867232
+x: 1.298792

geom/explicit_h/runtime/configs/loss_weights/tuned_weights.yaml

@@ -0,0 +1,15 @@
+# @package model.module.loss_weights
+# Canonical tuned weights for molecular training. Backbones can select this
+# preset explicitly, and task-specific experiments may override the group. The
+# package directive mounts the values under model.module.loss_weights.
+do_sub_a: 0.2
+sub_a_class: 0.03
+do_sub_e: 0.25
+sub_e_class: 0.5
+do_del: 0.4
+ins_rate: 0.15
+ins_state: 0.01
+ins_e: 0.4
+x: 2.0 # was 1.0
+c: 0.5
+ins_e_ii: 0.03

geom/explicit_h/runtime/configs/loss_weights/uniform.yaml

@@ -0,0 +1,15 @@
+# @package model.module.loss_weights
+# Uniform 1.0 weight per loss component. The key set matches
+# LightningModuleRates.LOSS_GROUPS exactly (the 11 actual loss terms
+# emitted by shared_step).
+do_sub_a: 1.0
+sub_a_class: 1.0
+do_sub_e: 1.0
+sub_e_class: 1.0
+do_del: 1.0
+ins_rate: 1.0
+ins_state: 1.0
+ins_e: 1.0
+ins_e_ii: 1.0
+x: 1.0
+c: 1.0

geom/explicit_h/runtime/configs/model/backbones/dit.yaml

@@ -0,0 +1,123 @@
+## Molecular DiT: the repository's first-class generative backbone.
+##
+## This intentionally keeps charge out of the state and adds neither an
+## equivariant lane nor an evolving dense pair state.  It retains the atom-count
+## condition: time plus current graph size disambiguates otherwise identical
+## local states and is especially useful while insertion/deletion is active.
+defaults:
+  - /optimizer: muon_aux_adam
+  - /loss_weights: tuned_weights
+  - _self_
+
+rotate: true
+n_augmentations: 1
+
+integrator: ${integrator.integrator}
+
+beta_sched:
+  _target_: torch.distributions.Beta
+  concentration1: 2.0
+  concentration0: 1.0
+
+linear_sched:
+  _target_: torch.distributions.Uniform
+  low: 0.0
+  high: 1.0
+
+sample_time_strategy: linear
+
+module:
+  _target_: morph.model.lightning.base.LightningModuleRates
+  insertion_head_mode: ${insertion_head.mode}
+  ins_noise_scale: ${data.ins_noise_scale}
+  ema_decay: 0.999
+  ema_warmup_steps: 5000
+  use_ema_for_eval: true
+  time_dist: ${if:${eq:${model.sample_time_strategy}, log}, ${model.beta_sched}, ${model.linear_sched}}
+  integrator: ${model.integrator}
+  vocab: ${data.vocab}
+  use_learnable_loss_weights: false
+  use_time_weights: false
+
+  model:
+    _target_: morph.model.backbones.dit.DiTBackboneWithHeads
+    _recursive_: false
+
+    atom_type_embedding_args: ${embeddings.atom_type_embedding}
+    edge_type_embedding_args: ${embeddings.edge_type_embedding}
+    time_embedding_args: ${embeddings.time_embedding}
+    node_count_embedding_args: ${embeddings.node_count_embedding}
+    # Sum the existing incident edge embeddings into each atom's input, excluding
+    # NO_BOND, through a zero-initialized projection. Uses the existing heads.
+    incident_bond_embedding: false
+    backbone_model_args: ${model.backbone}
+    heads_args: ${heads.heads}
+    insertion_head_args: ${insertion_head.model}
+    ins_edge_head_args: ${heads.ins_edge_head}
+
+head_node_dim: ${model.backbone.out_node_nf}
+head_edge_dim: ${model.backbone.out_edge_nf}
+
+# Shared adapters avoid seven separate full-width MLPs.  At the >150M preset,
+# node/insertion adapters widen while pair readouts remain compact.
+head_node_adapter_dim: 256
+head_edge_adapter_dim: 256
+head_insertion_adapter_dim: 256
+head_insertion_edge_dim: 256
+head_independent_adapters: true
+
+# Categorical biases are replaced at module construction with exact dataset
+# marginals.  There is no dataset-level transition-frequency statistic for the
+# binary/rate heads, so their defaults deliberately preserve the old neutral
+# predictions instead of inventing a corpus-independent edit prior.
+head_initialization:
+  output_std: 1.0e-3
+  coordinate_output_std: 0.0
+  do_del_probability: 0.5
+  do_sub_a_probability: 0.5
+  do_sub_e_probability: 0.5
+  ins_rate: 0.6931471805599453
+
+token_weighting:
+  type_loss_token_weights: uniform
+  weight_alpha: 0.4
+
+# Optional low-rank adaptation of the attention and SwiGLU projections. The
+# adapter residual is zero at initialization, so enabling this while loading a
+# pretrained checkpoint preserves its outputs exactly.
+lora:
+  enabled: false
+  rank: 8
+  alpha: 8.0
+  dropout: 0.0
+  targets: [qkv, proj, mlp.gate_value, mlp.out_proj]
+  freeze_base: true
+
+backbone:
+  _target_: morph.model.backbones.dit.DiTBackbone
+  _recursive_: false
+  d_input: ${embeddings.atom_type_embedding.out_dim}
+  d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+  d_model: 512
+  out_node_nf: 256
+  out_edge_nf: 256
+  num_layers: 8
+  nhead: 16
+  mlp_ratio: 4.0
+  ffn_hidden_dim: null
+  ffn_multiple_of: 64
+  qk_norm: true
+  coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+  center_coordinate_output: true
+  in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+  edge_pair_dim: 256
+  pair_bias_hidden_dim: 32
+  num_atom_types: ${len:${data.vocab.atom_tokens}}
+  lora: ${model.lora}
+  rbf_embedding_args:
+    _target_: morph.model.layers.embedding.RBFEmbedding
+    num_rbf: 128
+    rbf_dmax: 3.5
+    out_dim: 64
+    dropout: 0.0
+    trainable: false

geom/explicit_h/runtime/configs/model/finetune/lora_finetune.yaml

@@ -0,0 +1,13 @@
+# Parameter-efficient fine-tuning of the unconditional DiT.
+defaults:
+  - /model/backbones/dit@_here_
+  - _self_
+
+lora:
+  enabled: true
+
+module:
+  _target_: morph.model.finetune.base.LightningModuleRatesFinetune
+  pretrained_ckpt: ${oc.env:MORPH_PRETRAINED_CKPT}
+  pretrained_prefer_ema: true
+  freeze_pretrained_backbone: true

geom/explicit_h/runtime/configs/model/finetune/scaffold_finetune.yaml

@@ -0,0 +1,134 @@
+# DiT backbone augmented with Flamingo-style gated cross-attention for
+# scaffold-decoration fine-tuning.
+#
+# Conditioning (all derived from mol_t.keep_mask):
+#   - node_cond : per-atom scaffold membership            (always on)
+#   - edge_cond : per-edge scaffold class                 (always on)
+#   - cross-attn: scaffold subgraph encoded by a shared LightweightContextEncoder
+#                 (OPTIONAL — toggle via backbone.use_cross_attn):
+#                   * false → option 1: node+edge only
+#                   * true  → option 2: node+edge+cross-attn
+#
+# Inherits everything from configs/model/backbones/dit.yaml and overrides:
+#   - module._target_ → LightningModuleRatesScaffoldFinetune (loads a pretrained
+#       unconditional DiT checkpoint and adds scaffold loss/inference masking)
+#   - model._target_  → ScaffoldConditionalDiTBackboneWithHeads (auto-builds the
+#       node/edge cond and, when use_cross_attn, the scaffold cross-attn tokens)
+#   - backbone._target_ → ConditionalDiTBackbone (DiT + gated cross-attn blocks)
+#
+# Provides a ``pretrained_ckpt`` field for the lightning module to load.
+defaults:
+  - /model/backbones/dit@_here_
+  - _self_
+
+module:
+  _target_: morph.model.finetune.scaffold.LightningModuleRatesScaffoldFinetune
+  # Path to the unconditional DiT checkpoint to load weights from.
+  # ConditionalDiT keys are auto-retargeted (backbone.* → backbone.dit.*).
+  pretrained_ckpt: ${oc.env:MORPH_PRETRAINED_CKPT}
+  pretrained_prefer_ema: true
+  # If true, dispatches to ``self.model.freeze_pretrained()`` — for the
+  # cross-attn scaffold model this freezes the entire clean lane
+  # (embedding, ``backbone.dit``, all heads). Only the gated cross-attn
+  # blocks and edge gate plus the scaffold encoder train.
+  freeze_pretrained_backbone: true
+  optimizer_config:
+    optimizer:
+      lr: 1e-3
+    aux_optimizer:
+      lr: 1e-4
+  # Whether to exclude scaffold atoms from sub / x losses. These atoms are
+  # frozen at inference, so training them adds no useful signal.
+  exclude_scaffold_from_sub_loss: true
+  exclude_scaffold_from_x_loss: true
+  # Pin kept coordinates during integration and restore the input coordinate
+  # frame on output. To train flowing scaffold geometry, set this and the x-loss
+  # exclusion above to false.
+  freeze_keep_positions: true
+
+  # Override the model target to the scaffold-aware cross-attn wrapper.
+  model:
+    _target_: morph.model.finetune.scaffold.ScaffoldConditionalDiTBackboneWithHeads
+    _recursive_: false
+
+    atom_type_embedding_args: ${embeddings.atom_type_embedding}
+    edge_type_embedding_args: ${embeddings.edge_type_embedding}
+    time_embedding_args: ${embeddings.time_embedding}
+    node_count_embedding_args: ${embeddings.node_count_embedding}
+
+    backbone_model_args: ${model.backbone}
+
+    heads_args: ${heads.heads}
+    insertion_head_args: ${insertion_head.model}
+    ins_edge_head_args: ${heads.ins_edge_head}
+
+    # Cross-attention context encoder (option 2). Required because the backbone
+    # below has use_cross_attn: true. For option 1 (node+edge only), set
+    # backbone.use_cross_attn=false AND null this (the two must agree).
+    xattn_encoder_args: ${model.xattn_encoder}
+    # Geometric correspondence prior (query atom ↔ co-located scaffold token).
+    use_distance_bias: true
+
+# Cross-attention context encoder (option 2): encodes the scaffold subgraph of
+# mol_t (extracted via filter_nodes) into context tokens. ctx_dim must match
+# model.backbone.ctx_dim.
+# Node/edge conditioning is NOT configured here — the wrapper builds those directly
+# from backbone.d_model / backbone.in_edge_nf. Null this (and set
+# backbone.use_cross_attn=false) for node+edge-only scaffold conditioning.
+xattn_encoder:
+  _target_: morph.model.finetune.base.LightweightContextEncoder
+  num_atom_types: ${len:${data.vocab.atom_tokens}}
+  num_charges: ${len:${data.vocab.charge_tokens}}
+  atom_embed_dim: 64
+  charge_embed_dim: 16
+  pos_embed_dim: 32
+  ctx_dim: ${model.backbone.ctx_dim}
+  num_attn_layers: 4
+  nhead: 8
+  mlp_ratio: 4.0
+  dropout: 0.0
+  # Edge-based pair bias on the encoder's self-attention (RBF distance +
+  # edge-type of the scaffold's own bonds) so tokens mix with geometry/topology
+  # awareness, like the DiT backbone. num_edge_types/edge_embed_dim only used here.
+  use_pair_bias: true
+  num_edge_types: ${len:${data.vocab.edge_tokens}}
+  edge_embed_dim: 32
+  num_rbf: 32
+  rbf_dmax: 5.0
+  pair_bias_hidden_dim: 32
+
+# Override the backbone to use ConditionalDiTBackbone.
+backbone:
+  _target_: morph.model.backbones.conditional_dit.ConditionalDiTBackbone
+  _recursive_: false
+  d_input: ${embeddings.atom_type_embedding.out_dim}
+  d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+  d_model: 512
+  out_node_nf: 256
+  out_edge_nf: 256
+  num_layers: 8
+  nhead: 16
+  mlp_ratio: 4.0
+  ffn_hidden_dim: null
+  ffn_multiple_of: 64
+  qk_norm: true
+  coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+  center_coordinate_output: true
+  in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+  edge_pair_dim: 256
+  pair_bias_hidden_dim: 32
+  num_atom_types: ${len:${data.vocab.atom_tokens}}
+  lora: ${model.lora}
+  ctx_dim: 128
+  use_cross_attn: true   # cross-attention (cond_tokens) port
+  # Per-node conditioning: scaffold membership (which atoms to keep).
+  use_node_cond: true
+  # Per-edge conditioning: the scaffold encoder emits edge_cond at in_edge_nf.
+  use_edge_cond: true
+  rbf_embedding_args:
+    _target_: morph.model.layers.embedding.RBFEmbedding
+    num_rbf: 128
+    rbf_dmax: 3.5
+    out_dim: 64
+    dropout: 0.0
+    trainable: false

geom/explicit_h/runtime/configs/optimizer/muon_aux_adam.yaml

@@ -0,0 +1,39 @@
+# @package model.module.optimizer_config
+# Muon + auxiliary AdamW preset for the molecular DiT. Per-run overrides
+# go on ``model.module.optimizer_config``:
+#   model.module.optimizer_config.optimizer.lr=0.005 ...
+implementation: torch
+
+optimizer:
+  # no _target_ -> triggers Muon path in configure_optimizers
+  # The RTX 4090 Euler template selects native-arithmetic CUDA graph replay.
+  execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+  lr: 0.01
+  momentum: 0.95
+  weight_decay: 0.0
+  nesterov: true
+  ns_steps: 5
+  eps: 1e-7
+  adjust_lr_fn: original
+
+aux_optimizer:
+  lr: 3e-4
+  betas: [0.9, 0.95]
+  eps: 1e-10
+  weight_decay: 0.0
+
+# Lightning's two-optimizer path uses manual optimization, so clipping is
+# configured here and applied once to the full model gradient.
+gradient_clip_val: 1.0
+gradient_clip_algorithm: norm
+
+scheduler:
+  _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+  # configure_optimizers injects each native optimizer here.
+  warmup_steps: 2000        # Typically 2-5% of total_steps
+  total_steps: 200000        # The total number of training iterations
+  min_lr_fraction: 0.05      # Decays to 5% of each optimizer's peak LR
+  last_epoch: -1
+
+interval: step
+monitor: null

geom/explicit_h/runtime/configs/resume_geom.yaml

@@ -0,0 +1,417 @@
+representation: molecule
+seed: 42
+heads:
+  heads:
+    _target_: morph.model.layers.heads.MultiHeadModule
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    heads_configs:
+      graph_aggregation: mean
+      node_heads:
+        atom_type_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.atom_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        charge_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.charge_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        ins_rate_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_rate: ${model.head_initialization.ins_rate}
+        do_del_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_del_probability}
+        do_sub_a_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_sub_a_probability}
+      edge_heads:
+        edge_type_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: ${len:${data.vocab.edge_tokens}}
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+        do_sub_e_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+          initial_probability: ${model.head_initialization.do_sub_e_probability}
+  ins_edge_head:
+    _target_: morph.model.layers.heads.InsertionEdgeHead
+    _recursive_: false
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_edge_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    n_edge_types: ${len:${data.vocab.edge_tokens}}
+    n_atom_types: ${len:${data.vocab.atom_tokens}}
+    n_charge_types: ${len:${data.vocab.charge_tokens}}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+insertion_head:
+  mode: unimodal
+  dimension: 3
+  num_components: 1
+  model:
+    _target_: morph.model.layers.heads.UnimodalInsertionHead
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_adapter_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    N_a: ${len:${data.vocab.atom_tokens}}
+    N_c: ${len:${data.vocab.charge_tokens}}
+model:
+  module:
+    optimizer_config:
+      implementation: torch
+      optimizer:
+        execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+        lr: 0.01
+        momentum: 0.95
+        weight_decay: 0.0
+        nesterov: true
+        ns_steps: 5
+        eps: 1.0e-07
+        adjust_lr_fn: original
+      aux_optimizer:
+        lr: 0.0003
+        betas:
+        - 0.9
+        - 0.95
+        eps: 1.0e-10
+        weight_decay: 0.0
+      gradient_clip_val: 1.0
+      gradient_clip_algorithm: norm
+      scheduler:
+        _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+        warmup_steps: 2000
+        total_steps: 200000
+        min_lr_fraction: 0.05
+        last_epoch: -1
+      interval: step
+      monitor: null
+    loss_weights:
+      do_sub_a: 0.2
+      sub_a_class: 0.03
+      do_sub_e: 0.25
+      sub_e_class: 0.5
+      do_del: 0.4
+      ins_rate: 0.15
+      ins_state: 0.01
+      ins_e: 0.4
+      x: 2.0
+      c: 0.5
+      ins_e_ii: 0.03
+    _target_: morph.model.lightning.base.LightningModuleRates
+    insertion_head_mode: ${insertion_head.mode}
+    ins_noise_scale: ${data.ins_noise_scale}
+    ema_decay: 0.999
+    ema_warmup_steps: 5000
+    use_ema_for_eval: true
+    time_dist: ${if:${eq:${model.sample_time_strategy}, log}, ${model.beta_sched},
+      ${model.linear_sched}}
+    integrator: ${model.integrator}
+    vocab: ${data.vocab}
+    use_learnable_loss_weights: false
+    use_time_weights: false
+    model:
+      _target_: morph.model.backbones.dit.DiTBackboneWithHeads
+      _recursive_: false
+      atom_type_embedding_args: ${embeddings.atom_type_embedding}
+      edge_type_embedding_args: ${embeddings.edge_type_embedding}
+      time_embedding_args: ${embeddings.time_embedding}
+      node_count_embedding_args: ${embeddings.node_count_embedding}
+      incident_bond_embedding: false
+      backbone_model_args: ${model.backbone}
+      heads_args: ${heads.heads}
+      insertion_head_args: ${insertion_head.model}
+      ins_edge_head_args: ${heads.ins_edge_head}
+  rotate: true
+  n_augmentations: 1
+  integrator: ${integrator.integrator}
+  beta_sched:
+    _target_: torch.distributions.Beta
+    concentration1: 2.0
+    concentration0: 1.0
+  linear_sched:
+    _target_: torch.distributions.Uniform
+    low: 0.0
+    high: 1.0
+  sample_time_strategy: linear
+  head_node_dim: ${model.backbone.out_node_nf}
+  head_edge_dim: ${model.backbone.out_edge_nf}
+  head_node_adapter_dim: 256
+  head_edge_adapter_dim: 256
+  head_insertion_adapter_dim: 256
+  head_insertion_edge_dim: 256
+  head_independent_adapters: true
+  head_initialization:
+    output_std: 0.001
+    coordinate_output_std: 0.0
+    do_del_probability: 0.5
+    do_sub_a_probability: 0.5
+    do_sub_e_probability: 0.5
+    ins_rate: 0.6931471805599453
+  token_weighting:
+    type_loss_token_weights: uniform
+    weight_alpha: 0.4
+  lora:
+    enabled: false
+    rank: 8
+    alpha: 8.0
+    dropout: 0.0
+    targets:
+    - qkv
+    - proj
+    - mlp.gate_value
+    - mlp.out_proj
+    freeze_base: true
+  backbone:
+    _target_: morph.model.backbones.dit.DiTBackbone
+    _recursive_: false
+    d_input: ${embeddings.atom_type_embedding.out_dim}
+    d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+    d_model: 512
+    out_node_nf: 256
+    out_edge_nf: 256
+    num_layers: 8
+    nhead: 16
+    mlp_ratio: 4.0
+    ffn_hidden_dim: null
+    ffn_multiple_of: 64
+    qk_norm: true
+    coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+    center_coordinate_output: true
+    in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+    edge_pair_dim: 256
+    pair_bias_hidden_dim: 32
+    num_atom_types: ${len:${data.vocab.atom_tokens}}
+    lora: ${model.lora}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+schedule:
+  _target_: morph.flow_matching.schedules.SmoothstepSchedule
+  shift: 0.8
+trainer:
+  trainer:
+    deterministic: warn
+    devices: 2
+    accelerator: auto
+    strategy: ddp_find_unused_parameters_false
+    precision: bf16-mixed
+    max_epochs: -1
+    accumulate_grad_batches: 1
+    num_sanity_val_steps: 0
+    gradient_clip_val: null
+    gradient_clip_algorithm: norm
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+    limit_val_batches: null
+    limit_predict_batches: 4
+    max_time:
+      hours: 47
+      minutes: 15
+  do_train: true
+  do_validate: false
+  do_predict: false
+  checkpoint_path: /cluster/project/krause/frankem/morph-chem/outputs/geom_resume_e48_20261008/resume_epoch_048.ckpt
+  final_checkpoint_path: final.ckpt
+callbacks:
+  monitor_metric: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},val/validity,val/batch/pair_dist_l1}
+  monitor_metric_mode: ${if:${oc.eval:"'${representation}' in ('geometric_graph',
+    'molecule')"},max,min}
+  model_checkpoints:
+    save_top_k: 1
+    save_last: false
+    verbose: false
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link
+    verbose: false
+    save_on_train_epoch_end: true
+logging:
+  enabled: true
+  project: morph-${data.dataset_name}
+  name: ${oc.env:MORPH_RUN_ID}
+  save_dir: .
+  offline: false
+integrator:
+  integrator:
+    _target_: morph.flow_matching.integration.RateIntegrator
+    vocab: ${data.vocab}
+    distributions: None
+    insertion_head_mode: ${insertion_head.mode}
+    insertion_dimension: ${insertion_head.dimension}
+    insertion_num_components: ${insertion_head.num_components}
+    time_strategy: log
+    time_beta_alpha: 2.0
+    time_beta_beta: 1.0
+    num_integration_steps: 100
+    ins_noise_scale: ${data.ins_noise_scale}
+    renoise_insertions: true
+    exact_schedule_integration: true
+    del_schedule: ${data.interpolator.del_schedule}
+    ins_schedule: ${data.interpolator.ins_schedule}
+    sub_schedule: ${data.interpolator.sub_schedule}
+    sub_e_schedule: ${data.interpolator.sub_e_schedule}
+data:
+  root_path: /cluster/project/krause/frankem/morph-chem/data
+  preprocess_workers: 10
+  preprocess_chunk_size: 128
+  statistics_workers: ${data.preprocess_workers}
+  hydrogen_mode: explicit
+  cache_name: ${data.dataset_name}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}
+  source_dir: ${oc.env:GEOM_SOURCE_DIR,null}
+  max_atoms: 72
+  allowed_atom_symbols:
+  - H
+  - C
+  - 'N'
+  - O
+  - S
+  - F
+  - Cl
+  - Br
+  - P
+  - I
+  n_atoms_strategy: uniform
+  uniform_n_atoms_padding: 5
+  optimal_transport: equivariant
+  ins_noise_scale: 0.65
+  move_noise_scale: 0.5
+  allow_charged: true
+  preprocessing:
+    _target_: morph.dataset.core.preprocessing.Preprocessing
+    _recursive_: false
+    root: /cluster/project/krause/frankem/morph/data/geom
+    num_workers: ${data.statistics_workers}
+    train_dataset: ${data.datamodule.datasets.train}
+    atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+    edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+    charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+    distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+  vocab: null
+  interpolator:
+    _target_: morph.flow_matching.interpolation.Interpolator
+    distributions: null
+    optimal_transport: ${data.optimal_transport}
+    ins_noise_scale: ${data.ins_noise_scale}
+    move_noise_scale: ${data.move_noise_scale}
+    atom_type_noise_prob: 0.0
+    pre_align_iters: 2
+    c_move: 1.0
+    c_sub: 0.0
+    c_ins: 6.125
+    c_del: 6.125
+    del_schedule: ${schedule}
+    ins_schedule: ${schedule}
+    sub_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+    sub_e_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+  datamodule:
+    _target_: morph.dataset.core.lightning_datamodule.LightningDataModule
+    vocab: ${data.vocab}
+    distributions: null
+    representation: ${representation}
+    datasets:
+      train:
+        _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+      val:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+      test:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+    interpolator: ${data.interpolator}
+    rotate: ${model.rotate}
+    n_augmentations: ${model.n_augmentations}
+    num_workers:
+      train: 20
+      val: 0
+      test: 0
+    batch_size:
+      train: 320
+      val: 32
+      test: 32
+    n_atoms_strategy: ${data.n_atoms_strategy}
+    uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+    time_dist: ${model.module.time_dist}
+    train_prefetch_factor: 1
+  dataset_name: geom
+  dataset_target: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+embeddings:
+  atom_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.atom_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  edge_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.edge_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  time_embedding:
+    _target_: morph.model.layers.embedding.TimeEmbedding
+    embedding_dim: 128
+    out_dim: 128
+  node_count_embedding:
+    _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+    embedding_dim: 64
+    out_dim: 128
+    min_period: 128.0
+    max_period: 1024.0
+    max_count: 96.0
+hydra:
+  run:
+    dir: ${oc.env:MORPH_RUN_ROOT}/geom/geom/${now:%Y-%m-%d}/${oc.env:MORPH_RUN_ID}

geom/explicit_h/runtime/configs/schedule/beta.yaml

@@ -0,0 +1,3 @@
+_target_: morph.flow_matching.schedules.BetaSchedule
+k_alpha: 1.0
+k_beta: 1.5

geom/explicit_h/runtime/configs/schedule/cubic.yaml

@@ -0,0 +1 @@
+_target_: morph.flow_matching.schedules.CubicSchedule