PyTorch中能否将可训练参数作为指数使用？训练出现nan如何解决？

核心结论 #

PyTorch原生支持将可训练参数作为指数的运算，你遇到反向传播后参数变为nan的问题和算子支持无关，属于数值稳定性问题。

问题根因 #

torch.pow(T, x)对参数x的梯度计算公式为 d(T^x)/dx = T^x * ln(T)，以下场景都会触发nan或者梯度爆炸：

• T中存在小于等于0的元素，此时ln(T)在实数域无定义，直接返回nan
• T中存在极接近0的正数，ln(T)会输出绝对值极大的负数，引发梯度爆炸，参数更新后直接溢出为nan
• T中存在极大的正数，T^x计算时直接溢出为inf，后续传导变为nan

修复方案 #

• 预处理张量T，确保所有元素严格为正且不接近0：可以对T做截断T = T.clamp(min=1e-8)或者加极小正偏移T = T + 1e-8；如果T本身存在负数，需要先确认业务上负数作为底数是否合理，实数域内负数的非整数次幂无定义，这种情况可以先对T取绝对值再运算
• 合理初始化参数x，避免初始值过大导致T^x溢出，通常初始值设置在0~1区间即可适配大多数场景
• 增加梯度裁剪操作，反向传播后、优化器更新前对x的梯度做截断，避免梯度爆炸冲毁参数
• 可将指数运算改写为等价的对数形式，进一步提升数值稳定性：T^x = torch.exp(x * torch.log(T))，和直接调用torch.pow数学等价，更方便做中间数值校验

可运行示例代码 #

import torch

# 模拟输入张量T，确保全正
T = torch.randn(8, 8).abs()
# 截断避免数值过接近0
T = T.clamp(min=1e-8)

# 定义可训练参数x，初始值设为1.0
x = torch.nn.Parameter(torch.tensor(1.0))
optimizer = torch.optim.Adam([x], lr=1e-4)

# 前向传播
output = torch.pow(T, x)
# 等价对数形式写法，数值稳定性更高
# output = torch.exp(x * torch.log(T))

# 计算损失（此处为示例，替换为实际业务损失即可）
loss = output.mean()
# 反向传播
loss.backward()

# 梯度裁剪，避免梯度爆炸
torch.nn.utils.clip_grad_norm_([x], max_norm=1.0)
# 更新参数
optimizer.step()

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