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

TL;DR

FedAttn introduces a privacy-preserving distributed inference framework that integrates federated principles into self-attention, allowing $N$ participants to locally compute attention and periodically share KV matrices every $H$ Transformer blocks to collaboratively generate LLM outputs. The approach establishes a Federated Duality between FL and FedAttn, linking local contextual refinement to global model aggregation, and provides an analytical bound on error propagation across blocks and rounds. Theoretical results reveal trade-offs between response quality, communication, and computation, and guide design choices like the placement of global attention and the use of sparse attention and adaptive KV aggregation. Empirical validation on Qwen2.5 models with GSM8K demonstrates substantial efficiency gains and robust performance, highlighting FedAttn’s potential for scalable, edge-enabled collaborative LLM inference with strong privacy guarantees.

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.

Figures (10)

• Figure 1: Illustration of Transformer architecture and Attention map.
• Figure 2: Framework Overview of FedAttn. We exemplify FedAttn through three participants, each maintaining private input tokens while collaboratively computing attention through periodic KV exchange.
• Figure 3: Algorithmic Procedure of FedAttn. We illustrate FedAttn through a representative example of one communication round involving three participants, each with two local input tokens, executing Transformer-based LLM inference. For expository clarity, the notations for hidden representations, Queries, Keys, and Values are streamlined as vectors indexed by token IDs as subscripts.
• Figure 4: Illustration of a) Input Format and b) Input Segmentation Settings. We exemplify the input format with a 3-shot instructional examples and demonstrate 4 segmentation settings under 3 participants, each holding a segment as its local input sequence.
• Figure 5: Trade-off between response quality and communication cost. Evaluated with 4-shot prompting, greedy decoding, max 256 new tokens.

Theorems & Definitions (11)

• Theorem 1
• Remark 1: Lipschitz gain
• Remark 2: Error Injection
• Remark 3: Error Propagation
• Corollary 1
• Remark 4
• Remark 5: Diminishing gains in communication efficiency
• Theorem 2
• Remark 6: Where to Perform Global Attentions
• Theorem 3