address #39

Author: lucidrains

README.md

@@ -20,6 +20,8 @@ A fork with full Lightning + Hydra support is being maintained by <a href="https
 
 - <a href="https://github.com/gitabtion">Heng</a> for pointing out inconsistencies with the paper and pull requesting the solutions
 
+- <a href="https://github.com/luwei0917">Wei Lu</a> for catching a few erroneous hyperparameters
+
 - <a href="https://github.com/amorehead">Alex</a> for the PDB dataset preparation script!
 
 - <a href="https://github.com/patrick-kidger">Patrick</a> for <a href="https://docs.kidger.site/jaxtyping/">jaxtyping</a>, <a href="https://github.com/fferflo">Florian</a> for <a href="https://github.com/fferflo/einx">einx</a>, and of course, <a href="https://github.com/arogozhnikov">Alex</a> for <a href="https://einops.rocks/">einops</a>

alphafold3_pytorch/alphafold3.py

@@ -1455,6 +1455,7 @@ def __init__(
         attn_window_size = None,
         attn_pair_bias_kwargs: dict = dict(),
         attn_num_memory_kv = False,
+        trans_expansion_factor = 2,
         num_register_tokens = 0,
         serial = False,
         use_linear_attn = False,
@@ -1509,7 +1510,8 @@ def __init__(
             )
 
             transition = Transition(
-                dim = dim
+                dim = dim,
+                expansion_factor = trans_expansion_factor
             )
 
             conditionable_pair_bias = ConditionWrapper(
@@ -1959,7 +1961,7 @@ def __init__(
         sigma_data = 0.5,      # standard deviation of data distribution
         rho = 7,               # controls the sampling schedule
         P_mean = -1.2,         # mean of log-normal distribution from which noise is drawn for training
-        P_std = 1.2,           # standard deviation of log-normal distribution from which noise is drawn for training
+        P_std = 1.5,           # standard deviation of log-normal distribution from which noise is drawn for training
         S_churn = 80,          # parameters for stochastic sampling - depends on dataset, Table 5 in apper
         S_tmin = 0.05,
         S_tmax = 50,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.54"
+version = "0.1.55"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }