| """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']) |
| |
| 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=[] |
| |
| 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']) |
|
|
