[Box] Robust Training and Initialization of Deep Neural Networks: An Adaptive Basis Viewpoint

目录

• 概
• 主要内容
  • LSGD
  • Box 初始化
  • Box for Resnet
• 代码

Cyr E C, Gulian M, Patel R G, et al. Robust Training and Initialization of Deep Neural Networks: An Adaptive Basis Viewpoint.[J]. arXiv: Learning, 2019.

@article{cyr2019robust,
title={Robust Training and Initialization of Deep Neural Networks: An Adaptive Basis Viewpoint.},
author={Cyr, Eric C and Gulian, Mamikon and Patel, Ravi G and Perego, Mauro and Trask, Nathaniel},
journal={arXiv: Learning},
year={2019}}

概

这篇文章介绍了一种梯度下降的改进, 以及Box参数初始化方法.

主要内容

$$\tag{6} \arg \min_{\xi^L \xi^H} \sum_{k=1}^K \epsilon_k \|\mathcal{L}_k[u] - \sum_i \xi_i^L \mathcal{L}_k [\Phi_i(x, \xi^H)]\|^2_{\ell_2(\mathcal{X}_k)}.$$

LSGD

固定$\xi^H, \mathcal{X}_k$, 并令$\epsilon_k=1$, 则问题(6)退化为一个最小二乘问题

$$\arg \min_{\xi^L} \|A\xi^L -b\|^2_{\ell_2 (\mathcal{X})},$$

其中$b_i = \mathcal{L}[u](x_i)$, $A_{ij}=\mathcal{L} [\Phi_j (x_i, \xi^H)]$, $x_i \in \mathcal{X}, i=1,\ldots, N, j=1, \ldots, w$.

所以算法如下

Box 初始化

该算法期望使得feature-rich，但是我不知道这个rich从何而来.

假设第$l$层的输入为$x \in \mathbb{R}^{d_1}$, 输出为$y \in \mathbb{R}^{d_2}$, 则该层的权重矩阵$W \in \mathbb{R}^{d_2 \times d_1}$. 我们逐行地定义$W$:

1. 采样$p$, $p\sim U[0 ,1]^{d_1}$;
2. 采样$n$, $n \sim \mathcal{N}(0,I_{d_1})$, 并令$n=n/\|n\|$;
3. 求参数$k$使得

$$\max_{x \in [0, 1]^{d_1}} \sigma(k(x-p) \cdot n)=1.$$

4. $W$第$i$行$w_i=kn^T$, $b_i=-kp \cdot n$.

其中$\sigma$表示激活函数, 文中指的是ReLU.
求解参数$k$:

1. $p_{max} = \max (0, \mathrm{sign}(n))$;
2. $k=\frac{1}{(p_{max}-p) \cdot n}$

此$k$即为所需$k$, 只需证明$p_{max}$是最大化

$$(x - p)\cdot n, \quad x \in [0,1]^{d_1}$$

的解. 最大化上式, 可以分解为

$$\max_{x_i \in [0, 1]} x_in_i,$$

故$x_i = \max(0, \mathrm{sign}(n_i))$.

这个初始化有什么好处呢, 可以发现, 输入$x \in[0,1]^{d_1}$满足， 则输出$y \in [0, 1]^{d_2}$, 保证二者的"值域"范围一致, 以此类推整个网络节点值范围近似.

如果, 作者构建了一个2-2-2-2-2-2-2-2的网络, 可以发现, Xavier 和 Kaiming的初始化方法经过一定层数后, 就会塌缩在某个点, 而Box初始化方法能够缓解这一现象.

下面是文中列出的算法(与这里的符号有一点点不同, 另外$b$作者应该是遗漏了负号).

Box for Resnet

因为Resnet特殊的结构,

$$y=(W+I)x+b.$$

假设$x \in [0,m]^{d_1}$, 则:

1. 采样$p$, $p\sim U[0 ,m]^{d_1}$;
2. 采样$n$, $n \sim \mathcal{N}(0,I_{d_1})$, 并令$n=n/\|n\|$;
3. 求参数$k$使得

$$\max_{x \in [0, m]^{d_1}} \sigma(k(x-p) \cdot n)=\delta m.$$

4. $W$第$i$行$w_i=kn^T$, $b_i=-kp \cdot n$.

$$k=\frac{\delta m}{(mp_{max}-p) \cdot n}.$$

若第一层输入$x_i \in [0,1]$, 去$\delta=1/L$, 其中$L$为总的层数, 则

$$[0,1] \rightarrow [0,1+\frac{1}{L}] \rightarrow [0,(1+\frac{1}{L})^2] \rightarrow \cdots$$

代码

'''
initialization.py
'''
import torch
import torch.nn as nn
import warnings





def generate(size, m, delta):
    p = torch.rand(size) * m
    n = torch.randn(size)
    temp = 1 / torch.norm(n, p=2, dim=1, keepdim=True)
    n = temp * n
    pmax = nn.functional.relu(torch.sign(n)) * m
    temp = (pmax - p) * n
    k = (m * delta) / temp.sum(dim=1, keepdim=True)
    w = k * n
    b = -(w * p).sum(dim=1)
    return w, b

def box_init(module, m=1, delta=1):
    if isinstance(module, nn.Linear):
        w, b = generate(module.weight.shape, m, delta)
        try:
            module.weight.data = w
            module.bias.data = b
        except AttributeError as e:
            s = "Error: \n" + str(e) + "\n stops the initialization" \
                                       " for this module: {}".format(module)
            warnings.warn(s)

    elif isinstance(module, nn.Conv2d):
        outc, inc, h, w = module.weight.size()
        w, b = generate((outc, inc * h * w), m, delta)
        try:
            module.weight.data = w.reshape(module.weight.size())
            module.bias.data = b
        except AttributeError as e:
            s = "Error: \n" + str(e) + "\n stops the initialization" \
                                       " for this module: {}".format(module)
            warnings.warn(s)

    else:
        pass

"""config.py"""

nums = 10
layers = 6
method = "kaiming"  #box/xavier/kaiming
net = "Net"  #Net/ResNet

"""
测试
"""



import torch
import torch.nn as nn
import config
from initialization import box_init



class Net(nn.Module):

    def __init__(self, l):
        super(Net, self).__init__()

        self.linears = []
        for i in range(l):
            name = "linear" + str(i)
            self.__setattr__(name, nn.Sequential(nn.Linear(2, 2),
                                                 nn.ReLU()))
            self.linears.append(self.__getattr__(name))
        if config.method == 'box':
            self.box_init()
        elif config.method == "xavier":
            self.xavier_init()
        else:
            self.kaiming_init()

    def box_init(self):
        for module in self.modules():
            box_init(module)

    def xavier_init(self):
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                nn.init.xavier_normal_(module.weight)

    def kaiming_init(self):
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_normal_(module.weight)

    def forward(self, x):
        out = []
        temp = x
        for linear in self.linears:
            temp = linear(temp)
            out.append(temp)
        return out



class ResNet(nn.Module):

    def __init__(self, l):
        super(ResNet, self).__init__()

        self.linears = []
        for i in range(l):
            name = "linear" + str(i)
            self.__setattr__(name, nn.Sequential(nn.Linear(2, 2),
                                                 nn.ReLU()))
            self.linears.append(self.__getattr__(name))
        if config.method == 'box':
            self.box_init(l)
        elif config.method == "xavier":
            self.xavier_init()
        else:
            self.kaiming_init()

    def box_init(self, layers):
        delta = 1 / layers
        m = 1. + delta
        l = 0
        for module in self.modules():
            if isinstance(module, (nn.Linear)):
                if l == 0:
                    box_init(module, 1, 1)
                else:
                    box_init(module, m ** l, delta)
                l += 1

    def xavier_init(self):
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                nn.init.xavier_normal_(module.weight)

    def kaiming_init(self):
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_normal_(module.weight)

    def forward(self, x):
        out = []
        temp = x
        for linear in self.linears:
            temp = linear(temp) + temp
            out.append(temp)
        return out


if config.net == "Net":
    net = Net(config.layers)
else:
    net = ResNet(config.layers)

x = torch.linspace(0, 1, config.nums)
y = torch.linspace(0, 1, config.nums)

grid_x, grid_y = torch.meshgrid(x, y)

x = grid_x.flatten()
y = grid_y.flatten()
data = torch.stack((x, y), dim=1)
outs = net(data)


import  matplotlib.pyplot as plt


def axplot(x, y, ax):
    x = x.detach().numpy()
    y = y.detach().numpy()
    ax.scatter(x, y)

def plot(x, y, outs):
    fig, axs = plt.subplots(1, config.layers+1, sharey=True, figsize=(12, 2))
    axs[0].scatter(x, y)
    axs[0].set(title="layer0")
    for i in range(config.layers):
        ax = axs[i+1]
        out = outs[i]
        x = out[:, 0]
        y = out[:, 1]
        axplot(x, y, ax)
        ax.set(title="layer"+str(i+1))
    plt.tight_layout()
    plt.savefig("C:/Users/pkavs/Desktop/fig.png")
    #plt.show()
plot(x, y, outs)