This repository contains the implementation code for training large language models using knowledge graph-guided reinforcement learning, as described in our paper.

Knowledge Graph-Guided Reinforcement Learning for LLMs (arXiv:2601.15160)

Kansal, Yuval and Jha, Niraj K. Knowledge Graphs are Implicit Reward Models: Path-Derived Signals Enable Compositional Reasoning. arXiv preprint arXiv:2601.15160, 2026.

For detailed methodology, experimental results, and theoretical foundations, please refer to the paper.

Important: Our method was developed and evaluated using a knowledge graph derived from the paradigm presented in QA-GNN: Reasoning with Language Models and Knowledge Graphs for Question Answering (Yasunaga et al., NAACL 2021). While the overall algorithm (SFT → RL with path-derived rewards) is knowledge-graph agnostic, the following components are KG-dependent and may require adaptation for your use case:

• Data loading and preprocessing (data_loader.py, data_prep.py): Assumes specific schema (e.g., question_and_explanation, paths, category, source_concept, target_concept)
• Diversity-based filtering (create_filtered_dataset.py): Relies on KG metadata (categories, concepts, path patterns)
• Path alignment reward (rl_training.py): Expects paths field with KG structure

Users working with different knowledge graphs may need to modify these scripts to match their data schema and KG representation. The core training pipeline (LoRA SFT + GRPO RL) remains flexible and model-agnostic.

Keywords: reinforcement-learning · knowledge-graph · large-language-models · GRPO · LoRA · question-answering · compositional-reasoning

# Clone the repository
git clone <repository-url>
cd <repository-name>

# Install dependencies
pip install -r requirements.txt

• Python 3.11+
• PyTorch 2.0+
• CUDA 12.0+ (for GPU training)
• DeepSpeed
• Transformers
• TRL (Transformer Reinforcement Learning)
• PEFT (Parameter-Efficient Fine-Tuning)

.
├── data_loader.py               # Data loading utilities (placeholder for training data)
├── data_prep.py                 # Dataset preprocessing and train/test splitting
├── create_filtered_dataset.py   # Diversity-based filtering for SFT/RL splits
├── sft_training.py              # Supervised Fine-Tuning with LoRA
├── rl_training.py               # Reinforcement Learning (GRPO) training
├── configs/
│   ├── deepspeed_config.json    # DeepSpeed ZeRO-3 configuration
│   └── slurm_template.sh        # SLURM job submission template
├── requirements.txt             # Python dependencies
└── README.md                    # This file

Our training pipeline uses a diversity-based filtering approach to split data between SFT and RL:

Create a high-diversity subset (e.g., 5k examples) for RL training:

python create_filtered_dataset.py \
    --input_path <path-to-full-dataset> \
    --output_path <path-to-filtered-dataset> \
    --target_size 5000 \
    --min_per_category 2

This script uses stratified sampling to maximize:

• Category coverage: All categories represented
• Concept diversity: Rare source/target concepts prioritized
• Path pattern variety: Diverse knowledge graph paths
• Node coverage: Maximum unique nodes from KG

The filtered dataset (5k examples) is used for RL training, while the remaining examples (~19.6k) are used for SFT training.

The training scripts automatically handle data preprocessing. However, you can optionally preprocess datasets in advance:

For SFT training:

python data_prep.py \
    --input_path <path-to-sft-dataset> \
    --output_path <path-to-processed-sft-data> \
    --mode sft

For RL training:

python data_prep.py \
    --input_path <path-to-filtered-dataset> \
    --output_path <path-to-processed-rl-data> \
    --mode rl \
    --enable_thinking

Note: If you skip this step, the training scripts will automatically convert the data to the correct format during training.

Train a base model using LoRA for parameter-efficient fine-tuning:

# Single-node multi-GPU training
torchrun --nproc_per_node=8 sft_training.py \
    --model_name "Qwen/Qwen3-14B" \
    --dataset_path <path-to-training-data> \
    --output_dir ./sft_models/qwen3-14b-lora \
    --learning_rate 2e-4 \
    --num_train_epochs 20 \
    --deepspeed configs/deepspeed_config.json

# Or submit via SLURM
sbatch configs/slurm_template.sh

Continue training with RL using the SFT checkpoint:

torchrun --nproc_per_node=8 rl_training.py \
    --model_name "Qwen/Qwen3-14B" \
    --sft_checkpoint_path ./sft_models/qwen3-14b-lora/checkpoint-XXX \
    --dataset_path <path-to-rl-data> \
    --output_dir ./rl_models/qwen3-14b-grpo \
    --learning_rate 8e-6 \
    --num_train_epochs 2

• Diversity-Based Data Splitting: Sophisticated filtering that maximizes concept, category, and KG path coverage
• Parameter-Efficient Training: Uses LoRA for memory-efficient fine-tuning of large models
• Knowledge Graph Integration: Incorporates KG path information in reward functions
• Multi-Stage Training: SFT followed by RL for improved reasoning capabilities
• Distributed Training: Supports multi-GPU and multi-node training with DeepSpeed ZeRO-3
• Flexible Architecture: Easily adaptable to different model architectures (Qwen, LLaMA, etc.)

• Learning rate: 2e-4
• Batch size: 1 per device, gradient accumulation: 32
• LoRA rank: 16, alpha: 16
• Max sequence length: 2048 tokens
• Epochs: 20

• Learning rate: 8e-6
• Beta (KL penalty): 0.05
• Number of generations: 2
• Max prompt length: 896 tokens
• Max completion length: 896 tokens
• Epochs: 2

The RL training supports multiple reward signals. The default configuration uses Correctness and Path Alignment; two additional reward functions are available in the codebase and can be enabled as needed:

To enable the optional rewards, uncomment thinking_quality_reward_func and/or semantic_answer_similarity_reward_func in rl_training.py.

The training uses a two-stage data split:

1. Filtered Dataset (RL): 5k examples selected for maximum diversity

   • Ensures coverage of all categories
   • Prioritizes rare concepts and path patterns
   • Maintains long-tail coverage

2. Remaining Dataset (SFT): ~19.6k examples for supervised fine-tuning

   • Provides broad coverage and pattern learning
   • Builds strong base capabilities

This split ensures the RL stage focuses on diverse, challenging examples while the SFT stage provides comprehensive coverage.

Models are saved at regular intervals during training. To load a checkpoint:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-14B")

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "./sft_models/qwen3-14b-lora")

# For inference, merge adapters
model = model.merge_and_unload()

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-14B")

If you use this code in your research, please cite our paper:

@article{kansal2026knowledge,
  title={Knowledge Graphs are Implicit Reward Models: Path-Derived Signals Enable Compositional Reasoning},
  author={Kansal, Yuval and Jha, Niraj K},
  journal={arXiv preprint arXiv:2601.15160},
  year={2026}
}

For questions or issues, please open a GitHub issue or email yuvalkansal@princeton.edu for more information.

This project is licensed under the Princeton License - see the LICENSE file for details.

This work builds upon the following:

Knowledge Graph & QA:

• QA-GNN: Reasoning with Language Models and Knowledge Graphs for Question Answering (Yasunaga et al., NAACL 2021) — Our knowledge graph methodology is derived from this paradigm.

• Bottom-up Domain-specific Superintelligence: A Reliable Knowledge Graph is What We Need (Dedhia et al., 2025) — Our data curation pipeline is derived from the outlined pipeline.

Libraries & Frameworks:

• Transformers (Hugging Face)
• TRL — Transformer Reinforcement Learning (GRPO, SFT)
• PEFT (LoRA)
• DeepSpeed (ZeRO-3)