Knowledge Graph-Assisted LLM Post-Training for Enhanced Legal Reasoning

TL;DR

This paper introduces a knowledge graph-based method to improve large language models' legal reasoning by integrating domain-specific legal concepts, resulting in enhanced performance on legal reasoning benchmarks.

Contribution

The authors propose a novel KG-assisted post-training approach using the IRAC framework and legal case data to significantly boost LLMs' reasoning in high-stakes legal domains.

Findings

Models trained with KG outperform baselines on legal benchmarks.

70B DPO model achieves top scores on reasoning tasks.

KG-enhanced models demonstrate improved legal reasoning capabilities.

Abstract

LLM post-training has primarily relied on large text corpora and human feedback, without capturing the structure of domain knowledge. This has caused models to struggle dealing with complex reasoning tasks, especially for high-stakes professional domains. In Law, reasoning requires deep understanding of the relations between various legal concepts, a key component missing in current LLM post-training. In this paper, we propose a knowledge graph (KG)-assisted approach for enhancing LLMs' reasoning capability in Legal that is generalizable to other high-stakes domains. We model key legal concepts by following the \textbf{IRAC} (Issue, Rule, Analysis and Conclusion) framework, and construct a KG with 12K legal cases. We then produce training data using our IRAC KG, and conduct both Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) with three state-of-the-art (SOTA) LLMs (30B, 49B and 70B), varying architecture and base model family. Our post-trained models obtained better average performance on 4/5 diverse legal benchmarks (14 tasks) than baselines. In particular, our 70B DPO model achieved the best score on 4/6 reasoning tasks, among baselines and a 141B SOTA legal LLM, demonstrating the effectiveness of our KG for enhancing LLMs' legal reasoning capability.