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