梯度估计STE

文章目录

1. 1. 背景
2. 2. 应用
   1. 2.1. 恒等函数
   2. 2.2. SG函数
3. 3. 参考资料

背景

反向传播是现在训练模型的重要方法，但是在部分场景下，会遇到不可微分的函数，从而导致梯度传播失败。比如量化里的取整函数。因此，需要对梯度进行估计然后反向传播。

STE(Straight-Through Estimator)是2013年Yoshua Bengio等人针对梯度估计进行了研究，那篇论文提出了几种梯度估计的方法，并推导出了一些理论性质，然后通过实验证明，STE是效果最好的方法。

由于那篇论文很多篇幅在介绍较为复杂的估计方法，且理论推导也极为复杂，效果没有简洁的STE好，因此不对其进行详细介绍。

应用

恒等函数

STE在反向传播理论推导的时候，把不可微的原子函数（比如量化函数里有放缩和取整两部分，其中取整是不可微的原子函数）替换为恒等函数。

这种应用可以大大减少理论推导的难度，但是在代码里应用反向传播的时候不太方便，以pytorch框架为例，可能需要自己手写函数类 torch.autograd.Function，具体文档可以查看：PyTorch: Defining New autograd Functions — PyTorch Tutorials 2.5.0+cu124 documentation 。

SG函数

苏神在他的博客（VQ-VAE的简明介绍：量子化自编码器 - 科学空间|Scientific Spaces）中提出了一个函数 sg(stop gradient)，代表梯度反向传播终止的恒等函数。

所以原来的不可微的原子函数可以写为\(f(x)=x+sg(x'-x)\) ，在前向传播的时候和x'相同，在反向传播的时候梯度等于直接对 x反向传播梯度。

这个式子也能用于理论推导，但是不如视为恒等函数麻烦，但是在代码方面容易完成，可以使用pytorch的默认反向传播函数，不需要自定义 torch.autograd.Function（当然自定义也可以完成任务）。

BitNet/bitnet/bitlinear.py at main · kyegomez/BitNet 代码示例：

# w_quant 是 w 量化（不可微）后的值
# STE using detach
w_quant = w + (weight_quant(w) - w).detach()

这里也给出 torch.autograd.Function的一个示例(OneBit/transformers/src/transformers/models/bitnet.py at main · xuyuzhuang11/OneBit)：

class SignSTEFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.sign(input)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * (1.001 - torch.tanh(input) ** 2)
        # return grad_output * (1.01 - torch.tanh(input) ** 2)

参考资料

• [1308.3432] Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• PyTorch: Defining New autograd Functions — PyTorch Tutorials 2.5.0+cu124 documentation
• VQ-VAE的简明介绍：量子化自编码器 - 科学空间|Scientific Spaces
• kyegomez/BitNet: Implementation of "BitNet: Scaling 1-bit Transformers for Large Language Models" in pytorch
• OneBit/transformers/src/transformers/models/bitnet.py at main · xuyuzhuang11/OneBit