File size: 5,426 Bytes
3e62986
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
{
  "evoformer": {
    "max_relative_chain": 2,
    "max_relative_idx": 32,
    "msa_channel": 64,
    "msa_stack": {
      "msa_attention": {
        "num_head": 8
      },
      "msa_transition": {
        "num_intermediate_factor": 4,
        "use_glu_kernel": true
      },
      "num_layer": 4,
      "outer_product_mean": {
        "chunk_size": 128,
        "num_outer_channel": 32
      },
      "pair_attention": {
        "num_head": 4
      },
      "pair_transition": {
        "num_intermediate_factor": 4,
        "use_glu_kernel": true
      },
      "shard_transition_blocks": true,
      "triangle_multiplication_incoming": {
        "equation": "kjc,kic->ijc",
        "use_glu_kernel": true
      },
      "triangle_multiplication_outgoing": {
        "equation": "ikc,jkc->ijc",
        "use_glu_kernel": true
      }
    },
    "num_msa": 1024,
    "pair_channel": 128,
    "pairformer": {
      "block_remat": false,
      "num_layer": 48,
      "pair_attention": {
        "num_head": 4
      },
      "pair_transition": {
        "num_intermediate_factor": 4,
        "use_glu_kernel": true
      },
      "remat_block_size": 8,
      "shard_transition_blocks": true,
      "single_attention": {
        "key_dim": null,
        "num_head": 16,
        "value_dim": null
      },
      "single_transition": {
        "num_intermediate_factor": 4,
        "use_glu_kernel": true
      },
      "triangle_multiplication_incoming": {
        "equation": "kjc,kic->ijc",
        "use_glu_kernel": true
      },
      "triangle_multiplication_outgoing": {
        "equation": "ikc,jkc->ijc",
        "use_glu_kernel": true
      }
    },
    "per_atom_conditioning": {
      "atom_transformer": {
        "attention": {
          "key_dim": 128,
          "num_head": 4,
          "value_dim": 128
        },
        "num_blocks": 3,
        "num_intermediate_factor": 2
      },
      "per_atom_channels": 128,
      "per_atom_pair_channels": 16,
      "per_token_channels": 384
    },
    "seq_channel": 384,
    "template": {
      "dgram_features": {
        "max_bin": 50.75,
        "min_bin": 3.25,
        "num_bins": 39
      },
      "num_channels": 64,
      "template_stack": {
        "num_layer": 2,
        "pair_attention": {
          "num_head": 4
        },
        "pair_transition": {
          "num_intermediate_factor": 2,
          "use_glu_kernel": true
        },
        "shard_transition_blocks": true,
        "single_attention": null,
        "single_transition": null,
        "triangle_multiplication_incoming": {
          "equation": "kjc,kic->ijc",
          "use_glu_kernel": true
        },
        "triangle_multiplication_outgoing": {
          "equation": "ikc,jkc->ijc",
          "use_glu_kernel": true
        }
      }
    }
  },
  "global_config": {
    "bfloat16": "all",
    "final_init": "zeros",
    "flash_attention_implementation": "triton",
    "pair_attention_chunk_size": [
      [
        1536,
        128
      ],
      [
        null,
        32
      ]
    ],
    "pair_transition_shard_spec": [
      [
        2048,
        null
      ],
      [
        null,
        1024
      ]
    ]
  },
  "heads": {
    "confidence": {
      "dgram_features": {
        "max_bin": 50.75,
        "min_bin": 3.25,
        "num_bins": 39
      },
      "max_error_bin": 31.0,
      "no_embedding_prob": 0.2,
      "num_bins": 64,
      "num_plddt_bins": 50,
      "pae": {
        "max_error_bin": 31.0,
        "num_bins": 64
      },
      "pairformer": {
        "num_layer": 4,
        "pair_attention": {
          "num_head": 4
        },
        "pair_transition": {
          "num_intermediate_factor": 4,
          "use_glu_kernel": true
        },
        "shard_transition_blocks": true,
        "single_attention": {
          "key_dim": null,
          "num_head": 16,
          "value_dim": null
        },
        "single_transition": {
          "num_intermediate_factor": 4,
          "use_glu_kernel": true
        },
        "triangle_multiplication_incoming": {
          "equation": "kjc,kic->ijc",
          "use_glu_kernel": true
        },
        "triangle_multiplication_outgoing": {
          "equation": "ikc,jkc->ijc",
          "use_glu_kernel": true
        }
      }
    },
    "diffusion": {
      "atom_transformer": {
        "attention": {
          "key_dim": 128,
          "num_head": 4,
          "value_dim": 128
        },
        "num_blocks": 3,
        "num_intermediate_factor": 2
      },
      "conditioning": {
        "pair_channel": 128,
        "prob": 0.8,
        "seq_channel": 384
      },
      "eval": {
        "gamma_0": 0.8,
        "gamma_min": 1.0,
        "noise_scale": 1.003,
        "num_samples": 5,
        "step_scale": 1.5,
        "steps": 200
      },
      "eval_batch_dim_shard_size": 5,
      "eval_batch_size": 5,
      "per_atom_channels": 128,
      "per_atom_pair_channels": 16,
      "per_token_channels": 768,
      "transformer": {
        "attention": {
          "key_dim": null,
          "num_head": 16,
          "value_dim": null
        },
        "block_remat": false,
        "num_blocks": 24,
        "num_intermediate_factor": 2,
        "super_block_size": 4
      }
    },
    "distogram": {
      "first_break": 2.3125,
      "last_break": 21.6875,
      "num_bins": 64
    }
  },
  "num_recycles": 10,
  "return_distogram": false,
  "return_embeddings": false
}