个性化联邦学习中两类损失函数为何表现不同？

核心问题 #

版本2中全局模型参数的设备转移代码无效，导致正则项（prox_term）无法对客户端模型参数产生约束，最终无论lambda取值如何，都退化为纯本地训练。

具体原因 #

PyTorch的Tensor.to(device)方法不会修改原张量，而是返回一个位于目标设备上的新张量。你当前的循环：

for x in shared_model_parameters:
    x.to(device='cuda')

只是创建了临时的CUDA张量，但原列表shared_model_parameters中的元素依然是CPU上的张量。后续计算客户端模型参数与全局参数的差值时，虽然PyTorch会自动将CPU张量转移到CUDA，但这种隐式转换会导致正则项的梯度计算异常，无法正确约束客户端模型的更新方向。

修复方案 #

将全局参数列表中的每个元素替换为CUDA上的张量，用列表推导式实现：

def mtl_loss_fn(self, logits, labels, shared_model_parameters):
    sample_loss_fn = torch.nn.CrossEntropyLoss()
    mean_batch_term = sample_loss_fn(logits, labels)

    # 正确将全局参数转移到CUDA，替换原列表元素
    shared_model_parameters = [x.to(device='cuda') for x in shared_model_parameters]

    w_diff = torch.tensor(0., device=self.device)
    for w, w_t in zip(self.model.parameters(), shared_model_parameters):
        w_diff += torch.pow(torch.norm(w - w_t), 2)

    prox_term = 0.5 * self.lam * w_diff
    return mean_batch_term + prox_term

额外注意事项 #

服务器端传递全局参数时，建议传递不带梯度的副本，避免客户端的计算影响全局模型的参数状态，比如：

# 服务器端生成待传递的全局参数列表
shared_model_parameters = [p.detach().cpu() for p in global_model.parameters()]

内容的提问来源于stack exchange，提问作者raegan yang