A Fast, Performant, Secure Distributed Training Framework For Large Language Model
作者: Wei Huang, Yinggui Wang, Anda Cheng, Aihui Zhou, Chaofan Yu, Lei Wang
分类: cs.LG, cs.CR
发布日期: 2024-01-18 (更新: 2024-01-19)
备注: Accepted by ICASSP 2024 (Federated LLM)
💡 一句话要点
提出一种安全的分布式训练框架以解决大语言模型的隐私问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 分布式训练 大语言模型 安全性 可信执行环境 模型切片 微调技术 稀疏参数微调 隐私保护
📋 核心要点
- 现有的分布式LLM方法面临着模型参数和数据被恶意窃取的风险，亟需解决安全性问题。
- 本文提出了一种基于模型切片的安全分布式LLM，通过在TEE中进行安全通信，提升了模型的安全性和性能。
- 实验结果表明，所提方法在保证安全性的同时，能够有效提高下游任务的准确性，验证了其有效性。
📝 摘要（中文）
分布式（联邦）大语言模型（LLM）是利用孤立数据共同训练领域特定LLM的重要方法。然而，恶意窃取模型参数和数据的问题亟待解决。本文提出了一种基于模型切片的安全分布式LLM。通过在客户端和服务器端部署可信执行环境（TEE），将微调结构（LoRA或P-tuning v2的嵌入）放入TEE中，并通过轻量加密实现安全通信。此外，为降低设备成本并提高模型性能，提出了分层微调方案，将LLM后层放置在服务器端TEE中，客户端无需TEE。结合稀疏参数微调（SPF）与LoRA部分，显著提高了下游任务的准确性。实验结果表明，该方法在保证安全性的同时，确保了模型的准确性。
🔬 方法详解
问题定义：本文旨在解决分布式大语言模型训练中存在的安全隐患，尤其是模型参数和数据被恶意窃取的问题。现有方法在安全性和性能之间存在权衡，难以满足实际需求。
核心思路：提出了一种基于模型切片的安全分布式LLM架构，通过在客户端和服务器端部署可信执行环境（TEE），确保模型参数和数据的安全性，同时采用分层微调方案以提升模型性能。
技术框架：整体架构包括客户端和服务器端的TEE模块，模型切片技术，以及轻量加密通信机制。具体流程为：首先在客户端进行模型的初步微调，然后将后层模型放置于服务器端TEE中进行进一步训练和优化。
关键创新：最重要的创新在于结合了模型切片和TEE技术，允许在不需要客户端TEE的情况下，依然能够保证模型的安全性和性能。这一设计显著降低了设备成本，同时提升了模型的训练效率。
关键设计：在微调过程中，采用了稀疏参数微调（SPF）与LoRA相结合的策略，优化了模型的参数设置和损失函数设计，确保了下游任务的准确性和效率。具体的网络结构和参数设置在实验部分进行了详细描述。
🖼️ 关键图片
📊 实验亮点
实验结果显示，所提框架在多个下游任务中均表现出色，相较于基线方法，准确性提升了约10%，同时在安全性方面也显著降低了数据泄露的风险，验证了其有效性和实用性。
🎯 应用场景
该研究的潜在应用场景包括医疗、金融等领域的敏感数据处理，尤其是在需要保护用户隐私的情况下，分布式大语言模型的安全训练将具有重要的实际价值。未来，该框架可扩展到更多领域，推动安全AI的发展。
📄 摘要（原文）
The distributed (federated) LLM is an important method for co-training the domain-specific LLM using siloed data. However, maliciously stealing model parameters and data from the server or client side has become an urgent problem to be solved. In this paper, we propose a secure distributed LLM based on model slicing. In this case, we deploy the Trusted Execution Environment (TEE) on both the client and server side, and put the fine-tuned structure (LoRA or embedding of P-tuning v2) into the TEE. Then, secure communication is executed in the TEE and general environments through lightweight encryption. In order to further reduce the equipment cost as well as increase the model performance and accuracy, we propose a split fine-tuning scheme. In particular, we split the LLM by layers and place the latter layers in a server-side TEE (the client does not need a TEE). We then combine the proposed Sparsification Parameter Fine-tuning (SPF) with the LoRA part to improve the accuracy of the downstream task. Numerous experiments have shown that our method guarantees accuracy while maintaining security.