AI联邦学习与隐私计算完全教程

数据隐私法规日趋严格，跨机构数据协作需求却日益增长——联邦学习（Federated Learning, FL）正是解决这一矛盾的关键技术。它让多个参与方在不共享原始数据的前提下协作训练模型，同时结合差分隐私、安全聚合、同态加密等隐私计算技术，构建端到端的隐私保护机器学习流水线。

1. 联邦学习概述与动机

1.1 为什么需要联邦学习

传统的集中式机器学习要求将所有数据汇聚到一处进行训练。这在以下场景中面临严重挑战：

• 医疗数据：医院之间共享患者数据受HIPAA等法规严格限制
• 金融数据：银行间共享客户交易数据违反合规要求
• 移动设备数据：用户键盘输入、照片等敏感信息不应上传至云端
• 跨企业协作：竞对之间需要合作但不能暴露各自数据

联邦学习的核心思想：数据不动，模型动。每个参与方在本地训练模型，只上传模型参数更新（而非原始数据），由中央服务器聚合后分发全局模型。

1.2 联邦学习的基本流程

1. 服务器初始化全局模型
2. 选择参与本轮训练的客户端子集
3. 每个客户端下载全局模型
4. 各客户端用本地数据训练模型
5. 客户端上传模型更新（梯度或权重）
6. 服务器聚合所有更新，得到新的全局模型
7. 重复步骤2-6直至收敛

import torch
import torch.nn as nn
import numpy as np
from collections import OrderedDict

class SimpleFederatedServer:
    """联邦学习服务器的基本实现"""
    def __init__(self, global_model, n_clients, sample_fraction=0.3):
        self.global_model = global_model
        self.n_clients = n_clients
        self.sample_fraction = sample_fraction
    
    def select_clients(self):
        """随机选择参与本轮训练的客户端"""
        n_selected = max(1, int(self.n_clients * self.sample_fraction))
        return np.random.choice(self.n_clients, n_selected, replace=False)
    
    def aggregate(self, client_states, client_weights):
        """FedAvg聚合：加权平均客户端模型参数"""
        global_dict = OrderedDict()
        total_weight = sum(client_weights)
        
        for key in client_states[0].keys():
            global_dict[key] = sum(
                client_states[i][key].float() * (client_weights[i] / total_weight)
                for i in range(len(client_states))
            )
        
        self.global_model.load_state_dict(global_dict)
        return self.global_model.state_dict()
    
    def distribute(self):
        """分发全局模型给客户端"""
        return self.global_model.state_dict()

2. 横向联邦/纵向联邦/联邦迁移学习

2.1 横向联邦（Horizontal Federated Learning）

参与方拥有相同特征空间、不同样本。典型场景：不同地区的医院拥有相同类型的患者检查指标，但患者群体不同。

class HorizontalFLClient:
    """横向联邦学习客户端"""
    def __init__(self, model, local_data, client_id):
        self.model = model
        self.data = local_data  # 本地数据集
        self.client_id = client_id
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
    
    def train_local(self, global_state, epochs=5):
        """用本地数据训练模型"""
        self.model.load_state_dict(global_state)
        self.model.train()
        
        for epoch in range(epochs):
            for batch_x, batch_y in self.data:
                self.optimizer.zero_grad()
                output = self.model(batch_x)
                loss = nn.CrossEntropyLoss()(output, batch_y)
                loss.backward()
                self.optimizer.step()
        
        return self.model.state_dict(), len(self.data)

2.2 纵向联邦（Vertical Federated Learning）

参与方拥有相同样本、不同特征空间。典型场景：银行和电商平台拥有同一批用户，但银行掌握金融数据，电商掌握消费行为数据。

纵向联邦的核心挑战是需要在不暴露各方特征的前提下，联合计算模型输出。通常需要基于安全多方计算（MPC）或同态加密的协议。

class VerticalFLParty:
    """纵向联邦中的参与方"""
    def __init__(self, model_part, party_id, feature_indices):
        self.model = model_part  # 只负责部分特征的子模型
        self.party_id = party_id
        self.feature_indices = feature_indices
    
    def forward_bottom(self, x):
        """前向传播的底层：处理本方特征"""
        local_features = x[:, self.feature_indices]
        return self.model(local_features)
    
    def compute_local_gradient(self, upstream_grad, local_output):
        """计算本方参数的梯度（需要安全协议交换中间结果）"""
        local_grad = torch.autograd.grad(
            local_output, self.model.parameters(),
            grad_outputs=upstream_grad, retain_graph=True
        )
        return local_grad

2.3 联邦迁移学习

参与方的样本和特征空间都不同，通过迁移学习技术桥接各方的知识。典型场景：一方有大量标注数据但特征有限，另一方有丰富特征但标注稀缺。

3. 联邦平均算法（FedAvg）详解

FedAvg是最经典的联邦学习算法，由McMahan等人在2017年提出。核心思想简单但有效：在客户端进行多轮本地SGD训练，然后对模型参数进行加权平均。

3.1 完整实现

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from collections import OrderedDict

class FederatedNet(nn.Module):
    """简单的联邦学习网络"""
    def __init__(self, input_dim=784, hidden_dim=128, output_dim=10):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

def create_noniid_data(n_clients=10, samples_per_client=500):
    """模拟非IID数据分布"""
    # 假设每个客户端的标签分布不均匀
    client_data = []
    for i in range(n_clients):
        # 每个客户端偏好特定的几个标签
        preferred_labels = [(i * 2) % 10, (i * 2 + 1) % 10]
        # 生成偏向这些标签的数据...
        x = torch.randn(samples_per_client, 784)
        y = torch.randint(0, 10, (samples_per_client,))
        dataset = TensorDataset(x, y)
        loader = DataLoader(dataset, batch_size=32, shuffle=True)
        client_data.append(loader)
    return client_data

class FedAvgTrainer:
    """FedAvg完整训练流程"""
    def __init__(self, model_fn, n_clients, rounds=100,
                 local_epochs=5, sample_fraction=0.3, lr=0.01):
        self.global_model = model_fn()
        self.n_clients = n_clients
        self.rounds = rounds
        self.local_epochs = local_epochs
        self.sample_fraction = sample_fraction
        self.lr = lr
        self.history = {"round": [], "loss": []}
    
    def train_round(self, client_data, selected_clients):
        """执行一轮联邦训练"""
        global_state = self.global_model.state_dict()
        client_states = []
        client_sizes = []
        
        for client_id in selected_clients:
            client_model = FederatedNet()
            client_model.load_state_dict(global_state)
            optimizer = torch.optim.SGD(client_model.parameters(), lr=self.lr)
            client_model.train()
            
            total_loss = 0
            n_batches = 0
            for epoch in range(self.local_epochs):
                for batch_x, batch_y in client_data[client_id]:
                    optimizer.zero_grad()
                    output = client_model(batch_x)
                    loss = nn.CrossEntropyLoss()(output, batch_y)
                    loss.backward()
                    optimizer.step()
                    total_loss += loss.item()
                    n_batches += 1
            
            client_states.append(client_model.state_dict())
            client_sizes.append(len(client_data[client_id]))
        
        # 加权平均聚合
        aggregated = OrderedDict()
        total_size = sum(client_sizes)
        for key in global_state.keys():
            aggregated[key] = sum(
                client_states[i][key].float() * (client_sizes[i] / total_size)
                for i in range(len(client_states))
            )
        
        self.global_model.load_state_dict(aggregated)
        return total_loss / max(n_batches, 1)
    
    def fit(self, client_data, eval_fn=None):
        """完整的联邦训练循环"""
        for round_num in range(self.rounds):
            selected = np.random.choice(
                self.n_clients,
                max(1, int(self.n_clients * self.sample_fraction)),
                replace=False
            )
            loss = self.train_round(client_data, selected)
            self.history["round"].append(round_num)
            self.history["loss"].append(loss)
            
            if round_num % 10 == 0:
                print(f"Round {round_num}: avg_loss = {loss:.4f}")
                if eval_fn:
                    acc = eval_fn(self.global_model)
                    print(f"  → Global accuracy: {acc:.2%}")

3.2 FedAvg的局限性

• Non-IID数据：当各客户端数据分布差异很大时，FedAvg收敛困难
• 通信开销：每轮需传输完整模型参数
• 公平性：模型可能偏向数据量大的客户端

4. 通信优化与异步聚合

4.1 梯度压缩

通信是联邦学习的主要瓶颈。梯度压缩通过减少传输数据量来加速训练：

def topk_compress(tensor, compression_ratio=0.01):
    """Top-K梯度压缩：只传输最大的K%梯度"""
    k = max(1, int(tensor.numel() * compression_ratio))
    values, indices = torch.topk(tensor.abs().flatten(), k)
    
    compressed = torch.zeros_like(tensor.flatten())
    compressed[indices] = tensor.flatten()[indices]
    return compressed.reshape(tensor.shape), indices, values

def quantize_gradient(tensor, bits=8):
    """梯度量化：降低数值精度"""
    min_val = tensor.min()
    max_val = tensor.max()
    scale = (max_val - min_val) / (2**bits - 1)
    quantized = torch.round((tensor - min_val) / scale).to(torch.uint8)
    return quantized, min_val, scale

class CompressedClient:
    """带梯度压缩的联邦客户端"""
    def __init__(self, model, compression="topk", ratio=0.01):
        self.model = model
        self.compression = compression
        self.ratio = ratio
    
    def get_compressed_update(self, global_state):
        """返回压缩后的模型更新"""
        local_state = self.model.state_dict()
        compressed_update = {}
        
        for key in global_state.keys():
            delta = local_state[key] - global_state[key]
            if self.compression == "topk":
                compressed, indices, values = topk_compress(delta, self.ratio)
                compressed_update[key] = compressed
            elif self.compression == "quantize":
                quantized, min_val, scale = quantize_gradient(delta)
                compressed_update[key] = {"data": quantized, "min": min_val, "scale": scale}
        
        return compressed_update

4.2 异步聚合

在真实联邦场景中，客户端训练速度差异很大（设备性能不同、网络状况不同）。异步聚合不要求所有客户端同步完成：

class AsyncFedServer:
    """异步联邦学习服务器"""
    def __init__(self, global_model, staleness_weight="linear"):
        self.global_model = global_model
        self.global_version = 0
        self.staleness_weight = staleness_weight  # "linear", "polynomial", "constant"
    
    def get_staleness_weight(self, client_version):
        """根据陈旧度调整聚合权重"""
        staleness = self.global_version - client_version
        if self.staleness_weight == "linear":
            return max(0.1, 1.0 - 0.1 * staleness)
        elif self.staleness_weight == "polynomial":
            return 1.0 / (staleness + 1) ** 0.5
        else:
            return 1.0
    
    def async_aggregate(self, client_state, client_version, alpha=0.1):
        """异步聚合单个客户端的更新"""
        weight = self.get_staleness_weight(client_version)
        global_dict = self.global_model.state_dict()
        
        for key in global_dict.keys():
            # 加权移动平均
            global_dict[key] = (1 - alpha * weight) * global_dict[key] + \
                               alpha * weight * client_state[key]
        
        self.global_model.load_state_dict(global_dict)
        self.global_version += 1

4.3 FedProx

FedProx在本地损失函数中添加近端项，约束本地模型不要偏离全局模型太远，有效缓解Non-IID问题：

class FedProxClient:
    def __init__(self, model, mu=0.01):
        self.model = model
        self.mu = mu  # 近端项系数
    
    def train_local(self, global_state, data_loader, epochs=5):
        self.model.load_state_dict(global_state)
        self.model.train()
        optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
        
        for epoch in range(epochs):
            for batch_x, batch_y in data_loader:
                optimizer.zero_grad()
                output = self.model(batch_x)
                task_loss = nn.CrossEntropyLoss()(output, batch_y)
                
                # FedProx近端项：惩罚与全局模型的偏离
                prox_term = 0.0
                for name, param in self.model.named_parameters():
                    global_param = global_state[name]
                    prox_term += torch.sum((param - global_param) ** 2)
                
                total_loss = task_loss + (self.mu / 2) * prox_term
                total_loss.backward()
                optimizer.step()
        
        return self.model.state_dict()

5. 差分隐私与安全聚合

5.1 差分隐私基础

差分隐私（Differential Privacy, DP）提供严格的数学隐私保证：无论是否包含某条特定记录，查询结果的分布几乎不变。

定义（ε-差分隐私）： 对于相邻数据集 D 和 D'（仅差一条记录），随机化机制 M 满足：

P[M(D) ∈ S] ≤ e^ε × P[M(D') ∈ S]

ε 越小，隐私保护越强。

5.2 本地差分隐私

在联邦学习中，每个客户端在上传模型更新前，先对梯度添加噪声：

class DPClient:
    """差分隐私联邦学习客户端"""
    def __init__(self, model, epsilon=1.0, delta=1e-5, max_grad_norm=1.0):
        self.model = model
        self.epsilon = epsilon
        self.delta = delta
        self.max_grad_norm = max_grad_norm
    
    def clip_gradients(self):
        """梯度裁剪：限制每条样本的贡献"""
        total_norm = 0
        for p in self.model.parameters():
            if p.grad is not None:
                total_norm += p.grad.data.norm(2).item() ** 2
        total_norm = total_norm ** 0.5
        
        clip_coef = min(1.0, self.max_grad_norm / (total_norm + 1e-6))
        for p in self.model.parameters():
            if p.grad is not None:
                p.grad.data.mul_(clip_coef)
    
    def add_noise(self):
        """添加高斯噪声实现差分隐私"""
        # 根据隐私预算计算噪声标准差
        sigma = self.max_grad_norm * np.sqrt(2 * np.log(1.25 / self.delta)) / self.epsilon
        
        for p in self.model.parameters():
            if p.grad is not None:
                noise = torch.normal(0, sigma, size=p.grad.shape)
                p.grad.data.add_(noise)
    
    def get_noisy_update(self, global_state, data_loader, epochs=1):
        """训练并返回加噪后的模型更新"""
        self.model.load_state_dict(global_state)
        self.model.train()
        optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
        
        for epoch in range(epochs):
            for batch_x, batch_y in data_loader:
                optimizer.zero_grad()
                loss = nn.CrossEntropyLoss()(self.model(batch_x), batch_y)
                loss.backward()
                self.clip_gradients()   # 先裁剪
                self.add_noise()        # 再加噪
                optimizer.step()
        
        return self.model.state_dict()

5.3 安全聚合协议

安全聚合确保服务器只能看到聚合后的模型更新，而无法获知任何单个客户端的更新内容。核心思想是使用秘密共享：

import hashlib

class SecureAggregator:
    """简化的安全聚合实现（基于秘密共享的思想）"""
    def __init__(self, n_clients):
        self.n_clients = n_clients
    
    def generate_mask(self, client_id, round_id, peer_id, seed):
        """为每对客户端生成配对的随机掩码"""
        # 双方使用相同的种子生成相同的掩码
        key = f"{min(client_id,peer_id)}_{max(client_id,peer_id)}_{round_id}_{seed}"
        np.random.seed(int(hashlib.md5(key.encode()).hexdigest()[:8], 16))
        return np.random.randn()  # 实际中应为模型参数维度的向量
    
    def masked_upload(self, client_id, model_update, round_id, all_clients, seed):
        """客户端添加掩码后上传"""
        masked_update = {}
        for key, param in model_update.items():
            masked = param.clone()
            for peer_id in all_clients:
                if peer_id != client_id:
                    # 与peer配对的掩码：client_id < peer时加，否则减
                    mask = torch.randn_like(param)
                    if client_id < peer_id:
                        masked += mask
                    else:
                        masked -= mask
            masked_update[key] = masked
        return masked_update
    
    def secure_aggregate(self, masked_updates):
        """安全聚合：所有掩码相互抵消，得到真实聚合结果"""
        aggregated = {}
        for key in masked_updates[0].keys():
            # 掩码两两配对抵消，只剩真实梯度的和
            aggregated[key] = sum(update[key] for update in masked_updates.values()) / len(masked_updates)
        return aggregated

6. 同态加密在FL中的应用

6.1 同态加密概述

同态加密（Homomorphic Encryption, HE）允许在密文上直接进行计算，解密后结果与在明文上计算一致。这对于纵向联邦学习至关重要。

• 部分同态加密（PHE）：支持一种运算（加法或乘法），如Paillier加密
• 全同态加密（FHE）：支持任意运算，但计算开销极大

6.2 基于Paillier的加密聚合

# 使用PySyft或TenSEAL等库进行同态加密
# 这里展示概念性实现

class PaillierEncryptedAggregation:
    """基于Paillier加密的安全聚合"""
    def __init__(self, key_size=2048):
        # 实际中应使用专业加密库（如phe、tenseal）
        self.key_size = key_size
    
    def encrypt_model_update(self, model_state, public_key):
        """加密模型参数（概念示意）"""
        encrypted_state = {}
        for key, param in model_state.items():
            # 实际实现中，对每个参数值进行Paillier加密
            # encrypted_state[key] = public_key.encrypt(param.numpy())
            encrypted_state[key] = param  # 简化示意
        return encrypted_state
    
    def aggregate_encrypted(self, encrypted_updates):
        """在密文上进行聚合（Paillier的加法同态性质）"""
        aggregated = {}
        for key in encrypted_updates[0].keys():
            # Paillier: E(a) * E(b) = E(a + b)
            # 即密文相乘等于明文相加
            aggregated[key] = sum(update[key] for update in encrypted_updates)
        return aggregated
    
    def decrypt_aggregated(self, aggregated_state, private_key):
        """解密聚合结果"""
        decrypted = {}
        for key, enc_param in aggregated_state.items():
            # decrypted[key] = private_key.decrypt(enc_param)
            decrypted[key] = enc_param  # 简化示意
        return decrypted

6.3 使用TenSEAL进行加密计算

# TenSEAL是Microsoft SEAL的Python绑定
# 实际代码示例（需要安装tenseal）

def tenseal_example():
    """
    import tenseal as ts
    
    # 创建加密上下文
    context = ts.context(
        ts.SCHEME_TYPE.CKKS,
        poly_modulus_degree=8192,
        coeff_mod_bit_sizes=[60, 40, 40, 60]
    )
    context.global_scale = 2**40
    context.generate_galois_keys()
    
    # 加密客户端的模型梯度
    client1_grad = ts.ckks_vector(context, [0.1, 0.2, 0.3])
    client2_grad = ts.ckks_vector(context, [0.4, 0.5, 0.6])
    
    # 密文上直接相加
    aggregated = client1_grad + client2_grad
    
    # 解密得到聚合结果
    result = aggregated.decrypt()  # [0.5, 0.7, 0.9]
    """
    pass

7. 联邦学习框架

7.1 Flower

Flower是最流行的开源联邦学习框架，支持任意ML框架（PyTorch、TensorFlow、JAX），提供高度灵活的API。

# 使用Flower实现联邦学习
import flwr as fl
from collections import OrderedDict

class FlowerClient(fl.client.NumPyClient):
    """Flower联邦学习客户端"""
    def __init__(self, model, train_loader, test_loader):
        self.model = model
        self.train_loader = train_loader
        self.test_loader = test_loader
    
    def get_parameters(self, config):
        """返回模型参数"""
        return [val.cpu().numpy() for val in self.model.state_dict().values()]
    
    def set_parameters(self, parameters):
        """设置模型参数"""
        params_dict = zip(self.model.state_dict().keys(), parameters)
        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
        self.model.load_state_dict(state_dict)
    
    def fit(self, parameters, config):
        """本地训练"""
        self.set_parameters(parameters)
        self.model.train()
        optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01, momentum=0.9)
        
        for epoch in range(5):
            for batch_x, batch_y in self.train_loader:
                optimizer.zero_grad()
                loss = nn.CrossEntropyLoss()(self.model(batch_x), batch_y)
                loss.backward()
                optimizer.step()
        
        return self.get_parameters(config), len(self.train_loader.dataset), {}
    
    def evaluate(self, parameters, config):
        """本地评估"""
        self.set_parameters(parameters)
        self.model.eval()
        correct, total, total_loss = 0, 0, 0.0
        
        with torch.no_grad():
            for batch_x, batch_y in self.test_loader:
                output = self.model(batch_x)
                total_loss += nn.CrossEntropyLoss()(output, batch_y).item() * len(batch_y)
                correct += (output.argmax(1) == batch_y).sum().item()
                total += len(batch_y)
        
        return total_loss / total, total, {"accuracy": correct / total}

def start_flower_server(n_rounds=50):
    """启动Flower服务器"""
    strategy = fl.server.strategy.FedAvg(
        fraction_fit=0.3,           # 每轮选择30%的客户端
        fraction_evaluate=0.2,      # 20%客户端参与评估
        min_fit_clients=2,          # 最少需要2个客户端训练
        min_available_clients=2,    # 最少需要2个客户端在线
    )
    
    fl.server.start_server(
        server_address="0.0.0.0:8080",
        strategy=strategy,
        config=fl.server.ServerConfig(num_rounds=n_rounds),
    )

7.2 FATE（Federated AI Technology Enabler）

FATE是微众银行开源的企业级联邦学习平台，特别针对金融场景优化，支持横向联邦、纵向联邦和联邦迁移学习。

FATE的核心组件：

• FATE-Flow：联邦学习任务调度引擎
• FATE-Board：可视化训练监控
• EggRoll：分布式计算引擎
• FATE-Client：Python SDK

# FATE纵向联邦学习的Pipeline配置（概念示意）
from pipeline.component import HeteroLR, Intersection, Reader

def fate_vertical_lr_pipeline():
    """
    # 读取数据
    reader_guest = Reader(name="guest_reader")
    reader_host = Reader(name="host_reader")
    
    # 样本对齐（隐私集合求交PSI）
    intersection = Intersection(name="intersection")
    
    # 纵向逻辑回归
    hetero_lr = HeteroLR(
        name="hetero_lr",
        learning_rate=0.01,
        max_iter=30,
        batch_size=256,
        optimizer="sgd",
        tol=1e-4
    )
    
    # Pipeline组装
    pipeline = PipeLine()
    pipeline.add_component(reader_guest)
    pipeline.add_component(reader_host)
    pipeline.add_component(intersection, data=Data(
        data=[reader_guest.output.data, reader_host.output.data]
    ))
    pipeline.add_component(hetero_lr, data=Data(train_data=intersection.output.data))
    """
    pass

7.3 PySyft

PySyft是OpenMined开发的隐私保护机器学习库，支持联邦学习、差分隐私和安全计算：

# PySyft的使用示例（概念示意）
def pysyft_example():
    """
    import syft as sy
    import torch
    
    # 创建虚拟联邦节点
    alice = sy.VirtualMachine(name="alice")
    bob = sy.VirtualMachine(name="bob")
    
    alice_client = alice.get_root_client()
    bob_client = bob.get_root_client()
    
    # 将数据发送到远程节点
    alice_data = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
    bob_data = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
    
    alice_ptr = alice_data.send(alice_client)
    bob_ptr = bob_data.send(bob_client)
    
    # 远程计算（数据不离开节点）
    result = alice_ptr + bob_ptr  # 安全计算
    """
    pass

8. 联邦大模型微调

8.1 核心挑战

将联邦学习应用于大语言模型微调面临特殊挑战：

• 模型规模巨大：7B参数的模型约28GB（FP32），通信开销难以承受
• 本地计算资源有限：边缘设备无法完整训练大模型
• 数据异质性更强：NLP任务的Non-IID问题更严重

8.2 FedLoRA方案

结合LoRA（低秩适配器）的联邦微调，大幅降低通信和计算成本：

import torch
import torch.nn as nn

class LoRALayer(nn.Module):
    """LoRA低秩适配层"""
    def __init__(self, original_layer, rank=8, alpha=16):
        super().__init__()
        self.original = original_layer
        self.original.weight.requires_grad = False  # 冻结原始权重
        
        in_features = original_layer.in_features
        out_features = original_layer.out_features
        
        self.lora_A = nn.Parameter(torch.randn(in_features, rank) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))
        self.scaling = alpha / rank
    
    def forward(self, x):
        original_out = self.original(x)
        lora_out = (x @ self.lora_A @ self.lora_B) * self.scaling
        return original_out + lora_out

class FedLoRAClient:
    """联邦LoRA微调客户端"""
    def __init__(self, model_with_lora, local_data):
        self.model = model_with_lora
        self.data = local_data
        # 只优化LoRA参数
        self.optimizer = torch.optim.Adam(
            [p for n, p in self.model.named_parameters() if 'lora' in n],
            lr=1e-3
        )
    
    def get_lora_params_only(self):
        """只返回LoRA参数用于联邦通信"""
        return {n: p for n, p in self.model.named_parameters() if 'lora' in n}
    
    def set_lora_params(self, lora_state):
        """只更新LoRA参数"""
        model_dict = self.model.state_dict()
        model_dict.update(lora_state)
        self.model.load_state_dict(model_dict)

class FedLoRAServer:
    """联邦LoRA服务器"""
    def __init__(self, base_model_fn, rank=8):
        self.base_model = base_model_fn()
        self.lora_modules = {}
        self.rank = rank
        self._inject_lora()
    
    def _inject_lora(self):
        """在base model的关键层注入LoRA"""
        for name, module in self.base_model.named_modules():
            if isinstance(module, nn.Linear) and module.out_features > 64:
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]
                parent = dict(self.base_model.named_modules())[parent_name]
                setattr(parent, child_name, LoRALayer(module, rank=self.rank))
    
    def aggregate_lora(self, client_lora_states, client_weights):
        """只聚合LoRA参数"""
        aggregated = {}
        total_weight = sum(client_weights)
        
        for key in client_lora_states[0].keys():
            aggregated[key] = sum(
                state[key].float() * (w / total_weight)
                for state, w in zip(client_lora_states, client_weights)
            )
        
        return aggregated

8.3 通信优化策略

• 只传输LoRA参数：对于7B模型，LoRA参数通常只有原始模型的0.1%-1%
• 梯度压缩：对LoRA梯度进行Top-K或量化压缩
• 分层聚合：不同层使用不同的聚合频率

9. 可信执行环境（TEE）

9.1 TEE概述

可信执行环境（Trusted Execution Environment）是硬件级别的安全隔离区域，即使操作系统或hypervisor被攻破，TEE内的数据和代码仍然安全。

主流TEE技术：

• Intel SGX：CPU级别的enclave
• ARM TrustZone：移动端TEE
• AMD SEV：虚拟机级别的加密

9.2 TEE在联邦学习中的应用

TEE可以在服务器端安全聚合时提供额外保护：客户端将加密的模型更新发送到TEE中，TEE内部解密并聚合，确保服务器运维人员也无法窥探原始更新。

class TEESecureAggregation:
    """基于TEE的安全聚合（概念示意）"""
    def __init__(self):
        # 实际中需要在SGX enclave内运行
        self.enclave_initialized = False
    
    def init_enclave(self):
        """
        实际实现中：
        1. 加载SGX enclave
        2. 进行远程认证（Remote Attestation）
        3. 建立安全通道
        """
        self.enclave_initialized = True
    
    def secure_aggregate_in_enclave(self, encrypted_updates):
        """
        在TEE内执行聚合：
        1. 接收各方加密的模型更新
        2. 在enclave内解密
        3. 执行聚合计算
        4. 加密聚合结果
        5. 返回加密结果
        """
        if not self.enclave_initialized:
            raise RuntimeError("Enclave not initialized")
        
        # 模拟TEE内的安全聚合
        aggregated = {}
        for key in encrypted_updates[0].keys():
            # 这些计算在TEE内完成，外部不可见
            aggregated[key] = sum(u[key] for u in encrypted_updates) / len(encrypted_updates)
        
        return aggregated

9.3 硬件安全与远程认证

远程认证（Remote Attestation）让客户端验证：TEE中的代码确实是你期望的聚合逻辑，而非恶意代码。这是建立信任链的关键步骤。

10. 实战案例：医疗数据联邦学习

10.1 场景设定

多家医院合作训练一个胸部X光疾病分类模型。每家医院有自己的患者数据，受隐私法规约束不能直接共享。

import torch
import torch.nn as nn
from torchvision import models

class ChestXrayModel(nn.Module):
    """胸部X光分类模型"""
    def __init__(self, n_classes=5):
        super().__init__()
        self.backbone = models.resnet18(pretrained=True)
        self.backbone.fc = nn.Linear(512, n_classes)
    
    def forward(self, x):
        return self.backbone(x)

class HospitalClient:
    """医院联邦学习客户端"""
    def __init__(self, hospital_id, train_loader, test_loader, device="cpu"):
        self.hospital_id = hospital_id
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.device = device
        self.model = ChestXrayModel().to(device)
    
    def train(self, global_state, epochs=3, lr=1e-4):
        """本地训练（含差分隐私保护）"""
        self.model.load_state_dict(global_state)
        self.model.train()
        optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
        
        for epoch in range(epochs):
            for images, labels in self.train_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                optimizer.zero_grad()
                outputs = self.model(images)
                loss = nn.CrossEntropyLoss()(outputs, labels)
                loss.backward()
                
                # 梯度裁剪（差分隐私）
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                
                # 添加高斯噪声
                for param in self.model.parameters():
                    if param.grad is not None:
                        noise = torch.normal(0, 0.01, size=param.grad.shape, device=self.device)
                        param.grad.data.add_(noise)
                
                optimizer.step()
        
        return self.model.state_dict(), len(self.train_loader.dataset)
    
    def evaluate(self, global_state):
        """本地评估"""
        self.model.load_state_dict(global_state)
        self.model.eval()
        correct, total = 0, 0
        
        with torch.no_grad():
            for images, labels in self.test_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                outputs = self.model(images)
                predicted = outputs.argmax(dim=1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        return correct / total

def run_federated_medical_training():
    """模拟医院间联邦学习"""
    # 假设有5家医院参与
    n_hospitals = 5
    n_rounds = 50
    global_model = ChestXrayModel()
    
    clients = []
    for i in range(n_hospitals):
        # 实际中每家医院有自己的本地数据
        train_loader = create_hospital_data(hospital_id=i)
        test_loader = create_hospital_test_data(hospital_id=i)
        clients.append(HospitalClient(i, train_loader, test_loader))
    
    print("开始联邦训练...")
    for round_num in range(n_rounds):
        # 选择参与本轮的医院
        selected = np.random.choice(n_hospitals, size=3, replace=False)
        
        client_states = []
        client_sizes = []
        for idx in selected:
            state, size = clients[idx].train(global_model.state_dict())
            client_states.append(state)
            client_sizes.append(size)
        
        # 聚合
        aggregated = {}
        total_size = sum(client_sizes)
        for key in client_states[0].keys():
            aggregated[key] = sum(
                s[key].float() * (sz / total_size)
                for s, sz in zip(client_states, client_sizes)
            )
        global_model.load_state_dict(aggregated)
        
        # 评估
        if round_num % 10 == 0:
            accuracies = [c.evaluate(global_model.state_dict()) for c in clients]
            print(f"Round {round_num}: 各医院准确率 = {[f'{a:.2%}' for a in accuracies]}")
            print(f"  平均准确率: {np.mean(accuracies):.2%}")

10.2 关键设计决策

1. 隐私预算分配：每轮消耗部分ε，总预算在训练结束后审计
2. 模型架构选择：使用预训练的ResNet作为backbone，减少通信轮次
3. Non-IID处理：不同医院的疾病分布差异大，使用FedProx或个性化联邦学习
4. 合规性文档：记录每轮的隐私消耗和模型审计日志

11. 合规要求与挑战

11.1 全球隐私法规概览

11.2 联邦学习的隐私风险

尽管联邦学习不直接共享原始数据，仍存在多种隐私攻击：

1. 梯度反演攻击（Gradient Inversion）：从模型梯度中重建原始训练数据
2. 成员推理攻击（Membership Inference）：判断某条数据是否参与了训练
3. 模型反演攻击（Model Inversion）：从模型输出中推断敏感属性

def gradient_inversion_demo():
    """梯度反演攻击的概念演示"""
    # 假设攻击者获得了客户端上传的梯度
    model = nn.Linear(784, 10)
    
    # 模拟客户端的真实梯度
    real_data_x = torch.randn(1, 784)
    real_data_y = torch.tensor([3])
    output = model(real_data_x)
    loss = nn.CrossEntropyLoss()(output, real_data_y)
    loss.backward()
    real_grad = {n: p.grad.clone() for n, p in model.named_parameters()}
    
    # 攻击者尝试重建输入数据
    dummy_x = torch.randn(1, 784, requires_grad=True)
    dummy_y = real_data_y  # 假设标签已知
    
    optimizer = torch.optim.Adam([dummy_x], lr=0.1)
    for step in range(1000):
        optimizer.zero_grad()
        model.zero_grad()
        dummy_output = model(dummy_x)
        dummy_loss = nn.CrossEntropyLoss()(dummy_output, dummy_y)
        dummy_loss.backward()
        
        # 使虚拟梯度匹配真实梯度
        grad_loss = 0
        for name, param in model.named_parameters():
            grad_loss += ((param.grad - real_grad[name]) ** 2).sum()
        
        grad_loss.backward()
        optimizer.step()
    
    # dummy_x 在一定程度上逼近 real_data_x
    return dummy_x.detach()

11.3 防御策略

class PrivacyDefensePipeline:
    """多层隐私防御流水线"""
    def __init__(self, dp_epsilon=1.0, dp_delta=1e-5, clip_norm=1.0, 
                 use_secure_agg=True, use_tee=False):
        self.dp_epsilon = dp_epsilon
        self.dp_delta = dp_delta
        self.clip_norm = clip_norm
        self.use_secure_agg = use_secure_agg
        self.use_tee = use_tee
    
    def apply_defense(self, model_update):
        """应用多层防御"""
        # 第1层：梯度裁剪
        for key in model_update:
            norm = model_update[key].norm(2)
            if norm > self.clip_norm:
                model_update[key] = model_update[key] * (self.clip_norm / norm)
        
        # 第2层：差分隐私噪声
        sigma = self.clip_norm * np.sqrt(2 * np.log(1.25 / self.dp_delta)) / self.dp_epsilon
        for key in model_update:
            noise = torch.normal(0, sigma, size=model_update[key].shape)
            model_update[key] = model_update[key] + noise
        
        # 第3层：安全聚合（网络层实现）
        if self.use_secure_agg:
            model_update = apply_secure_aggregation_mask(model_update)
        
        # 第4层：TEE保护（可选）
        if self.use_tee:
            model_update = encrypt_for_tee(model_update)
        
        return model_update
    
    def audit_privacy_budget(self, rounds_completed):
        """审计隐私预算消耗"""
        # 使用高级组合定理计算总隐私消耗
        total_epsilon = self.dp_epsilon * np.sqrt(2 * rounds_completed * np.log(1 / self.dp_delta))
        return {
            "total_epsilon": total_epsilon,
            "rounds": rounds_completed,
            "per_round_epsilon": self.dp_epsilon,
            "remaining_budget": max(0, 10.0 - total_epsilon)  # 假设总预算为10
        }

11.4 实践建议

1. 隐私预算管理：在训练开始前设定总ε预算，监控每轮消耗
2. 模型审计：定期对模型进行隐私攻击测试，验证防御有效性
3. 合规文档：记录数据使用、隐私保护措施、模型访问控制等
4. 最小权限原则：客户端只能获取聚合后的全局模型，不能查看其他参与方信息
5. 数据生命周期管理：明确本地训练数据、模型更新的存储和销毁策略

总结

联邦学习与隐私计算正在从学术研究走向大规模工业落地。关键要点：

技术选型指南：

• 数据特征相同、样本不同 → 横向联邦（FedAvg/FedProx）
• 样本相同、特征不同 → 纵向联邦（需要MPC/HE配合）
• 大模型微调 → FedLoRA + 梯度压缩
• 高隐私要求 → 差分隐私 + 安全聚合 + TEE多层防御

框架选择：

• 快速原型 → Flower
• 金融企业级 → FATE
• 隐私计算研究 → PySyft

持续关注的方向：

• 个性化联邦学习（每个客户端获得定制化模型）
• 联邦学习中的公平性与激励机制
• 联邦大模型训练的可扩展性
• 隐私保护技术的效率提升（更快的HE、更小的TEE开销）

隐私不是功能的对立面——好的隐私保护设计能让AI系统更值得信赖、更容易获得用户和合作伙伴的接受。联邦学习正是构建这种信任的技术基石。