为何输入的梯度为零？三种修复方案原理探究

运行以下代码后，会出现input.grad全为0的情况：

import torch
import torch.nn as nn

DATASET_SIZE = 10
TUNING_EPOCH = 10
RANDOMIZATION = "rand"
shape = (10000,)
bsize = 1
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

input = eval(f"torch.{RANDOMIZATION}")(bsize, *shape).to(device).requires_grad_()
numel = input.numel()
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(numel, numel//2),
    nn.Tanh(),
    nn.Linear(numel//2, 1)
).to(device)
model_optimizer = torch.optim.Adam(model.parameters())
optimizer = torch.optim.Adam([input])
dataset = [(eval(f"torch.{RANDOMIZATION}_like")(input), (torch.rand(bsize,1)+8).to(device)) for i in range(DATASET_SIZE)]

model.train()
for e in range(TUNING_EPOCH):
    for x, y in dataset:
        model_optimizer.zero_grad()
        output = model(x)
        loss = torch.nn.MSELoss()(output, y)
        loss.backward(inputs=list(model.parameters()))
        model_optimizer.step()

model.eval()
optimizer.zero_grad()
output = model(input)
output.backward(inputs=input)
print("\ninput.grad", input.grad)

存在三种独立修复方案，各自解决零梯度问题的原因如下：

1. 设置TUNING_EPOCH = 0（不训练模型） #

模型未训练时，参数处于随机初始化状态，第一个Linear层的权重不会出现极端值，输入经过Linear层后的输出不会触发Tanh激活函数的饱和区（Tanh在输入绝对值过大时导数趋近于0）。此时梯度可以正常从输出层反向传递到输入，不会出现梯度消失导致input.grad全为0的情况。

2. 设置RANDOMIZATION = "randn"（更换输入分布） #

原代码使用torch.rand生成的输入是[0,1]区间的均匀分布，结合高维度（10000维）的输入，经过第一个Linear层后，输出值会因为权重初始化的特性累积成较大的绝对值，直接让Tanh进入饱和状态，导数趋近于0，导致梯度无法传递到输入。换成torch.randn（标准正态分布，均值0、方差1）后，输入分布更集中在0附近，经过Linear层后的输出不会轻易触发Tanh饱和，梯度得以正常传递。

3. 设置shape = (1000,)（缩小输入尺寸） #

输入维度从10000缩小到1000后，第一个Linear层的权重数量大幅减少，输入经过Linear层后的输出值不会因为高维度的累积效应变得极端，Tanh激活不会进入饱和区。同时，高维度场景下，梯度会因为链式法则的多次连乘快速衰减消失，缩小维度后梯度传递的“链路”更短，衰减程度降低，有效梯度能够保留并传递到输入。

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