Skip to content

分割学习与协同推理

分割学习(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联邦学习
通信内容中间激活模型参数/梯度
通信量较小较大
计算分配各方分摊各方独立
隐私风险激活泄露梯度泄露
适用场景模型大、各方算力弱数据多、各方有算力

分割学习优势

  1. 客户端只需存储部分模型,降低内存需求
  2. 通信量小于传输完整模型参数
  3. 适合模型大但终端算力弱的场景
  4. 可与联邦学习结合(SplitFed)

激活泄露

研究表明,通过中间激活可以部分重建输入数据。防御措施:1) 添加噪声;2) 加密传输;3) 使用更大切割层;4) 限制激活精度。

相关资源

最近更新