🧬 BioReason-Pro
Advancing Protein Function Prediction with
Multimodal Biological Reasoning

Abstract

Protein function annotation is fundamental to understanding biological mechanisms, designing therapeutics, and advancing biomedical research. Current computational methods either rely on shallow sequence similarity or treat function prediction as isolated classification tasks, failing to capture the integrative reasoning across sequence, structure, domains, and interactions that expert biologists perform to infer function. We introduce BioReason-Pro, the first multimodal reasoning large language model (LLM) for protein function prediction that integrates protein embeddings with biological context to generate structured reasoning traces. A key input into BioReason-Pro is the set of GO term predictions made by GO-GPT, our autoregressive transformer that captures hierarchical and cross-aspect dependencies of GO terms. BioReason-Pro is trained via supervised fine-tuning on synthetic reasoning traces generated by GPT-5 for over 130K proteins and further optimized through reinforcement learning. It achieves 73.6% F_max on GO term prediction and an LLM judge score of 8/10 on functional summaries, substantially outperforming previous methods. Evaluations with human protein experts show that BioReason-Pro annotations are preferred over ground truth UniProt annotations in 79% of cases. Remarkably, BioReason-Pro de novo predicted experimentally confirmed binding partners with per-residue attention localizing to the exact contact residues resolved in cryo-EM structures of those complexes. Together, GO-GPT and BioReason-Pro establish a framework for protein function prediction that combines precise ontology modeling with interpretable biological reasoning.

Web Interface

Try BioReason-Pro directly through our web-based inference server:

🔗 bioreason.net

Precomputed predictions for 223,000+ proteins (including the Human Protein Atlas) are available at bioreason.net.

The full catalogue is also available for download as a HuggingFace dataset: wanglab/protein_catalogue.

Datasets

The datasets used to train and evaluate BioReason-Pro comprise 133,492 proteins across 3,135 organisms curated from UniProt with experimental GO annotations, InterPro domains, STRING protein-protein interactions, and PDB structures. Temporal holdout follows the CAFA framework. Everything is on our HuggingFace collection:

The training and test sets load directly through datasets. Structures and GO embeddings are fetched with scripts/download_assets.py — see Training.

Note on coverage: the SFT set is 124,367 of the 133,492 curated proteins. The remaining 9,125 were excluded. 6,574 of those have no InterPro annotation (verified: every protein without one was dropped). The reason for the remaining 2,551 is not recorded in the release.

Checkpoints

Model weights are available on our HuggingFace collection:

Installation

Prerequisites

Training additionally requires a CUDA GPU at import time — train_protein_llm.py imports Unsloth unconditionally, which raises on a CPU-only machine. Inference (predict.py, eval.py) does not import Unsloth.

1. Install

git clone https://github.com/bowang-lab/BioReason-Pro.git
cd BioReason-Pro

# ESM must be installed with --no-deps: it pins a transformers version that conflicts with vllm
pip install esm --no-deps

# Everything else (torch, vllm, transformers, unsloth, lightning, ...)
pip install -e .

# flash-attn last: it builds against the already-installed torch
pip install flash-attn --no-build-isolation --no-cache-dir

Order matters. Installing esm without --no-deps, or flash-attn before torch, will fail.

2. Authenticate with HuggingFace

The protein encoder esm3_sm_open_v1 is a gated model. You must accept its licence once at EvolutionaryScale/esm3-sm-open-v1, then log in — otherwise model loading fails with a 401 and no other explanation:

hf auth login          # or: export HF_TOKEN=hf_...

The datasets, checkpoints, structures, and GO embeddings are all public and need no special access.

3. Check it worked

python -c "import torch, unsloth, esm, vllm; print('ok', torch.cuda.is_available())"

Inference

predict.py is a single-file pipeline that runs InterPro, GO-GPT, and BioReason-Pro to predict protein function from sequence. It requires a single GPU (A100 80GB recommended). You can use both the SFT model (more hypothesis, more hallucinations) or the RL model (more accurate, less mechnistically deep).

Quick Start

python predict.py --input examples/test_proteins.tsv --output results.tsv --model_type rl

Input Format

A tab-separated file with a header row and three required columns:

Example (proteins.tsv):

protein_id	organism	sequence
P51864	Homo sapiens (Human)	MDCRKMVRFSYSVIWIMAISKAFELGLVA...
P0A9K3	Escherichia coli (strain K12)	MNIDTREITLEPADNARLLSLCGPFDDNI...

Output Format

A tab-separated file with the same columns as the input plus additional results:

How It Works

The pipeline runs three sequential stages on a single GPU:

1. InterPro (CPU/network) — Queries the EBI InterProScan online API to annotate protein domains. Runs in parallel using all available CPU threads. This stage dominates wall-clock time and is highly variable: measured 94 s and 633 s for two proteins submitted together, depending on EBI queue load. A run that appears stuck here is usually just waiting. Each job gives up after 30 minutes rather than hanging forever.
2. GO-GPT (GPU) — Loads the GO-GPT model to predict Gene Ontology terms. One protein at a time. Unloads from GPU after completion.
3. BioReason-Pro (GPU) — Downloads the selected model checkpoint from HuggingFace, loads it, and generates functional annotations in batches of --batch_size.

Only stage 3 is batched. Stages 1 and 2 are per-protein, so raising --batch_size speeds up the last stage but not the run as a whole. All three stages checkpoint to disk, so --resume picks up where a previous run stopped.

Model checkpoints and GO embeddings are automatically downloaded from HuggingFace — no manual setup required.

Options

--model_type {sft,rl}   Model checkpoint to use (default: rl)
--resume                Resume from checkpoints / skip completed proteins
--batch_size N          Batch size for BioReason-Pro inference (default: 16)
--max_new_tokens N      Maximum tokens to generate (default: 5000)
--temperature F         Sampling temperature (default: 0.0)
--top_p F               Top-p sampling (default: 0.95)

Supported Organisms

BioReason-Pro supports 200+ organisms. See organism_list.txt for the full list. Common examples:

• Homo sapiens (Human)
• Mus musculus (Mouse)
• Escherichia coli (strain K12)
• Saccharomyces cerevisiae (strain ATCC 204508 / S288c) (Baker's yeast)
• Arabidopsis thaliana (Mouse-ear cress)
• Drosophila melanogaster (Fruit fly)

Organism names must match the format in organism_list.txt exactly. Unsupported organisms will still run but GO-GPT predictions may be less accurate.

Training

Supervised fine-tuning, end to end, after Installation:

python scripts/download_assets.py --dest /data/bioreason   # 1. get the data
$EDITOR scripts/sh_train_protein_qwen_staged.sh            # 2. fill in 4 paths
sbatch scripts/sh_train_protein_qwen_staged.sh             # 3. train

The released model was trained on 2 nodes x 4 GPUs; a single 80 GB GPU is enough to run the pipeline at small batch size. Each step is detailed below.

1. Download the assets

Two things training needs are too large for git and live on the Hub. One command fetches both and lays them out exactly as the scripts expect:

python scripts/download_assets.py --dest /data/bioreason

# -> /data/bioreason/go_embeddings   (GO_EMBEDDINGS_PATH)  177 MB down, 338 MB on disk
# -> /data/bioreason/structures      (STRUCTURE_DIR)       ~34 GB down, ~150 GB on disk

The download is resumable — re-run it if interrupted. To fetch only the required asset, add --skip-structures; training then runs sequence-only and will not reproduce the released checkpoint.

All ~370k structures in the shards are extracted, so the local copy is a complete mirror of the published set. Only ~123k are referenced by the released datasets — pass --referenced-only if you would rather save ~90 GB of disk.

Verify at any time:

python scripts/download_assets.py --dest /data/bioreason --verify

Why verify matters. Missing structure files are handled silently: the collate function substitutes empty coordinates, so a wrong STRUCTURE_DIR degrades the model to sequence-only without raising anything. Training prints a loud warning when fewer than 90% of sampled structures resolve, and --verify checks coverage against the released datasets directly.

The GO embeddings can also be rebuilt from scratch instead of downloaded, though downloading is preferred — a different model revision yields different vectors:

python -m bioreason2.utils.go_embed \
    --output_dir /data/bioreason/go_embeddings --batch_size 32 --device cuda

2. Edit one file

Everything is driven by scripts/sh_train_protein_qwen_staged.sh. It is the only file you need to change. Open it and fill in the Paths block near the top:

Then uncomment the #SBATCH header at the top and adjust it for your scheduler.

3. Launch

sbatch scripts/sh_train_protein_qwen_staged.sh
# or, on a single machine (srun is used only when SLURM_JOB_ID is set):
bash scripts/sh_train_protein_qwen_staged.sh

The script preflights every path and exits immediately with an explanation if one is unset or missing, rather than failing an hour into model loading. Training data (wanglab/bioreason-pro-sft-reasoning-data, 117,002 train / 7,365 validation) is pulled from HuggingFace automatically.

To sanity-check the whole path on a 50-sample slice before committing a long run, set --debug True in the script's BASE_COMMAND. That takes about ten minutes on one H100.

Training stages

--training_stage 1 trains only the protein projector and GO encoder with the LLM frozen; --training_stage 2 fine-tunes the whole stack with LoRA. The released checkpoint was produced by Stage 2 alone, starting from randomly-initialised projectors, which is what the script does by default. To run the warm-up first and initialise Stage 2 from its weights:

RUN_STAGE1=true sbatch scripts/sh_train_protein_qwen_staged.sh

Key hyperparameters

Convert the checkpoint

Training writes a PyTorch Lightning .ckpt; eval.py and predict.py need a HuggingFace directory. Convert before evaluating:

bash scripts/sh_convert_unsloth_ddp_to_hf_ckpt.sh /path/to/last.ckpt /path/to/output-hf

The hyperparameters inside that script must match the run that produced the checkpoint (they default to the released SFT settings, including LORA_DROPOUT=0).

Evaluation

Two ways to run the model, depending on whether its inputs already exist:

Use eval.py to reproduce benchmark numbers, predict.py for proteins of your own. The rest of this section is eval.py: convert (above), generate predictions, then score them.

1. Generate. Edit the Paths block of scripts/sh_eval.sh (GO_EMBEDDINGS_PATH, DATASET_CACHE_DIR, STRUCTURE_DIR, and MODEL_PATH — either your own checkpoint or a released one) and run it. It writes one JSON per protein into EVALS_PATH, and is resumable: re-running skips proteins that already have a result.

bash scripts/sh_eval.sh

2. Score. CAFA Fmax and IA-weighted Fmax over that directory:

bash evals/run_cafa_eval.sh <evals_path> [output_dir]

To score a released checkpoint without training anything, point MODEL_PATH at a snapshot of wanglab/bioreason-pro-sft or wanglab/bioreason-pro-rl.

Throughput. --batch_size is how many sequences vLLM decodes concurrently and is the main lever; --gpu_memory_utilization is vLLM's share of the GPU only, since ESM3, the GO encoder and the batch's prompt embeddings are allocated outside it. Shard across GPUs with NUM_CHUNKS/CHUNK_ID; each shard writes into the same directory and skips what is done.

Two things that will silently give you wrong numbers. The GO encoder flags at evaluation must match training exactly (512 / 3 / 8 / 200 / 2560, --unified_go_encoder True, --protein_embedding_layer 37) or the projector loads mismatched weights. And --max_new_tokens must be large enough for the model to finish reasoning and emit its GO summary — the default of 3000 is sized for this. Truncating the generation produces parseable-but-empty predictions and an Fmax near zero, which looks like a bad model rather than a bad setting.

Key Contributions

• First multimodal reasoning LLM for protein function: BioReason-Pro deeply integrates ESM3 protein embeddings, a GO graph encoder, and biological context within a unified LLM to generate structured reasoning traces and functional annotations.

• Autoregressive GO term prediction (GO-GPT): A novel autoregressive transformer that treats Gene Ontology prediction as a sequence generation task, capturing hierarchical and cross-aspect dependencies that discriminative methods miss, achieving state-of-the-art weighted F_max of 0.65–0.70.

• Expert-level functional reasoning: Human protein experts preferred BioReason-Pro annotations over curated UniProt entries in 79% of evaluated cases, with an LLM judge score of 8.03/10 across five evaluation axes.

• De novo structural predictions: BioReason-Pro predicted experimentally validated binding partners (e.g., SBP2 for eEFSec) with per-residue attention localizing to the exact contact interfaces resolved in cryo-EM structures.

• Structural reasoning beyond domain transfer: The model performs contextual architectural reasoning that overrides misleading superfamily-level annotations, as demonstrated for CFAP61's non-enzymatic scaffolding role.

• Broad-scale release: All model weights, training code, curated datasets, a web interface, and precomputed predictions for 223,000+ proteins including the Human Protein Atlas are publicly available.

Citation

If you find this work useful, please cite our papers:

@article {Fallahpour2026.03.19.712954,
	author = {Fallahpour, Adibvafa and Seyed-Ahmadi, Arman and Idehpour, Parsa and Ibrahim, Omar and Gupta, Purav and Naimer, Jack and Zhu, Kevin and Shah, Arnav and Ma, Shihao and Adduri, Abhinav and G{\"u}loglu, Talu and Liu, Nuo and Cui, Haotian and Jain, Arihant and de Castro, Max and Fallahpour, Amirfaham and Cembellin-Prieto, Antonio and Stiles, John S. and Nem{\v c}ko, Filip and Nevue, Alexander A. and Moon, Hyungseok C. and Sosnick, Lucas and Markham, Olivia and Duan, Haonan and Lee, Michelle Y. Y. and Salvador, Andrea F. M. and Maddison, Chris J. and Thaiss, Christoph A. and Ricci-Tam, Chiara and Plosky, Brian S. and Burke, Dave P. and Hsu, Patrick D. and Goodarzi, Hani and Wang, Bo},
	title = {BioReason-Pro: Advancing Protein Function Prediction with Multimodal Biological Reasoning},
	elocation-id = {2026.03.19.712954},
	year = {2026},
	doi = {10.64898/2026.03.19.712954},
	publisher = {Cold Spring Harbor Laboratory},
	URL = {https://www.biorxiv.org/content/early/2026/03/20/2026.03.19.712954},
	eprint = {https://www.biorxiv.org/content/early/2026/03/20/2026.03.19.712954.full.pdf},
	journal = {bioRxiv}
}

@misc{fallahpour2025bioreasonincentivizingmultimodalbiological,
      title={BioReason: Incentivizing Multimodal Biological Reasoning within a DNA-LLM Model}, 
      author={Adibvafa Fallahpour and Andrew Magnuson and Purav Gupta and Shihao Ma and Jack Naimer and Arnav Shah and Haonan Duan and Omar Ibrahim and Hani Goodarzi and Chris J. Maddison and Bo Wang},
      year={2025},
      eprint={2505.23579},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2505.23579}, 
}

Authors

• Adibvafa Fallahpour¹²³⁴⁵ * ([email protected])
• Arman Seyed-Ahmadi²⁵ *
• Parsa Idehpour¹⁵⁹ *
• Omar Ibrahim²⁵ *
• Purav Gupta³⁴⁵ *
• Jack Naimer⁵ ¹⁰
• Kevin Zhu⁵⁸
• Arnav Shah³⁴⁵
• Shihao Ma²³⁴⁵
• Abhinav Adduri¹⁵
• Talu Güloglu⁴⁵ ¹¹
• Nuo Liu¹
• Haotian Cui¹³
• Arihant Jain¹⁹
• Max de Castro⁹
• Amirfaham Fallahpour⁴
• Antonio Cembellin-Prieto¹
• John S. Stiles¹
• Filip Nemčko¹
• Alexander A. Nevue¹
• Hyungseok C. Moon¹
• Lucas Sosnick¹⁶
• Olivia Markham¹²
• Haonan Duan³⁴
• Michelle Y. Y. Lee¹⁶
• Andrea F. M. Salvador¹⁶
• Chris J. Maddison³⁴
• Christoph A. Thaiss¹⁶
• Chiara Ricci-Tam¹
• Brian S. Plosky¹
• Dave P. Burke¹
• Patrick D. Hsu¹⁸
• Hani Goodarzi†‡¹⁷ ([email protected])
• Bo Wang†‡²³⁴ ¹³ ([email protected])

¹ Arc Institute ² University Health Network ³ Vector Institute ⁴ University of Toronto ⁵ Core Contributor
⁶ Stanford University ⁷ University of California, San Francisco ⁸ University of California, Berkeley
⁹ University of Pennsylvania ¹⁰ EPFL ¹¹ ETH Zürich ¹² Cohere ¹³ Xaira Therapeutics

* Equal contribution. The order of authors is not a reflection of their relative contributions.
† These authors, listed alphabetically, jointly supervised this work.
‡ Corresponding authors

---

Made with ❤️ at Arc Institute, University of Toronto, Vector Institute, and University Health Network