FuhaiLiAiLab/GALAX

1---2language: en3license: mit4tags:5  - graph-ml6  - bioinformatics7  - precision-medicine8  - explainable-ai9  - reinforcement-learning10datasets:11  - FuhaiLiAiLab/Target-QA12library_name: transformers13pipeline_tag: text-generation14model-index:15  - name: GALAX16    results:17      - task:18          type: text-generation19          name: Target Prioritization20        dataset:21          name: Target-QA22          type: FuhaiLiAiLab/Target-QA23        metrics:24          - type: precision25            value: 0.547226          - type: recall27            value: 0.533228          - type: hit@1029            value: 0.881530          - type: hit@531            value: 0.924932---33 34# GALAX: Graph-Augmented Language Model for Explainable Reinforcement-Guided Subgraph Reasoning in Precision Medicine35 36<div align="center">37  <img src="https://github.com/FuhaiLiAiLab/GALAX/blob/main/Figures/GALAX-logo.png?raw=true" width="40%" alt="GALAX" />38</div>39 40<div align="center" style="line-height: 1;">41  <!-- GitHub -->42  <a href="https://github.com/FuhaiLiAiLab/GALAX" target="_blank" style="margin: 2px;">43    <img alt="GitHub" src="https://img.shields.io/badge/GitHub-GALAX%20Code-181717?logo=github&logoColor=white" style="display: inline-block; vertical-align: middle;"/>44  </a>45 46  <!-- Hugging Face Model -->47  <a href="https://huggingface.co/FuhaiLiAiLab/GALAX" target="_blank" style="margin: 2px;">48    <img alt="Hugging Face Model" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-GALAX%20Model-ffc107?color=ffc107&logoColor=white" style="display: inline-block; vertical-align: middle;"/>49  </a>50 51  <!-- Hugging Face Dataset -->52  <a href="https://huggingface.co/datasets/FuhaiLiAiLab/Target-QA" target="_blank" style="margin: 2px;">53    <img alt="Hugging Face Dataset" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Target--QA%20Dataset-ff6f61?color=ffc107&logoColor=white" style="display: inline-block; vertical-align: middle;"/>54  </a>55</div>56 57<div align="center" style="line-height: 1;">58  <!-- arXiv -->59  <a href="https://arxiv.org/abs/2509.20935" target="_blank" style="margin: 2px;">60    <img alt="arXiv" src="https://img.shields.io/badge/arXiv-GALAX%20Paper-b31b1b?logo=arxiv&logoColor=white" style="display: inline-block; vertical-align: middle;"/>61  </a>62 63  <!-- License -->64  <a href="LICENSE" style="margin: 2px;">65    <img alt="License" src="https://img.shields.io/badge/License-MIT-0a4d92?logo=open-source-initiative&logoColor=white" style="display: inline-block; vertical-align: middle;"/>66  </a>67</div>68 69---70 71## 🧩 Model Overview72 73![GALAX Overall Architecture](https://github.com/FuhaiLiAiLab/GALAX/blob/main/Figures/Figure3.png?raw=true)74 75**GALAX** is a graph-augmented language model designed for explainable target prioritization in precision medicine. It combines three key components:  76- **LLaMA3-8B-Instruct** as the language backbone, further adapted with the BioMedGraphica corpus and fine-tuned on Target-QA.  77- **Graph Attention Network (GAT)** pretrained on integrated multi-omics data and BioMedGraphica knowledge graphs.  78- **A reinforcement-guided subgraph generator** that enables interpretable reasoning by constructing biologically meaningful subgraphs from multi-omics and knowledge graph signals.  79 80By jointly leveraging **multi-omics features**, **protein–protein interactions**, and **disease–target associations**, GALAX provides an interpretable framework for **CRISPR target prioritization** across diverse cancer cell lines. To support benchmarking and reproducibility, we also introduce the **[Target-QA dataset](https://huggingface.co/datasets/FuhaiLiAiLab/Target-QA)**.81 82---83 84## 🚀 How to Use85 86```python87from transformers import AutoModelForCausalLM, AutoTokenizer88from huggingface_hub import snapshot_download89import os, torch90 91# 1. Load GALAX language model92model_id = "FuhaiLiAiLab/GALAX"93tokenizer = AutoTokenizer.from_pretrained(model_id)94lm_model = AutoModelForCausalLM.from_pretrained(95    model_id,96    device_map="auto",97    torch_dtype="auto"98)99 100# 2. Access graph foundation model101repo_path = snapshot_download(model_id)102combined_model_path = os.path.join(repo_path, "best_combined_model.pt")103device = "cuda" if torch.cuda.is_available() else "cpu"104best_combined_model = torch.load(combined_model_path, map_location=device)105```106 107---108 109## ⚙️ Experimental Setup110 111- **Backbone LM:** LLaMA3-8B-Instruct (QA-tuned).  112- **Graph Encoder:** BioBERT-v1.1 embeddings + GAT with edge masking.  113- **Training:** Adam optimizer on 2× NVIDIA H100 (80GB).  114- **Top features per omics modality:** K = 10.  115- **Subgraph rollout depth:** L = 5, candidate nodes η = 20.  116- **Evaluation:** Precision, Recall, F1, Jaccard, Hit@5, Hit@10.  117 118---119 120 121## 📊 Results122 123GALAX consistently outperforms baselines and ablation variants.124 125- **Overall Precision:** 0.5472  126- **Overall Recall:** 0.5332  127- **Hit@10:** 0.8815  128- **Hit@5:** 0.9249  129 130**Table 1. Precision and Recall across datasets**131 132| Model                  | Overall Precision ↑ | Overall Recall ↑ | LUAD Precision ↑ | LUAD Recall ↑ | BRCA Precision ↑ | BRCA Recall ↑ |133|-------------------------|---------------------|------------------|------------------|---------------|------------------|---------------|134| M2T                     | 0.0016              | 0.0011           | 0.0020           | 0.0014        | 0.0000           | 0.0000        |135| GAT                     | 0.0006 ± 0.0000     | 0.0006 ± 0.0000  | 0.0000 ± 0.0000  | 0.0000 ± 0.0000 | 0.0033 ± 0.0000 | 0.0033 ± 0.0000 |136| L3 + Omics              | 0.0071 ± 0.0032     | 0.0013 ± 0.0002  | 0.0079 ± 0.0137  | 0.0005 ± 0.0008 | 0.0020 ± 0.0035 | 0.0017 ± 0.0029 |137| L3 + Omics + KG         | 0.0125 ± 0.0032     | 0.0029 ± 0.0003  | 0.0014 ± 0.0025  | 0.0010 ± 0.0017 | 0.0073 ± 0.0068 | 0.0033 ± 0.0029 |138| L3-FT(Med) + Omics      | 0.0179 ± 0.0045     | 0.0133 ± 0.0064  | 0.0091 ± 0.0018  | 0.0105 ± 0.0044 | 0.0110 ± 0.0086 | 0.0106 ± 0.0075 |139| L3-FT(Med) + Omics + KG | 0.0158 ± 0.0030     | 0.0058 ± 0.0011  | 0.0081 ± 0.0071  | 0.0024 ± 0.0017 | 0.0149 ± 0.0057 | 0.0050 ± 0.0000 |140| L3-FT(QA) + Omics       | 0.5250 ± 0.0282     | 0.4959 ± 0.0435  | 0.5201 ± 0.0408  | 0.4905 ± 0.0532 | 0.5074 ± 0.0498 | 0.4856 ± 0.0570 |141| L3-FT(QA) + Omics + KG  | 0.5185 ± 0.0240     | 0.4908 ± 0.0402  | 0.5214 ± 0.0242  | 0.4952 ± 0.0432 | 0.4856 ± 0.0395 | 0.4656 ± 0.0436 |142| G-Retriever + pre-GAT   | 0.4763 ± 0.0004     | 0.3929 ± 0.0063  | 0.4642 ± 0.0181  | 0.3881 ± 0.0264 | 0.4414 ± 0.0099 | 0.3772 ± 0.0010 |143| **GALAX**               | **0.5472 ± 0.0053** | **0.5332 ± 0.0031** | **0.5345 ± 0.0185** | **0.5157 ± 0.0043** | **0.5608 ± 0.0031** | **0.5533 ± 0.0033** |144 145**Table 2. Hit@10 and Hit@5 across datasets**146 147| Model                  | Overall Hit@10 ↑ | Overall Hit@5 ↑ | LUAD Hit@10 ↑ | LUAD Hit@5 ↑ | BRCA Hit@10 ↑ | BRCA Hit@5 ↑ |148|-------------------------|------------------|-----------------|---------------|--------------|---------------|--------------|149| M2T                     | 0.0029           | 0.0000          | 0.0000        | 0.0000       | 0.0000        | 0.0000       |150| GAT                     | 0.0000 ± 0.0000  | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 |151| L3 + Omics              | 0.0021 ± 0.0037  | 0.0032 ± 0.0055 | 0.0048 ± 0.0082 | 0.0095 ± 0.0165 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 |152| L3 + Omics + KG         | 0.0122 ± 0.0033  | 0.0085 ± 0.0037 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 | 0.0056 ± 0.0096 | 0.0111 ± 0.0192 |153| L3-FT(Med) + Omics      | 0.0122 ± 0.0072  | 0.0116 ± 0.0097 | 0.0000 ± 0.0000 | 0.0000 ± 0.0000 | 0.0111 ± 0.0192 | 0.0000 ± 0.0000 |154| L3-FT(Med) + Omics + KG | 0.0132 ± 0.0040  | 0.0106 ± 0.0048 | 0.0048 ± 0.0082 | 0.0095 ± 0.0165 | 0.0111 ± 0.0192 | 0.0000 ± 0.0000 |155| L3-FT(QA) + Omics       | 0.8693 ± 0.0157  | 0.8889 ± 0.0168 | 0.8667 ± 0.0218 | 0.8476 ± 0.0165 | 0.8389 ± 0.0096 | 0.8889 ± 0.0509 |156| L3-FT(QA) + Omics + KG  | 0.8529 ± 0.0153  | 0.8794 ± 0.0114 | 0.8048 ± 0.0541 | 0.7905 ± 0.0436 | 0.8222 ± 0.0347 | 0.8778 ± 0.0192 |157| G-Retriever + pre-GAT   | 0.8550 ± 0.0046  | 0.8804 ± 0.0037 | 0.8524 ± 0.0165 | 0.8857 ± 0.0000 | **0.8667 ± 0.0000** | 0.8667 ± 0.0000 |158| **GALAX**               | **0.8815 ± 0.0033** | **0.9249 ± 0.0048** | **0.8810 ± 0.0082** | **0.9238 ± 0.0436** | 0.8500 ± 0.0441 | **0.8889 ± 0.0839** |159 160---161 162## 🔬 Intended Uses163 164- **Research use only**  165- Benchmarking **graph-language foundation models** in target priorization166- Target prioritization in **cancer biology**167 168---169 170## 📜 Citation171 172If you use this model, please cite:173 174```bibtex175@article{zhang2025galax,176  title     = {GALAX: Graph-Augmented Language Model for Explainable Reinforcement-Guided Subgraph Reasoning in Precision Medicine},177  author    = {Zhang, Heming and Huang, Di and Li, Wenyu and Province, Michael and Chen, Yixin and Payne, Philip and Li, Fuhai},178  journal   = {arXiv preprint arXiv:2509.20935},179  year      = {2025},180  doi       = {10.48550/arXiv.2509.20935},181  url       = {https://arxiv.org/abs/2509.20935}182}183