PyTorch实现简单线性回归无法收敛至正确形态的问题排查

我正在学习PyTorch，尝试自行实现简单线性回归。可视化结果中各元素含义：

• 红色：训练集
• 蓝色：测试集
• 橙色：每次迭代后的预测结果
• 黄色：最后一批数据的预测结果

观察发现橙色预测点试图穿过训练集红点中心，但斜率始终无法收敛。已尝试调整学习率、增加训练轮数、修改随机初始化方式、使用L1损失（绝对值损失），咨询ChatGPT 4o但得到幻觉回答。

感觉参数a和b存在类似囚徒困境的关系：a需增大的同时b需减小，怀疑当前处于局部最小值，请问如何跳出？求相关思路、资源链接或论坛推荐。

附代码：

import torch
import matplotlib.pyplot as plt

X = torch.linspace(2, 5, 50)
y = X * 3 - 10 + torch.rand(50)

X_train, X_test = X[:30], X[30:]
y_train, y_test = y[:30], y[30:]
plt.scatter(X_train, y_train, color="red")
plt.scatter(X_test, y_test, color="blue")

torch.manual_seed(43)
a = torch.tensor(6., requires_grad=True)
b = torch.tensor(-20., requires_grad=True)

for i in range(10):
    loss = (a * X_train + b - y_train).pow(2).mean()
    print(loss, a, b)
    y_predict = a * X + b
    plt.scatter(X, y_predict.detach(), color="orange")

    loss.backward()
    # break
    with torch.no_grad():
        a -= .03 * a.grad
        b -= .01 * b.grad
    a.grad.zero_()
    b.grad.zero_()

loss = (a * X_train + b - y_train).abs().mean()
display("loss", loss)

y_predict = a * X_test + b
plt.scatter(X_test, y_predict.detach(), color="yellow")

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