Pin exact evaluated starting checkpoint, objective equivalence tests, source and joint training protocol

preparation_complete.json

@@ -0,0 +1,31 @@
+{
+  "completed": true,
+  "time": 1793210639.778262,
+  "initialization": {
+    "initial_checkpoint_repository": "Queimo/rem3di-pcqm-epoch2-cosmo",
+    "initial_checkpoint_path": "epoch_10",
+    "initial_checkpoint_revision": "9266a41fbe3b36afe0615e0c3ba232ca74bdeef7",
+    "initial_hashes": {
+      "encoder.pth": "97e616a16e8317717784929495ba3a030bfcd821038be6e05324624b17c409c0",
+      "atomic_preprocessor.pth": "9683883267a902773b4f6777eedb366befd016f4f7083bed62a1bd35ae113e3d",
+      "geometric_preprocessor.pth": "4da0843646d06fcd2ff5b13e1fd4fb4b22c8c0456d0507c59df01e665c4d6d7b",
+      "cosmo_head.pth": "1a43270e60f00c9428099e0e37ab62f77a2c219ea611d4f05194bc1abf69b0a7"
+    },
+    "frozen_foundation_sha256": "fab8b8713c832f31a2a853aaa22fd638be8a369cbf5095e6b3e982a18d10e93a",
+    "verified_against_bigsoldb_protocol": true,
+    "bigsoldb_protocol_sha256": "dcd893dc297eef5d50b225367f295cdded69a6a9cd9b53c8f8f63c241f7e371b",
+    "denoising_decoder": {
+      "repository": "Queimo/rem3di-pcqm500k-corrected",
+      "revision": "2f5c3b40b36259a13adc8412689254054bb93288",
+      "path": "runs/corrected_20261027/evaluation/paper_audit/author_pcqm500k_5ep/epoch_02/decoder.pth",
+      "sha256": "bffdd09bae9bcb2d863165546dea090a10be06d8a52b89e4f651e0f1c80ae7c4"
+    },
+    "explanation": "All encoder, atomic/geometric preprocessor and COSMO head weights come from the exact evaluated COSMO10 checkpoint. Only its omitted denoising decoder is restored from PCQM epoch2."
+  },
+  "data": {
+    "repository": "Queimo/combisolv-qm-rem3di-multitask",
+    "revision": "de87254deae3bbdc031e3c2185df7c3ad63945ac",
+    "verified": true,
+    "files": 23
+  }
+}

prepared_data_hub.json

@@ -0,0 +1,6 @@
+{
+  "repository": "Queimo/combisolv-qm-rem3di-multitask",
+  "revision": "de87254deae3bbdc031e3c2185df7c3ad63945ac",
+  "verified": true,
+  "files": 23
+}

smoke_validation.json

@@ -0,0 +1,83 @@
+{
+  "passed": true,
+  "all_head_gradients": {
+    "encoder.pth": true,
+    "atomic_preprocessor.pth": true,
+    "geometric_preprocessor.pth": true,
+    "cosmo_head.pth": true,
+    "denoising_decoder.pth": true,
+    "combisolv_head.pth": true
+  },
+  "encoder_updated": true,
+  "task_weights_updated": true,
+  "optimizer_and_rng_resume_next_step_exact": true,
+  "normalization_buffers_unchanged": true,
+  "initial_bigsoldb_embedding_match": {
+    "molecules": 1334,
+    "max_absolute": 4.184246063232422e-05,
+    "mean_absolute": 1.621990719513633e-07
+  },
+  "real_batch_loss": {
+    "denoising": 0.0400693416595459,
+    "denoising_raw": 43.59544372558594,
+    "cosmo": 0.020166773349046707,
+    "cov_cosmo": 1.5453895684913732e-05,
+    "descriptor_std_cosmo": 0.11399243026971817,
+    "cosmo_sigma": 0.05140712857246399,
+    "cosmo_area": 0.0041954368352890015,
+    "cosmo_volume": 0.004897752311080694,
+    "combisolv": 1.3595600128173828,
+    "cov_combisolv": 0.0001902484946185723,
+    "descriptor_std_combisolv": 0.18535533547401428,
+    "total": 0.4212467670440674,
+    "gradient_norm": 2.3515171745204047,
+    "combisolv_rmse_kcal_mol": 5.220246867294003,
+    "log_variance_rem3di_denoising": -0.5009999871253967,
+    "precision_rem3di_denoising": 0.8251854181289673,
+    "log_variance_cosmo": -0.5009999871253967,
+    "precision_cosmo": 0.8251854181289673,
+    "log_variance_combisolv": -0.49900001287460327,
+    "precision_combisolv": 0.8235366940498352
+  },
+  "seconds_per_joint_step": 2.9659399400006805,
+  "steady_seconds_per_joint_step": 0.5084904360010114,
+  "smoke_cosmo_molecules": 64,
+  "smoke_qm_pairs": 256,
+  "peak_gpu_GB": 13.077587456,
+  "config": {
+    "epochs": 5,
+    "seed": 42026,
+    "learning_rate_backbone": 1e-05,
+    "learning_rate_cosmo": 1e-05,
+    "deterministic_training": true,
+    "learning_rate_denoiser": 1e-05,
+    "learning_rate_combisolv": 0.0003,
+    "learning_rate_uncertainty": 0.001,
+    "weight_decay": 0.0001,
+    "gradient_clip": 1.0,
+    "cosmo_batch_size": 96,
+    "cosmo_atom_cap": 4000,
+    "qm_pair_batch_size": 256,
+    "qm_graph_batch_size": 64,
+    "qm_atom_cap": 3600,
+    "noise_sigma": 0.3,
+    "denoising_reduction": "Original atom_denoising_loss / 1088 embedding channels; raw loss also logged",
+    "task_order": [
+      "rem3di_denoising",
+      "cosmo",
+      "combisolv"
+    ],
+    "uncertainty_init": -0.5,
+    "uncertainty_objective": "sum(L_t/(2*exp(s_t)) + s_t/2); s is learned log variance; no weight normalization in objective",
+    "decorrelation_weight": 8.0,
+    "decorrelation": "Monroe off_diag_cov_loss: sum(offdiag(sample covariance)^2)/(D*(D-1)); average COSMO and unique CombiSolv molecular batches; outside UW",
+    "monroe_revision": "de075eabe98e4adbbedbc86ca9e877f4bfd9b629",
+    "temperature_input": "Temperature_K/298",
+    "qm_head": "concat(solute640, solvent640, T/298) -> Linear1281x256 -> SiLU -> Linear256x1",
+    "epoch_definition": "Maximum of full COSMO and full CombiSolv training-loader lengths. Cycle shorter loader. Every epoch covers every retained training row from each corpus.",
+    "initial_checkpoint": "Exact PCQMep2COSMOep10 used in existing BigSolDB evaluation",
+    "denoiser_initialization": "PCQM epoch2 decoder only; encoder/preprocessors/COSMO head all from COSMO10",
+    "frozen": "MACE-POLAR foundation only; full float32 raw features cached; learned atomic/geometric preprocessors trainable",
+    "evaluation": "Identical BigSolDB2.0 author test IDs, shared baseline geometries, concatenated solute+solvent+T, Kumo Large four estimators seed42 full75731-row context and all1281 columns; evaluate every epoch without test-based selection"
+  }
+}

source/multitask_data.py

@@ -194,7 +194,8 @@ def prepare_tables():
 
 def prepare_geometries():
     path=DATA/'combisolv_geometries.parquet'
-    if path.exists():return
+    if path.exists():
+        record_geometry_manifest(path);return
     molecules=pd.read_parquet(DATA/'combisolv_molecules.parquet')
     status('generating_combisolv_3d',total=len(molecules))
     result=[]
@@ -206,8 +207,18 @@ def prepare_geometries():
     assert not failures,failures[:10]
     assert [r['smiles'] for r in result]==molecules.smiles.tolist()
     pd.DataFrame(result).to_parquet(path,index=False,compression='zstd')
-    write(DATA/'geometry_manifest.json',dict(rows=len(result),sha256=sha(path),seed=42026,
-        algorithm='RDKit ETKDGv3 explicit H; deterministic retries with random coordinates; MMFF94s or UFF optimization',
+    record_geometry_manifest(path)
+
+
+def record_geometry_manifest(path):
+    from rdkit import rdBase
+    frame=pd.read_parquet(path,columns=['smiles','geometry_method'])
+    molecules=pd.read_parquet(DATA/'combisolv_molecules.parquet',columns=['smiles'])
+    assert frame.smiles.tolist()==molecules.smiles.tolist()
+    write(DATA/'geometry_manifest.json',dict(rows=len(frame),sha256=sha(path),seed=42026,
+        rdkit_version=rdBase.rdkitVersion,
+        algorithm='RDKit ETKDGv3 explicit H, deterministic random-coordinate retries, generic 3D distance geometry fallback; MMFF94s or UFF where parameterized',
+        methods={str(k):int(v) for k,v in frame.geometry_method.value_counts().items()},
         no_2d_fallback=True,missing=0))
 
 
@@ -252,7 +263,10 @@ def cache_mace(name,frame):
     mace=build_author_mace_runtime().cuda().eval().requires_grad_(False)
     captures=[]
     hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda m,a:captures.append(a[0].detach()))
-    started=time.time();plan=list(batches(np.arange(len(frame)),lengths,max_size=64,atom_cap=3600))
+    started=time.time()
+    if name=='bigsoldb':
+        plan=list(common.batch_indices(np.arange(len(frame)),lengths,max_size=24,atom_cap=2400))
+    else:plan=list(batches(np.arange(len(frame)),lengths,max_size=64,atom_cap=3600))
     status('caching_frozen_mace',corpus=name,total=len(frame),atoms=int(ptr[-1]),bytes=int(ptr[-1])*4096*4,batches=len(plan))
     with torch.inference_mode():
         for bi,ids in enumerate(plan,1):
@@ -276,9 +290,55 @@ def cache_mace(name,frame):
 def publish_data():
     marker=OUT/'prepared_data_hub.json'
     if marker.exists():return
+    pairs=pd.read_parquet(DATA/'combisolv_pairs.parquet')
+    for split,frame in pairs.groupby('split',sort=False):
+        frame.to_parquet(DATA/f'combisolv_{split}.parquet',index=False,compression='zstd')
+    manifest=json.loads((DATA/'manifest.json').read_text())
+    manifest['files']={p.name:sha(p) for p in DATA.iterdir() if p.is_file() and p.name!='manifest.json'}
+    write(DATA/'manifest.json',manifest)
+    header='''---
+license: cc-by-4.0
+task_categories:
+- tabular-regression
+configs:
+- config_name: raw_combisolv_qm
+  default: true
+  data_files:
+  - split: records
+    path: raw/combisolv_qm.parquet
+- config_name: prepared_pairs
+  data_files:
+  - split: train
+    path: prepared/combisolv_train.parquet
+  - split: validation
+    path: prepared/combisolv_validation.parquet
+  - split: benchmark_overlap
+    path: prepared/combisolv_benchmark_overlap.parquet
+  - split: auxiliary_holdout
+    path: prepared/combisolv_auxiliary_holdout.parquet
+- config_name: combisolv_geometries
+  data_files:
+  - split: molecules
+    path: prepared/combisolv_geometries.parquet
+- config_name: cosmo_pool
+  data_files:
+  - split: records
+    path: prepared/cosmo-*.parquet
+---
+'''
+    card=header+'\n# CombiSolv-QM and COSMO multi-task continuation data\n\n'
+    card+='The unchanged original one-million-row CombiSolv-QM table is preserved under `raw/`, together with source archive/CSV/Parquet checksums. The target is signed solvation Gibbs free energy at 298 K, in kcal/mol, using 1 mol/L reference concentrations. It is distinct from experimental solubility.\n\n'
+    card+=f"Prepared CombiSolv memberships: `{json.dumps(manifest['combisolv_splits'],sort_keys=True)}`. All {manifest['unique_qm_molecules']:,} unique structures have explicit-H coordinates; supplied InChI defines identity, and original SMILES/InChI mappings are retained. Solute/solvent indices refer to row order in the geometry and molecule tables.\n\n"
+    card+='Geometries use seeded RDKit ETKDGv3, random-coordinate retries and generic 3D distance geometry for strained/unusual structures, followed by MMFF94s or UFF when parameterized. There is no 2D fallback. Geometry methods, software version, coverage and checksums are recorded.\n\n'
+    card+=f"The COSMO pool retains the source training/validation records and adds an explicit `joint_split` membership: `{json.dumps(manifest['cosmo_splits'],sort_keys=True)}`. Filter that column for training. The complete original COSMO dataset and its existing ESOL/FreeSolv quarantine remain available in [Queimo/pubchemqc-cosmo-pm6](https://huggingface.co/datasets/{COSMO_REPO}) at revision `{COSMO_REV}`.\n\n"
+    card+='BigSolDB2.0 held-out solute connectivity groups are excluded from all new objectives and both CombiSolv input roles. CombiSolv validation solutes are excluded across objectives. Existing COSMO validation groups are excluded from CombiSolv training. Raw rows remain available with explicit quarantine membership; earlier exposure in the starting checkpoint cannot be undone. ESOL/FreeSolv quarantine is preserved for COSMO, but is not claimed for CombiSolv.\n\n'
+    card+='CombiSolv target mean/std is fitted on its retained training rows only; signed energies are not log transformed. The existing COSMO log1p/ln transform and normalization are preserved so the initialized output head retains its meaning.\n\n'
+    card+='[SolProp v1.2 source](https://zenodo.org/records/5970538) · [Vermeire and Green (2021)](https://doi.org/10.1016/j.cej.2021.129307) · [Vermeire, Chung and Green (2022)](https://doi.org/10.1021/jacs.2c01768). Please cite the original data authors.\n'
+    card_path=OUT/'dataset_README.md';card_path.write_text(card)
     files=[p for p in DATA.iterdir() if p.is_file()]
     commit=HfApi().create_commit(repo_id=DATA_REPO,repo_type='dataset',
-        operations=[CommitOperationAdd(path_in_repo='prepared/'+p.name,path_or_fileobj=str(p)) for p in files],
+        operations=[CommitOperationAdd(path_in_repo='prepared/'+p.name,path_or_fileobj=str(p)) for p in files]+[
+            CommitOperationAdd(path_in_repo='README.md',path_or_fileobj=str(card_path))],
         commit_message='Add reproducible joint splits, all geometries, label transforms and benchmark exclusions')
     remote={p.rfilename:p for p in HfApi().dataset_info(DATA_REPO,revision=commit.oid,files_metadata=True).siblings}
     for p in files:
@@ -300,14 +360,11 @@ def main():
     graphs=pd.read_parquet(WORK/'bigsoldb_evaluation/features/shared_graphs.parquet')
     vocab=NODE_FEAT_LIST_ONE_HOT['atomic_number'];col=list(NODE_FEAT_LIST_ONE_HOT).index('atomic_number')
     ptable=Chem.GetPeriodicTable()
-    graph_atoms=pd.DataFrame(dict(smiles=graphs.smiles,positions=graphs.pos_rdkit,
+    graph_atoms=pd.DataFrame(dict(smiles=graphs.smiles,positions=[np.stack(p).astype(np.float32).tolist() for p in graphs.pos_rdkit],
         atomic_numbers=[[ptable.GetAtomicNumber(vocab[int(i)]) for i in np.stack(codes)[:,col]] for codes in graphs.node_codes],
         charge=[int(round(float(np.stack(v)[:,5].sum()))) for v in graphs.node_float]))
     graph_atoms.to_parquet(DATA/'bigsoldb_fixed_geometries.parquet',index=False,compression='zstd')
     cache_mace('bigsoldb',graph_atoms)
-    manifest=json.loads((DATA/'manifest.json').read_text())
-    manifest['files']={p.name:sha(p) for p in DATA.iterdir() if p.is_file() and p.name!='manifest.json'}
-    write(DATA/'manifest.json',manifest)
     publish_data()
     write(OUT/'preparation_complete.json',dict(completed=True,time=time.time(),initialization=json.loads((OUT/'initialization.json').read_text()),data=json.loads((OUT/'prepared_data_hub.json').read_text())))
     status('preparation_complete')

source/multitask_evaluate.py

@@ -21,7 +21,7 @@ def embeddings(epoch):
     model=MultiTask(checkpoint);model.eval().requires_grad_(False)
     cache=RawCache('bigsoldb');X=np.empty((len(geometry),640),dtype=np.float32)
     with torch.inference_mode():
-        for ids in batches(np.arange(len(geometry)),cache.lengths,max_size=24,atom_cap=2400):
+        for ids in common.batch_indices(np.arange(len(geometry)),cache.lengths,max_size=24,atom_cap=2400):
             _,_,z=model.encode(cache.sample(ids));X[ids]=z.float().cpu().numpy()
     assert np.isfinite(X).all()
     np.savez_compressed(path,X=X,smiles=geometry.smiles.to_numpy(dtype=str),checkpoint_sha256=json.dumps(hashes))

source/multitask_train.py

@@ -174,13 +174,15 @@ def smoke(data):
     else:check_monroe_losses()
     model=MultiTask();optimizer=optimizer_for(model)
     scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:1.)
-    model.eval();baseline=RawCache('bigsoldb');ids=list(range(24))
+    model.eval();baseline=RawCache('bigsoldb')
     with torch.no_grad():
-        _,_,ours=model.encode(baseline.sample(ids))
-        ref=np.load(WORK/'bigsoldb_evaluation/features/cosmo_epoch10.npz')['X'][ids]
-        difference=np.abs(ours.cpu().numpy()-ref)
-        assert np.allclose(ours.cpu().numpy(),ref,atol=2e-4,rtol=2e-4),(float(difference.max()),float(difference.mean()))
-    initial_difference=dict(max_absolute=float(difference.max()),mean_absolute=float(difference.mean()))
+        ref=np.load(WORK/'bigsoldb_evaluation/features/cosmo_epoch10.npz')['X']
+        ours=np.empty_like(ref)
+        for ids in common.batch_indices(np.arange(len(ref)),baseline.lengths,max_size=24,atom_cap=2400):
+            _,_,z=model.encode(baseline.sample(ids));ours[ids]=z.cpu().numpy()
+        difference=np.abs(ours-ref)
+        assert np.allclose(ours,ref,atol=2e-4,rtol=2e-4),(float(difference.max()),float(difference.mean()))
+    initial_difference=dict(molecules=len(ref),max_absolute=float(difference.max()),mean_absolute=float(difference.mean()))
     frozen={n:b.cpu().clone() for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers()}
     cp,qp,_=data.plans(1);ci=cp[10];qi=qp[10]
     initial=next(model.encoder.parameters()).detach().clone()
@@ -195,7 +197,9 @@ def smoke(data):
     expected_state={n:v.detach().cpu().clone() for n,v in model.state_dict().items()}
     model.load_parts(temp);state=torch.load(temp/'training_state.pt',map_location='cpu',weights_only=False)
     optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state)
+    replay_started=time.perf_counter()
     actual=update(model,optimizer,data,ci,qi)
+    steady_duration=time.perf_counter()-replay_started
     assert all(np.isclose(actual[k],expected[k],atol=1e-7,rtol=1e-7) for k in expected),(expected,actual)
     assert all(torch.equal(v.detach().cpu(),expected_state[n]) for n,v in model.state_dict().items()),'Resume changed next-step parameters'
     for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen[n])
@@ -203,6 +207,7 @@ def smoke(data):
     result=dict(passed=True,all_head_gradients=gradients,encoder_updated=True,task_weights_updated=True,
         optimizer_and_rng_resume_next_step_exact=True,normalization_buffers_unchanged=True,
         initial_bigsoldb_embedding_match=initial_difference,real_batch_loss=first,seconds_per_joint_step=duration,
+        steady_seconds_per_joint_step=steady_duration,
         smoke_cosmo_molecules=len(ci),smoke_qm_pairs=len(qi),peak_gpu_GB=peak,config=CONFIG)
     write(OUT/'smoke_validation.json',result);shutil.rmtree(temp)
     status('smoke_passed',seconds_per_joint_step=duration,peak_gpu_GB=peak,initial_embedding_difference=initial_difference)
@@ -219,7 +224,8 @@ def main():
     torch.manual_seed(CONFIG['seed']);torch.cuda.manual_seed_all(CONFIG['seed']);np.random.seed(CONFIG['seed']);random.seed(CONFIG['seed'])
     data=TrainingData()
     if args.smoke:smoke(data);return
-    assert json.loads((OUT/'smoke_validation.json').read_text())['passed']
+    smoke_result=json.loads((OUT/'smoke_validation.json').read_text())
+    assert smoke_result['passed'] and smoke_result['config']==CONFIG
     checkpoints=sorted((OUT/'checkpoints').glob('epoch_[0-9][0-9]')) if (OUT/'checkpoints').exists() else []
     model=MultiTask(checkpoints[-1] if checkpoints else None);optimizer=optimizer_for(model)
     scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:.1+.9*.5*(1+math.cos(math.pi*e/CONFIG['epochs'])))

source_hashes.json

@@ -1,8 +1,8 @@
 {
-  "multitask_data.py": "9e6a6926c427bba8b93f2225042a53e9d4a5cf0ee9bd2ee603e5d8acb284683c",
+  "multitask_data.py": "7ba9accbbd9a0d9c23720946cb0ebd2e410853aa7951dc58c9198a4f2586c616",
   "multitask_model.py": "1ffe262b355bc911e3ea35c9e6844d40ab5b6dafb5a11a90eba6565361083b05",
-  "multitask_train.py": "9cfe0b327476e8cc4b076a30d211b428f77a51fd21139785174301b2503f43f9",
-  "multitask_evaluate.py": "ab5352483e7d927755a8fed2fec150400ac71de804cbe93b1b34b039c1186ba1",
+  "multitask_train.py": "00983d0060b26f148eb67b021899d9a73d84231aaf821406b2a4d9f1c36d568d",
+  "multitask_evaluate.py": "ea5132b5bd4efe97924009bbe94f6d1f7ff4504f993b37689c0ce7ef93dea9c6",
   "multitask_driver.py": "ade6f9106b3be1d488e39db711119d43e07457ebe9425196b460c87440205ad7",
   "eval_common.py": "e694f4e3dc4d62f2567b55a3e3b4da9855686c2977ba3b182094c86b1fb8eaa5",
   "train_cosmo.py": "a2ce3a8bfcd7dd8eafb108151cf7ccce78b3e576b07df83c2b436fde931dd866",