Spectral Normalization 譜歸一化-原理及實現

//

一、譜范數及其計算方法

見我的這篇blog 譜范數求解方法-奇異值分解&冪迭代法

//

二、譜歸一化提出背景

譜歸一化由論文《Spectral Normalization For Generative Adversarial Networks》論文鏈接 提出。

原生 GAN 的目標函數等價於優化生成數據的分布和真實數據的分布之間的 J-S 散度 (Jensen–Shannon Divergence)。
而由於二者間幾乎不可能有不可忽略的重疊，所以無論它們相距多遠JS散度都是常數log2，最終導致生成器的梯度（近似）為0，梯度消失。
也就是說判別器訓練越好，生成器梯度消失越嚴重。

WGAN使用性質優良的 Wasserstein distance 代替原生 GAN 中的 J-S 散度。 然后利用KR對偶原理將 Wasserstein distance的求解問題轉換為求解最優的利普希茨連續函數的問題。 為了使得判別器 D 滿足利普希茨連續性，作者使用“梯度裁剪”將過大的參數直接裁剪到一個閾值以下。

“梯度裁剪”技術從每層神經網絡的參數矩陣的譜范數角度，引入利普希茨連續性約束，使神經網絡對輸入擾動具有較好的非敏感性，從而使訓練過程更穩定，更容易收斂。(深度學習模型存在“對抗攻擊樣本”，比如圖片只改變一個像素就給出完全不一樣的分類結果，這就是模型對輸入過於敏感的案例。)

我們可以這樣理解：局部最小點附近如果是平坦（flatness）的話（斜率有約束），那么其泛化的性能將較好，反之，若是不平坦（sharpness）的話，稍微一點變動，將產生較大變化，則其泛化性能就不好，也就不穩定。

Spectral Norm使用一種更優雅的方式使得判別器 D 滿足利普希茨連續性，限制了函數變化的劇烈程度，從而使模型更穩定。

//

三、Lipschitz 連續性

Lipschitz 條件限制的是函數變化的劇烈程度，即函數的最大梯度。

K-Lipschitz表示函數的最大梯度為K，K稱為Lipschitz constant（Lipschitz常量）。例如 y = sinx的最大斜率為1，所以它是是 1-Lipschitz的。

那么：

關鍵理論： 對矩陣A除以其譜范數可以使其具有 1-Lipschitz continuity（證明）。

//

四、GAN的譜歸一化原理

對GAN做Spectral Norm，實際就是要使得判別器D滿足1-Lipschitz條件。

又由於判別器D省略各個層加上的bias后，這種多層神經網絡實際上是多個復合函數嵌套的操作。
最常見的嵌套是：一層卷積，一層激活函數，再一層卷積，再一層激活函數，這樣層層包裹起來。而激活函數通常選取的 ReLU，Leaky ReLU 都是 1-Lipschitz 的，我們只需要保證卷積的部分是 1-Lipschitz continuous 的（如果有Linear層也要保證Linear層是1-Lipschitz continuous 的），就可以保證整個神經網絡都是 1-Lipschitz continuous 的。

那么怎樣保證卷積部分是 1-Lipschitz continuous 的呢？
圖像上每個位置的卷積操作，正好可以看成是一個矩陣乘法。因此，我們只需要約束各層卷積核的參數W ，使它是 1-Lipschitz continuous 的，就可以保證卷積的部分是 1-Lipschitz continuous 的，從而滿足整個神經網絡的 1-Lipschitz continuity。

因此，具體做法為：
對判別器D中的每一層卷積的參數矩陣W做譜歸一化，即：
step 1.對W進行SVD (實現時使用冪迭代法近似代替SVD，減少計算成本)，得到W的最大奇異值。
step 2.在每一次更新W之后都除以W的最大奇異值，從而使其滿足 1-Lipschitz continuity。

注意:
判別器 D 使用了 Spectral norm 之后，就不能使用 BatchNorm (或者其它 Norm) 了。 原因也很簡單，因為 Batch norm 的“除方差”和“乘以縮放因子”這兩個操作很明顯會破壞判別器的 Lipschitz 連續性。

//

五、GAN的譜歸一化實現

google用tensorflow實現了譜歸一化函數鏈接

pytorch中有實現好的譜歸一化函數torch.nn.utils.spectral_norm（官方文檔）（github）

import torch.nn as nn
import torch

# 對線性層做譜歸一化
sn_module = nn.utils.spectral_norm(nn.Linear(20,40))
# 驗證譜歸一化后的線性層是否滿足1-Lipschitz continuity
print(torch.linalg.norm(sn_module.weight,2))     # tensor(1.3898)  為啥不是1.000呢？

但是官方文檔中建議使用新版譜歸一化函數torch.nn.utils.parametrizations.spectral_norm（官方文檔）
不過目前看到的幾乎還是用torch.nn.utils.spectral_norm的。

用法：

# #############################################################################################
# 用法1 ref: https://blog.csdn.net/qq_37950002/article/details/115592633
# #############################################################################################
import torch
import torch.nn as nn

class TestModule(nn.Module):
    def __init__(self):
        super(TestModule,self).__init__()
        self.layer1 = nn.Conv2d(16,32,3,1)
        self.layer2 = nn.Linear(32,10)
        self.layer3 = nn.Linear(32,10)
 
    def forward(self,x):
        x = self.layer1(x)
        x = self.layer2(x)

model = TestModule()

def add_sn(m):
    for name, layer in m.named_children():
        m.add_module(name, add_sn(layer))
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            return nn.utils.spectral_norm(m)
        else:
            return m

my_model = add_sn(model)

# #############################################################################################
# 用法2 ref: https://github.com/Vbansal21/Custom_Architecture/EATS/models/v2_discriminator.py
# #############################################################################################
...
# 直接在卷積/線性層外面套一層nn.utils.spectral_norm()
nn.Sequential(
                nn.ReflectionPad1d(7),
                nn.utils.spectral_norm(nn.Conv1d(1, 16, kernel_size=15)),
                nn.LeakyReLU(0.2, True),
            ),
...

參考 :
Spectral Normalization 譜歸一化
令人拍案叫絕的Wasserstein GAN
pytorch的spectral_norm的使用
Thanks 😄