FuhaiLiAiLab/GALAX

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

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

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

🔬 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:

@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}
}