pytorch实现简单的straight-through estimator(STE)
现在深度学习中一般我们学习的参数都是连续的,因为这样在反向传播的时候才可以对梯度进行更新。但是有的时候我们也会遇到参数是离散的情况,这样就没有办法进行反向传播了,比如二值神经网络。本文中讲解了如何用pytorch
对二值化的参数进行梯度更新。
Question
提示
: 保证公式能正常显示请使用google chromeSTE
核心的思想就是我们的参数初始化的时候就是float
这样的连续值,当我们forward
的时候就将原来的连续的参数映射到{-1,, 1}带入到网络进行计算,这样就可以计算网络的输出。然后backward
的时候直接对原来float
的参数进行更新,而不是对二值化的参数更新。这样可以完成对整个网络的更新了。
首先我们对上面问题进行一下数学的讲解。
- 我们希望参数的范围是$r \in \mathbb{R}$
- 我们可以得到二值化的参数 $q = Sign(r)$, $Sign$函数可以参考
torch.sign
函数, 可以理解为取符号函数 backward
的过程中对$q$求梯度可得 $\frac{\partial loss}{\partial q}$- 对于$\frac{\partial q}{\partial r} = 0$, 所以可以得出 $\frac{\partial loss}{\partial r} = 0$, 这样的话我们就无法完成对参数的更新,因为每次
loss
对r
梯度都是0 - 所以
backward
的过程我们需要修改$\frac{\partial q}{\partial r}$这部分才可以使梯度继续更新下去,所以对$\frac{\partial loss}{\partial r}$进行如下修改: $\frac{\partial q}{\partial r} = \frac{\partial loss}{\partial q} * 1_{|r| \leq 1}$, 其中
$1_{|r| \leq 1}$ 可以看作$Htanh(x) = Clip(x, -1, 1) = max(-1, min(1, x))$对$x$的求导过程, 也就是是说:
$$\frac{\partial loss}{\partial r} = \frac{\partial loss}{\partial q} \frac{\partial Htanh}{\partial r}$$
Example
torch.sign
首先我们验证一下使用torch.sign
会是参数的梯度基本上都是0:
demo
我们需要重写sign
这个函数,就好像写一个激活函数一样。先看一下代码, github源码:LBSign.py
接下来我们做一下测试main.py
然后我们发现有梯度了
explain
接下来我们对代码就行一下解释pytorch文档链接:
- forward中的参数
ctx
是保存的上下文信息,input
是输入 - backward中的参数
ctx
是保存的上下文信息,grad_output
可以理解成 $\frac{\partial loss}{\partial q}$这一步的梯度信息,我们需要做的就是让
$$gradoutput * \frac{\partial Htanh}{\partial r}$$ 而不是让pytorch
继续默认的 $$gradoutput * \frac{\partial q}{\partial r}$$
但是我们可以从上面的公式可以看出函数$Htanh$对$x$求导是1, 当$x \in [-1, 1]$,所以程序就可以化简成保留原来的梯度就行了,然后裁剪到其他范围的。
reference
Custom binarization layer with straight through estimator gives error