LEXA: Legal Case Retrieval via Graph Contrastive Learning with Contextualised LLM Embeddings

The paper proposes LEXA, an enhanced legal case retrieval model that extends CaseGNN by integrating edge-updated graph attention, graph contrastive learning with augmentation, and contextualized LLM embeddings to fully leverage structural information and significantly outperform state-of-the-art methods.

The Gist

Imagine you are a lawyer trying to win a case. You need to find a past court decision (a "precedent") that is almost exactly like your current situation. In the old days, you might have searched through dusty libraries using keywords like "car accident" or "theft." If the words matched, you found a case. But what if the words were different, even though the situation was the same? You'd miss the perfect match.

This is the problem LEXA solves. It's a new AI system designed to find the right legal cases, not just by matching words, but by understanding the story and the relationships inside them.

Here is how LEXA works, explained through simple analogies:

1. The Old Way vs. The New Way

• The Old Way (Lexical Models): Imagine trying to find a friend in a crowd by only looking at their shirt color. If they wear a red shirt, you find them. But if they wear a blue shirt, even if it's the same person, you miss them. This is how older legal search tools worked; they just looked for matching words.
• The Middle Way (Language Models): Now, imagine you can recognize your friend's face. That's better! But you still don't know who they are with, what they are doing, or why they are there. You see the person, but not the context.
• The LEXA Way (Graph + Context): LEXA doesn't just see the person; it sees the entire social network. It knows who the friend is talking to, what they are arguing about, and how they are connected to others. It builds a "web" of relationships for every legal case.

2. Building the "Legal Web" (The Graph)

Legal cases are full of characters (the plaintiff, the defendant, the judge) and actions (stole, signed, drove).

• The Nodes (The Characters): LEXA turns every person or thing in the case into a dot on a map.
• The Edges (The Relationships): It draws lines between the dots to show how they are connected (e.g., "The Defendant stole from the Victim").
• The Problem with the Previous Version (CaseGNN): The authors' previous tool, CaseGNN, was good at looking at the dots (the people), but it treated the lines (the relationships) as static, unchangeable ropes. It never updated its understanding of the relationship as it looked deeper into the case.

3. The Three Magic Upgrades in LEXA

LEXA fixes the old tool with three specific superpowers:

A. The "Dynamic Rope" (EUGAT Layer)

In the old tool, the lines connecting the dots were fixed. In LEXA, the lines are dynamic.

• Analogy: Imagine the lines are made of smart rubber. As the AI looks at the people (dots), it realizes the relationship between them might be more complex than it thought. The "smart rubber" stretches and changes shape to reflect that new understanding.
• Result: The AI updates not just who the people are, but how they relate to each other, creating a much richer picture of the case.

B. The "Study Buddy" (Graph Contrastive Learning)

Training AI on legal data is hard because there aren't many labeled examples (it's expensive to get lawyers to label data).

• Analogy: Imagine you are studying for a test, but you only have one textbook. To learn better, you create "practice versions" of the textbook. You tear out a few pages (Edge Dropping) or blur some words (Feature Masking) and ask yourself, "Is this still the same story?"
• Result: By forcing the AI to recognize that a "blurred" version of a case is still the same case, it learns the core essence of the story rather than just memorizing specific words. This makes it much smarter and more robust.

C. The "Expert Translator" (LLM Embeddings)

Legal language is tricky. A word like "consideration" means something very specific in law, but something totally different in everyday life.

• Analogy: Instead of using a dictionary to translate words, LEXA hires a legal scholar (a Large Language Model) to read the case first. This scholar writes a summary that captures the nuance, the intent, and the hidden meaning before the AI even starts building the web.
• Result: The dots and lines in the web are now painted with "legal paint" instead of just "text paint." The AI understands the spirit of the law, not just the letters.

4. The Result: Finding the Needle in the Haystack

When you ask LEXA, "Find me a case where a driver was speeding but wasn't charged because of a medical emergency," it doesn't just look for the words "speeding" and "medical."

It looks at the web:

1. It sees the connection between the driver and the emergency.
2. It understands the relationship between the speed and the lack of a charge.
3. It compares this complex web to millions of other webs to find the one that matches the story, not just the words.

Summary

LEXA is like upgrading from a keyword search engine to a legal detective.

• It builds a 3D map of every case (Graph).
• It lets the relationships on that map evolve and update as it learns (EUGAT).
• It practices with distorted versions of cases to become a master of the core truth (Contrastive Learning).
• It uses a legal expert to translate the text into deep meaning before starting (LLM).

The result? It finds the right legal precedents significantly faster and more accurately than any previous system, helping lawyers and judges make better decisions.

Technical Summary

Here is a detailed technical summary of the paper "LEXA: Legal Case Retrieval via Graph Contrastive Learning with Contextualised LLM Embeddings."

1. Problem Statement

Legal Case Retrieval (LCR) is a specialized information retrieval task aimed at identifying relevant legal precedents for a given query case. While existing methods have evolved from traditional lexical models (e.g., BM25) to neural language models (LMs) and Large Language Models (LLMs), they suffer from three critical limitations:

1. Neglect of Structural Information: Most LM-based approaches rely solely on raw text, failing to capture the intricate structural relationships between legal entities (e.g., parties, evidence, criminal activities) that are crucial for legal reasoning.
2. Under-utilization of Edge Information: Previous graph-based approaches (specifically the authors' prior work, CaseGNN) update node features but keep edge features static during training. This ignores the rich, dynamic relational information inherent in legal connections.
3. Insufficient Training Signals & Context: Legal datasets are scarce and expensive to annotate. Furthermore, standard LMs often lack deep, domain-specific contextual understanding of legal semantics when encoding graph nodes and edges.

2. Methodology: The LEXA Framework

LEXA (Legal Case Retrieval via Graph Contrastive Learning with Contextualised LLM Embeddings) is an enhanced framework extending CaseGNN. It addresses the above limitations through four core components:

A. Contextualized Feature Initialization with LLMs

Instead of using generic embeddings, LEXA employs a fine-tuned Large Language Model (LEXA-8B, based on Qwen3-Embedding-8B) to generate initial features for the Text-Attributed Case Graph (TACG).

• Prompt Engineering: Specific prompts are used to guide the LLM in extracting legal facts and issues.
• Encoding: The LLM encodes relation triplets (head entity, relation, tail entity) and global nodes (representing overall legal facts/issues) into dense vector representations. This ensures the graph nodes and edges start with rich, domain-specific contextual semantics.

B. Edge-Updated Graph Attention Layer (EUGAT)

LEXA introduces a novel graph neural network layer called EUGAT to replace standard GNN layers.

• Dual Update Mechanism: Unlike previous models that only update node features, EUGAT simultaneously updates both node features and edge features in every GNN layer.
• Mechanism: It utilizes a multi-head attention mechanism where the attention weights and the resulting representations are computed based on the concatenation of the target node, neighbor nodes, and the connecting edge.
• Benefit: This allows the model to iteratively refine the semantic meaning of legal relationships (edges) as node representations evolve, leading to a more comprehensive understanding of case structure.

C. Graph Augmentation and Contrastive Learning

To overcome the scarcity of labeled legal data, LEXA employs Graph Contrastive Learning (GCL) with specific augmentation strategies.

• Augmentation: The model generates augmented views of the case graphs using Edge Dropping (randomly removing edges) and Feature Masking (masking node/edge attributes).
• Objective Function: The training objective pulls the representations of the query case and positive (relevant) cases closer, while pushing away negative samples.
  • Negative Sampling: Includes Easy Negatives (randomly sampled) and Hard Negatives (high BM25 similarity but irrelevant).
  • Augmented Views: The contrastive loss is applied not just to original graphs but also to their augmented versions, forcing the model to learn robust representations invariant to structural noise.

D. Graph Construction

Legal cases are transformed into two separate graphs:

1. Legal Fact Graph: Captures "who, when, what, where, why" elements.
2. Legal Issue Graph: Captures the legal disputes requiring resolution.
   Both graphs include a virtual global node connected to all other nodes to facilitate global information propagation.

3. Key Contributions

1. LEXA Framework: A state-of-the-art LCR model that integrates structural graph modeling with contextualized LLM embeddings.
2. EUGAT Layer: A novel GNN architecture that dynamically updates both node and edge features, fully leveraging the relational topology of legal cases.
3. Enhanced Contrastive Learning: A graph contrastive learning objective augmented with graph augmentation strategies (edge dropping, feature masking) to provide additional training signals in low-data regimes.
4. Contextualized Embeddings: The use of a fine-tuned, prompt-guided LLM (LEXA-8B) to initialize graph features, capturing deep legal semantics that standard encoders miss.

4. Experimental Results

The model was evaluated on two benchmark datasets: COLIEE 2022 and COLIEE 2023.

• Performance: LEXA achieved State-of-the-Art (SOTA) performance across all metrics (Precision@5, Recall@5, Micro/Macro F1, MRR@5, MAP, NDCG@5).
  • On COLIEE 2022, LEXA outperformed the previous best graph-based model (CaseLink) by a significant margin (e.g., NDCG@5 improved from ~70.3% to 79.3%).
  • On COLIEE 2023, LEXA similarly surpassed all baselines, including strong LLM embeddings and other graph models.
• Ablation Studies:
  • Removing the EUGAT layer (reverting to standard EdgeGAT) caused a performance drop, confirming the importance of updating edge features.
  • Removing Graph Contrastive Learning (GCL) reduced ranking metrics, proving the value of augmented training signals.
  • Using un-tuned LLMs for feature initialization resulted in lower performance, highlighting the necessity of domain-specific fine-tuning.
  • Edge Dropping was found to be the most effective augmentation strategy compared to feature masking.
• Sensitivity Analysis: The model performed best with a temperature coefficient ($\tau$) of 0.1 and a small number of easy negative samples (1), indicating a balance between smoothness and discriminative power.

5. Significance

• Bridging Structure and Semantics: LEXA successfully demonstrates that combining the structural reasoning of Graph Neural Networks with the semantic depth of Large Language Models yields superior legal retrieval capabilities compared to using either approach alone.
• Data Efficiency: By leveraging graph contrastive learning and augmentation, LEXA mitigates the challenge of limited labeled legal data, making it a practical solution for domains where annotation is costly.
• Dynamic Relationship Modeling: The introduction of EUGAT establishes a new paradigm for legal graph modeling, proving that legal relationships (edges) are not static but should be dynamically refined during the learning process to capture evolving case contexts.
• Practical Impact: The significant improvement in retrieval accuracy aids legal practitioners (judges, lawyers) in locating precedents more efficiently and provides accessible legal resources for individuals lacking financial means for legal counsel.