PyTorch优化器致损失上升、训练步非最优问题求助

训练过程中遇到异常：优化器的更新步骤不仅没帮助模型收敛，反而和损失函数的预期作用完全相反。为了测试模型的基本拟合能力，我把数据集所有标签设为固定值0.86，验证模型能否稳定预测这个常量。但无论使用MSELoss还是L1Loss，训练时模型输出在短暂接近0.86后，会逐渐向1收敛，同时损失持续上升。

模型定义 #

class DualBertForClassification(nn.Module):
    def __init__(self, bert_model_a, bert_model_b):
        super(DualBertForClassification, self).__init__()
        
        self.bert_model_wt = bert_model_a
        self.bert_model_mutant = bert_model_b
        self.layer_1 = nn.Linear(1024, 512)
        self.layer_2 = nn.Linear(512, 128)
        self.layer_3 = nn.Linear(128, 16)
        self.layer_4 = nn.Linear(16, 1)

    def forward(self, x):
        x_a = x[0]  
        x_b = x[1]
        x = torch.cat(
            (
                self.bert_model_wt(**x_a).last_hidden_state, # [batch_sz,sequence_sz,1024]
                self.bert_model_mutant(**x_b).last_hidden_state 
            ), 
            1
        )  # [batch_sz, 2*sequence_sz, 1024, ]
        x = torch.tanh(self.layer_1(x)) # [batch_sz, 2*sequence_sz, 512, ]
        x = torch.tanh(self.layer_2(x)) # [batch_sz, 2*sequence_sz, 128, ]
        x = torch.tanh(self.layer_3(x)) # [batch_sz, 2*sequence_sz, 16, ]
        x = torch.tanh(self.layer_4(x)) # [batch_sz, 2*sequence_sz, 1, ]
        x = torch.mean(x, dim = 1)  # [batch_sz, 1]
        return x

bert_model_a与bert_model_b为相同的预训练BERT模型，输出尺寸如代码注释所示。

训练循环代码 #

model1 = BertModel.from_pretrained(model_name)
model2 = BertModel.from_pretrained(model_name)
model = DualBertForClassification(model1, model2)

optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)

loss_fct = nn.L1Loss(reduction="sum") # 也试过nn.MSELoss(reduction="sum")
model.train()
for index, (input_a, input_b) in enumerate(zip(wt_inputs, alt_inputs)):
    label_val = torch.tensor([[0.86]], dtype = torch.float32)
    model_output = model((input_a, input_b))
    loss = loss_fct(model_output, label_val)
    print(loss, model_output, label_val)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

我怀疑问题出在autograd对两个BERT模型的处理逻辑上，但无法解释损失持续上升的具体原因。已经尝试修改forward方法末尾的输出层激活函数，问题依然存在，恳请帮忙排查解决。

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