{
  "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
    ]
  }
}