Federated Inference for Heterogeneous LLM Communication and Collaboration

This position paper introduces \texttt{FedRefine}, a novel federated inference framework that enables heterogeneous on-device Large Language Models to collaboratively enhance performance through privacy-preserving KV cache communication while adhering to task-oriented QoS and system constraints.

The Gist

Imagine you have a group of friends, each with a different level of expertise. One is a brilliant coder, another is a history buff, and a third is a creative writer. None of them has all the answers, but together, they could solve almost any problem.

The problem is, if they try to solve a puzzle together by shouting their thoughts out loud (sending text back and forth), it takes forever. They have to stop, think, write down a sentence, send it, wait for the other person to read it, and then start over. Plus, shouting their thoughts might reveal private secrets they didn't mean to share.

This paper proposes a new way for these "friends" (which are actually AI models running on your phone or laptop) to work together. They call this system FedRefine.

Here is the simple breakdown of how it works, using some fun analogies:

1. The Old Way: The "Text Message" Bottleneck

Usually, if two AIs want to collaborate, they talk like humans. AI A sends a sentence to AI B. AI B reads it, thinks, and sends a reply.

• The Problem: This is like trying to build a house by mailing one brick at a time. It's slow. Also, the "bricks" (the words) might contain private information about you that you don't want to send over the internet.

2. The New Way: The "Brain-to-Brain" Connection

Instead of sending words, FedRefine lets the AIs share their internal thought process directly. In AI terms, this is called the KV Cache (Key-Value Cache).

• The Analogy: Imagine you are reading a book. The "KV Cache" is like the mental bookmark and the summary of everything you've read so far in your head.
• How it works: Instead of AI A saying, "I am thinking about a cat," it hands AI B its "mental bookmark" (the KV Cache). AI B can instantly pick up exactly where AI A left off, without needing to read the whole story again.
• The Benefit: It's incredibly fast because they skip the "re-reading" part. It's also more private because they aren't sending the actual words (which might be sensitive); they are just sharing the structure of the thought.

3. The "Translator" Problem (Heterogeneity)

Here is the tricky part: Your phone might have a small, fast AI, while your friend's laptop has a huge, powerful AI. They speak different "languages" internally.

• The Solution: The paper introduces a Fuser (like a universal translator).
• The Analogy: Think of the Fuser as a specialized interpreter who can take the "mental bookmark" from the small AI and translate it perfectly so the big AI can understand it, and vice versa. This allows a tiny phone AI and a giant server AI to work together seamlessly, even though they are built differently.

4. The "Secret Code" (Privacy)

You might worry: "If I share my thoughts, isn't that private?"

• The Fix: Before the AI shares its "mental bookmark," it rewrites the question into a "secret code" (rephrased input).
• The Analogy: Imagine you want to ask a friend for help with a math problem, but you don't want them to know why you need the answer (e.g., you're taking a test). You rewrite the question as a generic puzzle. The friend solves the puzzle using their knowledge, and you get the answer, but they never knew the original context. FedRefine does this automatically.

5. The Results: Faster and Smarter

The authors tested this with different AI models (like Qwen and Llama).

• Accuracy: When the AIs worked together using this "brain-to-brain" method, they got much smarter answers than working alone. Adding more friends to the group made the answers even better.
• Speed: Even though the "mental bookmarks" are bigger files than simple text, the system is much faster overall because it skips the slow "reading and re-reading" delays.
• Privacy: They found that using the "secret code" (rephrasing) only made the answers slightly less accurate, but it kept your data safe.

The Big Picture

This paper suggests a future where your phone doesn't have to rely on a giant, slow cloud server to be smart. Instead, your phone can team up with other devices nearby. They can swap their "mental bookmarks" instantly, translate them for each other, and solve complex problems together, all while keeping your personal data hidden in your own pocket.

It's like turning a group of solo artists into a super-band, where they can hear each other's thoughts instantly without shouting, creating a masterpiece that none of them could make alone.

Technical Summary

1. Problem Statement

The paper addresses the limitations of on-device Large Language Models (LLMs), which often suffer from compromised accuracy and speed compared to full-scale cloud models. While fully offloading inference to the cloud is scalable, it neglects on-device capabilities and raises privacy concerns. Existing collaborative approaches face three critical challenges:

• Latency: Traditional token-to-token (T2T) communication requires transmitting text, which forces the receiving model to rebuild its Key-Value (KV) cache from scratch, incurring significant pre-fill delays.
• Privacy: Transmitting raw input/output tokens can reveal sensitive user data.
• Heterogeneity: Different LLMs have varying architectures, making the exchange of semantic information or knowledge difficult without a unifying mechanism.

2. Methodology: FedRefine Framework

The authors propose FedRefine, a novel federated inference framework that enables heterogeneous LLMs to collaborate via KV Cache-to-Cache (C2C) communication rather than token transmission. The framework is built on two core concepts:

A. From Self-Refinement to C2C

• Self-Refinement: LLMs iteratively improve their own outputs but are limited by internal knowledge.
• C2C Communication: Instead of sending text tokens, devices transmit their internal KV Caches.
  • Mechanism: A "Transmitter" LLM ($M_1$) sends its KV Cache ($C$) to a "Receiver" LLM ($M_2$).
  • Fuser Network: A pre-trained Fuser ($F_{12}$) bridges the gap between the heterogeneous architectures of $M_1$ and $M_2$. It projects the sender's KV cache into a format compatible with the receiver.
  • Inference: The receiver generates the next token ($t_{k+1}$) by concatenating its own model state with the fused KV cache from the sender, effectively skipping the pre-fill delay.
  • Formula: $t_{k+1} = P(t_k | C(F_{12}, M_1) \circ C(M_2))$, where $\circ$ denotes sequence-wise concatenation.

B. Bidirectional Collaboration (Co-C2C)

• The framework extends to bidirectional communication, where two LLMs act simultaneously as transmitters and receivers.
• Two fusers ($F_{12}$ and $F_{21}$) are trained to bridge the models in both directions, enabling mutual refinement and a fairer, incentive-compatible collaboration paradigm.

C. Federated Refinement (FedRefine)

• Multi-LLM System: In a system with $N$ heterogeneous LLMs, a central server maintains a library of pre-trained fusers for all possible bidirectional pairs ($F_{ij}, F_{ji}$).
• Privacy Preservation: To prevent intent leakage, the receiver model rephrases the original input tokens before distribution. The collaboration occurs using these rephrased tokens and the shared KV caches, ensuring private tokens remain local.
• Scalability: A gating network allows each LLM to dynamically select which fusers or models to collaborate with based on the task.

3. Key Contributions

1. New Paradigm (FedRefine): Introduces a federated inference framework that shifts collaboration from text-based (T2T) to cache-based (C2C) communication, specifically designed for heterogeneous LLMs.
2. Architecture-Agnostic Bridging: Utilizes pre-trained "Fuser" networks (MLPs) to project KV caches between different model architectures, solving the heterogeneity problem.
3. Privacy-Preserving Mechanism: Combines KV cache sharing with input rephrasing to protect user privacy while maintaining high inference quality.
4. Bidirectional Refinement: Proposes a Co-C2C approach allowing mutual refinement between devices, moving beyond unidirectional knowledge transfer.

4. Experimental Results

The authors evaluated FedRefine using a heterogeneous system comprising a receiver (Qwen3-0.6B) and four transmitter models (Qwen2.5 variants and Llama-3.2-1B).

• Accuracy Improvement:
  • Federated inference significantly outperformed standalone baseline models.
  • With four sharing models, the non-private KV collaboration improved accuracy by 21.2% over the independent baseline.
  • The privacy-preserving (rephrased) KV model showed only a 3% accuracy drop compared to the non-private version, proving the privacy strategy is effective without major performance loss.
  • C2C vs. T2T: C2C collaboration outperformed T2T by approximately 15% in accuracy when all models participated.
• Latency:
  • Although C2C requires transmitting larger data packets (88 KB per token vs. 16 bytes for T2T), the total latency of C2C is significantly lower than T2T. This is because C2C eliminates the pre-fill delay required to rebuild the KV cache from text tokens.
  • The privacy-preserving C2C method incurred slight latency due to query rewriting but remained far superior to T2T in total time.
• Scalability: Accuracy showed a clear upward trend as the number of participating "sharer" clients increased.

5. Significance and Future Directions

• Efficiency: FedRefine demonstrates that leveraging on-device inference capabilities through cache sharing is a scalable alternative to full cloud offloading, offering high accuracy with low latency.
• Privacy: It establishes a viable path for collaborative AI that respects data privacy without sacrificing the benefits of collective intelligence.
• Future Research: The paper highlights several avenues for future work:
  • Iterative Local Refinement: Designing multi-step refinement loops using cache communication.
  • Continuous Global Federation: Moving toward system-wide continuous refinement rather than single-task collaboration.
  • Multi-Modal Support: Adapting C2C strategies for multi-modal LLMs (image, audio, text).
  • Prompt Engineering: Developing specific prompt strategies for orchestrating cache-based federated inference.

In conclusion, FedRefine offers a robust, privacy-preserving, and efficient paradigm for heterogeneous LLM collaboration, fundamentally changing how edge devices can share knowledge via internal model states rather than raw data.