分割学习与协同推理
分割学习(Split Learning)将模型切割为多段,分布在不同参与方上执行。前向传播在各方间依次传递中间激活,反向传播依次传递梯度。
SplitNN 基本流程
参与方A 参与方B
┌──────────┐ ┌──────────┐
│ 输入数据 │ │ 标签数据 │
│ │ │ │
│ 前半模型 │──激活──→ │ 后半模型 │
│ (底模型) │←─梯度── │ (顶模型) │
└──────────┘ └──────────┘python
class SplitNNClient:
"""分割学习客户端(持有前半模型和数据)"""
def __init__(self, bottom_model):
self.bottom_model = bottom_model
def forward(self, x):
"""前向传播到切割层"""
with torch.no_grad():
activation = self.bottom_model(x)
return activation
def backward(self, x, grad_from_server):
"""接收服务端梯度,反向传播"""
activation = self.bottom_model(x)
activation.backward(grad_from_server)
class SplitNNServer:
"""分割学习服务端(持有后半模型和标签)"""
def __init__(self, top_model):
self.top_model = top_model
def forward(self, activation):
"""从切割层继续前向传播"""
activation.requires_grad = True
output = self.top_model(activation)
return output, activation
def backward(self, output, labels, activation):
"""计算损失并反向传播"""
loss = F.cross_entropy(output, labels)
loss.backward()
return activation.grad # 发送给客户端安全分割学习
中间激活可能泄露输入信息:
python
class SecureSplitNN(SplitNNClient):
"""安全分割学习"""
def forward(self, x):
activation = self.bottom_model(x)
# 方法1:添加噪声
noise = torch.randn_like(activation) * 0.01
noisy_activation = activation + noise
# 方法2:加密传输
encrypted = self.encrypt(noisy_activation)
return encrypted
def encrypt(self, tensor):
"""同态加密中间激活"""
# 使用 Paillier 或 CKKS
return tensor # 简化示例分割推理
模型训练完成后,推理也可以分割执行:
python
class SplitInference:
"""分割推理"""
def __init__(self, client_model, server_model):
self.client_model = client_model
self.server_model = server_model
def predict(self, x):
# 客户端计算前半
activation = self.client_model(x)
# 发送激活到服务端(可选加密)
encrypted_activation = self.encrypt_if_needed(activation)
# 服务端计算后半
result = self.server_model(encrypted_activation)
return result对比
| 特性 | SplitNN | 联邦学习 |
|---|---|---|
| 通信内容 | 中间激活 | 模型参数/梯度 |
| 通信量 | 较小 | 较大 |
| 计算分配 | 各方分摊 | 各方独立 |
| 隐私风险 | 激活泄露 | 梯度泄露 |
| 适用场景 | 模型大、各方算力弱 | 数据多、各方有算力 |
分割学习优势
- 客户端只需存储部分模型,降低内存需求
- 通信量小于传输完整模型参数
- 适合模型大但终端算力弱的场景
- 可与联邦学习结合(SplitFed)
激活泄露
研究表明,通过中间激活可以部分重建输入数据。防御措施:1) 添加噪声;2) 加密传输;3) 使用更大切割层;4) 限制激活精度。