QizhiPei commited on
Commit
cd73a83
Β·
verified Β·
1 Parent(s): c2c0ca1

Create README.md

Browse files
Files changed (1) hide show
  1. README.md +161 -0
README.md ADDED
@@ -0,0 +1,161 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ language:
4
+ - en
5
+ tags:
6
+ - biology
7
+ - chemistry
8
+ - molecule
9
+ - protein
10
+ - multimodal
11
+ - foundation-model
12
+ - pretrained
13
+ pipeline_tag: text-generation
14
+ base_model: Qwen/Qwen3-4B
15
+ library_name: transformers
16
+ ---
17
+
18
+ # BioMatrix-4B-Base
19
+
20
+ **BioMatrix** is a multimodal biological foundation model that natively integrates **1D sequences**, **3D structures**, and **natural language** for both **molecules** and **proteins** within a single decoder-only architecture.
21
+
22
+ This is the **4B-parameter Base model**, obtained via **multimodal continual pretraining** of Qwen3-4B on 304.4 billion tokens spanning text, molecular and protein 1D/3D data, and cross-modal corpora. This base checkpoint is intended for further fine-tuning on downstream tasks. For an instruction-tuned model ready for inference, see [BioMatrix-4B-SFT](https://huggingface.co/QizhiPei/BioMatrix-4B-SFT).
23
+
24
+ - πŸ“„ **Paper**: [BioMatrix: Towards a Comprehensive Biological Foundation Model Spanning the Modality Matrix of Sequences, Structures, and Language](https://arxiv.org/abs/xxxx.xxxxx)
25
+ - πŸ’» **Code**: [https://github.com/QizhiPei/biomatrix](https://github.com/QizhiPei/biomatrix)
26
+ - πŸ€— **Model & Data Collection**: [https://huggingface.co/collections/QizhiPei/biomatrix](https://huggingface.co/collections/QizhiPei/biomatrix)
27
+
28
+ ## Model Overview
29
+
30
+ BioMatrix maps **all biological modalities into a shared discrete token space** via a unified tokenization scheme:
31
+
32
+ - **Molecular 1D sequences** (both SMILES and SELFIES notations)
33
+ - **Molecular 3D structures** (via MolStrucTok with branch-decoupled decoder)
34
+ - **Protein 1D sequences** (residue-level tokens)
35
+ - **Protein 3D structures** (via GCP-VQVAE backbone tokenizer)
36
+ - **Natural language** (inherited from Qwen3 tokenizer)
37
+
38
+ All modalities are consumed and produced uniformly under a **single next-token prediction objective**β€”without external encoders, projection adapters, or modality-specific output heads.
39
+
40
+ | Model | Molecule 1D | Molecule 3D | Protein 1D | Protein 3D | Natural Language |
41
+ |-------|:-----------:|:-----------:|:----------:|:----------:|:----------------:|
42
+ | ESM3 | βœ— | βœ— | βœ“ | βœ“ | βœ“ |
43
+ | 3D-MoLM | βœ“ | βœ“ | βœ— | βœ— | βœ“ |
44
+ | AlphaFold3 | βœ“ | βœ“ | βœ“ | βœ“ | βœ— |
45
+ | BioT5/BioT5+ | βœ“ | βœ— | βœ“ | βœ— | βœ“ |
46
+ | BioMedGPT | βœ“ | βœ— | βœ“ | βœ— | βœ“ |
47
+ | **BioMatrix** | **βœ“** | **βœ“** | **βœ“** | **βœ“** | **βœ“** |
48
+
49
+ ## Model Details
50
+
51
+ - **Base Architecture**: Qwen3-4B
52
+ - **Parameters**: 4B
53
+ - **Training Stage**: Multimodal Continual Pretraining only (not instruction-tuned)
54
+ - **Training Tokens**: 304.4B
55
+ - **Context Length**: 8,192 tokens
56
+ - **Tokenizer**: Extended Qwen3 vocabulary with:
57
+ - 11,294 joint molecular 3D tokens (composed from SELFIES atom Γ— MolStrucTok codes)
58
+ - 4,096 protein 3D tokens (GCP-VQVAE codebook)
59
+ - 26 protein 1D tokens (amino acids + non-standard/unknown)
60
+ - SELFIES atom tokens and modality-specific control tokens
61
+
62
+ ### Embedding Initialization
63
+
64
+ New vocabulary entries are initialized via a **description-based scheme**: each new token is grounded in the pretrained Qwen3 embedding space by averaging the embeddings of the subword tokens of a short natural-language description (e.g., `<A_W>` β†’ "Tryptophan"), plus a small isotropic Gaussian perturbation to break symmetry. This provides a more stable starting point than random initialization.
65
+
66
+ ## Pretraining Corpus (304.4B tokens)
67
+
68
+ | Category | Tokens | Sources |
69
+ |----------|--------|---------|
70
+ | **Text** (105.3B) | General: 25.6B | FineWeb-Edu |
71
+ | | Scientific: 79.7B | FineFineWeb (biology/chemistry/medical/health), PubMed Full Articles |
72
+ | **Molecule** (73.7B) | 1D: 36.0B | PubChem, MolTextNet |
73
+ | | 3D: 17.6B | PubChem, PCQM4Mv2, PubChemQC |
74
+ | | Other: 24.0B | (text descriptions, properties, IUPAC names) |
75
+ | **Protein** (77.4B) | 1D: 17.1B | UniRef50 |
76
+ | | 3D: 38.5B | RCSB PDB, AlphaFold DB |
77
+ | | Other: 19.5B | Swiss-Prot, TrEMBL annotations |
78
+ | | Other (additional): 2.9B | |
79
+ | **Cross-entity** (48.0B) | Interleaved Text: 17.1B | PubMed, bioRxiv, S2ORC, USPTO |
80
+ | | 3D: 11.4B | CrossDocked, PPIRef |
81
+ | | Other: 19.5B | BindingDB, STITCH, jglaser, AlphaSeq |
82
+
83
+ ### Training Configuration
84
+
85
+ - **Framework**: LLaMA-Factory
86
+ - **Hardware**: 64 NVIDIA H100 GPUs
87
+ - **Global Batch Size**: 1,024
88
+ - **Maximum Sequence Length**: 8,192 tokens
89
+ - **Optimizer**: AdamW
90
+ - **Peak Learning Rate**: 2.0 Γ— 10⁻⁴ (cosine schedule)
91
+ - **Warmup Steps**: 2,000
92
+ - **Total Steps**: ~36.4K (1 epoch over the full 304.4B-token corpus)
93
+
94
+ ## Intended Use
95
+
96
+ This **Base model is not instruction-tuned**. It is suitable for:
97
+
98
+ - **Further fine-tuning** on custom biological tasks
99
+ - **Continued pretraining** on domain-specific corpora
100
+ - **Research on representation learning** across biomolecular modalities
101
+ - **Embedding extraction** for downstream classification/regression tasks
102
+
103
+ For ready-to-use instruction-following capabilities (e.g., molecule captioning, protein design, property prediction), please use the [SFT variant](https://huggingface.co/QizhiPei/BioMatrix-4B-SFT).
104
+
105
+ ## Quick Start
106
+
107
+ ```python
108
+ from transformers import AutoModelForCausalLM, AutoTokenizer
109
+
110
+ model_name = "QizhiPei/BioMatrix-4B-Base"
111
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
112
+ model = AutoModelForCausalLM.from_pretrained(
113
+ model_name,
114
+ torch_dtype="auto",
115
+ device_map="auto",
116
+ trust_remote_code=True
117
+ )
118
+
119
+ # Example: Continue a SMILES sequence
120
+ prompt = "<|mol_smi_start|>CC(=O)"
121
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
122
+ outputs = model.generate(**inputs, max_new_tokens=64)
123
+ print(tokenizer.decode(outputs[0], skip_special_tokens=False))
124
+ ```
125
+
126
+ ## Modality Wrapping
127
+
128
+ When constructing inputs, biomolecular content must be wrapped with the corresponding control tokens:
129
+
130
+ | Modality | Wrapping Example |
131
+ |----------|------------------|
132
+ | Molecule SMILES | `<\|mol_smi_start\|>CC#CC#N<\|mol_smi_end\|>` |
133
+ | Molecule SELFIES | `<\|mol_sfi_start\|>[C][#C][C][#N]<\|mol_sfi_end\|>` |
134
+ | Molecule 3D | `<\|mol_3d_start\|>[H_3][C_0][#C_6]...<\|mol_3d_end\|>` |
135
+ | Protein 1D | `<\|prot_aa_start\|><A_M><A_R><A_A>...<\|prot_aa_end\|>` |
136
+ | Protein 3D | `<\|prot_3d_start\|><S_4012><S_153>...<\|prot_3d_end\|>` |
137
+
138
+ Natural language text is left unwrapped and serves as the default carrier modality.
139
+
140
+ ## Limitations
141
+
142
+ - This model is **not instruction-tuned** and is unlikely to follow natural-language instructions out-of-the-box. Use the SFT variant for instruction-following.
143
+ - Molecular and protein 3D structures are tokenized in **disjoint geometric reference frames**, so the model cannot natively represent biomolecular complexes (e.g., docking poses).
144
+ - Heavy domain specialization may erode some general-purpose language capabilities of the underlying Qwen3 backbone.
145
+ - Coverage is limited to **small molecules and proteins**; nucleic acids, carbohydrates, and lipids are not currently supported.
146
+
147
+ ## Citation
148
+
149
+ If you find BioMatrix useful, please cite:
150
+
151
+ ```bibtex
152
+ @article{pei2026biomatrix,
153
+ title={BioMatrix: Towards a Comprehensive Biological Foundation Model Spanning the Modality Matrix of Sequences, Structures, and Language},
154
+ author={Pei, Qizhi and Zhou, Zhimeng and Duan, Yi and Zhao, Yiyang and He, Liang and Hsieh, Chang-Yu and He, Conghui and Yan, Rui and Wu, Lijun},
155
+ year={2026}
156
+ }
157
+ ```
158
+
159
+ ## License
160
+
161
+ This model is released under the Apache 2.0 license. The base model (Qwen3-4B) is subject to its own license terms.