Stanford 2 :Neural Network Compression for Noisy Storage Devices

本文是作者上一篇文章的延续，原文标题

Neural Network Compression for Noisy Storage Devices

背景

在上一篇文章中，作者研究了先压缩再coding的方法来提升NN的存储效率并对抗噪声。本文中，作者exploit联合压缩和coding算法。

作者考虑一个多输入单输出的神经网络，训练数据从 $p_{\text{data}}(x,y)$ 中采样得到。注意这里分布unknown只知道data，而我们训练神经网络的目的就是学习到这个联合分布，更准确的说是学习到 $p_w(y|x)$.

本文研究的内容仍然是把一个神经网络的参数压缩，编码，存储 (噪声 $\hat{w}=w+ϵ$)，解码，重构。原始参数 $w$ 与重构出的参数 $g(\hat{w})$之间的距离有两个衡量标准：

1. MSE: ${\min }_g\Vert g(\hat{w})-w{\Vert }_2$;
2. KL divergence: ${\min }_g{\mathbb{E}}_x[D_{\text{KL}}(p_w(y|x)||p_{g(\hat{w})}(y|x))]$

技术

我们现在把cell想象成信道，输入是 $w$, 输出是 $\hat{w}=w+ϵ$. 那么直观的两种对抗噪声方式是(见上一篇文章)
a) 多次使用信道，即repetition code;
b) 提高scaling factor。

本文作者使用了连各种对抗噪声的方式：

1. sparsity driven. 这些就是上一篇文章中的那些方法，比如保护sign，比如adaptive mapping，再比如adaptive redundancy.
   跟上篇文章一样，一些 typical observations:

2. Sensitivity-driven protection.
   以下都是本文的新技术，但实际上提升不大。比如这里说的sensitivity-driven protection 实际上就是identify出哪些weights是sensitive的，然后对他们用更多的cell保护而已。实验中它的improvement最多是2.8%。而且这是在1cell/weight的情况下，如果我们提升number of cells 这种方法几乎没有增益。但是我们还是写一写。

作者大概的思路就是想看看什么样的weights对KL divergence影响最大。假设weights的改变是
$$g(\hat{w})=w+{\delta }_w$$, 那么KL divergence可以用二阶泰勒展开近似为

$${\mathbb{E}}_x[D_{\text{KL}}(p_w(y|x)||p_{w+{\delta }_w}(y|x))]\approx {\delta }_w^{\top }\mathbf{F}{\delta }_w$$

其中 $\mathbf{F}$ 是 fisher information matrix：
$$\mathbf{F}={\mathbb{E}}_{x,y}[{\nabla }_w\log p_w(y|x){\nabla }_w^{\top }\log p_w(y|x)]$$

忽略掉 off-diagonal elements, KL 的近似变为
$${\delta }_w^{\top }\mathbf{F}{\delta }_w\approx {\mathbb{E}}_{x,y}\sum _{j=1}^d{\left({\delta }_{w_j}\frac{\partial \log p_w(y|x)}{\partial w_j}\right)}^2=\sum _{j=1}^d{\delta }_{w_j}^2{\mathbb{E}}_{x,y}{\left(\frac{\partial \log p_w(y|x)}{\partial w_j}\right)}^2$$
其中 $d$ 是 weights的个数。可以看到，不同的weights $w_j$, 它上的扰动 ${\delta }_{w_j}$对KL 的影响取决于后面那个 expectation项。作者把这一项称为 sensitivity of $w_j$ 并对所有的weights都算出一个sensitivity，然后根据sensitivity的大小来分配cell 个数。 over。

唯一的一个问题是如何计算 ${\mathbb{E}}_{x,y}{\left(\frac{\partial \log p_w(y|x)}{\partial w_j}\right)}^2$.

作者说就是一个在所有training data上算出的一个gradient，这是很显然的。但是这个代价函数岂不是maximizing $\log p_w(y|x)$？并不是很清楚maximizing一个prob distribution是啥意思。

Robust Training and Distillation

除了以上coding技术之外，作者还提出两个NN训练方法以加强最后的模型对噪声的robustness。这应该就是作者所说的joint optimization了。这两个方法本身并不新鲜，在Zhou 2020 那篇里面就有了，作者应该是做了一些调整，意义不大不讲了。