DL中常用的三种K-Lipschitz技术

在进入到正题前，首先来了解下什么是 K-Lipschitz 以及它在 DL 中能起到什么作用

Lipschitz continuity

利普希茨连续。满足如下性质的任意连续函数 $f$ 称为 K-Lipschitz：
$$\Vert f(x_1)-f(x_2)\Vert \le K\Vert x_1-x_2\Vert ,\forall x_1,x_2\in \text{dom}f$$
这里的 $\Vert \cdot \Vert$ 常用2-范。直观上看，Lipschitz 条件限制了函数变化的剧烈程度。在DL中，由于 Lipschitz continous 的函数的梯度上界被限制，因此函数会更平滑。因此利用梯度下降进行参数更新时，参数的变化不会太大/剧烈从而降低梯度爆炸的发生概率，使模型的更新更稳定。$K$ 称为 Lipschitz constant。

那么在 DL 中，有哪些方法可用于将 $f$ 限制在 K-Lipschitz 空间中？这里介绍三种：weight clipping、gradient penality、spectral normalization，一般而言，第三种的效果最佳

Weight clipping

由 Wasserstein GAN 提出。
在利用 gradient descent 进行参数更新后，在对所有参数进行如下操作：
$$w=\{\begin{array}{ll}c,& \text{if }w>c\\ -c,& \text{if }w<-c\\ w,& \text{otherwise}\end{array}$$
其中 $c$ 是人为设定的阈值。注意，Weight cplipping 并无法保证 $f$ 位于 1-Lipschitz，而只能保证其是 K-Lipschitz的（K具体无法确定）

Gradient penalty

由 Improved Training of Wasserstein GANs 提出。
理论支持：一个可微函数 $f$ 是 1-Lipschitz 当且仅当它对所有的输入 $x$ 均满足 $\Vert {\nabla }_{x}f(x)\Vert \le 1$，即，
$$f\in \text{1-Lipschitz}⟺\Vert {\nabla }_{x}f(x)\Vert \le 1,\forall x\in \text{dom}f$$
在具体实现时，即在 Objective function 中添加如下正则项：
$$\underset{\theta }{\min }\{\mathcal{L}(x,\theta )+\lambda \max (0,\Vert {\nabla }_{x}f(x)-1)\Vert \}$$
公式中的 $\mathcal{L}$ 即为 Loss/Objective function，而 $f$ 为 Score function。注意，优化该目标函数后，所解出的 $f$ 并无法保证一定满足 $\Vert {\nabla }_{x}f(x)\Vert \le 1$，但 $f$ 会偏向具有该属性

Spectral Normalization

谱归一化，由SN-GAN提出，是目前三种方法中效果最优的方法。

下面简要介绍其非严格的理论推导，主要来自知乎，再添加上自己的一些理解

理论推导

对于复合函数，存在如下定理：
$$\Vert f\circ g{\Vert }_{Lip}\le \Vert f{\Vert }_{Lip}\cdot \Vert g{\Vert }_{Lip}$$
neutral network 正是由多个复合函数嵌套而成，最常见的嵌套方式如下：$f(g(f(g(\cdots ))))$，其中 f 表示激活函数，$g$ 表示卷积操作（以CNN为例）。而 $f$ 常选取 LeakyRelu，Relu，Sigmoid，Tanh，而它们均为 1-Lipschitz。因此 $\Vert f\circ g{\Vert }_{Lip}\le \Vert f{\Vert }_{Lip}\cdot \Vert g{\Vert }_{Lip}=\Vert g{\Vert }_{Lip}$，故要使得复合函数 $f\circ g$ 为 1-Lipschitz，即需保证卷积操作 $g$ 是 1-Lipschitz，就可以保证整个网络都是 1-Lipschitz continous 的。

在图像上每个位置的卷积操作，正好就是做如下“局部区域“的变换：
$$\Vert {\text{unfold}}_{raw}(M)\cdot {\text{unfold}}_{col}(x)\Vert =y$$
其中 $x\in R^{f\times f}$ 为 local receptive field，$M\in R^{f\times f}$ 为卷积核，$y$ 为对应位置的卷积输出，${\text{unfold}}_{raw}(\cdot )$ 将 $\cdot$ 按行展开成行向量，${\text{unfold}}_{col}(\cdot )$ 将 $\cdot$ 按列展开成列向量。因此，只需保证 ${\text{unfold}}_{raw}(M)$ 是 1-Lipschitz，就可以使得整个 network 是 1-Lipschitz。

对任意矩阵 $A$ (${\text{unfold}}_{raw}(M)$ 是 $A$ 的一个具例)，存在如下定理：
$$\begin{array}{ll}& A\in \text{K-Lipschitz},\forall A:R^n\to R^m/\forall A\in R^{m\times n}\\ ⟺& \Vert Ax\Vert \le K\Vert x\Vert ,\forall x\in R^n\\ ⟺& ⟨Ax,Ax⟩\le K^2⟨x,x⟩\\ ⟺& x^T(A^{T}A-K^{2}I)x\le 0\\ ⟺& ⟨(A^{T}A-K^{2}I)x,x⟩\le 0\end{array}\text{\hspace{1em}(1)}$$
因 $A^{T}A$ 为实对称矩阵，故不同特征值对应的特征向量两两正交，又因 $A^{T}A$ 为半正定阵，故其所有特征值非负。不妨假设 $A^{T}A\in R^{n\times n}$ 的特征向量为 $v_i\text{for}i=1,2,\cdots ,n$，各自对应的特征值为 ${\lambda }_i\text{for}i=1,2,\cdots ,n$。因为 $v_i$ 两两互相正交（不严谨，因为不一定有 $n$ 个），不妨令 $v_i$ 已经单位化，则它们构成 $R^n$ 空间的一组单位正交基。因此 $x={\sum }_{i=1}^n{x}_i{v}_i,\text{ for }\forall x\in R^n$，则续 (1)：
$$\begin{array}{ll}& ⟨(A^{T}A-K^{2}I)x,x⟩\le 0\\ ⟺& ⟨(A^{T}A-K^{2}I){\sum }_{i=1}^n{x}_i{v}_i,{\sum }_{i=1}^n{x}_i{v}_i⟩\le 0\\ ⟺& {\sum }_{i=1}^n{\sum }_{j=1}^n{x}_i{x}_j⟨(A^{T}A-K^{2}I)v_i,v_j⟩\le 0\\ ⟺& {\sum }_{i=1}^n{\sum }_{j=1}^n{x}_i{x}_j{\left[(A^{T}A-K^{2}I)v_i\right]}^T{v}_j\le 0\end{array}\text{\hspace{1em}(2)}$$
又
$$\begin{array}{l}{\left[(A^{T}A-K^{2}I)v_i\right]}^T{v}_j={v_i}^T(A^{T}A-K^{2}I)v_j={v_i}^T{A}^{T}A{v}_j-K^2{v_i}^T{v}_j\\ =\{\begin{array}{ll}0,& \text{if }i≠j\\ ({\lambda }_i-K^2)⟨v_i,v_i⟩={\lambda }_i-K^2,& \text{if }i=j\end{array}\end{array}$$
故续 (2)
$$\sum _{i=1}^n\sum _{j=1}^n{x}_i{x}_j{\left[(A^{T}A-K^{2}I)v_i\right]}^T{v}_j\le 0⟺\sum _{i=1}^n{x}_i^2({\lambda }_i-K^2)\le 0$$
即
$$A\in \text{K-Lipschitz},\forall A:R^n\to R^m/\forall A\in R^{m\times n}⟺\sum _{i=1}^n{x}_i^2(K^2-{\lambda }_i)\ge 0\text{\hspace{1em}(3)}$$
其中 ${\lambda }_i$ 为 $A^{T}A$ 的第 $i$ 个特征值。
不妨令 $K^2={\max }_i({\lambda }_i)$，则必有 ${\sum }_{i=1}^n{x}_i^2(K^2-{\lambda }_i)\ge 0$ 成立，因此：
$$K^2=\underset{i}{\max }({\lambda }_i)\Rightarrow \left(\sum _{i=1}^n{x}_i^2(K^2-{\lambda }_i)\ge 0⟺A\in \text{K-Lipschitz},\forall A:R^n\to R^m/\forall A\in R^{m\times n}\right)\text{\hspace{1em}(4)}$$
而要令 $A$ 从 K-Lipschitz 变为 1-Lipschitz，仅需对 $A$ 作如下缩放即可：$A:=\frac{A}{K}$。即，$\text{矩阵 A 除以它的 spectral norm}（\Vert A{\Vert }_2=\sqrt{A^{T}A的最大特征值}）\text{可以使其具有1-Lipschitz continuity}$。

于是问题的关键就转变成如何求解 $A^{T}A$ 的最大特征值 ${\lambda }_1$ 了（不妨令 ${\lambda }_1\ge {\lambda }_2\ge \cdots \ge {\lambda }_n$），最经典的算法为 Power iteration

Power iteration

Power iteration 是用于近似计算方阵最大特征值和其对应特征向量的常用方法，其具体步骤如下：

1. 令 $B=A^{T}A\in R^{n\times n}$ ，假设 $B$ 是一个 $n\times n$ 的满秩方阵，其单位特征向量为 $v_1,v_2,\cdots ,v_n\in R^n$，对应特征值为 ${\lambda }_1,{\lambda }_2,\cdots ,{\lambda }_n$（因 $B$ 实对称，故 $v_1,\cdots ,v_n$相互正交）。那么对任意向量 $x\in R^n$ 均可写成 $x={\sum }_{i=1}^n{x}_i{v}_i$，有：
   $$\begin{array}{lll}Bx& =& B{\sum }_{i=1}^n{x}_i{v}_i\\ & =& {\sum }_{i=1}^n{x}_i{\lambda }_i{v}_i\end{array}$$
   经过 $k$ 次迭代后：
   $$B^{k}x=\sum _{i=1}^n{x}_i{\lambda }_i^k{v}_i={\lambda }_1^k{x}_1{v}_1+\sum _{i=2}^n{x}_i{\left(\frac{{\lambda }_i}{{\lambda }_1}\right)}^k{v}_i\text{\hspace{1em}(5)}$$
   不妨假定 ${\lambda }_1>{\lambda }_2>\cdots >{\lambda }_n$（不考虑特征值相等的情况，因为这在实际中很少见），因此，$\underset{k\to \infty }{\lim }{\left(\frac{{\lambda }_i}{{\lambda }_1}\right)}^k=0$，故式 (5) 可进一步化为：
   $$B^{k}x\approx {\lambda }_1^k{x}_1{v}_1\text{\hspace{1em}(6)}$$
   即经过 $k$ 次迭代后，我们将得到特征向量 $v_1$ 的线性缩放，只要将 $B^{k}x$ 归一化就可得到主特征向量 $v_1$，进而再利用 $B{v}_1={\lambda }_1{v}_1$ 解得 ${\lambda }_1$。

上述为 Power iteration 的理论部分，而伪代码如下：

u, v = 随机初始化, None
for k iteration:
	v = A^T u / |A^T u|_2  # (1)
	u = A v / |A v|_2      # (2)
sqrt_lambda_1 = u^T A v    # (3)

其中 $A:R^n\to R^m/A\in R^{m\times n}$，$v\in R^n,u\in R^m$，其中 $n,m$ 分别属于输入，输出空间的维度

伪代码中的 (1) 等价于公式 (6) 的原因：
$v=\frac{A^{T}u}{\Vert A^{T}u{\Vert }_2}=A^T\frac{Av}{\Vert Av{\Vert }_2}/{\Vert A^{T}u\Vert }_2={\alpha }_i(A^{T}A)v={\alpha }_{i}Bv$，该公式在迭代了 $k$ 次后，与上面的式 (6) 仅多乘了一个常数系数 $\alpha ={\prod }_{i=1}^k{\alpha }_i$，这并不影响主特征向量 $v_1$ 的方向，因此对 $v$ 归一化后即可得到 $v_1$。而在实际上，在每步迭代后，都会对 $v$ 进行归一化，因此 $v_1=v$

伪代码中的 (3) 可求解 ${\lambda }_1$ 的原因：
$$\begin{array}{ll}\because & A^{T}Av={\lambda }_{1}v,且\Vert v{\Vert }_2=1\\ \therefore & v^T{A}^{T}Av={\lambda }_1{v}^{T}v={\lambda }_1\\ 即& \sqrt{{\lambda }_1}=\Vert Av{\Vert }_2\\ 又& u=\frac{Av}{\Vert Av{\Vert }_2}\to u^{T}u=\frac{u^{T}Av}{\Vert Av{\Vert }_2}\to 1=\frac{u^{T}Av}{\Vert Av{\Vert }_2}\\ \therefore & \sqrt{{\lambda }_1}=u^{T}Av\end{array}$$

Tensorflow 实现

def spectral_norm(W, is_training, iteration=1):
	'''
	W: [f, f, in_c, out_c]
	'''
	shape = W.get_shape().as_list()    # [f, f, in_c, out_c]
	W = tf.reshape(W, [-1, shape[-1]]) # [N, out_c], 其中 N=f*f*in_c
									   # reshape，即理论推导里所说的展开操作——以通道为单位将 M 展开
	u = tf.get_variable(shape=[1, shape[-1]],
						trainable=False,
						initializer=...)
	# power iteration
	u_norm, v_norm = u, None
	for k in range(iteration):
		v_norm = tf.matmul(u_norm, W, transpose_b=True)  # [1, N]
		v_norm = tf.math.l2_normalize(v_norm)

		u_norm = tf.matmul(v_norm, W)  # [1, out_c]
		u_norm = tf.math.l1_normalize(u_norm)
	lambda_sqrt = tf.matmul(v_norm, W)
	lambda_sqrt = tf.matmul(u_norm, lambda_sqrt, transpose_b=True)
	# spectral norm
	W_sn = W / lambda_sqrt
	# update estimated 1st singular vector while training
	with tf.control_dependencies([tf.cond(is_training,
										  lambda: u.assign(u_norm),
										  lambda: u.assign(u))]):
		W_norm = tf.reshape(W_norm, shape)
	return W_norm