Papers
Topics
Authors
Recent
Search
2000 character limit reached

Three Birds, One Stone: Solving the Communication-Memory-Privacy Trilemma in LLM Fine-tuning Over Wireless Networks with Zeroth-Order Optimization

Published 14 Apr 2026 in cs.DC | (2604.12401v1)

Abstract: Federated Learning (FL) offers a promising pathway for collaboratively fine-tuning LLMs at the edge; however, this paradigm faces a critical bottleneck: the prohibitive communication and memory overheads incurred by exchanging high-dimensional gradients. Furthermore, recent studies reveal that user training data can still be recovered from these local gradients, undermining the core privacy promise of FL. In this paper, we address this trilemma of communication, memory, and privacy by proposing pAirZero, a novel framework that synergizes Zeroth-Order (ZO) optimization with Over-the-Air (OTA) computation. Uniquely, pAirZero enables resource-constrained devices to submit their local gradient with only bit-level communication loads while participating in federated fine-tuning of LLMs with inference-level memory costs. This approach not only eliminates the high memory requirements needed for LLM fine-tuning but also alleviates the strict synchronization requirements that plague conventional OTA methods. We further formulate a rigorous optimization model to adaptively determine the optimal transmit power and noise levels, ensuring consistent privacy protection regardless of channel conditions. Numerical experiments demonstrate the superiority of pAirZero in enabling secure, efficient LLM fine-tuning over wireless networks, with only 25% peak memory cost on OPT-125M and communication load orders of magnitude lower than conventional methods.

Summary

  • The paper introduces pAirZero, a framework that leverages zeroth-order optimization and OTA computation to drastically reduce communication load.
  • It achieves a 75% reduction in peak memory usage for models like OPT-125M while maintaining statistically equivalent performance under strict privacy constraints.
  • The approach embeds differential privacy at the physical layer, ensuring linear convergence and resource efficiency in federated edge learning.

Solving the Communication-Memory-Privacy Trilemma in Edge LLM Fine-tuning with pAirZero

Introduction and Motivation

As the deployment of LLMs shifts from centralized data centers to edge devices, the challenge of federated fine-tuning of these models becomes pronounced. Standard federated learning (FL) approaches encounter three major barriers on edge hardware: the communication bottleneck due to massive gradient payloads, the memory wall from backpropagation requirements, and inherent privacy leakage arising from gradient inversion attacks. The paper "Three Birds, One Stone: Solving the Communication-Memory-Privacy Trilemma in LLM Fine-tuning Over Wireless Networks with Zeroth-Order Optimization" (2604.12401) proposes pAirZero, a framework that integrates Zeroth-Order (ZO) optimization and over-the-air (OTA) computation, embedding differential privacy into the communication pipeline itself. The approach fundamentally reforms the computational, communication, and privacy aspects of resource-constrained federated LLM tuning.

System Model and pAirZero Workflow

The system considered consists of KK single-antenna clients connected to a single-antenna edge server over a fading wireless channel. Clients collaboratively fine-tune a global LLM without sharing raw data. pAirZero leverages ZO optimization, requiring only scalar gradient projections (rather than full gradients) for each iteration, minimizing communication and memory overhead per device.

The workflow is as follows:

  • Each client receives a shared random seed and computes the gradient projection along a pseudo-random direction via finite-difference estimation (SPSA variant).
  • Clients transmit only the scalar projection (potentially quantized to a single bit) over the wireless channel using either analog or sign-based digital OTA aggregation.
  • Artificial and channel noise are incorporated during uplink aggregation to provide calibrated differential privacy at the physical layer.
  • The aggregated value is broadcast by the edge server; clients update local models using the pseudo-random direction and received projection. Figure 1

    Figure 1: The workflow of pAirZero, integrating ZO optimization with analog/digital OTA and in-channel privacy.

This paradigm shifts memory requirements from backpropagation-level to inference-level and decouples per-iteration upload from model dimensionality and client count, while differentially private aggregation emerges from stochastic wireless superposition.

Zeroth-Order Optimization and Communication Efficiency

By relying on ZO optimization, pAirZero circumvents the need to store intermediate activations required by backpropagation in first-order (FO) methods. The memory cost per device for OPT-125M is reduced from approximately $600$ MB (SGD) or $955$ MB (Adam) to around $250$ MB, which is only modestly larger than model storage. Further, communication load per iteration drops orders of magnitude: whereas SGD requires uploading the entire gradient ($238.88$ MB for OPT-125M), pAirZero requires only one scalar (16 bits with analog OTA, a single bit with sign-based digital OTA), with upload cost independent of model size.

Privacy-Embedded Wireless Transmission

Traditional privacy mechanisms (DP, secure aggregation) operate at upper protocol layers and introduce additional computational and communication costs. pAirZero, instead, achieves differential privacy "for free" by leveraging physical-layer channel and artificially injected noise directly onto the transmitted scalars. The framework adaptively tunes transmit power and noise scale per iteration to guarantee (ϵ,δ)(\epsilon, \delta)-DP, regardless of changing channel SNR, and presents a closed-form optimization for these parameters. Privacy applies to the scalar projection rather than the entire gradient vector, reducing resource burden on clients.

Theoretical Guarantees: Convergence and Privacy

Rigorous convergence analysis is provided under mild assumptions (gradient Lipschitz, PL condition, local effective rank for the Hessian). Both analog and sign-based OTA support linear convergence to a neighborhood of the optimum, with the size of the neighborhood determined by the effective noise and channel gains. A unique aspect is the derivation of the optimality gap bounds, showing that, even with aggressive communication and memory reductions and in-channel DP, pAirZero matches the asymptotic convergence behavior traditionally only possible with FO optimization.

Empirical Evaluation: Robustness and Efficiency

Extensive experiments are conducted using the OPT-125M model on SST-2 and SQuAD, representing both classification and extractive QA tasks, under varying SNR and privacy regimes. The results conclusively demonstrate:

  • Statistically equivalent test performance to full-precision, noise-free (i.e., ideal) aggregation even under strong privacy guarantees (ϵ=5\epsilon=5, δ=0.01\delta=0.01).
  • Magnitudes-of-order reduction in communication cost and 75%75\% reduction in peak memory compared to FO baselines.
  • Consistent performance with digital OTA (Sign-pAirZero) across SNR; analog OTA suffers in low SNR due to aggressive noise injection necessary for privacy.
  • Solution-based power allocation (where per-iteration power/noise are adaptively optimized) strongly outperforms static allocation schemes. Figure 2

    Figure 2: Main results on OPT-125M with SST2 and SQuAD tasks, showing pAirZero matches ideal performance and vastly outperforms baselines in communication/memory efficiency.

    Figure 3

    Figure 3: Performance comparison of optimized (solution-based) versus static power allocation in analog modulation.

Fine-Grained Privacy Analysis and Gradient Projection Distribution

Empirical studies show the sign-reversal probability of the scalar gradient projection remains well-controlled (<0.5<0.5), justifying the practical efficiency of sign-based OTA under stochastic conditions. The distribution of gradient projections is near-symmetric and centered, supporting theoretical assumptions for privacy/utility trade-off calculations. Figure 4

Figure 4: Inherent sign-reversing probability density for scalar gradient projections.

Figure 5

Figure 5: Sign-reversing probability as a function of projected gradient value.

Figure 6

Figure 6: Distribution of scalar gradient projections across iterations, illustrating clipping threshold selection.

Figure 7

Figure 7: Distribution of gradient projections during fine-tuning, with $600$0 falling within $600$1.

Practical and Theoretical Implications

pAirZero establishes a scalable and resource-efficient solution to edge LLM fine-tuning, supporting indifferentiable objectives and massive client populations. By decoupling per-iteration communication and memory costs from ambient model dimension and client count, it facilitates practical LLM personalization, domain adaptation, and continual learning at the edge. Furthermore, the demonstrated in-channel DP mechanism, integrated at the physical layer, provides a robust privacy guarantee under real-world wireless constraints, and avoids both the computational burden and accuracy degradation typical of traditional DP.

The convergence and gap-optimality analysis supports the extension of this approach to larger models and more challenging communication settings, potentially inspiring further integration of wireless communication theory and federated optimization.

Conclusion

The pAirZero framework proposed in "Three Birds, One Stone: Solving the Communication-Memory-Privacy Trilemma in LLM Fine-tuning Over Wireless Networks with Zeroth-Order Optimization" (2604.12401) provides a principled resolution to the central challenges of federated LLM tuning at the edge. By synthesizing zeroth-order optimization, OTA computation, and in-channel differential privacy, it achieves significant reductions in communication and memory requirements without sacrificing privacy or utility. These advancements pave the way for practical, secure, and scalable on-device adaptation of large foundation models across diverse and bandwidth-constrained settings.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.