Balanced and efficient federated inference of large language models

Abstract

Large language models have shown significant performance across various tasks. This paper focuses on federated LLM inference, in which the context generated by silos requires being protected, resulting in both inefficiency and imbalance in communication overhead. To address this issue, this paper proposes BEFI, enabling Balanced and Efficient Federated LLM Inference equipped with carefully designed mechanisms for information sharing among silos. For important tokens, we propose a fine-grained, fixed-size cache sharing mechanism that enables direct and balanced sharing of the KV cache of these tokens. For less critical tokens, we propose a context-free intermediate state sharing mechanism, which allows for the sharing of tokens of arbitrary length with constant communication overhead. Finally, we evaluate the effectiveness of BEFI across various LLMs and privacy mechanisms, demonstrating that BEFI reduces communication overheads by up to 96.0%, while maintaining balanced communication.

Graphical abstract

Keywords

large language model / inference optimization / federated learning

References

[1] Zhu W, Liu H, Dong Q, Xu J, Huang S, Kong L, Chen J, Li L. Multilingual machine translation with large language models: empirical results and analysis. In: Proceedings of Findings of the Association for Computational Linguistics: NAACL 2024. 2024, 2765−2781

[2] He Z, Liang T, Jiao W, Zhang Z, Yang Y, Wang R, Tu Z, Shi S, Wang X . Exploring human-like translation strategy with large language models. Transactions of the Association for Computational Linguistics, 2024, 12: 229–246

[3] Wu J, Yang S, Zhan R, Yuan Y, Chao L S, Wong D F . A survey on LLM-generated text detection: necessity, methods, and future directions. Computational Linguistics, 2025, 51( 1): 275–338

[4] 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

[5] Ali M, Naeem F, Tariq M, Kaddoum G . Federated learning for privacy preservation in smart healthcare systems: a comprehensive survey. IEEE Journal of Biomedical and Health Informatics, 2023, 27( 2): 778–789

[6] Rocco di Torrepadula F, Fisichella M, di Martino S, Mazzocca N . Fedflow: a personalized federated learning framework for passenger flow prediction. Machine Learning, 2025, 114( 7): 166

[7] Voigt P, von dem Bussche A. The EU General Data Protection Regulation (GDPR): A Practical Guide. Cham: Springer, 2017

[8] Dwork C. Differential privacy. In: Proceedings of the 33rd International Conference on Automata, Languages and Programming - Volume Part II. 2006, 1−12

[9] Pope R, Douglas S, Chowdhery A, Devlin J, Bradbury J, Levskaya A, Heek J, Xiao K, Agrawal S, Dean J. Efficiently scaling transformer inference. In: Proceedings of the 6th MLSys Conference. 2023

[10] Mohtashami A, Jaggi M. Landmark attention: random-access infinite context length for transformers. 2023, arXiv preprint arXiv: 2305.16300

[11] Coppersmith D, Winograd S . Matrix multiplication via arithmetic progressions. Journal of Symbolic Computation, 1990, 9( 3): 251–280

[12] Bläser M . Fast matrix multiplication. Theory of Computing Library, Graduate Surveys, 2013, 5: 1–60

[13] Beltagy I, Peters M E, Cohan A. Longformer: the long-document transformer. 2020, arXiv preprint arXiv: 2004.05150

[14] Wang Y, Li S, Tong Y, Wei S, Zhou Z . Efficient shapley-based data valuation for federated trajectories. Frontiers of Computer Science, 2026, 20( 9): 2009621

[15] Zaheer M, Guruganesh G, Dubey A, Ainslie J, Alberti C, Ontanon S, Pham P, Ravula A, Wang Q, Yang L, Ahmed A. Big bird: transformers for longer sequences. In: Proceedings of the 34th International Conference on Neural Information Processing Systems. 2020

[16] Zhang Z, Sheng Y, Zhou T, Chen T, Zheng L, Cai R, Song Z, Tian Y, Ré C, Barrett C, Wang Z, Chen B. H₂O: heavy-hitter oracle for efficient generative inference of large language models. In: Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023, 1506

[17] Liu Z, Desai A, Liao F, Wang W, Xie V, Xu Z, Kyrillidis A, Shrivastava A. Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time. In: Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023, 2279

[18] Wang H, Zhang Z, Han S. SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning. In: The the 27th IEEE International Symposium on High-Performance Computer Architecture. 2021, 97-110

[19] Roy A, Saffar M, Vaswani A, Grangier D . Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 2021, 9: 53–68

[20] Tay Y, Bahri D, Yang L, Metzler D, Juan D C. Sparse sinkhorn attention. In: Proceedings of the 37th International Conference on Machine Learning. 2020, 9438−9447

[21] Touvron H, Martin L, Stone K, Albert P, Almahairi A, Babaei Y, Bashlykov N, Batra S, Bhargava P, Bhosale S, Bikel D, Blecher L, Canton Ferrer C, Chen M, Cucurull G, Esiobu D, Fernandes J, Fu J, Fu W, Fuller B, Gao C, Goswami V, Goyal N, Hartshorn A, Hosseini S, Hou R, Inan H, Kardas M, Kerkez V, Khabsa M, Kloumann I, Korenev A, Koura P S, Lachaux M, Lavril T, Lee J, Liskovich D, Lu Y, Mao Y, Martinet X, Mihaylov T, Mishra P, Molybog I, Nie Y, Poulton A, Reizenstein J, Rungta R, Saladi K, Schelten A, Silva R, Smith E M, Subramanian R, Tan X E, Tang B, Taylor R, Williams A, Kuan J X, Xu P, Yan Z, Zarov I, Zhang Y, Fan A, Kambadur M, Narang S, Rodriguez A, Stojnic R, Edunov S, Scialom T. Llama 2: open foundation and fine-tuned chat models. 2023, arXiv preprint arXiv: 2307.09288

[22] Jiang A Q, Sablayrolles A, Mensch A, Bamford C, Chaplot D S, de Las Casas D, Bressand F, Lengyel G, Lample G, Saulnier L, Lavaud L R, Lachaux M A, Stock P, Le Scao T, Lavril T, Wang T, Lacroix T, El Sayed W. Mistral 7B. 2023, arXiv preprint arXiv: 2310.06825

[23] Almazrouei E, Alobeidli H, Alshamsi A, Cappelli A, Cojocaru R, Debbah M, Goffinet E, Hesslow D, Launay J, Malartic Q, Mazzotta D, Noune B, Pannier B, Penedo G. The falcon series of open language models. 2023, arXiv preprint arXiv: 2311.16867

[24] Merity S, Xiong C, Bradbury J, Socher R. Pointer sentinel mixture models. In: Proceedings of the 5th International Conference on Learning Representations. 2017

[25] Rae J W, Potapenko A, Jayakumar S M, Hillier C, Lillicrap T P. Compressive transformers for long-range sequence modelling. In: Proceedings of the 8th International Conference on Learning Representations. 2020

[26] Bai Y, Lv X, Zhang J, Lyu H, Tang J, Huang Z, Du Z, Liu X, Zeng A, Hou L, Dong Y, Tang J, Li J. LongBench: a bilingual, multitask benchmark for long context understanding. In: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics. 2024, 3119−3137

[27] Lyu L, He X, Li Y. Differentially private representation for NLP: formal guarantee and an empirical study on privacy and fairness. In: Proceedings of Findings of the Association for Computational Linguistics. 2020, 2355−2365

[28] Meehan C, Mrini K, Chaudhuri K. Sentence-level privacy for document embeddings. In: Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics. 2022, 3367−3380

[29] Yue X, Du M, Wang T, Li Y, Sun H, Chow S S M. Differential privacy for text analytics via natural text sanitization. In: Proceedings of Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021. 2021, 3853−3866

[30] Chanyaswad T, Dytso A, Poor H V, Mittal P. MVG mechanism: differential privacy under matrix-valued query. In: Proceedings of 2018 ACM SIGSAC Conference on Computer and Communications Security. 2018, 230−246

[31] Abadi M, Chu A, Goodfellow I, McMahan H B, Mironov I, Talwar K, Zhang L. Deep learning with differential privacy. In: Proceedings of 2016 ACM SIGSAC Conference on Computer and Communications Security. 2016, 308−318

[32] Dwork C. Differential privacy: a survey of results. In: Proceedings of the 5th International Conference on Theory and Applications of Models of Computation. 2008, 1−19

[33] Du M, Yue X, Chow S S M, Wang T, Huang C, Sun H. DP-forward: fine-tuning and inference on language models with differential privacy in forward pass. In: Proceedings of 2023 ACM SIGSAC Conference on Computer and Communications Security. 2023, 2665−2679

[34] Zhao P, Shen L, Fan R, Li Q, Wu H, Wu J, Liu Z. Learning with user-level local differential privacy. 2024, arXiv preprint arXiv: 2405.17079

[35] Xu Z, Collins M, Wang Y, Panait L, Oh S, Augenstein S, Liu T, Schroff F, McMahan H B. Learning to generate image embeddings with user-level differential privacy. In: Proceedings of 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023, 7969−7980

[36] Chatzikokolakis K, Andrés M E, Bordenabe N E, Palamidessi C. Broadening the scope of differential privacy using metrics. In: Proceedings of International Symposium on Privacy Enhancing Technologies Symposium. 2013, 82−102

[37] Liang Y, Yi K. Concentrated geo-privacy. In: Proceedings of 2023 ACM SIGSAC Conference on Computer and Communications Security. 2023, 1934−1948

[38] Ji T, Li P, Yilmaz E, Ayday E, Ye Y, Sun J . Differentially private binary- and matrix-valued data query: an XOR mechanism. Proceedings of the VLDB Endowment, 2021, 14( 5): 849–862

[39] Blocki J, Blum A, Datta A, Sheffet O. The johnson-lindenstrauss transform itself preserves differential privacy. In: Proceedings of 2012 IEEE 53rd Annual Symposium on Foundations of Computer Science. 2012, 410−419

[40] Blum A, Roth A. Fast private data release algorithms for sparse queries. In: Proceedings of the 16th International Workshop on Approximation Algorithms for Combinatorial Optimization. 2013, 395−410

[41] Upadhyay J. Randomness efficient fast-Johnson-lindenstrauss transform with applications in differential privacy and compressed sensing. 2015, arXiv preprint arXiv: 1410.2470

[42] Dawid A P . Some matrix-variate distribution theory: notational considerations and a Bayesian application. Biometrika, 1981, 68( 1): 265–274

[43] Yang J, Xiang L, Chen R, Li W, Li B . Differential privacy for tensor-valued queries. IEEE Transactions on Information Forensics and Security, 2022, 17: 152–164

[44] Han S, Buyukates B, Hu Z, Jin H, Jin W, Sun L, Wang X, Wu W, Xie C, Yao Y, Zhang K, Zhang Q, Zhang Y, Joe-Wong C, Avestimehr S, He C. FedSecurity: a benchmark for attacks and defenses in federated learning and federated LLMs. In: Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024, 5070−5081

[45] Zhang Z, Zhang J, Huang J, Qu L, Zhang H, Wang Q, Zhou X, Xu Z. FewFedPIT: towards privacy-preserving and few-shot federated instruction tuning. 2024, arXiv preprint arXiv: 2403.06131

[46] Wu C, Li X, Wang J. Vulnerabilities of foundation model integrated federated learning under adversarial threats. 2025, arXiv preprint arXiv: 2401.10375v1

[47] Sun Y, Li Z, Li Y, Ding B. Improving LoRA in privacy-preserving federated learning. In: Proceedings of the 12th International Conference on Learning Representations. 2024

[48] Qin Z, Chen D, Qian B, Ding B, Li Y, Deng S. Federated full-parameter tuning of billion-sized language models with communication cost under 18 kilobytes. In: Proceedings of the 41st International Conference on Machine Learning. 2024, 1686

[49] Khalid U, Iqbal H, Vahidian S, Hua J, Chen C. CEFHRI: a communication efficient federated learning framework for recognizing industrial human-robot interaction. In: Proceedings of 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems. 2023, 10141−10148

[50] Charles Z, Mitchell N, Pillutla K, Reneer M, Garrett Z. Towards federated foundation models: Scalable dataset pipelines for group-structured learning. In: Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023, 1402

[51] Liu Z, Li D, Fernandez-Marques J, Laskaridis S, Gao Y, Dudziak Ł, Li S Z, Hu S X, Hospedales T M. Federated learning for inference at anytime and anywhere. 2022, arXiv preprint arXiv: 2212.04084

[52] Ma T, Hoang T N, Chen J. Federated inference through aligning local representations and learning a consensus graph. In: Proceedings of the 10th International Conference on Learning Representations. 2022

[53] Kovaleva O, Romanov A, Rogers A, Rumshisky A. Revealing the dark secrets of BERT. In: Proceedings of 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing. 2019, 4365−4374

[54] Child R, Gray S, Radford A, Sutskever I. Generating long sequences with sparse transformers. 2019, arXiv preprint arXiv: 1904.10509

[55] Xiao G, Tian Y, Chen B, Han S, Lewis M. Efficient streaming language models with attention sinks. In: Proceedings of the 12th International Conference on Learning Representations. 2024

[56] Lee W, Lee J, Seo J, Sim J. InfiniGen: efficient generative inference of large language models with dynamic KV cache management. In: Proceedings of the 18th USENIX Symposium on Operating Systems Design and Implementation. 2024, 155−172

[57] Xu Y, Jie Z, Dong H, Wang L, Lu X, Zhou A, Saha A, Xiong C, Sahoo D. ThinK: thinner key cache by query-driven pruning. In: Proceedings of the 13th International Conference on Learning Representations. 2025

[58] Tang H, Lin Y, Lin J, Han Q, Ke D, Hong S, Yao Y, Wang G. RazorAttention: efficient KV cache compression through retrieval heads. In: Proceedings of the 13th International Conference on Learning Representations. 2025

[59] Yang D, Han X, Gao Y, Hu Y, Zhang S, Zhao H. Pyramidinfer: pyramid KV cache compression for high-throughput LLM inference. In: Proceedings of Findings of the Association for Computational Linguistics: ACL 2024. 2024, 3258−3270

[60] Ge S, Zhang Y, Liu L, Zhang M, Han J, Gao J. Model tells you what to discard: Adaptive KV cache compression for LLMs. In: Proceedings of the 12th International Conference on Learning Representations. 2024

[61] Adnan M, Arunkumar A, Jain G, Nair P J, Soloveychik I, Kamath P. Keyformer: KV cache reduction through key tokens selection for efficient generative inference. In: Proceedings of the 7th MLSys Conference. 2024

[62] Gao B, Pavel L. On the properties of the softmax function with application in game theory and reinforcement learning. 2018, arXiv preprint arXiv: 1704.00805

Rights & permissions

Higher Education Press