update readme

README.md

@@ -31,7 +31,7 @@ Here we briefly introduce the details of pre-training of AIDO.RAGProtein-16B. Ma
 
 ### Data
 
-**UniRef50/Uniclust30 MSA dataset**: We utilized sequences from UniRef50 as queries to search for homologous sequences in UniClust30, subsequently constructing multiple sequence alignments (MSAs). UniRef50 comprises a total of 53.6 million sequences. Using HHblits, we searched all sequences, identifying over 25 homologous sequences for 23.7 million of them. This dataset was directly used as the training set, referred to as `HHblits_MSA`. The remaining 29.9 million sequences were input into MSA Retriever, resulting in 7.7 million sequences with more than 25 homologous sequences. This dataset was designated as `Retriever_MSA`. During training, RAGPLM randomly sampled from the two datasets with probabilities of 0.75 and 0.25
+**UniRef50/Uniclust30 MSA dataset**: We utilized sequences from UniRef50 as queries to search for homologous sequences in UniClust30, subsequently constructing multiple sequence alignments (MSAs). UniRef50 comprises a total of 53.6 million sequences. Using HHblits, we searched all sequences, identifying over 25 homologous sequences for 23.7 million of them. This dataset was directly used as the training set, referred to as `HHblits_MSA`. The remaining 29.9 million sequences were input into MSA Retriever, resulting in 7.7 million sequences with more than 25 homologous sequences. This dataset was designated as `Retriever_MSA`. During training, RAGPLM randomly sampled from the two datasets with probabilities of 0.75 and 0.25. Refer to AIDO.Protein-RAG-3B paper ([link](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1)) for more information.
 
 **AlphaFold Database MSA & Structure dataset**: We downloaded all the structural data from the AlphaFold Database and only kept the structures where the amino acid ratio of pLDDT>70 was greater than 40%. Then we used `mmseqs` to cluster the remaining sequences with `seq id=0.5`, and retained a representative sequence for each class. Final we get 46.9 million sequence/structure pairs. For each structure, we used [genbio-ai/AIDO.StructureTokenizer](https://huggingface.co/genbio-ai/AIDO.StructureTokenizer) to obtain the corresponding structure tokens and structure embedding. And used [MSA Retriever](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1) to obtain the MSA corresponding to the sequence.
 
@@ -45,34 +45,45 @@ Same training data with [AIDO.Protein-16B](https://huggingface.co/genbio-ai/AIDO
 
 #### (2) UniRef50/Uniclust30 MSA finetuning
 
-We used UniRef50/Uniclust30 MSA dataset to finetune the model from stage (1). Refer [AIDO.RAGPLM](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1) for more information.
+We used UniRef50/Uniclust30 MSA dataset to finetune the model from stage (1). Refer to AIDO.Protein-RAG-3B paper ([link](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1)) for more information.
 
 #### (3) AFDB MSA & Structure tokens finetuning:
 
 We fine-tuned a pretrained masked language model using MSA data by concatenating the query sequence with homologous sequences. The input structure embedding (hidden dimension 384) is linearly mapped to 2304 and then added to the corresponding embedding of the query sequence tokens. 
 
-**Mask of sequences**: We introduced several modifications to the standard BERT masking strategy: (1) We randomly sampled `0.05×L` span positions from a query sequence of length `L`, with span lengths following a geometric distribution (`p=0.2`), and capped the maximum length at 10. Our experiments revealed that this settings lead to an average of 15% of the query tokens were masked. (2) To prevent information leakage, when a residue was selected, all residues at the same index across all sequences (the column of the MSA matrix) were also masked. (3) When a column of MSA was selected for masking, the entire column was replaced with the `<MASK>` token in 80% of cases, with random amino acids in 10% of cases, and remained unchanged in the remaining 10% of cases. 
+**Ssequence masking**: We introduced several modifications to the standard BERT masking strategy: (1) We randomly sampled `0.05×L` span positions from a query sequence of length `L`, with span lengths following a geometric distribution (`p=0.2`), and capped the maximum length at 10. Our experiments revealed that this settings lead to an average of 15% of the query tokens were masked. (2) To prevent information leakage, when a residue was selected, all residues at the same index across all sequences (the column of the MSA matrix) were also masked. (3) When a column of MSA was selected for masking, the entire column was replaced with the `<MASK>` token in 80% of cases, with random amino acids in 10% of cases, and remained unchanged in the remaining 10% of cases. 
 
-**Mask of structure**: In 20% of the cases, we randomly replaced the structure embedding with 0; in 80% of the cases, we randomly sampled a certain number of amino acids using the BetaLinear30 distribution and masked their structure embedding. The BetaLinear30 distribution is defined as a combination of 20% of the [0, 1] uniform distribution and 80% of the Beta(3, 9) Beta distribution.
+**Structure masking**: In 20% of the cases, we randomly replaced the structure embedding with 0; in 80% of the cases, we randomly sampled the number of amino acids using the BetaLinear30 distribution and replaced the structure embedding with 0. The BetaLinear30 distribution is defined as a combination of 20% of the Uniform(0, 1) distribution and 80% of the Beta(3, 9) distribution.
 
-**Positional embedding**: To help the model distinguish which tokens are from the same chain and which tokens have the same residue index, we use [2D rotary position embedding](https://arxiv.org/abs/2406.05347) to encode the tokens.
+**Positional embedding**: To help the model distinguish which tokens are from the same chain and which tokens have the same residue index, we use [2D rotary position embedding](https://arxiv.org/abs/2406.05347) to encode the tokens. Refer to AIDO.Protein-RAG-3B paper ([link](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1)) for more information.
 
-**Loss**: The loss function consists of a sequence loss function and a structure loss function (weights are 1.0 and 0.01 respectively). The sequence loss function is the CrossEntropy function that recovers the masked sequence tokens, and the structure loss function is the CrossEntropy function that predicts each masked structure token.
+**Loss**: The loss function consists of a sequence loss function and a structure loss function (weights are 1.0 and 0.01 respectively). The sequence loss function is the CrossEntropy of recovering the masked sequence tokens, and the structure loss function is the CrossEntropy of predicting the masked structure tokens.
 
 | Hyper-params                | (1) 1D -> 2D finetuning | (2) UniRef50/Uniclust30 MSA finetuning | (3) AFDB MSA & Structure tokens finetuning |
 | --------------------------- | :---------------------: | :------------------------------------: | :----------------------------------------: |
+| Initialized parameters      |   AIDO.Protein-16B      |       Stage (1)                        |                      Stage (2)             |
 | Data                        |   ColabFoldDB, UniRef   |       HHblits_MSA, Retriever_MSA       |        AFDB MSA & Structure tokens         |
 | Global Batch Size           |           512           |                  256                   |                    256                     |
 | Sequence length             |          2048           |                 12800                  |                   12800                    |
 | Per Device Micro Batch Size |            1            |                   1                    |                     1                      |
 | Precision                   |     Mixed FP32-FP16     |            Mixed FP32-FP16             |              Mixed FP32-FP16               |
-| LR                |       [5e-6,5e-5]       |              [1e-6, 1e-5]              |                    1e-5                    |
-| Num Tokens        |       10 billion        |              100 billion               |                 80 billion                 |
+| LR                          |       [5e-6,5e-5]       |              [1e-6, 1e-5]              |                    1e-5                    |
+| Num Tokens                  |       10 billion        |              100 billion               |                 80 billion                 |
 
 ### Tokenization
 
 We encode protein sequence with single amino acid resolution with 44 vocabularies, where 24 tokens represent amino acid types and 20 are special tokens. Sequences were also suffixed with a `[SEP]` token as hooks for downstream tasks. 
 
+## Results
+
+### Supervised downstream tasks
+
+<center><img src="supervised_tasks.png" alt="supervised_tasks" style="width:90%; height:auto;" /></center>
+
+### Supervised DMS fitness score prediction of 25 samples
+
+<center><img src="supervised_dms.png" alt="supervised_dms" style="width:90%; height:auto;" /></center>
+
 ## How to Use
 
 ### Build any downstream models from this backbone with ModelGenerator
@@ -93,7 +104,7 @@ import torch
 from modelgenerator.tasks import Embed
 model = Embed.from_config({"model.backbone": "aido_ragprotein_16b"}).eval()
 model.backbone.max_length = 12800
-data = torch.load("ModelGenerator/experiments/AIDO.RAGPLM/examples.pt", 'cpu')[0]
+data = torch.load("examples.pt", 'cpu')[0]
 transformed_batch = model.transform(data)
 with torch.no_grad():
     embedding = model(transformed_batch)
@@ -108,7 +119,7 @@ import torch
 from modelgenerator.tasks import SequenceClassification
 model = SequenceClassification.from_config({"model.backbone": "aido_ragprotein_16b", "model.n_classes": 2}).eval()
 model.backbone.max_length = 12800
-data = torch.load("ModelGenerator/experiments/AIDO.RAGPLM/examples.pt", 'cpu')[0]
+data = torch.load("examples.pt", 'cpu')[0]
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)
@@ -124,7 +135,7 @@ import torch
 from modelgenerator.tasks import TokenClassification
 model = TokenClassification.from_config({"model.backbone": "aido_ragprotein_16b", "model.n_classes": 3}).eval()
 model.backbone.max_length = 12800
-data = torch.load("ModelGenerator/experiments/AIDO.RAGPLM/examples.pt", 'cpu')[0]
+data = torch.load("examples.pt", 'cpu')[0]
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)
@@ -139,7 +150,7 @@ print(torch.argmax(logits, dim=-1))
 from modelgenerator.tasks import SequenceRegression
 model = SequenceRegression.from_config({"model.backbone": "aido_ragprotein_16b"}).eval()
 model.backbone.max_length = 12800
-data = torch.load("ModelGenerator/experiments/AIDO.RAGPLM/examples.pt", 'cpu')[0]
+data = torch.load("examples.pt", 'cpu')[0]
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)

examples.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e60409786b001317d5150ab440521fa1f9a6a90ca18f4666e27740b4e6a75aa5
+size 5485326