File size: 9,192 Bytes
b42a228
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bd29d7c
af4d93a
e5610d1
66a97d5
e5610d1
10ce31a
e5610d1
10ce31a
 
 
e5610d1
10ce31a
 
 
 
e5610d1
10ce31a
 
 
 
e5610d1
 
 
 
10ce31a
 
0bfea06
e5610d1
 
10ce31a
e5610d1
10ce31a
e5610d1
 
af4d93a
 
 
 
 
66a97d5
bd29d7c
1649a92
 
 
 
af4d93a
1649a92
af4d93a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10ce31a
af4d93a
 
10ce31a
 
af4d93a
 
bd29d7c
 
 
 
 
 
 
 
 
 
 
 
 
 
af4d93a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10ce31a
 
af4d93a
 
 
 
 
 
 
 
 
10ce31a
 
 
 
 
 
 
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
---
language: en
license: mit
tags:
  - graph-ml
  - bioinformatics
  - precision-medicine
  - explainable-ai
  - reinforcement-learning
datasets:
  - FuhaiLiAiLab/Target-QA
library_name: transformers
pipeline_tag: text-generation
model-index:
  - name: GALAX
    results:
      - task:
          type: text-generation
          name: Target Prioritization
        dataset:
          name: Target-QA
          type: FuhaiLiAiLab/Target-QA
        metrics:
          - type: precision
            value: 0.5472
          - type: recall
            value: 0.5332
          - type: hit@10
            value: 0.8815
          - type: hit@5
            value: 0.9249
---

# GALAX: Graph-Augmented Language Model for Explainable Reinforcement-Guided Subgraph Reasoning in Precision Medicine

<div align="center">
  <img src="https://github.com/FuhaiLiAiLab/GALAX/blob/main/Figures/GALAX-logo.png?raw=true" width="40%" alt="GALAX" />
</div>

<div align="center" style="line-height: 1;">
  <!-- GitHub -->
  <a href="https://github.com/FuhaiLiAiLab/GALAX" target="_blank" style="margin: 2px;">
    <img alt="GitHub" src="https://img.shields.io/badge/GitHub-GALAX%20Code-181717?logo=github&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
  </a>

  <!-- Hugging Face Model -->
  <a href="https://huggingface.co/FuhaiLiAiLab/GALAX" target="_blank" style="margin: 2px;">
    <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;"/>
  </a>

  <!-- Hugging Face Dataset -->
  <a href="https://huggingface.co/datasets/FuhaiLiAiLab/Target-QA" target="_blank" style="margin: 2px;">
    <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;"/>
  </a>
</div>

<div align="center" style="line-height: 1;">
  <!-- arXiv -->
  <a href="https://arxiv.org/abs/2509.20935" target="_blank" style="margin: 2px;">
    <img alt="arXiv" src="https://img.shields.io/badge/arXiv-GALAX%20Paper-b31b1b?logo=arxiv&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
  </a>

  <!-- License -->
  <a href="LICENSE" style="margin: 2px;">
    <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;"/>
  </a>
</div>

---

## 🧩 Model Overview

![GALAX Overall Architecture](https://github.com/FuhaiLiAiLab/GALAX/blob/main/Figures/Figure3.png?raw=true)

**GALAX** is a graph-augmented language model designed for explainable target prioritization in precision medicine. It combines three key components:  
- **LLaMA3-8B-Instruct** as the language backbone, further adapted with the BioMedGraphica corpus and fine-tuned on Target-QA.  
- **Graph Attention Network (GAT)** pretrained on integrated multi-omics data and BioMedGraphica knowledge graphs.  
- **A reinforcement-guided subgraph generator** that enables interpretable reasoning by constructing biologically meaningful subgraphs from multi-omics and knowledge graph signals.  

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

---

## 🚀 How to Use

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import snapshot_download
import os, torch

# 1. Load GALAX language model
model_id = "FuhaiLiAiLab/GALAX"
tokenizer = AutoTokenizer.from_pretrained(model_id)
lm_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype="auto"
)

# 2. Access graph foundation model
repo_path = snapshot_download(model_id)
combined_model_path = os.path.join(repo_path, "best_combined_model.pt")
device = "cuda" if torch.cuda.is_available() else "cpu"
best_combined_model = torch.load(combined_model_path, map_location=device)
```

---

## ⚙️ Experimental Setup

- **Backbone LM:** LLaMA3-8B-Instruct (QA-tuned).  
- **Graph Encoder:** BioBERT-v1.1 embeddings + GAT with edge masking.  
- **Training:** Adam optimizer on 2× NVIDIA H100 (80GB).  
- **Top features per omics modality:** K = 10.  
- **Subgraph rollout depth:** L = 5, candidate nodes η = 20.  
- **Evaluation:** Precision, Recall, F1, Jaccard, Hit@5, Hit@10.  

---


## 📊 Results

GALAX consistently outperforms baselines and ablation variants.

- **Overall Precision:** 0.5472  
- **Overall Recall:** 0.5332  
- **Hit@10:** 0.8815  
- **Hit@5:** 0.9249  

**Table 1. Precision and Recall across datasets**

| Model                  | Overall Precision ↑ | Overall Recall ↑ | LUAD Precision ↑ | LUAD Recall ↑ | BRCA Precision ↑ | BRCA Recall ↑ |
|-------------------------|---------------------|------------------|------------------|---------------|------------------|---------------|
| M2T                     | 0.0016              | 0.0011           | 0.0020           | 0.0014        | 0.0000           | 0.0000        |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| **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** |

**Table 2. Hit@10 and Hit@5 across datasets**

| Model                  | Overall Hit@10 ↑ | Overall Hit@5 ↑ | LUAD Hit@10 ↑ | LUAD Hit@5 ↑ | BRCA Hit@10 ↑ | BRCA Hit@5 ↑ |
|-------------------------|------------------|-----------------|---------------|--------------|---------------|--------------|
| M2T                     | 0.0029           | 0.0000          | 0.0000        | 0.0000       | 0.0000        | 0.0000       |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| 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 |
| **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** |

---

## 🔬 Intended Uses

- **Research use only**  
- Benchmarking **graph-language foundation models** in target priorization
- Target prioritization in **cancer biology**

---

## 📜 Citation

If you use this model, please cite:

```bibtex
@article{zhang2025galax,
  title     = {GALAX: Graph-Augmented Language Model for Explainable Reinforcement-Guided Subgraph Reasoning in Precision Medicine},
  author    = {Zhang, Heming and Huang, Di and Li, Wenyu and Province, Michael and Chen, Yixin and Payne, Philip and Li, Fuhai},
  journal   = {arXiv preprint arXiv:2509.20935},
  year      = {2025},
  doi       = {10.48550/arXiv.2509.20935},
  url       = {https://arxiv.org/abs/2509.20935}
}