update readme3

README.md

@@ -4,11 +4,15 @@ license: other
 
 # AIDO.RAGProtein-16B 
 
-AIDO.RAGProtein-16B is a multimodal protein language model that integrates MSA and structural data based on the [AIDO.Protein-16B](https://huggingface.co/genbio-ai/AIDO.Protein-16B) model. Its training is divided into multiple stages, 100 billion tokens are trained on UniRef50/UniClust30 MSA data, and 80 billion tokens are trained on AlphaFold Database MSA and structural data.
+AIDO.RAGProtein-16B is a multimodal protein language model that integrates Multiple Sequence Alignment (MSA) and structural data, building upon the [AIDO.Protein-16B](https://huggingface.co/genbio-ai/AIDO.Protein-16B) foundation. The training process comprises three main stages:
+
+1. 2D RoPE encoding fine-tuning
+2. Initial training on 100 billion tokens from UniRef50/UniClust30 MSA data
+3. Subsequent training on 80 billion tokens from AlphaFold Database MSA and structural data
 
 ## Model Architecture Details
 
-AIDO.RAGProtein-16B is a transformer encoder-only architecture with the dense MLP layer in each transformer block replaced by a sparse MoE layer. It uses single amino acid tokenization and is optimized using a masked languange modeling (MLM) training objective. For each token, 2 experts will be selectively activated by the top-2 rounting mechiansim.
+AIDO.RAGProtein-16B employs a transformer encoder-only architecture featuring sparse Mixture-of-Experts (MoE) layers that replace dense MLP layers in each transformer block. Utilizing single amino acid tokenization and optimized through masked language modeling (MLM), the model activates 2 experts per token via top-2 routing mechanisms.
 
 <center><img src="proteinmoe_architecture.png" alt="An Overview of AIDO.Protein" style="width:70%; height:auto;" /></center>
 
@@ -27,37 +31,55 @@ More architecture details are shown below:
 
 ## Pre-training of AIDO.RAGProtein-16B
 
-Here we briefly introduce the details of pre-training of AIDO.RAGProtein-16B. Mainly divided into three stages: (1) 1D -> 2D finetuning; (2) UniRef50/Uniclust30 MSA finetuning; (3) AlphaFold Database MSA & Structure tokens finetuning.
+Here we briefly introduce the details of pre-training of AIDO.RAGProtein-16B. Mainly divided into three stages: (1) 1D -> 2D RoPE encoding finetuning; (2) UniRef50/Uniclust30 MSA finetuning; (3) AlphaFold Database MSA & Structure tokens finetuning.
 
 ### 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. 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.
+**AlphaFold Database MSA & Structure dataset**: We downloaded all structural data from the AlphaFold Database and kept only those where more than 40% of amino acids had a pLDDT score > 70. The remaining sequences were clustered using `mmseqs` (`seq id=0.5`), and one representative per cluster was retained, resulting in 46.9 million sequence/structure pairs. For each structure, we used [genbio-ai/AIDO.StructureTokenizer](https://huggingface.co/genbio-ai/AIDO.StructureTokenizer) to obtain structure tokens and embeddings. [MSA Retriever](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1) was used to obtain the corresponding MSA.
 
 ### Training Details
 
 Model training is divided into three stages: 
 
-#### (1) 1D -> 2D finetuning: 
+#### (1) 1D -> 2D RoPE Encoding Fine-tuning 
+
+Same training data as [AIDO.Protein-16B](https://huggingface.co/genbio-ai/AIDO.Protein-16B), but with [2D rotary position embedding](https://arxiv.org/abs/2406.05347) for token encoding.
+
+#### (2) UniRef50/UniClust30 MSA Fine-tuning
+
+The model from Stage 1 is further fine-tuned on the UniRef50/Uniclust30 MSA dataset. See the [AIDO.Protein-RAG-3B paper](https://www.biorxiv.org/content/10.1101/2024.12.02.626519v1) for more.
+
+#### (3) AlphaFold Database MSA & Structure Fine-tuning
+
+We fine-tuned the model with concatenated query and homologous sequences. Structure embeddings (dim = 384) are linearly mapped to 2304 and added to the query token embeddings.
+
+##### Sequence Masking
 
-Same training data with [AIDO.Protein-16B](https://huggingface.co/genbio-ai/AIDO.Protein-16B), but use [2D rotary position embedding](https://arxiv.org/abs/2406.05347) to encode the tokens;
+* Randomly sample `0.05 × L` span positions from a query of length `L`. Span lengths follow a geometric distribution (`p=0.2`), capped at length 10. On average, ~15% of query tokens are masked.
 
-#### (2) UniRef50/Uniclust30 MSA finetuning
+* When a residue is selected, its aligned residues across all sequences (MSA column) are also masked.
 
-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.
+* For masked MSA columns: 80% are replaced with `<MASK>`, 10% with random amino acids, and 10% left unchanged.
 
-#### (3) AFDB MSA & Structure tokens finetuning:
+##### Structure Masking
 
-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. 
+* In 20% of cases, structure embeddings are replaced with 0.
 
-**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. 
+* In 80% of cases, a number of amino acids is sampled using the BetaLinear30 distribution and corresponding embeddings are zeroed. (BetaLinear30 = 20% Uniform(0,1) + 80% Beta(3,9)).
 
-**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
 
-**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.
+We use [2D rotary position embedding](https://arxiv.org/abs/2406.05347) to help the model distinguish token chain identities and residue indices. See 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 of recovering the masked sequence tokens, and the structure loss function is the CrossEntropy of predicting the masked structure tokens.
+##### Loss Function
+
+Total loss is a weighted sum of sequence loss (weight 1.0) and structure loss (weight 0.01).
+
+* **Sequence loss**: CrossEntropy loss for masked token prediction.
+
+* **Structure loss**: CrossEntropy loss for masked structure token prediction.
 
 | Hyper-params                | (1) 1D -> 2D finetuning | (2) UniRef50/Uniclust30 MSA finetuning | (3) AFDB MSA & Structure tokens finetuning |
 | --------------------------- | :---------------------: | :------------------------------------: | :----------------------------------------: |
@@ -76,17 +98,17 @@ We encode protein sequence with single amino acid resolution with 44 vocabularie
 
 ## Results
 
-### Supervised downstream tasks
+### 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
+### Supervised DMS Fitness Score Prediction (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
+### Build Downstream Models Using ModelGenerator
 
 For more information, visit: [Model Generator](https://github.com/genbio-ai/modelgenerator)
 
@@ -95,7 +117,7 @@ mgen fit --model SequenceClassification --model.backbone aido_protein_rag_16b --
 mgen test --model SequenceClassification --model.backbone aido_protein_rag_16b --data SequenceClassificationDataModule --data.path <hf_or_local_path_to_your_dataset>
 ```
 
-### Or use directly in Python
+### Use Directly in Python
 
 #### Embedding
 
@@ -104,7 +126,12 @@ import torch
 from modelgenerator.tasks import Embed
 model = Embed.from_config({"model.backbone": "aido_protein_rag_16b"}).eval()
 model.backbone.max_length = 12800
-data = torch.load("examples.pt", 'cpu')[0]
+restypes = 'ARNDCQEGHILKMFPSTWYV'
+data = {
+    'sequences': [''.join(random.choice(restypes) for _ in range(50))],
+    'msa': [ [ ''.join(random.choice(restypes+'-') for _ in range(50)) for _ in range(25) ] ],
+    'str_emb': np.random.normal(size=(1, 50, 384))
+}
 transformed_batch = model.transform(data)
 with torch.no_grad():
     embedding = model(transformed_batch)
@@ -119,7 +146,12 @@ import torch
 from modelgenerator.tasks import SequenceClassification
 model = SequenceClassification.from_config({"model.backbone": "aido_protein_rag_16b", "model.n_classes": 2}).eval()
 model.backbone.max_length = 12800
-data = torch.load("examples.pt", 'cpu')[0]
+restypes = 'ARNDCQEGHILKMFPSTWYV'
+data = {
+    'sequences': [''.join(random.choice(restypes) for _ in range(50))],
+    'msa': [ [ ''.join(random.choice(restypes+'-') for _ in range(50)) for _ in range(25) ] ],
+    'str_emb': np.random.normal(size=(1, 50, 384))
+}
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)
@@ -135,7 +167,12 @@ import torch
 from modelgenerator.tasks import TokenClassification
 model = TokenClassification.from_config({"model.backbone": "aido_protein_rag_16b", "model.n_classes": 3}).eval()
 model.backbone.max_length = 12800
-data = torch.load("examples.pt", 'cpu')[0]
+restypes = 'ARNDCQEGHILKMFPSTWYV'
+data = {
+    'sequences': [''.join(random.choice(restypes) for _ in range(50))],
+    'msa': [ [ ''.join(random.choice(restypes+'-') for _ in range(50)) for _ in range(25) ] ],
+    'str_emb': np.random.normal(size=(1, 50, 384))
+}
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)
@@ -150,7 +187,12 @@ print(torch.argmax(logits, dim=-1))
 from modelgenerator.tasks import SequenceRegression
 model = SequenceRegression.from_config({"model.backbone": "aido_protein_rag_16b"}).eval()
 model.backbone.max_length = 12800
-data = torch.load("examples.pt", 'cpu')[0]
+restypes = 'ARNDCQEGHILKMFPSTWYV'
+data = {
+    'sequences': [''.join(random.choice(restypes) for _ in range(50))],
+    'msa': [ [ ''.join(random.choice(restypes+'-') for _ in range(50)) for _ in range(25) ] ],
+    'str_emb': np.random.normal(size=(1, 50, 384))
+}
 transformed_batch = model.transform(data)
 with torch.no_grad():
     logits = model(transformed_batch)

examples.pt

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