为什么反向传播后未使用的embedding向量被修改？部分微调如何实现？

问题成因 #

• 权重衰减（weight_decay）的直接影响：你配置的weight_decay属于L2正则项，它的更新逻辑独立于梯度计算，会对所有注册到优化器的参数统一施加衰减，更新规则为param = param * (1 - lr * weight_decay)，即使非目标索引的embedding梯度为0，也会被该规则修改，这是你给出的示例代码中问题的核心诱因。
• 带状态优化器的累积更新：如果你使用Adam、RMSprop、带动量的SGD这类带有历史状态缓存的优化器，即使某一步参数的梯度为0，优化器依然会基于之前累积的一阶动量、二阶矩等状态对参数执行更新。
• 优化器作用范围未做限制：你将完整的embedding矩阵参数注册到了优化器中，优化器默认会遍历所有传入的参数执行更新逻辑，不会自动过滤梯度为0的参数。

可行的解决方法 #

方法1：替换无额外更新逻辑的优化器 #

如果业务允许不使用自适应优化器和L2正则，可以直接替换为不带动量的普通SGD并移除weight_decay配置，普通SGD仅基于梯度执行更新，梯度为0的参数不会产生变更：

# 替换原优化器定义即可
optim = torch.optim.SGD(params=embeds.parameters(), lr=0.001)

该方案改造成本最低，但限制较多，无法使用自适应优化器的加速效果和L2正则的泛化效果。

方法2：拆分训练参数，仅注册目标向量到优化器 #

将不需要更新的embedding部分冻结，单独把目标索引对应的向量拆分为独立可训练参数，训练完成后再赋值回原矩阵，完全避免其他位置被修改，适配所有优化器和正则配置：

import torch
from torch import nn

embeds = nn.Embedding(10, 3)
indexes = torch.tensor([1,2,3])
# 冻结全量embedding梯度
embeds.weight.requires_grad = False
# 拆分目标向量为独立可训练参数
trainable_emb = nn.Parameter(embeds.weight[indexes].clone())
optim = torch.optim.Adam(params=[trainable_emb], lr=0.001, weight_decay=1e-6)
bce_loss = nn.BCELoss()

optim.zero_grad()
loss = bce_loss(trainable_emb.sum(dim=-1), torch.zeros(3))
loss.backward()
optim.step()
# 更新完成后回写原矩阵
embeds.weight.data[indexes] = trainable_emb.data

方法3：手动过滤非目标位置的梯度 #

PyTorch 1.8及以上版本的优化器会自动跳过梯度为None的参数，你可以在反向传播之后、优化器step之前，将非目标索引对应的梯度设置为None，即可避免这些位置被更新：

optim.zero_grad()
loss = bce_loss(embeds(indexes).sum(dim=-1), torch.zeros(3))
loss.backward()
# 新增：将非目标位置的梯度设为None
mask = torch.ones(embeds.weight.shape[0], dtype=torch.bool)
mask[indexes] = False
embeds.weight.grad[mask] = None
optim.step()

该方案不需要修改原有的前向逻辑，改造成本较低，也适配各类优化器。

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