| """Random-initialized author architecture, embedding denoising and COSMO heads.""" |
| from pathlib import Path |
| import sys |
| sys.path.insert(0,str(Path(__file__).resolve().parent)) |
| from scratch_data import * |
| from torch import nn |
| import yaml,math |
| 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=10,probes_per_epoch=5,seed=42026,architecture_source_commit='20df6fe9c61154e351d677cb0f77fd716d00beae', |
| initialization='Random learned Rem3Di weights; no published or PCQM/COSMO checkpoint loaded', |
| learning_rate_max=1e-4,weight_decay=1e-3,learning_rate_uncertainty=1e-3,gradient_clip=1., |
| noise_sigma=.3,denoising='Exact author atomic embedding denoising loss and gradient paths, on PCQM', |
| uncertainty='sum(L_t/(2*exp(s_t))+s_t/2); raw denoising loss; initial log variance ln(1088)-0.5 for denoising, -0.5 for COSMO', |
| uncertainty_initialization='Equivalent initial precision to author-denoising per channel; no modification of raw denoising function', |
| task_order=['rem3di_denoising','cosmo'],cosmo_head='640 -> 256 SiLU -> 155', |
| pcqm_batch_size=128,pcqm_atom_cap=5250,cosmo_batch_size=96,cosmo_atom_cap=4000, |
| epoch_definition='Maximum of complete PCQM and COSMO loader lengths; shorter loader cycles; all retained training molecules seen each epoch', |
| scheduler='OneCycleLR max_lr=1e-4, total optimizer steps computed from ten exact deterministic epoch plans', |
| decorrelation=False,combisolv=False,mace='Frozen MACE-POLAR-1-M; raw 4096-channel float32 cache', |
| normalization='Fresh train-only population invariant mean/std across PCQM and COSMO atoms', |
| probe='Frozen pooled 640; solute640+solvent640+Temperature_K; exact BigSolDB rows; full75731-row Kumo Large context, 4 estimators, seed42', |
| probe_schedule='ceil(k * epoch_steps / 5), k=1..5, each epoch; fifty test probes; no test-based selection', |
| checkpoint_policy='All fifty probe points include all weights, optimizer, scheduler and RNG; verified Hugging Face backup before probing') |
|
|
|
|
| def uncertainty(losses,log_variance): |
| return (losses/(2*log_variance.exp())+log_variance/2).sum() |
|
|
|
|
| class ScratchJoint(nn.Module): |
| def __init__(self,checkpoint=None,device='cuda'): |
| super().__init__() |
| checkpoint=Path(checkpoint) if checkpoint is not None else None |
| architecture_path=checkpoint/'post_training_architecture_config.yaml' if checkpoint is not None else OUT/'architecture.yaml' |
| architecture=yaml.safe_load(architecture_path.read_text()) |
| architecture['mace_config']=None |
| cfg=EncoderDecoderArchitectureConfig.model_validate(architecture) |
| |
| stats_path=checkpoint/'normalization.pt' if checkpoint is not None else DATA/'normalization.pt' |
| stats=torch.load(stats_path,map_location='cpu',weights_only=True) |
| self.pre=cfg._build_preprocessor() |
| self.pre.atomic_preprocessor.double();self.pre.geometric_preprocessor.float() |
| self.pre.atomic_preprocessor.invariant_normalization.set_stats(stats['mean'],stats['std'],dtype=torch.float64) |
| self.encoder=cfg._build_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.log_variance=nn.Parameter(torch.tensor([math.log(1088)-.5,-.5],dtype=torch.float32)) |
| self.noise=NoiseModule(ConstantSchedule(.3)) |
| if checkpoint is not None:self.load_parts(Path(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,'denoising_decoder.pth':self.denoiser, |
| 'cosmo_head.pth':self.cosmo_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): |
| clean=self.pre(sample);descriptor=self.encoder(clean) |
| z=descriptor if torch.is_tensor(descriptor) else descriptor.flat |
| return clean,descriptor,z |
|
|
| def denoising(self,sample): |
| clean,descriptor,z=self.encode(sample) |
| target=clean.preprocessed_atomic_embeddings.clone() |
| noisy=self.noise(target,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) |
| loss=atom_denoising_loss(target,reconstruction,clean.padding_mask,.3) |
| return loss,z,target |
|
|
| def cosmo(self,sample,target): |
| _,_,z=self.encode(sample);pred=self.cosmo_head(z) |
| loss,parts=loss_terms(pred,target) |
| return loss,z,pred,parts |
|
|
|
|
| def optimizer_for(model): |
| return torch.optim.AdamW([ |
| dict(params=list(model.pre.parameters())+list(model.encoder.parameters())+list(model.denoiser.parameters())+list(model.cosmo_head.parameters()),lr=CONFIG['learning_rate_max']), |
| dict(params=[model.log_variance],lr=CONFIG['learning_rate_uncertainty'],weight_decay=0.)],weight_decay=CONFIG['weight_decay']) |
|
|
