LexReward: A Taxonomy-Driven Reward Framework for Legal Language Models

LexReward. A Taxonomy-Driven Reward Framework for Legal Language Models. By Yida Cai and seven others, from Peking University, Northeastern University, and one other institution. Published on the thirtieth of September. It links to a code repository that is no longer there.

A legal language model needs a reward signal that captures more than whether the answer is correct. A legal response has quality along several dimensions at once. Existing reward methods tend to rely on coarse, holistic judgments, which offer little that is specific to law and little insight into why one response scored above another. A single overall verdict cannot say whether a response fell short in its wording, in the legal facts and statutes it handled, or in the reasoning that connects them.

LexReward replaces the single holistic judgment with a taxonomy. Legal response quality is broken into three complementary dimensions, and each is scored with its own rubric, setting out the evaluation criteria and the levels of quality.

The first dimension, Style, covers lexical and syntactic quality. The second, Element, assesses the legal subjects, facts, statutes and decisions a response deals with. The third, Chain, evaluates the legal reasoning itself: its order, completeness, correctness, and freedom from redundancy.

The three are meant to be read together: how a response is written, what legal content it contains, and how it reasons from that content to a conclusion.

The rewards these rubrics produce are then used to build pairwise preference data, in which one response is marked as better than another. That data serves two purposes. It feeds training with Direct Preference Optimization, and it trains reward models, which the authors call LexRM.

Experiments show the rubric-based rewards reliably separate legal responses of different quality, and DPO training on the preference data improves performance on all three dimensions. The learned reward models support downstream optimization as well. Each dimension-specific model improves a policy, through reinforcement learning, on its own dimension, and does so without needing reference answers at reward time. That last point matters for use beyond the training set, since a reward that needs no reference answer can score responses to questions nobody has answered yet. Analyses dimension by dimension further support the taxonomy and the way rewards are built from it.

Splitting the judgment is what buys both things the old rewards lacked. Domain specificity comes from rubrics written around legal elements and legal reasoning, and interpretability comes from knowing which of the three dimensions a response won or lost on.