GAN中的Spectral Normalization

GAN中的Spectral Normalization

  Spectral Normalization 出自 《Spectral Norm Regularization for Improving the Generalizability of Deep Learning》和《Spectral Normalization for Generative Adversarial Networks》，是为了解决GAN训练不稳定的问题，从“层参数”的角度用spectral normalization的方式施加regularization，从而使判别器D具备Lipschitz连续条件。

---

为什么要让D具备Lipschitz连续

为了防止判别器“放飞自我”

根据Wikipedia的定义，Lipschitz连续条件 如下。其意义在于使得 f 足够稳定，在输入发生少量变化时，输出不会有太巨大的变化。如果有图像A，修改少量像素得到图像B，输入判别器D后得到相差非常巨大的判别效果，那么判别器就是不稳定的，它对输入过于敏感。
$\frac{||f(x_1)-f(x_2)||}{||x_1-x_2||}\leq K, \forall x_1,x_2$
这里的 K 被称为 $f(x)$ 的 Lipschitz constant，K的最小值（上确界）被称为 $||f||_{Lip}$ ，称 $f(x)$ 满足 Lipschitz连续条件。

《Wasserstein GAN》 给出了衡量真实分布 $P_r$和生成分布 $P_g$的 Earth-Mover(EM) 距离：

$W(P_r, P_g)=\inf_{\gamma \in (P_r, P_g)}E_{(x,y)\sim \gamma}[||x-y||]$

由 Kantorovich-Rubinstein duality 可得EM距离的另一个形式，这里有用到 $||f||_{Lip}$：

$W(P_r, P_g)=\frac{1}{K} \sup_{||f||_{Lip}\leq K}E_{x\sim P_r}[f(x)]-E_{x\sim P_g}[f(x)]$

在 $||f||_{Lip}\leq 1$的条件下，求两个期望之差的上确界，就是EM距离，推导过程可参考这里。

---

矩阵范数

根据Wikipedia，矩阵的范数如下，它是右边的分式的上确界。

$||A||_p=\sup_{x\neq 0}\frac{||Ax||_p}{||x||_p}$

其中， $p=2$时被称为 Euclidean Norm 或 L2-Norm 或 Spectral Norm，它是矩阵的最大奇异值（或者最大特征值的开方），下面 $A^T$ 表示共轭转置。

$\sigma (A)=||A||_2=\sqrt{\lambda_{max}(A^TA)}$

矩阵范数的性质：

$||AB||_p\leq ||A||_p||B||_p$
$\sigma (AB)\leq \sigma (A)\sigma (B)$

---

Spectral Normalization

考虑采用了非线性激活函数的MLP网络构成的判别器D， $x^l=a^l(W^lx^{l-1}+b^l)$，把它的参数写成 $\theta=\{W^l, b^l\}_{l=1}^{L}$，使的判别器可以写成 $f_{\theta}(x^0)=x^L$。只考虑 $x$的一个很小的邻域，可以将判别器看作一个线性函数， $f_{\theta}(x)=W_{\theta,x}x+b_{\theta,x}$

$\frac{||f_{\theta}(x+\delta)-f_{\theta}(x)||_2}{||\delta||_2}=\frac{||W_{\theta,x}\delta||_2}{||\delta||_2}\leq \sigma(W_{\theta,x})=\sup_{\delta\neq 0}\frac{||W_{\theta,x}\delta||_2}{||\delta||_2}$

若激活函数使用 ReLU，可以把 $a^l(x^{l-1})$ 看作 $D^l_{\theta, x}x$，对角阵 $D^l_{\theta, x}$ 在 $x^{l-1}$ 非负的对应位置上为1，其他地方是0，这样 $\sigma(D^l_{\theta, x})=||D^l_{\theta, x}||_2\leq 1$。

$W_{\theta,x}=D^L_{\theta,x}W^{L}_{\theta,x}\cdots D^1_{\theta,x}W^1$

$\sigma(W_{\theta,x})\leq \sigma(D^L_{\theta,x})\sigma(W^{L})\cdots\sigma(D^1_{\theta,x})\sigma(W^1)\leq\prod_{l=1}^L\sigma(W^l)$

若对判别器D的每一层都做 Spectral Normalization：

$\hat W^l_{SN}=\frac{W^l}{\sigma(W^l)}$

$\sigma(\hat W^l_{SN})=\frac{\sigma(W)}{\sigma(W)}=1$

$\frac{||f_{\theta}(x+\delta)-f_{\theta}(x)||_2}{||\delta||_2}\leq\sigma(W_{\theta,x})\leq1$

---

Spectral Normalization 实现

Spectral Normalization实际上在做的事情，是将每层的参数矩阵除以自身的最大奇异值，本质上是一个逐层SVD的过程，但是真的去做SVD就太耗时了，所以采用幂迭代的方法求（参见Wikipedia）。

def spectral_norm(w, iteration=10, name="sn"):
	'''
	Ref: https://github.com/taki0112/Spectral_Normalization-Tensorflow/blob/65218e8cc6916d24b49504c337981548685e1be1/spectral_norm.py
	'''
	w_shape = w.shape.as_list() # [KH, KW, Cin, Cout] or [H, W]
	w = tf.reshape(w, [-1, w_shape[-1]]) # [KH*KW*Cin, Cout] or [H, W]

	u = tf.get_variable(name+"_u", [1, w_shape[-1]], initializer=tf.random_normal_initializer(), trainable=False)
	s = tf.get_variable(name+"_sigma", [1, ], initializer=tf.random_normal_initializer(), trainable=False)

	u_hat = u # [1, Cout] or [1, W]
	v_hat = None 

	for _ in range(iteration):
		v_hat = tf.nn.l2_normalize(tf.matmul(u_hat, tf.transpose(w))) # [1, KH*KW*Cin] or [1, H]
		u_hat = tf.nn.l2_normalize(tf.matmul(v_hat, w)) # [1, Cout] or [1, W]
		
	u_hat = tf.stop_gradient(u_hat)
	v_hat = tf.stop_gradient(v_hat)

	sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat)) # [1,1]
	sigma = tf.reshape(sigma, (1,))

	with tf.control_dependencies([u.assign(u_hat), s.assign(sigma)]):
		# ops here run after u.assign(u_hat)
		w_norm = w / sigma 
		w_norm = tf.reshape(w_norm, w_shape)
	
	return w_norm

若输入矩阵是全连接层的参数，尺寸为 $[H,W]$，则spectral_norm在效果上会直接对该二维矩阵求最大奇异值，但如果输入矩阵为卷积层的卷积核，其尺寸应该是 $[K_H, K_W, C_{in}, C_{out}]$，spectral_norm会先将该矩阵reshape成一个大小为 $[K_HK_W C_{in}, C_{out}]$的矩阵，再用迭代法对该二维矩阵求最大奇异值。不论哪种情况，求得的奇异值都是一个单值scalar。

---

带SN的卷积层

def conv2d(x, channel, k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02, name='conv2d'):
	with tf.variable_scope(name):
		w = tf.get_variable('w', [k_h, k_w, x.get_shape()[-1], channel], 
				initializer=tf.truncated_normal_initializer(stddev=stddev))
		w_sn = spectral_norm(w, iteration=3)

		conv = tf.nn.conv2d(x, filter=w_sn, strides=[1, d_h, d_w, 1], padding='VALID')

		biases = tf.get_variable('biases', [channel], initializer=tf.constant_initializer(0.0))
		conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())

		return conv

带SN的FC层

def dense(x, output_size, stddev=0.02, bias_start=0.0, activation=None, sn=False, reuse=False, name='dense'):
	shape = x.get_shape().as_list()
	with tf.variable_scope(name, reuse=reuse):
		W = tf.get_variable(
			'weights', [shape[1], output_size], 
			tf.float32, 
			tf.random_normal_initializer(stddev=stddev))
		bias = tf.get_variable(
			'biases', [output_size], 
			initializer=tf.constant_initializer(bias_start))
		if sn:
			W = spectral_norm(W, 20, name="sn")
		out = tf.matmul(x, W) + bias 
		if activation is not None:
			out = activation(out)
	
	return out

检验SpectralNorm

  用dense层检验，首先将W的初始值改为tf.ones_initializer()变成全1矩阵，然后用大小(3,3)的全1输入，在sn=True和sn=False的条件下分别输出结果。

x = tf.ones(shape=(3,3), dtype=tf.float32)
y1 = dense_ones(x, output_size=4, sn=False)
y2 = dense_ones(x, output_size=4, sn=True)
with tf.Session() as sess:
	sess.run(tf.global_variables_initializer())
	print sess.run(y1)
	print sess.run(y2)

  dense层的W是一个大小为(3,4)的全1矩阵，用matlab或np.linalg.svd可以算出它的最大奇异值为3.464，而3.0/3.464=0.8660，从下面的结果可见，在sn=True时，FC层输出的结果等于在sn=False时的结果除以FC层权值矩阵W的最大奇异值，该结果本质上是由于权值矩阵W自己除以最大奇异值导致的。
  因此，Spectral Norm的实现是正确的。

# y1
[[3. 3. 3. 3.]
 [3. 3. 3. 3.]
 [3. 3. 3. 3.]]
# y2
[[0.86602545 0.86602545 0.86602545 0.86602545]
 [0.86602545 0.86602545 0.86602545 0.86602545]
 [0.86602545 0.86602545 0.86602545 0.86602545]]

---

如何查看sigma的值

  有时候希望查看迭代法求得的最大奇异值是否正确，或者用它作为某种指示指标（比如用作GAN里判别器是否放飞自我的标准），希望以Tensor的形式获取它。只需在spectral_norm中专门声明一个variable用于存储它的值，在迭代法计算完之后把sigma赋给它即可。

def spectral_norm(...):
	...
	s = tf.get_variable("sn_sigma", [1, 1], initializer=tf.random_normal_initializer(), trainable=False)
	...
	with tf.control_dependencies([u.assign(u_hat), s.assign(sigma)]):
		...

  为了用sess.run获取它实时的值，需要先得到这个Tensor，一种简单暴力的方法是首先用tf.global_variables()查看所有variable的名字，找到该Tensor的名字，用tf.get_collection获取。要注意Tensor的名字往往前面带有很多variable_scope，就像文件路径一样，需要完整地传入tf.get_collection。

with tf.Session() as sess:
	sess.run(tf.global_variables_initializer())
	v = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, '你的前缀/sn_sigma:0')[0]
	sigma = sess.run(v, feed_dict={...})
	print sigma[0, 0]

---

control_dependencies的作用

  在我找到的所有Spectral Normalization的实现代码中，都有control_dependencies存在，其作用在于每次调用spectral_norm时用上一次调用中进行迭代得到的u来初始化本次调用的u，暂时不明白其必要性，如果为了更加精确而多次调用输入同一个权值矩阵，倒不如将iteration设的大一点（20以上就已经非常精确了），如果是不同的输入，则这样初始化没有意义，并不会变得更精确。
  这里演示一下control_dependencies、assign和identity的配合使用。
  在下面的代码中，control_dependencies保证了在每次out=tf.identity(u)前都会先将u_var更新为u的值，只有在最开始u_var的初始值为1.0。

u_var = tf.Variable(1.0)
u = u_var 

u = u * 3

with tf.control_dependencies([tf.assign(u_var, u, name='update_u')]):
    out = tf.identity(u)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for _ in range(3):
        print out.eval()
'''
3.0
9.0
27.0
'''

---

关于迭代次数

  迭代法计算最大奇异值，迭代次数达到25即可非常精确，但为了节省时间，一般不需要这么多。

---

Reference

Spectral Norm Regularization for Improving the Generalizability of Deep Learning
Spectral Normalization for Generative Adversarial Networks
BatchNormalization LSTMCell Tensorflow
spectral normalization Tensorflow
compare_gan spectral normalization Tensorflow
Wikipedia Matrix norm