FLAME: improving legal case retrieval through factor-aware graph modeling and mixture-of-experts

Abstract

Legal Case Retrieval (LCR) is a foundational task in intelligent legal systems that assists legal professionals in efficiently retrieving precedent cases from large-scale legal corpora. Unlike generic document retrieval, LCR demands deeper semantic modeling due to the complexity and domain-specific logic embedded in legal texts, where similarity depends crucially on the alignment of legal factors such as legal facts, dispute focus, and reasoning. While recent large language models (LLMs) have advanced legal document understanding, existing approaches generally treat legal documents as monolithic texts and overlook the rich internal structure of legal components and their complex juridical reasoning patterns. Moreover, they lack the flexibility to dynamically adjust the importance of legal factors based on specific case contexts. To address these challenges, we propose Factor-aware Legal cAse Matching with Expert Networks (FLAME), a novel framework that centers on legal factors as fundamental modeling units and dynamically identifies the most relevant ones between case pairs. FLAME comprises three core modules: (1) the Legal-Factor Extraction Module employs Multi-Agent Reinforcement Learning for instance-wise prompt selection and leverages LLMs to extract structured legal factors from raw case documents; (2) the Legal-Factor Interaction Module constructs a heterogeneous legal-factor graph using a Factor Graph Convolutional Network to model semantic relationships among factors; (3) the Legal-Factor Integration Module employs a Legal-Factor Aware Mixture-of-Experts mechanism to dynamically highlight salient legal factors through parallel expert assessment and gating-based weight determination. The framework is optimized through a principled two-stage process that first uses contrastive learning to train the core retrieval architecture, and then employs reinforcement learning for adaptive prompt refinement. Our approach enables interpretable and context-aware legal case retrieval by explicitly modeling factor-level semantics and interactions. Extensive experiments on two benchmark datasets demonstrate that FLAME significantly outperforms state-of-the-art baselines.

Graphical abstract

Keywords

legal case retrieval / reinforcement learning / graph neural networks / mixture-of-experts / prompt optimization

References

[1] Zatarain J M N . Artificial intelligence and legal analytics: New tools for law practice in the digital age. SCRIPT-ed, 2018, 15( 1): 156–161

[2] Lai J, Gan W, Wu J, Qi Z, Yu P S . Large language models in law: a survey. AI Open, 2024, 5: 181–196

[3] Robertson S E, Walker S. Some simple effective approximations to the 2-poisson model for probabilistic weighted retrieval. In: Proceedings of the 17th Annual International ACM SIGIR Conference on Research and Development in Information Retrieval. 1994, 232−241

[4] Salton G, Buckley C . Term-weighting approaches in automatic text retrieval. Information Processing & Management, 1988, 24( 5): 513–523

[5] Chalkidis I, Fergadiotis M, Malakasiotis P, Aletras N, Androutsopoulos I. LEGAL-BERT: The Muppets straight out of law school. In: Proceedings of Findings of the Association for Computational Linguistics: EMNLP 2020. 2020, 2898−2904

[6] Xiao C, Hu X, Liu Z, Tu C, Sun M . Lawformer: a pre-trained language model for Chinese legal long documents. AI Open, 2021, 2: 79–84

[7] Li H, Ai Q, Chen J, Dong Q, Wu Y, Liu Y, Chen C, Tian Q. SAILER: structure-aware pre-trained language model for legal case retrieval. In: Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2023, 1035−1044

[8] Zhou Z, Shi J X, Song P X, Yang X W, Jin Y X, Guo L Z, Li Y F. LawGPT: a Chinese legal knowledge-enhanced large language model. 2024, arXiv preprint arXiv: 2406.04614

[9] Shao Y, Mao J, Liu Y, Ma W, Satoh K, Zhang M, Ma S. BERTPLI: modeling paragraph-level interactions for legal case retrieval. In: Proceedings of the 29th International Joint Conference on Artificial Intelligence. 2020, 3501−3507

[10] Tang Y, Qiu R, Li X. Prompt-based effective input reformulation for legal case retrieval. In: Proceedings of the 34th Australasian Database Conference on Databases Theory and Applications. 2023, 87−100

[11] Deng C, Mao K, Dou Z. Learning interpretable legal case retrieval via knowledge-guided case reformulation. In: Proceedings of 2024 Conference on Empirical Methods in Natural Language Processing. 2024, 1253−1265

[12] Zhang K, Yang Z, Başar T. Multi-agent reinforcement learning: a selective overview of theories and algorithms. In: Vamvoudakis K G, Wan Y, Lewis F L, Cansever D, eds. Handbook of Reinforcement Learning and Control. Cham: Springer, 2021, 321−384

[13] Brown T B, Mann B, Ryder N, Subbiah M, Kaplan J, Dhariwal P, Neelakantan A, Shyam P, Sastry G, Askell A, Agarwal S, Herbert-Voss A, Krueger G, Henighan T, Child R, Ramesh A, Ziegler D, Wu J, Winter C, Hesse C, Chen M, Sigler E, Litwin M, Gray S, Chess B, Clark J, Berner C, McCandlish S, Radford A, Sutskever I, Amodei D. Language models are few-shot learners. In: Proceedings of the 34th International Conference on Neural Information Processing Systems. 2020, 159

[14] Deng M, Wang J, Hsieh C P, Wang Y, Guo H, Shu T, Song M, Xing E, Hu Z. RLPrompt: optimizing discrete text prompts with reinforcement learning. In: Proceedings of 2022 Conference on Empirical Methods in Natural Language Processing. 2022, 3369−3391

[15] Kwon M, Kim G, Kim J, Lee H, Kim J. StablePrompt: automatic prompt tuning using reinforcement learning for large language model. In: Proceedings of 2024 Conference on Empirical Methods in Natural Language Processing. 2024, 9868−9884

[16] Li L, Zhang Y, Chen L . Personalized prompt learning for explainable recommendation. ACM Transactions on Information Systems, 2023, 41( 4): 103

[17] Papadimitriou C H, Tamaki H, Raghavan P, Vempala S. Latent semantic indexing: a probabilistic analysis. In: Proceedings of the 17th ACM SIGACT-SIGMOD-SIGART Symposium on Principles of Database Systems. 1998, 159−168

[18] Henderson M, Al-Rfou R, Strope B, Sung Y H, Lukacs L, Guo R, Kumar S, Miklos B, Kurzweil R. Efficient natural language response suggestion for smart reply. 2017, arXiv preprint arXiv: 1705.00652

[19] Li H, Ai Q, Han X, Chen J, Dong Q, Liu Y. DELTA: Pre-train a discriminative encoder for legal case retrieval via structural word alignment. In: Proceedings of the 39th AAAI Conference on Artificial Intelligence. 2025, 27072–27080

[20] Gao C, Xiao C, Liu Z, Chen H, Liu Z, Sun M. Enhancing legal case retrieval via scaling high-quality synthetic query-candidate pairs. In: Proceedings of 2024 Conference on Empirical Methods in Natural Language Processing. 2024, 7086–7100

[21] Yan S, Zhang Z . SKYPER: Legal case retrieval via skeleton-aware hypergraph embedding in the hyperbolic space. Information Sciences, 2024, 682: 121162

[22] Su W, Hu Y, Xie A, Ai Q, Bing Q, Zheng N, Liu Y, Shen W, Liu Y. STARD: A Chinese statute retrieval dataset derived from real-life queries by non-professionals. In: Proceedings of Findings of the Association for Computational Linguistics: EMNLP 2024. 2024, 10658−10671

[23] Ma Y, Shao Y, Wu Y, Liu Y, Zhang R, Zhang M, Ma S. LeCaRD: A legal case retrieval dataset for Chinese law system. In: Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2021, 2342−2348

[24] Mentzingen H, António N, Bacao F . Automation of legal precedents retrieval: findings from a literature review. International Journal of Intelligent Systems, 2023, 2023( 1): 6660983

[25] Tang Y, Qiu R, Liu Y, Li X, Huang Z. CaseGNN: graph neural networks for legal case retrieval with text-attributed graphs. In: Proceedings of the 46th European Conference on Information Retrieval on Advances in Information Retrieval. 2024, 80–95

[26] Tang Y, Qiu R, Liu Y, Li X, Huang Z. CaseGNN++: graph contrastive learning for legal case retrieval with graph augmentation. 2024, arXiv preprint arXiv: 2405.11791

[27] Tang Y, Qiu R, Yin H, Li X, Huang Z. CaseLink: inductive graph learning for legal case retrieval. In: Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2024, 2199–2209

[28] Wang C, Liu Q, Wu R, Chen E, Liu C, Huang X, Huang Z. Confidence-aware matrix factorization for recommender systems. In: Proceedings of the 32nd AAAI Conference on Artificial Intelligence. 2018

[29] Wang C, Zhu H, Zhu C, Qin C, Xiong H. SetRank: a setwise Bayesian approach for collaborative ranking from implicit feedback. In: Proceedings of the 34th AAAI Conference on Artificial Intelligence. 2020, 6127−6136

[30] Zhu Y, Wang C, Zhang Q, Xiong H. Graph signal diffusion model for collaborative filtering. In: Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2024, 1380−1390

[31] Kipf T N, Welling M. Semi-supervised classification with graph convolutional networks. In: Proceedings of the 5th International Conference on Learning Representations. 2017

[32] Hamilton W L, Ying R, Leskovec J. Inductive representation learning on large graphs. In: Proceedings of the 31st International Conference on Neural Information Processing Systems. 2017, 1025−1035

[33] Cui H, Lu J, Ge Y, Yang C. How can graph neural networks help document retrieval: a case study on CORD19 with concept map generation. In: Proceedings of the 44th European Conference on IR Research. 2022, 75−83

[34] Li X, De Rijke M, Liu Y, Mao J, Ma W, Zhang M, Ma S. Learning better representations for neural information retrieval with graph information. In: Proceedings of the 29th ACM International Conference on Information & Knowledge Management. 2020, 795−804

[35] Yao L, Mao C, Luo Y. Graph convolutional networks for text classification. In: Proceedings of the 33rd AAAI Conference on Artificial Intelligence and 31st Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence. 2019, 905

[36] Ragesh R, Sellamanickam S, Iyer A, Bairi R, Lingam V. HeteGCN: heterogeneous graph convolutional networks for text classification. In: Proceedings of the 14th ACM International Conference on Web Search and Data Mining. 2021, 860−868

[37] Zhang K, Chen C, Wang Y, Tian Q, Bai L. CFGL-LCR: a counterfactual graph learning framework for legal case retrieval. In: Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2023, 3332−3341

[38] Yang J, Ma W, Zhang M, Zhou X, Liu Y, Ma S . LegalGNN: legal information enhanced graph neural network for recommendation. ACM Transactions on Information Systems (TOIS), 2022, 40( 2): 33

[39] Wu Z, Pan S, Chen F, Long G, Zhang C, Yu P S . A comprehensive survey on graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 2021, 32( 1): 4–24

[40] Wang C, Zhu H, Hao Q, Xiao K, Xiong H. Variable interval time sequence modeling for career trajectory prediction: deep collaborative perspective. In: Proceedings of the Web Conference 2021. 2021, 612−623

[41] Wang C, Zhu H, Wang P, Zhu C, Zhang X, Chen E, Xiong H . Personalized and explainable employee training course recommendations: A Bayesian variational approach. ACM Transactions on Information Systems (TOIS), 2022, 40( 4): 70

[42] Dai D, Deng C, Zhao C, Xu R X, Gao H, Chen D, Li J, Zeng W, Yu X, Wu Y, Xie Z, Li Y K, Huang P, Luo F, Ruan C, Sui Z, Liang W. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models. In: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024, 1280–1297

[43] Amato C. An introduction to centralized training for decentralized execution in cooperative multi-agent reinforcement learning. 2024, arXiv preprint arXiv: 2409.03052

[44] Mao W, Wu J, Chen W, Gao C, Wang X, He X . Reinforced prompt personalization for recommendation with large language models. ACM Transactions on Information Systems, 2025, 43( 3): 72

[45] Vaswani A, Shazeer N, Parmar N, Uszkoreit J, Jones L, Gomez A N, Kaiser Ł, Polosukhin I. Attention is all you need. In: Proceedings of the 31st International Conference on Neural Information Processing Systems. 2017, 6000−6010

[46] Van Den Oord A, Li Y, Vinyals O. Representation learning with contrastive predictive coding. 2018, arXiv preprint arXiv: 1807.03748

[47] Li H, Shao Y, Wu Y, Ai Q, Ma Y, Liu Y. LeCaRDv2: a large-scale Chinese legal case retrieval dataset. In: Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2024, 2251−2260

[48] Liu Y, Ott M, Goyal N, Du J, Joshi M, Chen D, Levy O, Lewis M, Zettlemoyer L, Stoyanov V. RoBERTa: a robustly optimized BERT pretraining approach. 2019, arXiv preprint arXiv: 1907.11692

[49] Chen J, Xiao S, Zhang P, Luo K, Lian D, Liu Z. M3-embedding: Multi-linguality, multi-functionality, multi-granularity text embeddings through self-knowledge distillation. In: Proceedings of Findings of the Association for Computational Linguistics: ACL 2024. 2024, 2318–2335

[50] Cui J, Ning M, Li Z, Chen B, Yan Y, Li H, Ling B, Tian Y, Yuan L. Chatlaw: a multi-agent collaborative legal assistant with knowledge graph enhanced mixture-of-experts large language model. 2024, arXiv preprint arXiv: 2306.16092v2

Rights & permissions

Higher Education Press