| { |
| "model_type": "rfd4_proteina", |
| "format_version": 1, |
| "weights_sha256": "5b549dd61adb2f6c6b5cb95bda13f40f5e31556b4e6c755edc0ca6f011c68fb0", |
| "config": { |
| "nn": { |
| "name": "local_latents_transformer", |
| "output_parameterization": { |
| "bb_ca": "v", |
| "local_latents": "v" |
| }, |
| "sequence_readout": { |
| "enabled": true, |
| "detach": true, |
| "num_classes": 147 |
| }, |
| "nlayers": 14, |
| "token_dim": 768, |
| "nheads": 12, |
| "parallel_mha_transition": false, |
| "token_type_emb_dim": 32, |
| "cath_emb_dim": 32, |
| "strict_feats": false, |
| "noise_align_chirality": true, |
| "noise_align_internal_coords": false, |
| "noise_align_chirality_translate": false, |
| "noise_align_internal_coords_translate": false, |
| "use_atom37_factory": true, |
| "feats_seq": [ |
| "xt_bb_ca", |
| "xt_local_latents", |
| "token_type", |
| "sequence_probabilities", |
| "cond_a37coors_nm", |
| "cond_a37coors_internal_nm", |
| "rasa_binned", |
| "secondary_structure", |
| "hbond_modernized", |
| "pistack", |
| "pication", |
| "saltbridge", |
| "hydrophobic", |
| "metal_coord", |
| "na_ss_state", |
| "hotspot", |
| "atom_charge", |
| "atomic_element", |
| "coord_group_seq", |
| "chirality", |
| "chain_centroid_xyz" |
| ], |
| "feats_cond_seq": [ |
| "time_emb_bb_ca", |
| "time_emb_local_latents", |
| "cath_code", |
| "recycle_remaining" |
| ], |
| "dim_cond": 256, |
| "idx_emb_dim": 256, |
| "t_emb_dim": 256, |
| "feats_pair_repr": [ |
| "xt_bb_ca_pair_dists", |
| "bond_mask", |
| "bond_order", |
| "is_aromatic", |
| "index_group_pair", |
| "molecule_idx_pair", |
| "molecule_entity_pair", |
| "token_bond_path_length", |
| "canonical_polymer_ordering", |
| "na_ss_pair" |
| ], |
| "feats_pair_cond": [ |
| "time_emb_bb_ca", |
| "time_emb_local_latents" |
| ], |
| "xt_pair_dist_dim": 30, |
| "xt_pair_dist_min": 0.1, |
| "xt_pair_dist_max": 3, |
| "seq_sep_dim": 64, |
| "pair_repr_dim": 384, |
| "update_pair_repr": true, |
| "update_pair_repr_every_n": 2, |
| "use_tri_mult": true, |
| "tri_mult_c": 256, |
| "tri_mult_bias": false, |
| "tri_mult_backend": "cuequivariance", |
| "recycle": { |
| "enabled": true, |
| "grad": "last", |
| "training_max_recycles": 9, |
| "condition_probability": 0.5 |
| }, |
| "checkpointing": "none", |
| "pair_write": "hadamard", |
| "use_qkln": true, |
| "attention_type": "torch", |
| "use_torch_compile": false, |
| "norm_type": "rmsnorm_torch_fused" |
| }, |
| "product_flowmatcher": { |
| "bb_ca": { |
| "zero_com_noise": true, |
| "guidance_enabled": true, |
| "dim": 3, |
| "stochastic_centering_mode": "x_0", |
| "stochastic_centering_scale": 0.2, |
| "stochastic_centering_scope": "per_group", |
| "stochastic_centering_group_keys": [ |
| "molecule_iid" |
| ] |
| }, |
| "local_latents": { |
| "zero_com_noise": false, |
| "guidance_enabled": true, |
| "dim": 32, |
| "stochastic_centering_mode": "none", |
| "stochastic_centering_scale": null |
| } |
| }, |
| "loss": { |
| "signed_volume_weight": 2.5, |
| "sequence_prediction": { |
| "enabled": true, |
| "t_min": 0.0, |
| "t_max": 1.0 |
| }, |
| "t_distribution": { |
| "bb_ca": { |
| "name": "mix_unif_beta", |
| "p1": 1.9, |
| "p2": 1.0, |
| "p3": 0.02 |
| }, |
| "local_latents": { |
| "name": "mix_unif_beta", |
| "p1": 1.0, |
| "p2": 1.5, |
| "p3": 0.2 |
| }, |
| "shared_groups": null |
| } |
| }, |
| "training": { |
| "exploration_k": 4, |
| "exploration_losses": [ |
| "bb_ca", |
| "local_latents" |
| ], |
| "n_recycle_max": 9, |
| "n_recycle_sampler": "poisson", |
| "n_recycle_lambda": 3.0, |
| "n_recycle_sync": true, |
| "n_recycle_val_loss": 3, |
| "motif_conditioning": false |
| }, |
| "opt": { |
| "lr": 5e-05, |
| "muon_warmup_steps": 0, |
| "adam_warmup_steps": 10000, |
| "gradient_clip_val": 1.0, |
| "max_epochs": 10000000, |
| "accumulate_grad_batches": 2, |
| "dist_strategy": "ddp", |
| "steps_per_epoch": 1000, |
| "val_every_n_epochs": 5, |
| "val_start_epoch": 0, |
| "compute_val_loss": false, |
| "foldseek_timeout_minutes": 25, |
| "muon_exclude_param_patterns": [ |
| "_factory\\.", |
| "^nn\\.(local_latents_linear|ca_linear|insertion_q_proj|insertion_k_proj)\\.", |
| "^nn\\.sequence_readout\\." |
| ], |
| "muon_adjust_lr_fn": "match_rms_adamw", |
| "muon_grouped": true, |
| "num_sanity_val_steps": 2, |
| "coordinate_descent_tuning": false |
| }, |
| "autoencoder_config": { |
| "nn_ae": { |
| "name": "transformer_ae_aa_rna_dna_glycan_ncaa", |
| "latent_z_dim": 32, |
| "allow_oxt": true, |
| "use_torch_compile": false, |
| "torch_compile_fullgraph": false, |
| "n_residue_types": 147, |
| "encoder": { |
| "nlayers": 12, |
| "token_dim": 768, |
| "nheads": 12, |
| "parallel_mha_transition": false, |
| "token_type_emb_dim": 32, |
| "strict_feats": false, |
| "feats_seq": [ |
| "token_type", |
| "x1_a37coors_nm", |
| "x1_a37coors_nm_rel", |
| "chain_idx_seq", |
| "incident_bond_type_counts", |
| "atom37_element_charge_embedding" |
| ], |
| "feats_cond_seq": null, |
| "dim_cond": 128, |
| "idx_emb_dim": 128, |
| "feats_pair_repr": [ |
| "token_bond_path_length", |
| "canonical_polymer_ordering", |
| "molecule_idx_pair", |
| "x1_bb_pair_dists_nm" |
| ], |
| "seq_sep_dim": 64, |
| "pair_repr_dim": 256, |
| "update_pair_repr": false, |
| "update_pair_repr_every_n": 3, |
| "use_tri_mult": false, |
| "use_attn_pair_bias": false, |
| "use_qkln": true, |
| "normalize_latent": false, |
| "latent_z_dim": 32, |
| "n_residue_types": 147, |
| "charge_emb_dim": 16 |
| }, |
| "decoder": { |
| "nlayers": 12, |
| "token_dim": 768, |
| "abs_coors": false, |
| "nheads": 12, |
| "token_type_emb_dim": 32, |
| "parallel_mha_transition": false, |
| "strict_feats": false, |
| "feats_seq": [ |
| "ca_coors_nm", |
| "z_latent_seq", |
| "token_type", |
| "cond_residue_type" |
| ], |
| "feats_cond_seq": null, |
| "dim_cond": 128, |
| "idx_emb_dim": 128, |
| "feats_pair_repr": [ |
| "token_bond_path_length", |
| "canonical_polymer_ordering", |
| "molecule_idx_pair", |
| "ca_coors_nm_pair_dists" |
| ], |
| "seq_sep_dim": 64, |
| "pair_repr_dim": 256, |
| "update_pair_repr": false, |
| "update_pair_repr_every_n": 3, |
| "use_tri_mult": false, |
| "use_attn_pair_bias": false, |
| "use_qkln": true, |
| "latent_z_dim": 32, |
| "n_residue_types": 147 |
| } |
| } |
| }, |
| "path_length_search_cutoffs": [ |
| 64 |
| ] |
| } |
| } |
|
|
