[Paper Review] FedJudge: Federated Legal Large Language Model

Linan Yue, Qi Liu+4 more|arXiv (Cornell University)|Sep 15, 2023

Artificial Intelligence in Law4 citations

TL;DR

FedJudge is the first federated learning framework for training legal large language models (LLMs) that addresses data privacy, computational overhead, and data distribution shifts. It uses parameter-efficient LoRA fine-tuning and continual learning to preserve global knowledge, achieving state-of-the-art performance on legal tasks while keeping raw data local.

ABSTRACT

Large Language Models (LLMs) have gained prominence in the field of Legal Intelligence, offering potential applications in assisting legal professionals and laymen. However, the centralized training of these Legal LLMs raises data privacy concerns, as legal data is distributed among various institutions containing sensitive individual information. This paper addresses this challenge by exploring the integration of Legal LLMs with Federated Learning (FL) methodologies. By employing FL, Legal LLMs can be fine-tuned locally on devices or clients, and their parameters are aggregated and distributed on a central server, ensuring data privacy without directly sharing raw data. However, computation and communication overheads hinder the full fine-tuning of LLMs under the FL setting. Moreover, the distribution shift of legal data reduces the effectiveness of FL methods. To this end, in this paper, we propose the first Federated Legal Large Language Model (FedJudge) framework, which fine-tunes Legal LLMs efficiently and effectively. Specifically, FedJudge utilizes parameter-efficient fine-tuning methods to update only a few additional parameters during the FL training. Besides, we explore the continual learning methods to preserve the global model's important parameters when training local clients to mitigate the problem of data shifts. Extensive experimental results on three real-world datasets clearly validate the effectiveness of FedJudge. Code is released at https://github.com/yuelinan/FedJudge.

Motivation & Objective

• To address data privacy concerns in training legal LLMs by avoiding raw data sharing across institutions like courts and legal firms.
• To reduce computational and communication overheads in fine-tuning large language models under federated learning settings.
• To mitigate performance degradation caused by data distribution shifts across heterogeneous legal clients (e.g., courts vs. legal consultancies).
• To develop a scalable, privacy-preserving framework for training high-quality legal LLMs using federated learning with parameter-efficient and continual learning techniques.

Proposed method

• Employs LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning, updating only trainable rank matrices while freezing the base LLM weights.
• Transfers only LoRA parameters between clients and server, drastically reducing communication overhead compared to full model parameter exchange.
• Applies continual learning with a regularization loss (Eq. 3) to prevent local models from forgetting important global knowledge during training.
• Uses FedAvg-style aggregation to update the global model using only aggregated LoRA parameters from clients.
• Introduces personalized client models (e.g., CL-Client) that maintain global knowledge while adapting to local data distributions.
• Employs a dual training strategy: global model training via federated aggregation and local model fine-tuning with knowledge preservation constraints.

Fig. 1 : (a) An example in the court client. (b) An example in the legal consultation client.

Experimental results

Research questions

• RQ1 Can federated learning be effectively applied to fine-tune large language models for legal tasks while preserving data privacy?
• RQ2 How can communication and computational overheads be minimized when fine-tuning legal LLMs in a federated setting?
• RQ3 To what extent does data distribution shift across legal institutions (e.g., courts vs. consultancies) degrade the performance of federated legal LLMs?
• RQ4 Can continual learning techniques improve the generalization and knowledge retention of federated legal LLMs under heterogeneous data distributions?

Key findings

• FedJudge-CL outperforms both the global model (FedJudge-Base) and personalized baselines (Base-Client) on most metrics, demonstrating the effectiveness of continual learning in mitigating distribution shifts.
• The personalized model CL-Client3 achieved the best performance on the legal reasoning task, correctly reasoning and answering a complex case, while Baichuan-7B and Center produced incorrect or irrelevant responses.
• Center-Client1 outperformed Center on its own data (by 4.87 BertScore), but failed on other clients’ data (e.g., 4.47 BLEU-N on Client1), highlighting the limitations of centralized training with heterogeneous data.
• FedJudge-CL achieved competitive results compared to the centralized model (Center), despite training without access to raw global data, proving the viability of federated learning in legal LLMs.
• The use of LoRA reduced communication and computational costs significantly, enabling efficient fine-tuning on resource-constrained clients.
• The continual learning constraint (Eq. 3) effectively reduced weight divergence between local and global models, preserving global knowledge and improving global model quality.

Fig. 2 : An overview of FedJudge, where represents trainable weights and denotes frozen LLMs.