File size: 9,163 Bytes
00a287b
ca2d66e
 
 
00a287b
ca2d66e
 
 
 
 
 
 
 
 
 
 
00a287b
 
ca2d66e
 
 
 
 
 
4a9a443
 
ca2d66e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74b1ffc
ca2d66e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
 
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
 
 
 
00a287b
ca2d66e
00a287b
ca2d66e
00a287b
ca2d66e
 
 
 
 
 
 
00a287b
ca2d66e
00a287b
ca2d66e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
---
license: apache-2.0
language:
- en
tags:
- biology
- chemistry
- molecule
- protein
- multimodal
- foundation-model
- drug-discovery
- protein-design
pipeline_tag: text-generation
base_model: QizhiPei/BioMatrix-4B-Base
library_name: transformers
---

# BioMatrix-4B-SFT

**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.

This is the **4B-parameter SFT (Supervised Fine-Tuned)** variant, instruction-tuned across 80 downstream biological tasks spanning 6 categories.

- 📄 **Paper**: [BioMatrix: Towards a Comprehensive Biological Foundation Model Spanning the Modality Matrix of Sequences, Structures, and Language](https://github.com/QizhiPei/BioMatrix/blob/main/biomatrix_tech_report.pdf)
- 💻 **Code**: [https://github.com/QizhiPei/BioMatrix](https://github.com/QizhiPei/BioMatrix)
- 🤗 **Model & Data Collection**: [https://huggingface.co/collections/QizhiPei/biomatrix](https://huggingface.co/collections/QizhiPei/biomatrix)

## Model Overview

BioMatrix closes the gap between native multimodality and broad entity coverage in biological foundation models. Unlike adapter-based approaches that bolt external encoders onto a language model, or prior native-tokenization models confined to a single entity type, BioMatrix maps **all modalities into a shared discrete token space** via a unified tokenization scheme:

- **Molecular 1D sequences** (both SMILES and SELFIES notations)
- **Molecular 3D structures** (via MolStrucTok with branch-decoupled decoder)
- **Protein 1D sequences** (residue-level tokens)
- **Protein 3D structures** (via GCP-VQVAE backbone tokenizer)
- **Natural language** (inherited from Qwen3 tokenizer)

All modalities are consumed and produced uniformly under a **single next-token prediction objective**—without external encoders, projection adapters, or modality-specific output heads.

| Model | Molecule 1D | Molecule 3D | Protein 1D | Protein 3D | Natural Language |
|-------|:-----------:|:-----------:|:----------:|:----------:|:----------------:|
| ESM3 | ✗ | ✗ | ✓ | ✓ | ✓ |
| 3D-MoLM | ✓ | ✓ | ✗ | ✗ | ✓ |
| AlphaFold3 | ✓ | ✓ | ✓ | ✓ | ✗ |
| BioT5/BioT5+ | ✓ | ✗ | ✓ | ✗ | ✓ |
| BioMedGPT | ✓ | ✗ | ✓ | ✗ | ✓ |
| NatureLM | ✓ | ✗ | ✓ | ✗ | ✓ |
| SciReasoner | ✓ | ✗ | ✓ | ✗ | ✓ |
| **BioMatrix** | **✓** | **✓** | **✓** | **✓** | **✓** |

## Model Details

- **Base Architecture**: Qwen3-4B-Base
- **Parameters**: 4B
- **Training Stages**:
  - **Continual Pretraining** on 304.4B tokens (general/scientific text, molecular & protein 1D/3D data, cross-modal interleaved corpora)
  - **Instruction Tuning** on a comprehensive suite of 80 downstream tasks across 6 categories
- **Context Length**: 8,192 tokens
- **Tokenizer**: Extended Qwen3 vocabulary with:
  - 11,294 joint molecular 3D tokens (composed from SELFIES atom × MolStrucTok codes)
  - 4,096 protein 3D tokens (GCP-VQVAE codebook)
  - 26 protein 1D tokens (amino acids + non-standard/unknown)
  - SELFIES atom tokens and modality-specific control tokens

## Pretraining Corpus (304.4B tokens)

| Category | Tokens | Sources |
|----------|--------|---------|
| **Text** | 105.3B | FineWeb-Edu, FineFineWeb (biology/chemistry/medical/health), PubMed Full Articles |
| **Molecule** | 73.7B | PubChem, PCQM4Mv2, PubChemQC, MolTextNet |
| **Protein** | 77.4B | UniRef50, RCSB PDB, Swiss-Prot, TrEMBL, AlphaFold DB |
| **Cross-entity** | 48.0B | Interleaved text (PubMed, bioRxiv, S2ORC, USPTO), Molecule–protein (BindingDB, STITCH, jglaser, CrossDocked), Protein–protein (AlphaSeq, PPIRef) |

## Performance Highlights

BioMatrix achieves **state-of-the-art or competitive performance on 77 out of 80 tasks**. Selected highlights for the 4B-SFT variant:

### Molecular Tasks
- **Unconditional 1D Generation** (GuacaMol): 0.998 validity, 1.000 uniqueness, 0.986 novelty
- **Name Conversion (I2S EM)**: 92.83% (vs. SciReasoner-8B: 84.40%)
- **Text-Based Molecule Generation (EM)**: 65.07% (vs. SciReasoner-8B: 48.00%)
- **MoleculeQA Total Accuracy**: 73.78% (vs. prior best MolCA-1.3B: 64.79%)
- **Property-Conditioned 3D Generation**: ~3-4× error reduction on QM9 electronic-structure targets

### Protein Tasks
- **Fold Type Prediction (Family level)**: 87.25% accuracy
- **Annotation Prediction (UniProtSeq Keywords F1)**: 91.26%
- **Inverse Folding AAR**: 75.50% (vs. DPLM-2-3B: 61.67%)
- **Sequence–Structure Co-generation**: scTM = 0.965, scRMSD = 2.80
- **Unconditional Backbone Generation**: scTM = 0.963 (joint frontier with RFDiffusion)

### Interaction Tasks
- **BindingDB Affinity (RMSE)**: 1.030 (new SOTA, surpasses prior literature SOTA of 1.340)
- **PDBBindv2020 3D Affinity**: RMSE = 1.260, Pearson = 0.737, MAE = 0.972

## Quick Start

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "QizhiPei/BioMatrix-4B-SFT"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)

# Example: Molecule captioning with SELFIES input
instruction = "I need a brief explanation of the molecule denoted in this SELFIES notation. <|mol_sfi_start|>[Te]<|mol_sfi_end|>"

messages = [
    {"role": "user", "content": instruction}
]
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=2048, do_sample=False)
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=False)
print(response)
```

## Modality Wrapping

When constructing prompts, biomolecular content must be wrapped with the corresponding control tokens:

| Modality | Wrapping Example |
|----------|------------------|
| Molecule SMILES | `<\|mol_smi_start\|>CC#CC#N<\|mol_smi_end\|>` |
| Molecule SELFIES | `<\|mol_sfi_start\|>[C][#C][C][#N]<\|mol_sfi_end\|>` |
| Molecule 3D | `<\|mol_3d_start\|>[H 3][C 0][#C 6]...<\|mol_3d_end\|>` |
| Protein 1D | `<\|prot_aa_start\|><A M><A R><A A>...<\|prot_aa_end\|>` |
| Protein 3D | `<\|prot_3d_start\|><S 4012><S 153><S 2091>...<\|prot_3d_end\|>` |

Natural language text is left unwrapped and serves as the default carrier modality.

## Supported Tasks

BioMatrix-4B-SFT was instruction-tuned across the following task categories:

**Molecule (1D)**: unconditional generation, name conversion, property prediction, captioning, text-based generation, forward/retrosynthesis, editing, optimization, customized generation, question answering

**Molecule (3D)**: unconditional generation, property-conditioned generation

**Protein (1D)**: sequence understanding, annotation prediction, knowledge mining, text-based design, unconditional generation

**Protein (3D)**: structure understanding, folding, inverse folding, sequence-structure co-generation, unconditional backbone generation

**Interaction**: molecule-protein binding affinity (1D & 3D), protein-protein interaction

> **Note on task-group variants**: As detailed in the paper, the released SFT model is trained on the union of all sub-task corpora with mild oversampling for small-data tasks. For best performance on specific benchmarks, please refer to the paper's task-group-specific variants.

## SMILES vs. SELFIES

BioMatrix supports both notations as parallel 1D molecular representations. Empirically:

- **SELFIES** excels on tasks requiring validity-by-construction (unconditional generation, property optimization)
- **SMILES** excels on tasks requiring surface-level structural anchoring (customized generation with atom/bond/functional-group constraints, forward synthesis, retrosynthesis)

See Section 9.2 of the paper for detailed analysis.

## Limitations

- Molecular and protein 3D structures are tokenized in **disjoint geometric reference frames**, so the model cannot natively represent biomolecular complexes (e.g., docking poses).
- Heavy domain specialization may erode some general-purpose language capabilities of the underlying Qwen3 backbone.
- Coverage is limited to **small molecules and proteins**; nucleic acids, carbohydrates, and lipids are not currently supported.
- Fine-grained 3D geometry (e.g., bond lengths) shows residual quantization error from finite codebooks; a lightweight post-hoc force-field refinement (e.g., MMFF) closes most of this gap.

## Citation

If you find BioMatrix useful, please cite:

```bibtex
@article{pei2026biomatrix,
  title={BioMatrix: Towards a Comprehensive Biological Foundation Model Spanning the Modality Matrix of Sequences, Structures, and Language},
  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},
  year={2026}
}
```

## License

This model is released under the Apache 2.0 license. The base model (Qwen3-4B-Base) is subject to its own license terms.