File size: 8,323 Bytes
31dee3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Three task heads with Monroe's exact UW objective and covariance penalty."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from multitask_data import *
from torch import nn
import yaml
from threedscriptors.configuration.architecture_config import EncoderDecoderArchitectureConfig
from threedscriptors.data_handling.sample import PreprocessedSample
from threedscriptors.training.noise_scheduler import NoiseModule,ConstantSchedule
from threedscriptors.training.pretraining import atom_denoising_loss
from train_cosmo import loss_terms,batches
CONFIG=dict(epochs=5,seed=42026,learning_rate_backbone=1e-5,learning_rate_cosmo=1e-5,
deterministic_training=True,
learning_rate_denoiser=1e-5,learning_rate_combisolv=3e-4,learning_rate_uncertainty=1e-3,
weight_decay=1e-4,gradient_clip=1.,cosmo_batch_size=96,cosmo_atom_cap=4000,
qm_pair_batch_size=256,qm_graph_batch_size=64,qm_atom_cap=3600,
noise_sigma=.3,denoising_reduction='Original atom_denoising_loss / 1088 embedding channels; raw loss also logged',
task_order=['rem3di_denoising','cosmo','combisolv'],uncertainty_init=-.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.,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=MONROE_REV,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')
def off_diag_cov_loss(emb):
"""Verbatim numerical operations from pinned Monroe train/loss.py."""
B,D=emb.shape
if B<2:return emb.new_tensor(0.)
emb=emb.float();emb=emb-emb.mean(dim=0)
cov=(emb.T@emb)/(B-1)
return (cov.pow(2).sum()-cov.diagonal().pow(2).sum())/(D*(D-1))
def uncertainty_objective(losses,log_variance):
precision=1./(2*log_variance.exp())
return (losses*precision+log_variance/2).sum()
class MultiTask(nn.Module):
def __init__(self,checkpoint=None,device='cuda'):
super().__init__()
if checkpoint is None:
initial,decoder_path=initialize()
architecture_path=common.ORIGINAL/'post_training_architecture_config.yaml'
else:
checkpoint=Path(checkpoint)
architecture_path=checkpoint/'post_training_architecture_config.yaml'
architecture=yaml.safe_load(architecture_path.read_text())
architecture['mace_config']=None
cfg=EncoderDecoderArchitectureConfig.model_validate(architecture)
self.pre=cfg._build_preprocessor();self.encoder=cfg._build_encoder()
self.pre.atomic_preprocessor.double();self.pre.geometric_preprocessor.float();self.encoder.float()
self.denoiser=cfg.decoder_config.build().float()
self.cosmo_head=nn.Sequential(nn.Linear(640,256),nn.SiLU(),nn.Linear(256,155))
self.qm_head=nn.Sequential(nn.Linear(1281,256),nn.SiLU(),nn.Linear(256,1))
self.log_variance=nn.Parameter(torch.full((3,),CONFIG['uncertainty_init']))
self.noise=NoiseModule(ConstantSchedule(CONFIG['noise_sigma']))
if checkpoint is None:
for name,module in [('encoder.pth',self.encoder),('atomic_preprocessor.pth',self.pre.atomic_preprocessor),
('geometric_preprocessor.pth',self.pre.geometric_preprocessor),('cosmo_head.pth',self.cosmo_head)]:
module.load_state_dict(torch.load(initial/name,map_location='cpu',weights_only=True),strict=True)
self.denoiser.load_state_dict(torch.load(decoder_path,map_location='cpu',weights_only=True),strict=True)
else:self.load_parts(checkpoint)
self.to(device)
@property
def parts(self):
return {'encoder.pth':self.encoder,'atomic_preprocessor.pth':self.pre.atomic_preprocessor,
'geometric_preprocessor.pth':self.pre.geometric_preprocessor,'cosmo_head.pth':self.cosmo_head,
'denoising_decoder.pth':self.denoiser,'combisolv_head.pth':self.qm_head}
def load_parts(self,path):
manifest=json.loads((path/'manifest.json').read_text())
for name,module in self.parts.items():
assert sha(path/name)==manifest['files'][name]['sha256'],name
module.load_state_dict(torch.load(path/name,map_location='cpu',weights_only=True),strict=True)
assert sha(path/'task_weighting.pth')==manifest['files']['task_weighting.pth']['sha256']
state=torch.load(path/'task_weighting.pth',map_location='cpu',weights_only=True)
with torch.no_grad():self.log_variance.copy_(state['log_variance'])
def encode(self,sample):
processed=self.pre(sample);descriptor=self.encoder(processed)
z=descriptor if torch.is_tensor(descriptor) else descriptor.flat
return processed,descriptor,z
def cosmo_denoising(self,sample,target,denoise=True):
clean,descriptor,z=self.encode(sample)
pred=self.cosmo_head(z);cosmo,parts=loss_terms(pred,target)
denoising=cosmo.new_zeros(())
if denoise:
clean_atomic=clean.preprocessed_atomic_embeddings.clone()
noisy=self.noise(clean_atomic,clean.padding_mask)
corrupted=PreprocessedSample(preprocessed_atomic_embeddings=noisy,
initial_pair_representation=clean.initial_pair_representation.detach(),
geometrical_encoding=clean.geometrical_encoding.detach(),padding_mask=clean.padding_mask,pair_mask=clean.pair_mask)
reconstruction=self.denoiser(corrupted,descriptor)
raw=atom_denoising_loss(clean_atomic,reconstruction,clean.padding_mask,CONFIG['noise_sigma'])
# Per-channel reduction changes only a constant, leaving the denoising target/noise/gradient paths intact.
denoising=raw/clean_atomic.shape[-1]
return torch.stack([denoising,cosmo]),off_diag_cov_loss(z),z,pred,parts
def qm(self,cache,pair_ids,pairs,targets):
a=pairs.solute_index.to_numpy()[pair_ids];b=pairs.solvent_index.to_numpy()[pair_ids]
unique,inverse=np.unique(np.r_[a,b],return_inverse=True)
zs=[];order=[]
# Unique molecules once per pair batch: shared solvents are not counted repeatedly in covariance.
for ids in batches(unique,cache.lengths,max_size=CONFIG['qm_graph_batch_size'],atom_cap=CONFIG['qm_atom_cap']):
_,_,z=self.encode(cache.sample(ids));zs.append(z);order.extend(ids)
z=torch.cat(zs);lookup={int(i):j for j,i in enumerate(order)}
ai=torch.tensor([lookup[int(i)] for i in a],device='cuda');bi=torch.tensor([lookup[int(i)] for i in b],device='cuda')
temperature=torch.as_tensor(pairs.temperature_K.to_numpy()[pair_ids]/298.,dtype=torch.float32,device='cuda')[:,None]
pred=self.qm_head(torch.cat([z[ai],z[bi],temperature],dim=1))[:,0]
loss=(pred-targets[pair_ids]).square().mean()
return loss,off_diag_cov_loss(z),pred,z
def optimizer_for(model):
return torch.optim.AdamW([
dict(params=list(model.pre.parameters())+list(model.encoder.parameters()),lr=CONFIG['learning_rate_backbone']),
dict(params=model.cosmo_head.parameters(),lr=CONFIG['learning_rate_cosmo']),
dict(params=model.denoiser.parameters(),lr=CONFIG['learning_rate_denoiser']),
dict(params=model.qm_head.parameters(),lr=CONFIG['learning_rate_combisolv']),
dict(params=[model.log_variance],lr=CONFIG['learning_rate_uncertainty'],weight_decay=0.)],weight_decay=CONFIG['weight_decay'])
|