PyTorch中含输出二阶导数的Loss不下降问题排查求助

我希望按如下形式计算损失：

损失由两部分组成：第一部分是样本x₁对应的预测值$\hat{u}{bc}$与精确值$u{bc}$的损失；第二部分是样本x₂对应的输出二阶导数预测值$\hat{u}''{r}$与精确值$u''{r}$的损失，x₁和x₂为不同样本。

我尝试通过以下代码实现：

# forward pass to calculate the first loss component
u_bc_pred = self.forward(self.X_u)
loss_bcs = self.loss_fnc(u_bc_pred, self.Y_u)

# the second loss component involves the second derivative of output
u_r_pred = self.forward(self.X_r)
u_x = torch.autograd.grad(u_r_pred, self.X_r, torch.ones_like(u_r_pred), retain_graph=True, create_graph=True)[0] / self.sigma_x
u_xx = torch.autograd.grad(u_x, self.X_r, torch.ones_like(u_x), retain_graph=True, create_graph=False)[0] / self.sigma_x
loss_res = self.loss_fnc(u_xx, self.Y_r)

# Total loss
loss = loss_res + loss_bcs
self.loss_log.append(loss.data)

# backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()

但代码运行后，涉及u_xx的第二个损失项并未下降。该网络用于学习标量函数，是单输入单输出结构，可视为[1, n, n, 1]的层结构。恳请帮忙排查实现含解导数的损失时是否存在错误。

以下是几个可能导致二阶导数损失不下降的原因及对应修正方案：

1. 二阶导数计算的梯度图被切断 #

计算u_xx时设置create_graph=False，会直接切断u_xx到网络参数的梯度传播链路，导致这部分损失无法反向更新模型参数。
修正：将u_xx计算中的create_graph改为True，保留梯度图以支持反向传播：

u_xx = torch.autograd.grad(u_x, self.X_r, torch.ones_like(u_x), retain_graph=True, create_graph=True)[0] / self.sigma_x

2. 梯度计算的权重逻辑不匹配 #

你使用torch.ones_like(u_r_pred)作为grad_outputs参数，若损失函数默认是平均（比如MSELoss()），全1的权重会让导数损失的求和逻辑与损失函数的平均逻辑冲突，导致梯度尺度异常。
修正：改用与损失函数匹配的权重，比如：

u_x = torch.autograd.grad(u_r_pred, self.X_r, torch.ones_like(u_r_pred)/u_r_pred.size(0), retain_graph=True, create_graph=True)[0] / self.sigma_x

或者直接省略grad_outputs参数（默认就是全1向量），再配合损失函数的求和模式（MSELoss(reduction='sum')）。

3. 导数的尺度归一化不匹配 #

你对预测的导数除以了self.sigma_x，但需确认精确值self.Y_r是否做了相同的归一化处理。如果只有预测导数做了缩放，标签未同步调整，会导致损失计算的尺度完全错位，模型无法有效优化。

4. 两个损失项的量级失衡 #

若loss_bcs的初始量级远大于loss_res，模型会优先拟合第一部分损失，导致二阶导数损失几乎得不到优化资源。
修正：给两个损失项添加权重系数平衡量级：

lambda_res = 10.0  # 根据实际情况调整
lambda_bcs = 1.0
loss = lambda_res * loss_res + lambda_bcs * loss_bcs

5. 网络结构的表达能力不足 #

单输入单输出的[1, n, n, 1]结构对二阶导数的拟合能力有限，尤其是使用ReLU这类二阶导数几乎为0的激活函数时。
优化建议：

• 增加隐藏层的神经元数量
• 替换激活函数为二阶导数更平滑的类型（比如Swish、GELU）

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