Three Birds, One Stone: Solving the Communication-Memory-Privacy Trilemma in LLM Fine-tuning Over Wireless Networks with Zeroth-Air Optimization
摘要¶
联邦学习（Federated Learning, FL）为在边缘设备上协同微调大语言模型（LLM）提供了有前景的路径，但面临三大瓶颈：通信开销（高维梯度交换）、内存开销（LLM 微调需大量显存）、隐私泄露（用户训练数据可从梯度中恢复）。本文提出 pAirZero，创新性地结合零阶（Zeroth-Order, ZO）优化与过空（Over-the-Air, OTA）计算，同时解决这三大问题。pAirZero 使资源受限设备仅需比特级通信负载即可提交本地梯度，以推理级内存成本参与 LLM 联邦微调，既消除了 LLM 微调的高内存需求，又缓解了传统 OTA 方法的严格同步要求。
核心贡献¶
- 通信-内存-隐私三难困境的系统性解决：pAirZero 首次在同一框架下同时优化通信量、内存占用和隐私保护三个指标。
- 零阶优化 + 过空计算的创新结合：利用 ZO 梯度估计替代高阶优化器，仅需前向传播（无需反向传播），将内存成本降至推理级别。
- 自适应最优化传输功率与噪声水平：建立严格的优化模型，根据信道条件自适应确定最优发射功率和噪声水平，确保跨信道条件的一致性隐私保护。
- 极低通信与内存开销：在 OPT-125M 上仅需 25% 峰值内存成本，通信负载比传统方法低数个量级。
为什么重要¶
联邦学习是端侧/边缘设备上进行个性化记忆系统训练的核心范式。然而传统 FL 方法要求设备存储完整模型参数、保留梯度历史用于同步，且面临隐私泄露风险。pAirZero 代表了面向资源受限设备的记忆系统轻量化训练的重要突破：
- 内存压缩：无需存储优化器状态和梯度，仅需推理级内存，使手机、可穿戴设备也能微调自己的个性化记忆模块。
- 通信效率：比特级通信负载，摆脱了对高带宽网络的依赖，对移动端记忆系统意义重大。
- 隐私保障：即使在无线信道条件差的情况下也能保证梯度不可逆性泄露，解决了个性化记忆系统中用户数据隐私的核心顾虑。
技术细节¶
核心方法：pAirZero¶
pAirZero 的关键设计是利用 Zeroth-Order (ZO) 优化进行梯度估计。传统 FL 需要每个设备存储完整模型梯度（内存密集），而 ZO 优化通过黑盒函数评估估计梯度：
$$ \hat{\nabla}f(x) = \frac{1}{q}\sum_{i=1}^{q}\frac{f(x+\mu u_i) - f(x)}{\mu}u_i^T $$
其中 $u_i$ 为随机方向向量，$\mu$ 为平滑参数。这使得设备只需进行前向传播（无需反向传播），内存需求大幅降低。
过空（OTA）计算¶
过空计算允许多个设备同时在空中汇聚梯度信号，而非传统的正交时分复用通信。通过精心设计的噪声注入和功率控制，即使在信道条件受限的情况下也能保证差分隐私级别的梯度安全性。
自适应优化¶
pAirZero 建立了端到端优化模型，联合优化：
- 发射功率（transmit power）
- 噪声注入水平（noise level）
- 设备选择策略（device selection）
目标是在给定信道条件约束下最大化隐私保护效用，同时最小化通信和内存开销。
实验结果¶
| 指标 | pAirZero | 传统 FL |
|---|---|---|
| 峰值内存（OPT-125M） | 25% | 100% |
| 通信负载 | ~1 bit/参数 | 全梯度交换 |
| 隐私保护 | 信道自适应 | 固定 |
与移动端/端侧记忆的相关性¶
- 端侧个性化训练：25% 峰值内存使手机/手表级别的设备也能微调自己的记忆模块
- 隐私保护的个性化记忆：用户数据永不离开本地设备，梯度也无法被窃取
- 低带宽适应：比特级通信使记忆系统在弱网环境下仍可工作
- 多设备协同：多个边缘设备可协同训练共享的记忆知识而不暴露私有记忆
参考文献¶
- 原论文：Zhijie Cai et al., "Three Birds, One Stone: Solving the Communication-Memory-Privacy Trilemma in LLM Fine-tuning Over Wireless Networks with Zeroth-Air Optimization", arXiv:2604.12401