add shrink and perturb fn

Author: lucidrains

README.md

@@ -438,3 +438,15 @@ docker run -v .:/data --gpus all -it af3
     year    = {2023}
 }
 ```
+
+```bibtex
+@article{Ash2019OnTD,
+    title   = {On the Difficulty of Warm-Starting Neural Network Training},
+    author  = {Jordan T. Ash and Ryan P. Adams},
+    journal = {ArXiv},
+    year    = {2019},
+    volume  = {abs/1910.08475},
+    url     = {https://api.semanticscholar.org/CorpusID:204788802}
+}
+```
+

alphafold3_pytorch/alphafold3.py

@@ -5415,6 +5415,20 @@ def init_and_load(
         alphafold3.load(path)
         return alphafold3
 
+    def shrink_and_perturb_(
+        self,
+        shrink_factor = 0.5,
+        perturb_factor = 0.01
+    ):
+        # Shrink & Perturb - Ash et al. https://arxiv.org/abs/1910.08475
+        assert 0. <= shrink_factor <= 1.
+
+        for p in self.parameters():
+            noise = torch.randn_like(p.data)
+            p.data.mul_(1. - shrink_factor).add_(noise * 0.1)
+
+        return self
+
     @typecheck
     def forward(
         self,

pyproject.toml

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