梯度累积训练效果与批量训练不一致的问题排查求助

作为个人练手项目，我正在尝试实现一个Transformer模型来预测一组数字的最小值。用常规批量训练时效果很不错，但切换到梯度累积策略后就出问题了——模型完全学不到正确的最小值，只会输出最小值的均值。我已经试过很多调整手段（比如把损失除以累积步数、调整学习率、打印梯度查看情况等等），但都没有任何改观。另外我也确认过，我的模型里没有批量归一化层（就是普通的Transformer结构后面接了一个简单的线性层，全程没有用任何归一化操作）。

下面是我的训练循环代码：

model = MinValuePredictor().cuda()
model.train()
loss_fn = nn.MSELoss()
accum_steps = 8 # 要测试大批量训练的话可以把这个值改成1
optimizer = optim.AdamW(model.parameters(), lr=0.000001)
# 训练循环
input_count = 10
epochs = 50000
batch_size = 1
optimizer.zero_grad()

for epoch in range(epochs):
    # 生成数据
    inputs, targets = generate_data(batch_size, input_count)

    # 前向传播
    predictions = model(inputs).squeeze(1)

    loss = loss_fn(predictions, targets) / accum_steps
    if epoch % 100 == 0:
        print(f"Difference: {predictions.item() - targets.item()}")

    # 反向传播
    loss.backward()
    loss.retain_graph = False

    # 每8个epoch执行一次优化器更新并清空梯度
    if (epoch + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()  # 清除累积的梯度

    if (epoch + 1) % 50 * accum_steps == 0:
        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}")


# 测试模型
test_input, test_target = generate_data(batch_size=1)
test_output = model(test_input).item()
print(f"Test Input: {test_input}")
print(f"Predicted Min: {test_output:.4f}, Actual Min: {test_target.item():.4f}")

我的随机数生成逻辑如下：

# 最大值固定为30
def generate_data(batch_size=64, n=10):
    min_values = 10 + (5 * torch.rand(batch_size, 1).cuda())  # 每个样本随机生成最小值
    max_values = 30 * torch.ones(batch_size, 1).cuda()  # 固定最大值
    input_tensors = (min_values + (max_values - min_values) * torch.rand(batch_size, n).cuda())
    targets = torch.min(input_tensors, dim=1).values
    return input_tensors, targets

就像我之前说的，当批量大小设为8时，不管累积步数是多少，代码都能正常工作；但当批量大小设为1时，无论累积步数设成多少，模型的表现都和批量大小、累积步数都是1的情况一模一样。有没有大佬能帮我看看我到底哪里写错了？

备注：内容来源于stack exchange，提问作者Davide