FedProxy: Federated Fine-Tuning of LLMs via Proxy SLMs and Heterogeneity-Aware Fusion

1. Problem Statement

The paper addresses the "trilemma" of Federated Learning (FL) for Large Language Models (LLMs), which involves balancing three conflicting objectives:

1. Intellectual Property (IP) Protection: Preventing the server's proprietary LLM from being fully exposed to clients.
2. Client Privacy: Ensuring client data remains local and private.
3. Performance on Heterogeneous Data: Maintaining high model performance despite non-IID (statistically heterogeneous) data distributions across clients.

Limitations of Existing Methods:
Current state-of-the-art approaches like Offsite-Tuning (OT) and its federated variants (FedOT, FedBiOT) attempt to solve this by training only lightweight adapters on the client side while keeping the backbone frozen and obfuscated. However, the authors identify a fundamental bottleneck: these lightweight adapters lack sufficient representation capacity, leading to significant performance gaps compared to centralized fine-tuning, especially in heterogeneous settings.

2. Methodology: The FedProxy Framework

FedProxy introduces a surrogate-driven optimization paradigm. Instead of training small adapters, it compresses the proprietary LLM into a unified, powerful Proxy Small Language Model (SLM). This Proxy SLM acts as a high-fidelity surrogate for collaborative fine-tuning. The framework operates in three distinct stages:

Stage 1: Efficient Representation via Server-Guided Compression

• Goal: Create a resource-friendly Proxy SLM ($f_\phi$) from the proprietary LLM ($f_\theta$) without exposing the full model IP.
• Mechanism: The server uses a public dataset ($D_{pub}$) to perform structured pruning on the LLM.
• Metric: It utilizes the Block Influence (BI) metric to quantify the importance of each transformer block based on its impact on hidden state representations.
• Result: Blocks with the highest BI scores are retained to form a compact Proxy SLM ($|\phi| \ll |\theta|$), which is then distributed to all clients. This ensures the proxy retains the backbone's structural knowledge while fitting client resource constraints.

Stage 2: Robust Optimization via Interference-Mitigating Aggregation

• Goal: Address data heterogeneity and parameter conflicts during federated training.
• Mechanism: A three-step aggregation strategy per communication round:
  1. Server-Side Analysis: The server computes Task Vectors (updates) and calculates pairwise cosine similarities to derive:
     • Heterogeneity Coefficient ($h_k$): Measures how much a client deviates from the group.
     • Adaptive Weight ($w_k$): Prioritizes clients with strong consensus.
     • Parameter-Level Conflict Score ($C_d$): Quantifies disagreement in update signs for specific parameters.
  2. Client-Side Regularization (PCR): Clients incorporate conflict scores into their local loss function. They are penalized for diverging from the global model on parameters with high conflict, encouraging alignment on consensus parameters while allowing personalization on conflicting ones.
  3. Heterogeneity-Aware TIES-Merging (H-TIES): The server merges updates using a modified TIES-Merging approach:
     • Sparsification: Updates from highly heterogeneous clients are sparsified (retention rate reduced).
     • Scaling: Updates are weighted by consensus strength.
     • Conflict Resolution: For each parameter dimension, only updates conforming to the dominant sign (positive or negative) are averaged, effectively canceling out opposing gradients.

Stage 3: Effortless Fusion via Training-Free "Plug-in"

• Goal: Integrate the learned knowledge from the Proxy SLM back into the original LLM.
• Mechanism: Since the Proxy SLM is a direct sub-network of the LLM (architectural mapping $\phi \subset \theta$), the framework performs a direct parameter replacement.
• Process: The updated layers from the converged Proxy SLM overwrite the corresponding layers in the original LLM.
• Benefit: This is a training-free process that incurs zero additional inference latency on the massive LLM, unlike decoding-time fusion methods.

3. Key Contributions

1. FedProxy Framework: A holistic solution resolving the IP-Privacy-Performance trilemma by replacing weak adapters with a powerful, compressed Proxy SLM.
2. H-TIES and PCR: Novel algorithms for Heterogeneity-Aware Aggregation.
   • PCR guides conflict-aware local training.
   • H-TIES performs heterogeneity-aware server merging, effectively mitigating parameter interference in diverse federated ecosystems.
3. Theoretical Analysis: The paper provides a theoretical bound on the fusion error, decomposing it into proxy suboptimality, compression distortion, and subspace approximation, justifying the direct parameter replacement strategy.
4. Empirical Superiority: Extensive experiments demonstrating that FedProxy significantly outperforms OT-based methods and approaches centralized fine-tuning performance.

4. Experimental Results

The framework was evaluated on LLaMA2-7B and Mistral-7B across two scenarios: Homogeneous (IID) and Heterogeneous (Non-IID) tasks (QA and GLUE benchmarks).

• Performance Gains:
  • In Heterogeneous settings (50% compression), FedProxy achieved a QA accuracy of 0.6011 and GLUE accuracy of 0.7944 (LLaMA2-7B).
  • This represents a massive improvement over FedOT (e.g., +12.6% absolute improvement on QA).
  • FedProxy reached ~90-91% of the performance of centralized fine-tuning (CentSFT), whereas OT-based methods lagged significantly behind.
• Cost-Performance Trade-off:
  • While FedProxy requires higher iterative communication and client computation (due to training a larger Proxy SLM vs. small adapters), the initial communication cost is comparable to OT.
  • The performance gains justify the increased resource investment, closing the gap to centralized training.
• Ablation Studies:
  • Removing either PCR or H-TIES resulted in significant performance drops, proving both components are critical.
  • FedProxy outperformed even a stronger baseline where the OT emulator was also trained (FedOT-ET), confirming the architectural advantage of the Proxy SLM over adapter-based structures.
  • IP Protection: The initial compressed Proxy SLM performed significantly worse than the original LLM (Zero-Shot), confirming that the compressed model does not leak the full IP.

5. Significance and Impact

• New Benchmark: FedProxy establishes a new standard for secure, high-performance federated LLM adaptation, proving that high-fidelity adaptation is possible without full model exposure.
• Practical Deployment: The "plug-in" fusion mechanism allows for immediate capability upgrades to the global LLM without retraining or inference overhead, making it highly suitable for real-world industrial applications.
• Scalability: By addressing the performance bottleneck of existing OT methods, FedProxy makes federated fine-tuning of LLMs a viable solution for privacy-sensitive domains (e.g., finance, healthcare) where data cannot be centralized.

Limitations: The authors note that higher compression ratios (>50%) lead to performance degradation due to capacity loss. Additionally, the current aggregation strategy scales quadratically ($O(K^2)$) with the number of clients, suggesting a need for approximate nearest neighbor search or hierarchical aggregation for massive-scale deployments.