Aligning LLM s for T hai Legal Question Answering with Efficient Semantic-Similarity Rewards (2025)

Abstract: The Retrieval-Augmented Generation (RAG) systems’ performance on Thai legal question answering is still limited, especially for questions requiring extensive, complex legal reasoning. To address these limitations, we introduce a resource-efficient approach that aligns Large Language Models (LLMs) for improved citation accuracy and response quality using Group-Relative Policy Optimization (GRPO). Our proposed method leverages BGE-M3 embeddings as a cost-efficient semantic-similarity reward, significantly reducing computational expenses up to 2.5x compared to an LLM-based reward model. Experiments on the NitiBench benchmark demonstrate substantial improvements: GRPO achieves up to 90% citation-F1 gains relative to the base model and a 31% increase in joint quality metrics over instruction tuning. Crucially, our approach provides a practical and effective solution for enhancing legal LLMs in resource-constrained environments.

## introduction the ability to deliver accurate and grounded answers with relevant law citations is essential for reliable legal question answering. most legaldomain llm solutions @xcite @xcite adopt retrieval-augmented generation (rag) to reduce hallucinations by attaching retrieved legal documents as supporting context. however, retrieved documents are not always fully leveraged and can contain false positives @xcite . a common approach to mitigate this issue is to require llms to emit explicit citations during generation. yet, @xcite show that even when golden contexts are provided, strong proprietary models often fail to cite all relevant law sections at generation time. these findings highlight a key limitation of current llms in producing factually grounded responses, thereby undermining the reliability of downstream legal applications. ensuring accurate and well-cited responses thus remains a central open challenge in the legal domain. while instruction tuning can partially address this gap, it provides limited control over citation behavior since its objective is to maximize next-token likelihood rather than citation accuracy. this highlights the need for more targeted alignment that not only improves factuality but also enforces verifiable citation standards. to enable fine-grained control over citation accuracy, we frame thai legal qa as a citation-sensitive alignment challenge. building on recent advances in reinforcement learning with verifiable rewards (rlvr) @xcite and group-relative policy optimization (grpo) @xcite , we treat citation accuracy and response quality as a verifiable reward and directly align llms toward higher citation quality during alignment tuning. although computing rewards for citation is straightforward and inexpensive, evaluating response quality is far more costly, as it typically requires an llm-based reward model during training, leading to significant computational overhead and higher alignment costs. this raises our central research question: how can we affordably and effectively align llms for citation-sensitive legal qa in resource-constrained settings such as thai law? to address this, we investigate the following research questions: does thai-specific continued pretraining enhance the effectiveness of alignment strategies? • (rq3) rlvr vs. instruction tuning: how can rlvr enhance response quality and citation accuracy compared to instruction tuning? our experiments on the nitibench benchmark @xcite yield three key contributions. first, we propose a low-cost semantic reward proxy to compute response quality reward using general embedding models, achieving over 2.5× improvement in training efficiency while maintaining performance comparable to an expensive llm-based reward model. second, we show that thai-cpt models are more receptive to alignment signals than their general-purpose counterparts. finally, we demonstrate that grpo consistently outperforms instruction tuning in improving citation fidelity and response quality, particularly for in-domain tasks. together, these findings provide a recipe for building specialized legal llms that achieve efficient and accurate law grounding under resource constraints. these documents are meticulously organized into a multi-tiered structure, often including divisions such as book, title, chapter, division, and ultimately, the section. a section is the fundamental unit of law, articulating a specific rule, right, or obligation. accordingly, this work focuses on evaluating legal reasoning and citation performance at the section level, treating it as the primary unit for retrieval and grounding. enhancing llm legal citation performance a growing body of work seeks to make llms produce verifiable citations. citalaw @xcite adapts the alce benchmark @xcite to the legal domain, introduces a syllogismlevel citation metric, and supports both statutes and precedent cases. alce itself evaluates statementlevel grounding using an nli verifier, requiring every generated claim to be backed by retrieved evidence. @xcite compares citation accuracy across three retrieval regimes: 1) retriever-only, 2) llm query-rewrite, and 3) hybrid method. this work focuses on australian case-law and shows that task-specific instruction tuning yields the largest gains in improving citation accuracy. legalbench-rag (pipitone and alami, 2024) isolates the retriever's contribution by measuring precision over expert-annotated snippets while varying chunking and top-k, revealing a retrieval-quality ceiling on downstream citation f1. ## proposed studies we frame the thai legal question answering (qa) task as a citation-sensitive generative task, where the model must generate a free-form response and citations based on a user query and a set of retrieved legal documents. the response must be semantically informative, cite relevant statutes, and avoid hallucinations that reference unsupported claims. to align llms with these objectives, we proposed a framework based on grpo @xcite , which aligns model behavior through a carefully designed reward function that encourages response quality and citation accuracy. our proposed method is summarized in figure 1 . this rewardbased formulation enables us to study how different alignment strategies influence legal response quality under the constraints of low-resource legal domains. ## results ## reward composition analysis to understand reward contributions, we performed ablations on openthaigpt1.5-7b foot_8 (see ## error analysis to better understand the trade-offs between different response quality reward strategies, we performed a qualitative analysis on the outputs of openthaigpt1.5-7b grpo with semantic and cov/con reward. we highlight three cases as shown in ## efficiency of reward signal proxies to investigate the efficiency gains between semantic and llm-based rewards, we compared the training cost in gpu-hours of both methods. when using llm-based judge, qwen2.5-72b-instruct served with optimized inference stack, the training time cost up to 264 gpu-hours (@xmath085), achieving up to over 2.5× cost and time saving.additionally, semantic reward is also more memory efficient as training only requires one gpu, while the llm-based reward model requires one additional gpu for hosting the reward model. 10 this disparity arises because the lightweight bge-m3 calculation adds minimal latency to the rl loop, whereas the 72b judge model requires a dedicated gpu for inference, creating a significant training bottleneck. while a large judge may offer reward signals with higher fidelity to final evaluation metrics, its computational cost is a major barrier to online rl training. the strong in-domain 10 cost in usd was estimated based on a100 80gb pcie median rental cost of $0.82/hr via https://vast.ai/ pricing/gpu/a100-pcie accessed april 2025. ## conclusion this work demonstrates that grpo can be used to align llms toward better legal question answering and usually outperforms instruction tuning. the model aligned with grpo yields significant citation accuracy gains on in-domain tasks and improves generalization on complex legal reasoning for thai-aligned models, a setting where instruction tuning consistently degrades performance. ablations of reward functions reveal a key tradeoff: an efficient semantic reward is cost-effective for in-domain tasks but loses effectiveness on complex generalization. while an llm-based reward model provides a more accurate signal, it does so at over 2.5× the computational cost. our ablations confirm that using both citation and answer quality rewards is necessary for the best outcomes. ultimately, our findings show that grpo is a highly effective approach for specialized domains, but its success depends on the careful synergy of the rl algorithm, base model capabilities, and reward design. future work should focus on creating reward mechanisms that are both accurate enough for complex reasoning and efficient enough for practical training.