Technical Summary: Efficient Zeroth-Order Federated Finetuning of Language Models on Resource-Constrained Devices

1. Problem Statement

Federated Learning (FL) offers a privacy-preserving framework for finetuning Large Language Models (LLMs) across distributed data sources. However, deploying these models on resource-constrained edge devices faces significant hurdles:

• Memory Constraints: First-order methods (e.g., standard backpropagation) and parameter-efficient techniques like LoRA require storing intermediate activations for gradient computation. This results in memory footprints (e.g., ~4 GB for RoBERTa-large) that exceed the capabilities of most edge devices.
• Communication Overhead: Transmitting full model parameters or even low-rank adaptation matrices (LoRA) incurs substantial upload costs, particularly problematic in low-bandwidth wireless environments.
• Zeroth-Order (ZO) Limitations: While Zeroth-Order Optimization (ZO) methods like MeZO eliminate the need for backpropagation by using forward passes and pseudo-random perturbations, they suffer from slow convergence and high computational demand. To achieve stable gradient estimates, existing ZO-FL methods require a large number of perturbations per training step, leading to excessive Forward-Only computations (FLOPs).

There exists a critical trade-off: current ZO methods are memory-efficient but computationally expensive, while first-order methods are computationally efficient but memory-intensive.

2. Methodology: FedSPZO

The authors propose Federated Split-Perturbation Zeroth-order Optimization (FedSPZO), a novel framework designed to reduce the computational burden of ZO-FL while preserving its memory and communication advantages.

Core Mechanism: Split-Perturbation Strategy

FedSPZO exploits the sequential compositional structure of neural networks by splitting the model into two consecutive blocks, $f_1$ (the first, larger block) and $f_2$ (the second, smaller block), such that $y = f_2(\theta_2; f_1(\theta_1; B))$.

• Asymmetric Perturbation Budgets: Instead of applying the same number of perturbations to the entire network, FedSPZO assigns a higher perturbation budget to the second, smaller block ($f_2$).
  • Block 1 ($\theta_1$): Perturbed with $P_1$ directions.
  • Block 2 ($\theta_2$): Perturbed with $P_2$ directions, where $P_2 = 2 \times P_1 \times P_s$.
• Activation Reuse: The key efficiency gain comes from reusing intermediate activations.
  1. The client perturbs $\theta_1$ in a single direction ($+z_1$) and performs a forward pass through $f_1$ to generate intermediate output $+y_l$.
  2. This single $+y_l$ is then used as the input for $P_s$ distinct perturbations of $\theta_2$ (both positive and negative directions).
  3. The process is repeated for the negative direction of $\theta_1$ ($-z_1$), generating $-y_l$, which is reused for another set of $P_s$ perturbations of $\theta_2$.
• Gradient Estimation:
  • The gradient for $\theta_2$ is estimated using the differences in losses from the $P_2$ perturbations, all sharing the same input from $f_1$.
  • The gradient for $\theta_1$ is estimated by averaging the loss differences across the multiple perturbations of $f_2$, effectively reducing the noise introduced by the stochasticity of $f_2$'s perturbations.

Communication Protocol

FedSPZO maintains a communication-efficient design:

• Client Side: Clients perform local training for $K$ steps. They do not upload model parameters. Instead, they upload only scalar gradient values ($G_1, G_2$) and random seeds ($S_1, S_2$) used to regenerate perturbation vectors.
• Server Side: The server reconstructs the exact client models by regenerating the perturbation vectors from the received seeds and applying the scalar updates. This eliminates the need for clients to transmit model weights, reducing upload overhead by orders of magnitude compared to parameter-based FL.

3. Key Contributions

1. FedSPZO Framework: A novel ZO-FL method that partitions the model into two blocks and allocates a larger perturbation budget to the smaller, deeper block. This leverages intermediate activation reuse to update the full network with fewer forward evaluations.
2. Computational Efficiency: The method achieves up to a 3× reduction in computation compared to state-of-the-art ZO-FL baselines (specifically Fang et al., 2022, and Li et al., 2025) while retaining the memory and communication benefits of ZO.
3. Comprehensive Evaluation: Extensive experiments on RoBERTa-large, OPT-1.3B, and LLaMA-3-3.2B models across multiple datasets (SST-2, RTE, WiC, MultiRC, SQuAD, BoolQ) demonstrate that FedSPZO reduces total GFLOPs significantly while maintaining accuracy comparable to first-order backpropagation methods (with minor degradation).
4. Ablation Analysis: The authors conducted studies on model splitting, perturbation budgets, and independent block-wise gradient computation, validating that the split-perturbation strategy is superior to independent block estimation or uniform perturbation schemes.

4. Experimental Results

The evaluation compares FedSPZO against first-order methods (FedAvg, FedAvg+LoRA) and other ZO-FL methods (FedZO, DecomFL).

• Memory Efficiency: FedSPZO maintains an inference-like memory footprint. For RoBERTa-large with context length 256, it requires ~1.75 GB, whereas FedAvg requires ~4.9 GB and LoRA requires ~3.9 GB. This makes FedSPZO feasible for edge devices where LoRA and standard FL are infeasible.
• Communication Overhead: FedSPZO reduces upload costs by more than three orders of magnitude compared to LoRA. While LoRA uploads low-rank matrices (~100s of MBs), FedSPZO uploads only scalar gradients and seeds (kilobytes to megabytes).
• Computational Cost:
  • Compared to FedAvg(LoRA): FedSPZO requires higher total computation (approx. 17× more) due to the slower convergence rate inherent to ZO methods.
  • Compared to Other ZO Methods: FedSPZO is significantly more efficient. It reduces computation by 2.1× to 3× compared to DecomFL and FedZO across various models and datasets.
• Accuracy: FedSPZO achieves accuracy comparable to FedAvg(LoRA), with only minor degradations (e.g., <1% on SST-2, ~4% on RTE for RoBERTa). It performs slightly worse than full FedAvg (backpropagation), which serves as an upper bound, but the gap is considered minor in the context of resource constraints.
• Training Time: In a simulated edge environment (Jetson Orin Nano), FedSPZO (2.6 hours for RoBERTa/SST-2) is slower than LoRA (0.45 hours) but drastically faster than DecomFL (6 hours) and FedZO (166 hours).

5. Significance and Claims

The paper positions FedSPZO as a pragmatic trade-off for resource-constrained environments.

• For Edge Devices: It is presented as a preferable choice when memory and communication bandwidth are the primary bottlenecks, scenarios where LoRA and standard backpropagation fail.
• Balancing Act: The authors acknowledge that FedSPZO does not eliminate the computational cost of ZO but significantly mitigates it. It offers a viable path to finetuning LLMs on devices where storing activations for backpropagation is impossible and transmitting full model updates is too costly.
• Future Outlook: The authors note that while FedSPZO improves computational efficiency over existing ZO methods, it still lags behind first-order methods in raw speed. Future work aims to explore ZO on inference-only accelerators with lower precision to further bridge this gap.

The paper concludes that FedSPZO successfully decouples the high memory/communication costs of first-order FL from the high computational costs of traditional ZO, enabling LLM finetuning in scenarios previously deemed infeasible.