[论文解读] Federated Full-Parameter Tuning of Billion-Sized Language Models with Communication Cost under 18 Kilobytes

Zhen Qin, Daoyuan Chen+4 more|arXiv (Cornell University)|Dec 11, 2023

DNA and Biological Computing被引用 8

一句话总结

论文提出了 FedKSeed，一种零阶联邦微调方法，通过使用固定的小种子池和标量梯度，以及非均匀种子采样增强 FedKSeed-Pro，在极低的通信量（每轮不足 18 KB）的情况下实现对十亿级 LLM 的全参数微调。

ABSTRACT

Pre-trained large language models (LLMs) need fine-tuning to improve their responsiveness to natural language instructions. Federated learning offers a way to fine-tune LLMs using the abundant data on end devices without compromising data privacy. Most existing federated fine-tuning methods for LLMs rely on parameter-efficient fine-tuning techniques, which may not reach the performance height possible with full-parameter tuning. However, federated full-parameter tuning of LLMs is a non-trivial problem due to the immense communication cost. This work introduces FedKSeed that employs zeroth-order optimization with a finite set of random seeds. It significantly reduces transmission requirements between the server and clients to just a few random seeds and scalar gradients, amounting to only a few thousand bytes, making federated full-parameter tuning of billion-sized LLMs possible on devices. Building on it, we develop a strategy enabling probability-differentiated seed sampling, prioritizing perturbations with greater impact on model accuracy. Experiments across six scenarios with various LLMs, datasets and data partitions demonstrate that our approach outperforms existing federated LLM fine-tuning methods in both communication efficiency and new task generalization.

研究动机与目标

• 证明在联邦环境中以极低通信量实现十亿级大型语言模型的全参数微调的可行性。
• 通过避免全模型参数交换和反向传播所带来的内存使用，减少服务器-客户端传输。
• 开发基于种子扰动的策略，在联邦、非独立同分布数据下维持或提升准确性。
• 对零阶优化框架中种子重用进行理论与经验分析。

提出的方法

• 使用零阶优化（ZOO）结合有限数量的随机种子来生成全参数微调的扰动。
• 在服务器端维护 K 个种子池和一个标量梯度累积器，以在不传输完整参数的情况下重构最新全局模型。
• 客户端使用与所选种子相关联的两点梯度估计进行本地更新，并仅报告种子-梯度历史而非完整模型更新。
• 在服务器端对各种种子进行标量梯度聚合以更新全局扰动投影，并通过基于种子的更新规则在本地重构最新模型。
• 通过按估计梯度大小对种子进行加权，引入非均匀种子采样（FedKSeed-Pro），以提升效率和准确性。

实验结果

研究问题

• RQ1 十亿级大型语言模型的全参数微调是否能实现足够的通信效率，以便实际部署在联邦场景？
• RQ2 将种子池限定为 K 个扰动对比无限种子的方法，对收敛性和准确性有何影响？
• RQ3 非均匀种子采样是否提升了 FedKSeed-Pro 的准确性并降低所需的种子基数？
• RQ4 ZOO 基于联邦学习中种子重用的理论关系与收敛性保证之间是什么关系？
• RQ5 在不同的 LLM、数据集和数据划分设置下，FedKSeed 与 FedKSeed-Pro 相比基于 PEFT 的联邦微调基线在准确性与通信/内存效率方面的表现如何？

主要发现

• FedKSeed 通过将扰动编码为种子并仅传输标量梯度，在每轮通信低于 18 千字节的情况下实现全参数微调。
• 收敛性分析表明，在提出的框架下，种子多样性不会影响收敛性，从而实现种子重用而不牺牲理论保证。
• 非均匀种子采样（FedKSeed-Pro）持续提升 Rouge-L 分数并降低所需的种子基数，在各场景平均相比基线实现最高可达 7.26% 的相对提升。
• 在六种场景、模型规模为 1B–3B 的实验中，FedKSeed 与 FedKSeed-Pro 在准确性与通信/内存效率方面均优于基于 PEFT 的联邦微调基线。
• FedKSeed-Pro 在较小的 K（如 1024–4096）下也能实现更高的准确性，在若干设定中还优于 FedKSeed 的同步效率。