深入浅出Sampling-Argmax

说来惭愧，《Localization with Sampling-Argmax》这篇NeurIPS 2021的论文我非常喜欢，但是写的笔记很多小伙伴反馈说还是看不懂，我自己也感觉其中有很多地方当时理解不够到位，所以大家看不懂其实是因为我自己不够懂。于是最近又重新研读梳理了一下这篇论文的脉络，写一篇新的解读，将一些前置知识补齐，力求做到深入浅出。

由于故意写得比较浅显易懂，对于IPR有一定了解的同学可以跳过你熟知的部分。

0. Argmax的问题

要理解一篇论文，首先最重要的是要理解它试图解决一个什么样的问题，这往往也是论文作者跟读者之间最大的壁垒，因为站在作者的视角，很容易注意力全部集中在解释如何解决问题上，而忽视掉发现这个问题所需要的前置知识有哪些，理解“哪些前置知识是我具有而读者不一定具有的”，这是文章容易被理解的关键。

IPR方法由来已久，它的本质是对Argmax操作的光滑，将不可导的Argmax变为可导的Soft-Argmax，从而实现端到端的训练。

说到Argmax，我想从我们最熟悉的分类问题说起。

其实在分类问题中我们就已经开始用到了Argmax，只不过我们只在推理阶段使用：对于N个类的分类问题，模型会输出一个长度为N的向量，我们用Argmax来获取响应值最大的位置，作为模型分类结果的预测。

按照正常的思路，模型的训练和推理应当是一致的，因为只有这样才可以保证我们的训练目标和我们的推理目标绝对统一，一旦不一致，各种问题将层出不穷。

所以按照这个思路，我们最初的训练流程应该是：模型输出一个长度为N的向量，用Argmax获取最大值位置，然后损失函数负责计算和标注的差异。

这个流程为什么行不通呢，因为Argmax是不可导的，梯度流会在这里被截断，无法完成反向传播，所以模型无法训练。

而解决方案是，既然经过了Argmax之后的信息无法进行监督，所以我将监督的位置放到Argmax之前，在那个N维向量上做监督，让这个模型预测的N维向量去拟合一个one-hot的概率分布，从而绕开了Argmax不可导的问题。

而这放在姿态估计任务上，正是Heatmap-based方法的做法：

1. Soft-Argmax诞生

分类方法的思想来到姿态估计任务后，摇身一变成为了Heatmap-based方法，也确实表现不凡，但还是不可避免地有些水土不服：量化误差无可避免。

因为二维的Heatmap计算量庞大，Heatmap-based方法在使用时往往尺寸是小于输入图片的，这就导致Argmax的结果一定无法跟标注信息一一对应，这称为量化误差。

为了解决量化误差问题，大家开始动心思如何把监督信号挪到Argmax后面，而最直接的做法就是进行光滑化，用一个可导的函数来近似，于是Soft-Argmax应运而生。

分类的本质是什么？其实是把输出的向量看成一个离散概率分布，每一个像素是这个离散分布的样本，这个像素上的响应值是概率分量的值，而Soft-Argmax的公式正是求离散概率分布期望（均值）的公式。

用Soft-Argmax能取得跟Argmax近似的结果，最重要的是，这个公式是可导的，所以我们把Argmax换成Soft-Argmax后，监督信号可以移到流程的最后了（也就是我们最初设想的地方）。

最初提出这种方法的论文给Soft-Argmax起名叫DSNT：

2. Soft-Argmax的问题

Soft-Argmax是存在问题的，而问题产生的根源是它让监督信号跟模型学习的目标不一致了。

我们的模型一直在学习如何建模关键点位置的概率分布函数，学习的方式有两种：

一是学习概率分布函数的均值和方差，这一般是用在已经假设好分布函数的标准形状时，均值和方差只是对标准形状进行平移和拉伸。

二是学习概率分布函数的概率分量，也就是让模型把概率分布函数图像给“画”出来。

先是由于输出结果是一个一个的像素串成的向量，所以是离散的，所以模型画出来的图像只能是连续分布的离散采样，离散分布跟连续分布相比，在这一步就已经丢失了一部分信息。

接下来到了Soft-Argmax，在训练和推理时，我们都只关心离散概率分布的均值，而对分布长什么样毫不在意，这仿佛就是在说“我就想当县长夫人，谁是县长我不在乎”。

监督信号跟学习目标不一致，一个均值可以对应无数种分布，所以模型学习的目标不明确，效果大打折扣。

3. 形状约束

这个问题在DSNT论文中就已经被作者意识到了，而他提出的改进方法是，给模型学习的概率分布加上先验约束。

什么叫加先验约束，所谓先验，就是我们对于目标的一种假设，一种信念，我们相信目标符合某种分布，或相信目标的分布与我们已知的某种分布近似。

在分类问题中，我们让模型拟合one-hot向量，其实是在拟合类别分布（categorical distribution），到了Heatmap-based方法，我们又人工生成高斯热图，是在让模型拟合高斯分布。

中心极限定理表明，只要你的样本够大，即使群体不服从高斯分布，平均值的分布也会服从高斯分布，所以高斯分布成为万金油，分布假设中的常客——它也许不够好，但它一定不会错，这样一种下限的保障。

所以我们可以增加一个监督信号，在模型预测的离散概率分布上增加一个监督，让模型预测的离散分布拟合高斯分布，如此一来约束了分布的形状。

4. 先验分布的局限性和改进

就像我上面说的，高斯分布属于是一种下限的保障，真实的数据分布那么复杂，要想让模型学得够好，光是拟合高斯分布是不够的，于是我们开始思考如何让模型拟合数据的真实分布。

数据的真实分布是复杂的，复杂到我们写不出公式，写不出公式就没办法给模型做示范，没办法人工生成学习对象扔给模型说“照着这个形状学就行”。

不过大家也不是这么容易放弃的。虽然我不知道真实分布的形状是什么样，但我至少知道这种形状肯定是错的：

虽然我给不出正确答案，但好歹可以去掉一些错误答案嘛，这样对于模型的学习应该也是有帮助的，你先给我把这个性质学上：

• 均值所在的概率分量应该是最大的

5. 新的约束

于是我们现在思考的问题变成：如何在不知道真实分布形状的情况下，让概率分量集中在均值？

论文作者给出了他的思路：把“期望的误差”变成“误差的期望”。

在过去的传统做法中，我们都是先对概率分布求期望，再计算期望与GroundTruth的距离损失，公式可以写成：

而误差的期望则是改变顺序，先计算误差，再对误差求期望：

看上去只是计算顺序发生的变化，旧方法的问题在上面已经分析得足够了，新的方法有何不同呢？让我们仔细品一下这个公式：

其中 是模型预测的每个像素上的概率分量，而 里所有内容都是固定的，我举一个例子：假如在一个长度为4的离散概率分布上，目标为2，距离函数为L1距离，那么这个公式的结果会是：

又因为概率分量之和为1，也就是说，要降低这个损失，模型必须尽量多地把概率分量集中到0所在的位置上，也就是GT所在的位置的 越大越好。

因此，这个损失函数能够在完全不知道目标分布形状的情况下，约束分布形状，让分布呈现出“中间大周围小”的尖锐形状。

6. 新损失函数的问题

通过上面的例子我们明白，新的损失函数的确可以满足我们的要求，然而用这个损失函数训练的效果却非常差：普通的IPR能在COCO上取得64.5AP，而新损失函数不仅收敛慢，最后结果只有30.9AP。

问题出在哪里呢？ 的方差太大了。

在梯度估计方法中，有一类经典的梯度估计方法叫SFGE(Score Function Gradient Estimator)，其梯度形式与我们的新损失函数梯度形式相同。SFGE在满足求导几分可换的条件下，是一个无偏的估计，但是其方差：

当 中包含了很多与梯度无关的项时，方差将会显著变大，因此SFGE通常会通过减去一个baseline的方式来降低方差。

7. 采样方式进行梯度估计

要直接解决新损失函数的梯度问题是困难的，但是我们还有另一条路，也就是通过采样来进行梯度估计。

什么叫通过采样来进行梯度估计？

假如我们有一个长度为3的离散概率分布，概率分量为[0.7, 0.2, 0.1]，那么按照新损失函数计算为：

如果我们把 看成训练样本, 假如我们有 10 个训练样本, 其中7个是 个是 个 是 , 那么这 10 个训练样本的平均损失会是:

我们可以发现，这两种方法得到的损失是相等的，只要样本出现的概率，等于离散分布的概率分量，换句话说，只要能从这个离散概率分布中随机采样，我们就可以近似得到相同的梯度信息。

那怎么让采样具有随机性、满足特定的概率分布呢？很简单，我们可以按照[0.7, 0.2, 0.1]的比例划分一个转盘，然后随机转到哪块区域就生成一个对应区域的样本，写成代码的话可以是：

p = np.random.rand()  
  
if p < pi[0]:  
    y = 0  
elif p < pi[1]:  
    y = 1  
else:  
    y = 2 

随机性有了，但这种采样是不可导的，有什么办法可以让采样既可导又有随机性？

说白了，采样的意思就是“确定样本是谁”，要可导的关键是把模型输出的参数写进计算公式里，“这个样本是谁”要通过跟模型输出计算来得到，而不是大于小于号分支语句，于是这里要引出我们的重参数技巧。

8.重参数技巧

所谓重参数技巧(reparameterization trick)，又叫参数重整化，是变分推断的一个技巧，在著名的VAE中就有用到。

需要注意的是，近年比较火爆的RepVGG使用的技巧叫结构重参数化，两者都常被简称为重参数化，比较容易误导。

变分推断的重参数技巧，简单来说，就是把采样步骤移出计算图，先从一个标准分布中随机采样，然后通过运算来得到符合我们模型输出范围的结果。

以VAE为例, 模型预测一个高斯分布的均值和方差, 但我们不能从这两个参数决定的高斯分布 中直接计算出随机样本编号。所以我们可以先从一个标准正态分布中采样一个 , 然后计算

如此一来, 在整个计算图中 相当于一个常数, 因此可以反向传播。

以上是连续高斯分布的重参数技巧，而离散分布的重参数技巧可以用Gumbel-Softmax来解决。

我们可以不用管什么是Gumbel分布，不影响理解，只需要知道从它里面采样后通过一顿计算可以得到特定离散概率分布的近似样本，实现跟VAE中获取高斯分布采样结果一样的效果。

9. Gumbel-Max & Gumbel-Softmax

在讲Gumbel-Softmax之前，我们需要先了解一下它的不可导版本，Gumbel-Max来铺垫一下，它的式子为：

而OneHot(Argmax(*))操作是不可导的，于是我们可以用Softmax来近似替代，变为

通过调节平滑系数 , 我们就可以改变样本的形状, 系数越大, 采样出来的向量越接近于均匀 分布; 系数越小, 计算出来的向量约接近于one-hot向量, 比如可能是 , 如果说 代表了1个 , 它的贡献度是 ，那么 可以看成是 个 个 个 。

由于在训练初期模型预测的概率值也不太靠谱，所以我们可以用比较大的平滑系数，来生成一些比较均匀分布的向量，也就是每个样本的贡献度差不多。等到训练后期，再减小系数，最后基本是近似one-hot的向量。

在有了Gumbel-Softmax后，我们已经可以通过模型预测的离散概率分布函数，来获得无限多个服从该分布的样本了。也是说，我们可以通过：

来近似我们一开始的目标函数了：

又因为Gumbel-Softmax得到的样本毕竟只是高度近似，还不是完全的one-hot，所以需要通过加权来得到样本，这里为了避免混淆，我用 来表示Gumbel-Softmax生成的向量中的每个位置的分量：

写成代码会是：

hm = model(x)  
w_x = torch.arange(hm.shape[-1])  
loss = 0.  
for i in range(10):  
    g = gumbel_softmax(hm, tau=_tau)  
    y_hat = (g * w_x).sum()  
    loss += d(y_gt, y_hat)  
loss /= 10. 

10. 真实概率分布函数

其实故事到了这里已经比较完整了，但它有一个比较明显的缺点：上面的一整套流程都是在离散概率分布上进行的，但是真实的数据概率分布应该是连续空间上的，比如训练时的监督信息 ，更可能是一个小数而不是整数。

所以我们还可以尽量想办法通过离散分布来逆推连续空间里的概率密度函数。

香农采样定理标明，采样率必须要达到目标频率的两倍，目标函数才能被完美重建，所以我们离散分布重建连续函数肯定是不现实的，但我们可以退而求其次追求一个尽量近似的结果。

假如真实的概率分布函数长这样：

那么我们模型预测出来的离散分布其实长这样：

本文将真实概率分布看成连续空间上的混合分布，而混合分布密度函数等于其所有子分布的密度函数的加权和，写成公式为：

每一个 都是一个子分布的密度函数，为了简单其见，我们可以让所有子分步的标准密度函数相同，换句话说，各个子分步方差相同，均值不同，于是有：

用示意图表示为：

如此一来，我们就可以通过离散概率分布逆推重建一个近似的连续分布。而子分布的形状比较简单，论文中实验了三种最常见的分布：

• 均匀分布（矩形）
• 三角分布（三角形）
• 高斯分布（钟型）

11. 连续空间采样过程

将离散分布推广到连续空间上后，原来的采样过程也需要进一步完善了。

在上面的流程中，通过Gumbel-Softmax我们只能从离散分布中采样，样本也是离散值，而现在我们的密度函数是连续空间上的，我们希望能采样到连续空间上的值。

有了混合分布的假设后，其实后面的采样过程不难理解。混合分布的假设让我们把连续空间平分成了N个等宽的区间，所以一开始的离散采样的结果，相当于是在选择从哪一个子区间里采样。

又因为一个子区间对应了一个子分步，所以我们只需要在这个子分步上进行采样就行了。

而在已知形状的连续分布上采样，我们在VAE的例子里已经学过了，如果子分步是高斯分布，那么我们就从标准高斯分布中采样；如果子分步是均匀分布，我们就从标准均匀分布采样，然后通过对采样结果加上均值得到等价于子分步采样的结果。

给一个均匀分布的例子写成代码：

# 从均匀分布采样出(-0.5, 0.5)的随机数  
eps_x = torch.rand_like(w_x) - 0.5  
# 生成对应子分步的均值  
w_x = torch.arange(hm_x.shape[-1])  
# 给采样结果加上对应位置的偏移，此时的结果等价于从每个子分步采样一次的结果拼成的向量  
w_x = w_x + eps_x 

所以连续空间采样的过程有两步：

1. 通过Gumbel-Softmax离散采样确定样本来自哪个子区间
2. 从子区间对应的子分步进行采样

最后还别忘了，我们是通过连续采样来实现近似梯度估计的，所以要多次取样求平均，以获得损失函数的近似。

12. 尾声

总结一下，这篇论文通过简单调整公式，让我们可以对原本不知道形状不知道公式的概率分步进行形状约束，其中的数学原理是很有意思的。

而通过采样来进行梯度估计的方法同样很神奇，让我们可以获得良好的梯度。

最后用离散分布逆推重建连续分布让我们得到了一个近似的真实分布函数。

回顾我们一开始手里有的东西，只有目标点的标注而已，既不知道真实分布形状，也没法稳定训练添加约束，这些都被本文一步一步解决了。

希望这篇笔记可以帮助大家更好地理解Sampling-Argmax，如果读完让你有收获，欢迎点一个赞支持一下~

参考文章

关于Gumbel-Softmax解决离散输出不可反向求导的思考：https://zhuanlan.zhihu.com/p/516117090

梯度估计——初步：https://zhuanlan.zhihu.com/p/104991140