Federated Attention: A Distributed Paradigm for Collaborative LLM Inference over Edge Networks

TL;DR

Federated Attention (FedAttn) introduces a distributed LLM inference framework that enhances privacy, reduces communication, and improves computational efficiency by integrating federated learning principles into self-attention mechanisms at the edge.

Contribution

This work presents FedAttn, a novel distributed inference method for LLMs that combines federated learning with self-attention, enabling privacy-preserving, efficient collaborative inference.

Findings

Theoretical analysis of error propagation in FedAttn.

Trade-off characterization between response quality and efficiency.

Experimental validation showing scalability and optimization opportunities.

Abstract

Large language models (LLMs) are proliferating rapidly at the edge, delivering intelligent capabilities across diverse application scenarios. However, their practical deployment in collaborative scenarios confronts fundamental challenges: privacy vulnerabilities, communication overhead, and computational bottlenecks. To address these, we propose Federated Attention (FedAttn), which integrates the federated paradigm into the self-attention mechanism, creating a new distributed LLM inference framework that simultaneously achieves privacy protection, communication efficiency, and computational efficiency. FedAttn enables participants to perform local self-attention over their own token representations while periodically exchanging and aggregating Key-Value (KV) matrices across multiple Transformer blocks, collaboratively generating LLM responses without exposing private prompts. Further, we identify a structural duality between contextual representation refinement in FedAttn and parameter optimization in FL across private data, local computation, and global aggregation. This key insight provides a principled foundation for systematically porting federated optimization techniques to collaborative LLM inference. Building on this framework, we theoretically analyze how local self-attention computation within participants and heterogeneous token relevance among participants shape error propagation dynamics across Transformer blocks. Moreover, we characterize the fundamental trade-off between response quality and communication/computation efficiency, which is governed by the synchronization interval and the number of participants. Experimental results validate our theoretical analysis, and reveal significant optimization opportunities through sparse attention and adaptive KV aggregation, highlighting FedAttn's potential to deliver scalability and efficiency in real-world edge deployments.