WGAN 和 Wasserstein distance

在文献 [2] 中作者分析了 GAN [1] 难以训练的原因在于原生 GAN 的目标函数等价于优化生成数据的分布 $p_g$ 和真实数据的分布 $p_r$ 之间的 J-S 散度 (Jensen–Shannon Divergence)。

接着作者提出 WGAN [3]，使用性质优良的 Wasserstein distance 代替原生 GAN 中的 J-S 散度。然后利用KR对偶原理将 Wasserstein distance 的求解问题转换为求解最优的利普希茨连续函数的问题。为了使得判别器 D 满足利普希茨连续性，作者使用“梯度裁剪”将过大的参数直接裁剪到一个阈值以下。

本文要介绍的这篇文章 “Spectral normalization for generative adversarial network” (以下简称 Spectral Norm) 使用一种更优雅的方式使得判别器 D 满足利普希茨连续性。

为了说明 Spectral Norm背后的动机，我们从 Wasserstein distance (以下简称 W distance)开始说起。 W distance的定义为：

\begin{equation} \begin{split} \text{Was}(p_r, p_g) &= \inf_{\gamma \in \prod(p_r, p_g)} \int_{x, y} \gamma(x, y) \cdot \lVert x-y \rVert \\ &= \inf_{\gamma \in \prod(p_r, p_g)} \mathbb{E} [\lVert x - y\rVert] \end{split} \label{eq:wasserstein} \end{equation}

其中 $p_r, p_g$ 为两个概率分布，$\gamma(x, y)$ 为边缘分布分别为 $p_r, p_g$ 的联合概率分布，$\lVert x-y \rVert$ 为事件 $x, y$ 之间的某种距离。公式$\ref{eq:wasserstein}$ 的定义看起来比较难懂，具体 Wasserstein distance 的定义和求解也不是本文的重点，强烈推荐大家看我的另一篇文章来详细了解 Wasserstein distance。在这里，我们只需要了解 Wasserstein distance 是一种性质优良的衡量两个概率分布之间差异的方式。

Wasserstein distance 的求解并不那么直接，式 $\ref{eq:wasserstein}$ 所描述的问题一般可以通过线性规划 (Linear programming) 求解，因为问题中的求解目标以及约束条件均为线性。即便如此，在每一次训练的迭代去求解一个线性问题仍然是非常耗时的，而且我们并不需要得到最优传输，我们只需要得到最优传输对应的最小传输代价，以及该代价的导数。在文献 [3] 中作者使用了 K-R 对偶 (Kantorovich-Rubinstein Duality) 将 wasserstein distance 求解的问题转换为求解最优函数的问题。 K-R 对偶指出，当测度 $\lVert i-j \rVert = (i-j)^2$ 的时候，$\ref{eq:wasserstein}$ 式所对应的问题可以转换为寻找一个满足利普希茨连续 (Lipschitz continuity) 的函数 $f(x)$，使得下式取得最大值：

\begin{equation} \text{Was}(p_r, p_g) = \sup_{\lVert f \rVert\_{L \leq 1}} \mathbb{E}_{x \sim p_r} f(x) - \mathbb{E}_{y \sim p_g} f(y) \label{eq:dual-problem} \end{equation}

也就是说，为了求概率分布 $p_r$ 和 $p_g$ 之间的 Wasserstein distance，可以先从两个分布中分别采样出大量样本 $x, y$，然后找到一个满足普希茨连续性的函数 $f$，并求解 $\mathbb{E}_{x \sim p_r} f(x) - \mathbb{E}_{y \sim p_g} f(y)$ 最大值，此最大值即为 $p_r$ $p_g$ 之间的 Wasserstein distance。至于 $\ref{eq:dual-problem}$ 式的结论为什么成立可以参考这篇文章。本文关注另外一个问题，就是如何让函数 $f$ 满足“利普希茨连续性” (Lipschitz continuity)。

利普希茨连续性 (Lipschitz Continuity)

利普希茨连续性 (Lipschitz continuity) 的定义很简单，他要求函数 $f(x)$ 满足对任意 $x, y$ 都有:

\begin{equation} \frac{f(x) - f(y)}{x - y} \leq K \label{eq:lipschitz} \end{equation}

实际上式 $\ref{eq:lipschitz}$ 甚至都没有要求函数 $f(x)$ 可导。在 K-R 对偶中，要求 $k=1$，也即：

$$ \frac{f(x) - f(y)}{x - y} \leq 1 $$

在 WGAN[3] 中作者使用一个神经网络来逼近函数 $f(x)$，并且使用一个简单的 trick 使得函数 $f(x)$ 满足利普希茨连续性：在每一次更新网络参数之后，如果参数大于某个阈值，则强行将参数剪裁到这个阈值。这样这个神经网络的每一层的参数都不会大于这个阈值，总体上这个神经网络的就一定满足利普希茨连续性。

但这样的做法似乎有些暴力，因为这样强行更改网络参数可能会影响到网络的收敛。于是在文献[4]中作者提出了 Spectral Norm，对网络参数进行谱归一化使得网络满足利普希茨连续条件。

矩阵的奇异值分解 (Singular Value Decomposition)

奇异值分解

为了方便大家理解 Spectral Norm ，我们先介绍一下 SVD 分解。

假设有 $m\times n$ 的矩阵 $A$，对矩阵 $A$ 存在这样一种分解：

\begin{equation} \mathbf{A} = \mathbf{U} \mathbf{\Sigma} \mathbf{V^T} \label{eq:svd} \end{equation}

当矩阵 $W$ 是一个实数阵的时候，

$\mathbf{U}$ 是一个 $m \times m$ 的正交矩阵 (orthogonal matrix)；
$\mathbf{\Sigma}$ 是一个 $m \times n$ 的对角阵，对角线上的元素为非负实数，这些数称为矩阵 $A$ 的“奇异值”；
$\mathbf{V}$ 是一个 $n \times n$ 的正交矩阵 (orthogonal matrix)。

矩阵 $A$ 的奇异值计算起来也不麻烦，只需要求得 $n\times n$ 的方阵 $A^TA$ 的特征值，然后把特征值开平方根，即为矩阵 $A$ 的奇异值，因为：

\begin{equation*} \begin{split} \mathbf{A}^T \mathbf{A} &= (\mathbf{U} \mathbf{\Sigma} \mathbf{V^T})^T \cdot (\mathbf{U} \mathbf{\Sigma} \mathbf{V^T}) \\ &= (\mathbf{V} \mathbf{\Sigma}^T \mathbf{U}^T)(\mathbf{U} \mathbf{\Sigma} \mathbf{V}^T) \\ &= \mathbf{V} (\mathbf{\Sigma}^T \mathbf{\Sigma}) \mathbf{V}^T。 \end{split} \end{equation*}

关于奇异值的计算并不是本文的重点，我们仔细研究一下式 $\ref{eq:svd}$ 所表达的几何意义。我们知道，每一个矩阵对应着一个线性变换。而正交矩阵对应的线性变换很特殊，叫旋转变换 (rotation)。而对角阵对应的变换，叫做拉伸变换 (scale)。也就是说式 $\ref{eq:svd}$ 把矩阵 $A$ 分解为三个变换：

旋转；
拉伸；
旋转。

这里就很有意思了。假设神经网络的某一层权重 $W$ 为一个 $m \times n$ 的矩阵，对于输入的特征 $x$，输出为:

\begin{equation} \begin{split} y &= W \cdot x \\ &= \mathbf{U} \cdot \mathbf{\Sigma} \cdot \mathbf{V^*} \cdot x \end{split} \label{eq:layer} \end{equation}

也相当于对输入做了三次变换。如果将输入视为一个向量的话，只有 $\Sigma$ 对应的拉伸变换才会改变输入的长度；其他两次旋转变换都不会。

矩阵的谱系数

矩阵 $W$ 的谱系数 (spectral norm) $\sigma(W)$ 定义为其 奇异值的最大值，其物理含义是: 对于任意的输入向量 $x$，经过 $W$ 的线性变换 对 $x$ 向量可能的最大的拉伸系数。

当一神经网络的参数矩阵的谱系数被限制小于1的时候，它对于任意的输入都不会有放大作用，因此自然也就满足利普希茨连续了。这里引入拉伸这个概念只是为了直观上容易理解，关于限制谱系数和函数导数的关系，在[4]有证明，这里不再赘述。

Spectral Normalization

介绍完 SVD 以及几何意义之后，Spectral Normalization的做法就很简单了：将神经网络的每一层的参数 $W$ 作 SVD 分解，然后将其最大的奇异值限定为1，具体地，在每一次更新 $W$ 之后都除以 $W$ 最大的奇异值。这样，每一层对输入$x$最大的拉伸系数不会超过 1。

经过 Spectral Norm 之后，神经网络的每一层 $g_l(x)$，都满足

$$ \frac{g_l(x) - g_l(y)}{x - y} \leq 1。 $$

对于整个神经网络 $f(x) = g_N(g_{N-1}(…g_1(x)…))$ 自然也就满足利普希茨连续性了。

然而，在每一次训练迭代中都对神经网络的每一层作 SVD 分解也是很不现实的，特别是当网络权重维度很大的时候。因此作者采用了一种叫做 power iteration[5] 的方法去迭代得到奇异值的近似解。

Power iteration 首先初始化一个随机的 $\hat{u}$ 然后根据以下公式迭代 $\hat{u}$ 和 $\hat{v}$：

\begin{equation} \hat{v} \leftarrow \frac{W^T\hat{u}}{\lVert W^T\hat{u} \rVert_2}, \ \\ \hat{u} \leftarrow \frac{W^T\hat{v}}{\lVert W^T\hat{v} \rVert_2}。 \label{eq:power-iter} \end{equation}

最后矩阵 $W$ 的最大奇异值 $\sigma(W)$ 可以根据 $\hat{u}$ 和 $\hat{v}$ 求得：

\begin{equation} \sigma(W) \approx \hat{u}W^T\hat{v}。 \label{eq:sigma_w} \end{equation}

在得到 $\sigma(W)$ 之后，网络每更新一次参数，都对参数 $W$ 进行谱归一化：

\begin{equation} W \leftarrow \frac{W}{\sigma(W)}。 \end{equation}

值得注意的是，$\ref{eq:power-iter}$ 式的迭代算法通常需要迭代多次才能比较好的逼近 $\sigma(W)$，但是Spectral Normal[4] 在神经网络的每一次更新过程中只进行一次 power iteration 迭代，仍然能取得不错的效果。这是因为神经网络的每次迭代中参数 $W$ 是在缓慢更新的，因此每两次迭代过程中可以认为 $W$ 的奇异值是近似相等的。因此虽然在网络的每一次迭代中，只进行一次power iteration，但是随着网络的训练，power iteration 对奇异值的估计会越来越准。这种做法其实十分巧妙，因为神经网络的优化本身也是个迭代的过程，因此对 $\sigma(W)$ 的逼近就可以利用神经网络的迭代逐渐来逼近，在网络参数的每一次更新过程中对 $\sigma(W)$ 的迭代不必太准确。

最后需要注意的一点是，判别器 D 使用了 Spectral norm 之后，就不能使用 BatchNorm (或者其它 Norm) 了。原因也很简单，因为 Batch norm 的“除方差”和“乘以缩放因子”这两个操作很明显会破坏判别器的 Lipschitz 连续性。

参考文献

Goodfellow, Ian, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. “Generative adversarial nets.” In Advances in neural information processing systems, pp. 2672-2680. 2014. [PDF]
Arjovsky, Martin, and Léon Bottou. “Towards principled methods for training generative adversarial networks.” arXiv preprint arXiv:1701.04862 (2017). [PDF]
Arjovsky, Martin, Soumith Chintala, and Léon Bottou. “Wasserstein generative adversarial networks.” In International Conference on Machine Learning, pp. 214-223. 2017. [PDF]
Miyato, Takeru, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. “Spectral normalization for generative adversarial networks.” arXiv preprint arXiv:1802.05957 (2018). [PDF]
Gulrajani, Ishaan, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C. Courville. “Improved training of wasserstein gans.” In Advances in Neural Information Processing Systems, pp. 5767-5777. 2017. [PDF]
Yoshida, Yuichi, and Takeru Miyato. “Spectral Norm Regularization for Improving the Generalizability of Deep Learning.” arXiv preprint arXiv:1705.10941 (2017).

感谢阅读🤗 本文内容谢绝任何形式的转载，如果您想和朋友分享本文内容，请分享本文链接 kaizhao.net/blog/spectral-norm。如果您发现文中的错误，或者有任何疑问，欢迎在下方留言交流 (留言功能基于 disqus，在中国大陆的读者可能需要一些技术手段才能连接🥲)。