Enhancing Legal LLMs through Metadata-Enriched RAG Pipelines and Direct Preference Optimization

1. Problem Statement

Large Language Models (LLMs) struggle with long, complex legal documents, often producing hallucinations (e.g., incorrect clauses, fabricated precedents, or wrong citations). While Retrieval-Augmented Generation (RAG) is a standard solution, it faces two critical failure modes in the legal domain:

1. Retrieval Errors: Due to high lexical redundancy and structural similarity in legal corpora, retrievers often select chunks from incorrect source documents (Document-Level Retrieval Mismatch).
2. Decoding Errors: Models frequently generate answers even when the retrieved context is insufficient, or conversely, they become overly conservative and refuse to answer when the context is sufficient (over-refusal).

These issues are exacerbated by the need for data privacy, which often necessitates using smaller, locally deployed models that lack the robustness of massive foundation models.

2. Methodology

The authors propose a dual-pronged approach addressing both the retrieval pipeline and the generation alignment:

A. Metadata-Enriched Hybrid RAG Pipeline

To improve retrieval accuracy, the authors developed a novel RAG strategy that moves beyond simple text chunking:

• Recursive Chunking: Utilizes RecursiveCharacterTextSplitter to respect legal document hierarchies (sections, subsections, paragraphs) rather than arbitrary character limits.
• Metadata Injection: Chunks are augmented with domain-specific metadata to create "context-aware" retrieval units.
  • Global/Local Summaries: A local window strategy groups 4 consecutive chunks and generates a concise summary (150–200 tokens) using a local LLM, which is appended to each chunk.
  • Document Details: Injection of document names, jurisdictions, and specific legal metadata (e.g., parties in a contract).
  • Token Allocation: Metadata occupies ~20–25% of the token space, balancing enrichment with core content preservation.
• Hybrid Retrieval: Combines Dense Embeddings (using thenlper/gte-large) with Sparse Lexical Retrieval (BM25) in a weighted hybrid approach ($\alpha=0.8$ for dense, $0.2$ for sparse). This ensures both semantic understanding and precise matching of legal citations and case names.

B. Direct Preference Optimization (DPO) for Safe Refusal

To address decoding errors, the authors apply DPO to a small instruction-tuned model (LLaMA 3.2, 1B) to teach selective refusal:

• Dataset Construction:
  • Set 1 (Correct Context): The model is trained to answer if the context is sufficient and to refuse only if it is not.
  • Set 2 (Incorrect Context): The model is trained to refuse answering when provided with irrelevant context.
• Alignment Objective: The model learns to prefer "answering" when context is sufficient and "refusing" (with a fixed phrase: "Given context is not sufficient to answer") when context is insufficient. This corrects the baseline model's tendency toward over-refusal or hallucination.

3. Key Contributions

1. Metadata-Enriched Chunking: Demonstrated that injecting local summaries and document metadata significantly improves retrieval precision in dense, cross-referential legal texts, addressing the Document Retrieval Mismatch (DRM) problem.
2. DPO for Legal Safety: Showed that DPO is superior to prompt engineering for enforcing selective refusal behavior in small models, effectively reducing over-refusal on valid contexts while increasing refusal on invalid ones.
3. Comprehensive Evaluation Framework: Introduced a rigorous evaluation protocol using Document Retrieval Mismatch (DRM) and Span Recall metrics across three diverse datasets (PrivacyQA, MAUD, Australian Legal QA).
4. End-to-End Validation: Proved that combining improved retrieval with DPO alignment yields superior end-to-end performance compared to optimizing either component in isolation.

4. Results

The experiments were conducted on three datasets: PrivacyQA (privacy policies), MAUD (merger & acquisition contracts), and Australian Legal QA (court judgments).

• RAG Performance (Retrieval):

  • MAUD Dataset: Showed the most dramatic improvement. Metadata-enhanced chunking increased Span Recall from ~25% to ~105% (at $k=64$) and reduced DRM from 87% to 13%.
  • Australian Legal QA: Achieved a 34.9 percentage point improvement in Span Recall at $k=16$ and an 18.7 percentage point reduction in DRM.
  • PrivacyQA: Showed minimal gains, suggesting metadata schemas must be tailored to specific document structures (privacy policies are less structured than contracts).

• DPO Performance (Alignment):

  • Refusal Behavior: The zero-shot baseline refused to answer 53.2% of the time on correct contexts (over-refusal). Post-DPO, this dropped to 1.5%. Conversely, refusal on incorrect contexts rose from 87.3% to 99.3%.
  • Answer Quality: On correct contexts, the DPO-aligned model achieved a higher BERTScore (F1) (0.9074) compared to the zero-shot baseline (0.8526), a statistically significant improvement.

• End-to-End Evaluation:

  • Combining Metadata-Enhanced RAG with DPO resulted in a mean F1 score of 0.8657, significantly outperforming the baseline (0.8479) with a 95% confidence interval of [0.0113, 0.0244].

5. Significance

This work addresses a critical gap in legal AI: the reliability of small, privacy-preserving models.

• Trust & Safety: By enabling models to reliably distinguish between sufficient and insufficient context, the system reduces the risk of hallucinated legal advice, which is vital for high-stakes legal applications.
• Scalability: The approach proves that high-performance legal AI does not strictly require massive models; rather, it requires better data engineering (metadata enrichment) and targeted alignment (DPO).
• Domain Adaptation: The paper highlights that "one-size-fits-all" RAG strategies fail in law. The success of metadata injection on structured contracts (MAUD) versus unstructured judgments suggests that future legal AI systems must be highly adaptable to document structure.

In conclusion, the paper establishes that a Metadata-Enriched Hybrid RAG pipeline coupled with DPO creates a robust, safe, and accurate framework for legal question answering, effectively mitigating hallucinations at both the retrieval and generation stages.