File size: 2,785 Bytes
3a8e4d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
_target_: models.simplefold.torch.architecture.FoldingDiT

hidden_size: 1536
num_heads: 24
atom_num_heads: 8
output_channels: 3
use_atom_mask: False
use_length_condition: True
esm_dropout_prob: 0.0
esm_model: esm2_3B

time_embedder:
  _target_: models.simplefold.torch.layers.TimestepEmbedder
  hidden_size: 1536
aminoacid_pos_embedder:
  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
  in_dim: 1
  embed_dim: 1536
  include_input: True
pos_embedder:
  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
  in_dim: 3
  include_input: True
  min_freq_log2: 0
  max_freq_log2: 12
  num_freqs: 128
  log_sampling: True

trunk:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 36
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 1536
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 1536
      num_heads: 24
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 1536
        num_heads: 24
        base: 100.0

atom_hidden_size_enc: 512
atom_n_queries_enc: 32
atom_n_keys_enc: 128
atom_encoder_transformer:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 3
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 512
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 512
      num_heads: 8
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 512
        num_heads: 8
        base: 100.0

atom_hidden_size_dec: 512
atom_n_queries_dec: 32
atom_n_keys_dec: 128
atom_decoder_transformer:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 3
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 512
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 512
      num_heads: 8
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 512
        num_heads: 8
        base: 100.0