| { |
| "schema_version": "0.1.0", |
| "model": { |
| "family": "boltz-2", |
| "weights_source": "boltz-community/boltz-2 (default)", |
| "scope": "single-sequence protein only (no MSA, no templates, no affinity head)", |
| "concrete_shapes": { |
| "B": 1, |
| "N_max": 46, |
| "A_max": 352, |
| "note": "Graphs exported with concrete shapes. Dynamic axes are a Phase 5+ follow-up." |
| }, |
| "hyperparams": { |
| "token_s": 384, |
| "token_z": 128, |
| "atoms_per_window_queries_W": 32, |
| "atoms_per_window_keys_H": 128, |
| "max_num_atoms_per_token": 23 |
| } |
| }, |
| "diffusion": { |
| "sigma_min": 0.0001, |
| "sigma_max": 160.0, |
| "sigma_data": 16.0, |
| "rho": 7.0, |
| "gamma_0": 0.8, |
| "gamma_min": 1.0, |
| "noise_scale": 1.003, |
| "step_scale": 1.5, |
| "alignment_reverse_diff": true, |
| "default_num_sampling_steps": 5, |
| "schedule_formula": "for i in range(num_sampling_steps): sigmas[i] = (sigma_max**(1/rho) + i/(num-1) * (sigma_min**(1/rho) - sigma_max**(1/rho)))**rho * sigma_data; then append 0.0 as final step." |
| }, |
| "graphs": { |
| "trunk": { |
| "files": { |
| "fp32": "fp32/trunk.onnx (+ trunk.onnx.data)", |
| "fp16": "fp16/trunk_fp16.onnx (+ trunk_fp16.onnx.data)", |
| "int8": "int8/trunk_int8.onnx (+ trunk_int8.onnx.data)" |
| }, |
| "purpose": "One recycling pass. JS owns the recycling loop \u2014 feed (s, z) back as (s_prev, z_prev) for the next iteration.", |
| "inputs": "78 feats tensors (see feats_spec below) + s_prev [B, N, token_s] + z_prev [B, N, N, token_z]. On iteration 0, s_prev and z_prev are zeros.", |
| "outputs": [ |
| { |
| "name": "s", |
| "shape": "[B, N, token_s]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "z", |
| "shape": "[B, N, N, token_z]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "pdistogram", |
| "shape": "[B, N, N, 1, 64]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "q", |
| "shape": "[B, A, 128]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "c", |
| "shape": "[B, A, 128]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "atom_enc_bias", |
| "shape": "[B, K, W, H, 12] K=A/W, W=atoms_per_window_queries_W, H=atoms_per_window_keys_H", |
| "dtype": "float32" |
| }, |
| { |
| "name": "atom_dec_bias", |
| "shape": "[B, K, W, H, 12]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "token_trans_bias", |
| "shape": "[B, N, N, token_s]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "s_inputs", |
| "shape": "[B, N, token_s]", |
| "dtype": "float32" |
| } |
| ] |
| }, |
| "diffusion_step": { |
| "files": { |
| "fp32": "fp32/diffusion_step.onnx (+ .data)", |
| "fp16": "fp16/diffusion_step_fp16.onnx (+ .data)", |
| "int8": "int8/diffusion_step_int8.onnx (+ .data)" |
| }, |
| "purpose": "One denoising step. JS owns the sampling loop.", |
| "inputs": "78 feats tensors + the 8 trunk-cached tensors (s, s_inputs, q, c, atom_enc_bias, atom_dec_bias, token_trans_bias) + x_noisy [B, A, 3] (current atom coords) + sigma [B] (1-D tensor with t_hat = sigma_tm * (1 + gamma); must NOT be a scalar).", |
| "outputs": [ |
| { |
| "name": "x_denoised", |
| "shape": "[B, A, 3]", |
| "dtype": "float32" |
| } |
| ], |
| "to_keys_note": "to_keys (used by the atom encoder) is NOT a graph input. The diffusion graph reconstructs it internally via get_indexing_matrix(K=A/W, W, H) at trace time." |
| }, |
| "confidence": { |
| "files": { |
| "fp32": "(local: phase5b_confidence/confidence.onnx) \u2014 not yet uploaded to HF", |
| "fp16": "(local: phase5b_confidence/confidence_fp16.onnx)", |
| "int8": "(local: phase5b_confidence/confidence_int8.onnx)" |
| }, |
| "purpose": "Confidence head \u2014 one forward pass after diffusion completes.", |
| "inputs": "78 feats tensors + s_inputs [B, N, token_s] + s [B, N, token_s] + z [B, N, N, token_z] + x_pred [B, A, 3] (final denoised atom coords).", |
| "outputs": [ |
| { |
| "name": "plddt_logits", |
| "shape": "[B, N, 50]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "pae_logits", |
| "shape": "[B, N, N, 64]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "pde_logits", |
| "shape": "[B, N, N, 64]", |
| "dtype": "float32" |
| }, |
| { |
| "name": "resolved_logits", |
| "shape": "[B, N, 2]", |
| "dtype": "float32" |
| } |
| ] |
| } |
| }, |
| "orchestration": { |
| "recycling_loop": "for i in range(recycling_steps + 1):\n s, z, pdistogram, q, c, aeb, adb, ttb, s_inputs = trunk(feats, s_prev, z_prev)\n s_prev = s; z_prev = z\n(s_prev, z_prev start as zeros on iteration 0.)", |
| "sampling_loop": "sigmas = build sigma schedule (length sampling_steps+1, last = 0)\ngammas = where(sigmas > gamma_min, gamma_0, 0.0)\natom_coords = sigmas[0] * randn([B, A, 3])\natom_coords_denoised = None\nfor step in range(sampling_steps):\n sigma_tm, sigma_t, gamma = sigmas[step], sigmas[step+1], gammas[step+1]\n R, tr = compute_random_augmentation(B) # Haar uniform rotation + translation\n atom_coords = (atom_coords - atom_coords.mean(axis=-2, keepdims=True)) @ R + tr\n if atom_coords_denoised is not None:\n atom_coords_denoised = (atom_coords_denoised - atom_coords_denoised.mean(axis=-2, keepdims=True)) @ R + tr\n t_hat = sigma_tm * (1 + gamma)\n noise_var = noise_scale**2 * (t_hat**2 - sigma_tm**2)\n eps = sqrt(noise_var) * randn([B, A, 3])\n atom_coords_noisy = atom_coords + eps\n atom_coords_denoised = diffusion_step(feats, s, s_inputs, q, c, aeb, adb, ttb,\n x_noisy=atom_coords_noisy,\n sigma=[t_hat]) # 1-D tensor\n if alignment_reverse_diff: # default True for Boltz-2\n atom_coords_noisy = weighted_rigid_align(atom_coords_noisy, atom_coords_denoised, atom_mask, atom_mask)\n denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat\n atom_coords = atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma", |
| "ca_extraction": "USE `token_to_center_atom`, NOT `token_to_rep_atom`!\nca_coords = einsum('bna,bad->bnd', feats['token_to_center_atom'].float(), atom_coords)\nshape: [B, N, 3]. Drop the batch dim for PDB writing.\nWHY: token_to_rep_atom points at C\u03b2 (distogram input). token_to_center_atom points at C\u03b1. Using rep_atom for C\u03b1 output writes C\u03b2 coords labeled 'CA' \u2014 produces a tangled rope when rendered. See P-9 in EXPORT_PLAN.md for the postmortem.", |
| "plddt_decoding": "Per-residue pLDDT from plddt_logits [B, N, 50]:\nprobs = softmax(plddt_logits, axis=-1)\nbin_width = 1.0 / 50 = 0.02\nbin_centers = [0.01, 0.03, 0.05, ..., 0.99] # 50 values, bin_width * (i + 0.5)\nplddt = sum(probs * bin_centers, axis=-1) # [B, N] in [0, 1]\nplddt_for_display = plddt * 100 # in [0, 100]" |
| }, |
| "feats_spec": { |
| "affinity_token_mask": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "asym_id": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "atom_backbone_feat": { |
| "shape": [ |
| 1, |
| 352, |
| 17 |
| ], |
| "dtype": "int64" |
| }, |
| "atom_pad_mask": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "float32" |
| }, |
| "atom_resolved_mask": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "bool" |
| }, |
| "atom_to_token": { |
| "shape": [ |
| 1, |
| 352, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "bfactor": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "float32" |
| }, |
| "chiral_atom_index": { |
| "shape": [ |
| 1, |
| 4, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "chiral_atom_orientations": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "chiral_reference_mask": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "connected_atom_index": { |
| "shape": [ |
| 1, |
| 2, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "connected_chain_index": { |
| "shape": [ |
| 1, |
| 2, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "contact_conditioning": { |
| "shape": [ |
| 1, |
| 46, |
| 46, |
| 5 |
| ], |
| "dtype": "int64" |
| }, |
| "contact_negation_mask": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "contact_pair_index": { |
| "shape": [ |
| 1, |
| 2, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "contact_threshold": { |
| "shape": [ |
| 1, |
| 46, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "contact_thresholds": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "float32" |
| }, |
| "contact_union_index": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "coords": { |
| "shape": [ |
| 1, |
| 1, |
| 352, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "cyclic_period": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "deletion_mean": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "deletion_value": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "disto_center": { |
| "shape": [ |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "disto_coords_ensemble": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "disto_target": { |
| "shape": [ |
| 1, |
| 46, |
| 46, |
| 1, |
| 64 |
| ], |
| "dtype": "float32" |
| }, |
| "ensemble_ref_idxs": { |
| "shape": [ |
| 1, |
| 1 |
| ], |
| "dtype": "int64" |
| }, |
| "entity_id": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "frame_resolved_mask": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "bool" |
| }, |
| "frames_idx": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "int64" |
| }, |
| "has_deletion": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "bool" |
| }, |
| "method_feature": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "modified": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "mol_type": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "msa": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "msa_mask": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "msa_paired": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "planar_bond_index": { |
| "shape": [ |
| 1, |
| 6, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "planar_ring_5_index": { |
| "shape": [ |
| 1, |
| 5, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "planar_ring_6_index": { |
| "shape": [ |
| 1, |
| 6, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "plddt": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "float32" |
| }, |
| "profile": { |
| "shape": [ |
| 1, |
| 46, |
| 33 |
| ], |
| "dtype": "float32" |
| }, |
| "query_to_template": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "r_set_to_rep_atom": { |
| "shape": [ |
| 1, |
| 46, |
| 352 |
| ], |
| "dtype": "int64" |
| }, |
| "rdkit_bounds_angle_mask": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "rdkit_bounds_bond_mask": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "rdkit_bounds_index": { |
| "shape": [ |
| 1, |
| 2, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "rdkit_lower_bounds": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "float32" |
| }, |
| "rdkit_upper_bounds": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "float32" |
| }, |
| "ref_atom_name_chars": { |
| "shape": [ |
| 1, |
| 352, |
| 4, |
| 64 |
| ], |
| "dtype": "int64" |
| }, |
| "ref_charge": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "float32" |
| }, |
| "ref_chirality": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "int64" |
| }, |
| "ref_element": { |
| "shape": [ |
| 1, |
| 352, |
| 128 |
| ], |
| "dtype": "int64" |
| }, |
| "ref_pos": { |
| "shape": [ |
| 1, |
| 352, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "ref_space_uid": { |
| "shape": [ |
| 1, |
| 352 |
| ], |
| "dtype": "int64" |
| }, |
| "res_type": { |
| "shape": [ |
| 1, |
| 46, |
| 33 |
| ], |
| "dtype": "int64" |
| }, |
| "residue_index": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "stereo_bond_index": { |
| "shape": [ |
| 1, |
| 4, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "stereo_bond_orientations": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "stereo_reference_mask": { |
| "shape": [ |
| 1, |
| 0 |
| ], |
| "dtype": "bool" |
| }, |
| "sym_id": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "symmetric_chain_index": { |
| "shape": [ |
| 1, |
| 2, |
| 0 |
| ], |
| "dtype": "int64" |
| }, |
| "template_ca": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "template_cb": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "template_frame_rot": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "template_frame_t": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 3 |
| ], |
| "dtype": "float32" |
| }, |
| "template_mask": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "template_mask_cb": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "template_mask_frame": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "template_restype": { |
| "shape": [ |
| 1, |
| 1, |
| 46, |
| 33 |
| ], |
| "dtype": "int64" |
| }, |
| "token_bonds": { |
| "shape": [ |
| 1, |
| 46, |
| 46, |
| 1 |
| ], |
| "dtype": "float32" |
| }, |
| "token_disto_mask": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "token_index": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "token_pad_mask": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "token_resolved_mask": { |
| "shape": [ |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| }, |
| "token_to_center_atom": { |
| "shape": [ |
| 1, |
| 46, |
| 352 |
| ], |
| "dtype": "int64" |
| }, |
| "token_to_rep_atom": { |
| "shape": [ |
| 1, |
| 46, |
| 352 |
| ], |
| "dtype": "int64" |
| }, |
| "type_bonds": { |
| "shape": [ |
| 1, |
| 46, |
| 46 |
| ], |
| "dtype": "int64" |
| }, |
| "visibility_ids": { |
| "shape": [ |
| 1, |
| 1, |
| 46 |
| ], |
| "dtype": "float32" |
| } |
| } |
| } |
