模拟min(x,y)的PyTorch神经网络无法收敛，求故障原因

问题背景 #

我尝试用一个简单神经网络模拟最小值函数f(x,y)=min(x,y)，输入限制为0至100的整数（已归一化到0-1区间）。我手动设置了一组能让损失为0的权重，但模型无论怎么调整学习率，都无法收敛到这组权重值，想请教问题出在哪。

模型实现代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

class myMinFun(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(2, 4)
        self.lin2 = nn.Linear(4, 1)

    def forward(self, x):
        x = self.lin1(x)
        x = F.relu(x)
        x = self.lin2(x)
        return x

model = myMinFun()
crit = torch.nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)
max_int = 100
min_int = 0

for epoch in range(100):
    x1 = np.random.randint(low=min_int, high=max_int, size=10).astype(np.float32) / max_int
    x2 = np.random.randint(low=min_int, high=max_int, size=10).astype(np.float32) / max_int
    x = torch.tensor(np.column_stack((x1, x2)))
    y = torch.tensor(np.minimum(x2,x1)).unsqueeze(1)
    y_h = model(x)
    loss = crit(y_h, y)
    print(loss.item())
    model.zero_grad()
    loss.backward()
    optimizer.step()

已知最优权重（手动设置可使损失为0） #

with torch.no_grad():
    model.lin1.weight = torch.nn.Parameter(torch.tensor([[1,1],[1,-1],[-1,1],[-1,-1]], dtype=torch.float32), requires_grad=True)
    model.lin1.bias = torch.nn.Parameter(torch.tensor([0,0,0,0], dtype=torch.float32), requires_grad=True)
    model.lin2.weight = torch.nn.Parameter(torch.tensor([[0.5,-0.5,-0.5,-0.5]], dtype=torch.float32), requires_grad=True)
    model.lin2.bias = torch.nn.Parameter(torch.tensor([0], dtype=torch.float32), requires_grad=True)

无法收敛的原因分析 #

• ReLU导致的神经元“死亡”，梯度阻断：你手动设置的第一层权重里，[1,-1]和[-1,1]这两个神经元，在一半的输入场景下会输出负数，经过ReLU后直接变为0。比如当x>y时，[-1,1]的输出是y-x（负数），ReLU后为0，这个神经元相当于“罢工”，梯度无法反向传播到它的权重，自然没法更新到目标值。
• SGD批量太小，更新噪声大：你每次训练只用10个样本，随机梯度的噪声非常大，权重更新方向不稳定，很难精准朝着目标权重收敛。即使调整学习率，也抵消不了样本量太小带来的波动。
• 初始化与目标权重差异过大：PyTorch默认的线性层权重初始化（比如Xavier初始化）和你手动设置的权重值差异极大，加上ReLU的非线性作用，模型一开始就处于一个梯度信号很弱的状态，没法朝着目标方向优化。
• MSE损失在分界点的梯度不连续：min(x,y)是分段线性函数，在x=y的分界点处，损失函数的梯度是不连续的。SGD处理这种不连续的损失曲面时，更新会出现跳变，难以收敛到精确的目标权重。

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