网络权重共享场景下：独立与联合优化器的差异及选型疑问

我想了解以下两种场景的副作用：

• netA与netB共享权重，但各自使用独立的优化器
• netA与netB共享权重，但使用同一个优化器

两种方案的结果十分相似，但并不完全相同。

示例代码 #

import torch
import torch.nn as nn

torch.manual_seed(0)

share_fc2 = True
batch_size = 4
channels = 2

training_data_A = torch.rand((batch_size, channels))
training_data_B = torch.rand((batch_size, channels))


class Net(torch.nn.Module):
    """Minimal network"""

    def __init__(self) -> None:
        super().__init__()
        self.fc1 = nn.Linear(channels, channels, bias=False)
        self.fc2 = nn.Linear(channels, channels, bias=False)

    def forward(self, x):
        return self.fc2(self.fc1(x))


netA = Net().requires_grad_()
netB = Net().requires_grad_()

if share_fc2:
    # replace fc2 by netA.fc2
    netB.fc2 = netA.fc2

lossA = netA(training_data_A).mean()
lossB = netA(training_data_A).mean()

lossA.backward()
lossB.backward()

方案1：独立优化器 #

# OPTION 1 (independent run, seed used)
optA = torch.optim.Adam(params=netA.parameters())  # contains shared fc2
optB = torch.optim.Adam(params=netB.parameters())  # contains shared fc2
optA.step()
optB.step()

print(list(netA.parameters()))
print(list(netB.parameters()))

输出 #

netA与netB的fc1参数数值不同；共享的fc2参数经过两次Adam更新后，数值与方案2存在明显差异。

方案2：同一优化器 #

opt = torch.optim.Adam(
    params=(
        list(netA.parameters())  # contains netA.fc1, netA/B.fc2
        + list(netB.parameters())[:1]  # contains netA.fc1
    )
)
opt.step()
print(list(netA.parameters()))
print(list(netB.parameters()))

输出 #

netA与netB的fc1参数数值一致；共享的fc2参数仅经过一次Adam更新，数值与方案1不同。

疑问 #

• 差异是否源于Adam的内部参数调整，因此可以忽略？
• 差异是否源于梯度计算与更新的数学逻辑？
• 两种方案中哪一种更优？

解答 #

1. 差异根源：Adam内部状态与更新次数 #

差异完全源于Adam的内部动量状态以及参数更新次数，和梯度计算逻辑无关。

• 方案1中，共享的fc2参数被两个独立Adam优化器管理：optA和optB各自维护一套一阶/二阶动量状态。调用两次step()后，fc2会被连续更新两次——第一次用optA的初始状态计算更新，第二次用optB的全新初始状态再次计算更新，两次更新的逻辑完全独立。
• 方案2中，Adam会自动对重复参数去重，因此fc2仅被一个优化器管理并更新一次；netA.fc1虽被重复加入参数列表，但也只会被更新一次，符合正常的优化逻辑。

两种方案的梯度计算是一致的：lossA.backward()和lossB.backward()会将梯度累加在共享参数的.grad属性中，最终fc2和fc1的梯度都是两次损失的梯度之和。

2. 差异是否可忽略？ #

不能直接忽略。Adam的两次连续更新会改变参数的收敛轨迹：两次小步更新和一次基于累加梯度的大步更新效果完全不同，尤其是训练初期动量状态尚未稳定时，差异会被放大，直接影响模型的收敛速度和最终精度。

3. 哪种方案更优？ #

方案2是更合理的推荐选择，原因如下：

• 逻辑一致性：共享参数仅应该被更新一次，多个优化器重复更新会破坏Adam动量累积的设计逻辑，导致参数更新混乱。
• 收敛稳定性：单个优化器统一维护所有参数的动量状态，符合Adam的优化逻辑，训练过程更稳定。
• 资源高效：多个优化器会占用额外内存存储各自的动量状态，单个优化器更节省计算资源。

如果需要对netA和netB的独立参数（如fc1）设置不同优化配置，可以通过参数分组实现，无需使用多个优化器：

opt = torch.optim.Adam([
    {"params": netA.fc1.parameters(), "lr": 1e-3},
    {"params": netB.fc1.parameters(), "lr": 2e-3},
    {"params": netA.fc2.parameters()}  # 共享参数统一配置
])

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