DeepPoly中卷积的纯矩阵实现
DeepPoly是一种证明神经网络在特定输入域上稳定的方法，对于普通的线性层可以达到完全的准确性，对于ReLU层使用线性边界近似。卷积层可以看成一种特殊的线性层，但是在实际实现中，为了性能考虑，需要额外处理。本文借用PyTorch的矩阵操作实现DeepPoly卷积。
线性层
我们先从相对简单的线性层出发，来解释一下DeepPoly的回传过程。
问题定义
对于线性层 $\vec x^{m+1}=W^m \vec x^m+\vec b^m$ ，已知
\[W_{l}^{m,m-1} \vec x^{m-1} + \vec b_{l}^{m,m-1} \leq \vec x^{m} \leq W_{u}^{m,m-1} \vec x^{m-1} + \vec b_{u}^{m,m-1}\]求 $\vec x^{m+1}$ 的范围。
理论分析
$\vec x^{m+1}$ 可以表示成
\[W_{L}^{m+1,m-1} \vec x^{m-1} + \vec b_{L}^{m+1,m-1} \leq \vec x^{m+1} \leq W_{U}^{m+1,m-1} \vec x^{m-1} + \vec b_{U}^{m+1,m-1}\]的形式，与输入域的形式类似，因此，只需要求出参数 $W_{l}^{i+1,i-1}, \vec b_{l}^{i+1,i-1} , W_{u}^{i+1,i-1}, b_{u}^{i+1,i-1}$ 就可以合并两层。通过代入，可以得到参数的计算公式：
\[W_{L}^{m+1,m-1}[x, z] = \sum_y W^{m}[x, y] (W_L^{m, m-1}[y, z] [W^{m}[x, y] > 0] + W_U^{m, m-1}[y, z] [W^{m}[x, y] < 0])\] \[W_{U}^{m+1,m-1}[x, z] = \sum_y W^{m}[x, y] (W_U^{m, m-1}[y, z] [W^{m}[x, y] > 0] + W_L^{m, m-1}[y, z] [W^{m}[x, y] < 0])\]直观理解上，就是如果该层的权值为正，下限就取上一层的下限，上限就取上一层的上限；如果为负，就反之。
具体实现
在具体实现上，为了批量化，需要使用gather函数来根据正负选取权值：
def backsubsitute(self, lowerW: Tensor, lowerB: Tensor, upperW: Tensor, upperB: Tensor):
new_sz = lowerW.shape[0]
pre_sz, base_sz = self.lowerW.shape
W = torch.stack([self.lowerW, self.upperW], dim=1) # [p,2,q]
W = W.expand((new_sz, -1, -1, -1)) # [o,p,2,q]
B = torch.stack([self.lowerB, self.upperB], dim=1) # [p,2]
B = B.expand((new_sz, -1, -1)) # [o,p,2]
idx_neg = (lowerW < 0).long() # [o,p]
idxW = idx_neg.unsqueeze(-1).unsqueeze(-1).expand((-1, -1, -1, base_sz)) # [o,p,1,q]
idxB = idx_neg.unsqueeze(-1) # [o,p,1]
lW = torch.gather(W, dim=2, index=idxW).squeeze(2) # [o,p,1,q] -> [o,p,q]
lB = torch.gather(B, dim=2, index=idxB) # [o,p,1]
lowerW = lowerW.unsqueeze(1) # [o,p] -> [o,1,p]
newLowerW = torch.bmm(lowerW, lW).squeeze(1) # [o,1,p]x[o,p,q]=[o,1,q] -> [o,q]
newLowerB = torch.bmm(lowerW, lB).view(-1) # [o,1,p]x[o,p,1]=[o,1,1] -> [o]
newLowerB += lowerB
idx_pos = (upperW > 0).long()
idxW = idx_pos.unsqueeze(-1).unsqueeze(-1).expand((-1, -1, -1, base_sz))
idxB = idx_pos.unsqueeze(-1)
uW = torch.gather(W, dim=2, index=idxW).squeeze(2)
uB = torch.gather(B, dim=2, index=idxB)
upperW = upperW.unsqueeze(1)
newUpperW = torch.bmm(upperW, uW).squeeze(1)
newUpperB = torch.bmm(upperW, uB).view(-1)
newUpperB += upperB
return self.preLayer.backsubsitute(newLowerW, newLowerB, newUpperW, newUpperB)
ReLU层也是用线性上下限证明的，因此可以直接复用，区别在于层本身的上下限要改成DeepPoly中的
\[\lambda \vec x^{m-1} \leq \vec x^m \leq A \vec x^{m-1} + B\]卷积层
DeepPoly中卷积层的主要矛盾在于，普通的卷积是共享权值的，但是对于DeepPoly的卷积ReLU层，每个像素的 $\lambda$ 不一定相等，导致每一层的卷积核不一定相等，因此需要分别存储处理才能达到最高的准确度。
常见的卷积操作为
\[X^{m+1}_{uvw} = \sum_{xyz} W^{m+1, m}_{uxyz} X^m_{x,v+y,w+z}\]权值只有在通道上是分离的，而在长宽上都是共享的。在这里，我们就需要把长宽也分离出来：
\[X^{m+1}_{uvw} = \sum_{xyz} W^{m+1, m}_{uvwxyz} X^m_{x,v+y,w+z}\]代入 $X^m_{xyz}$ ，得到 \(X^{m+1}_{uvw},X^{m-1}_{uvw}\) 之间的关系：
\[X^{m+1}_{ijk} = \sum_{uvw} W^{m+1, m}_{ijkuvw} \sum_{xyz} W^{m, m-1}_{uvwxyz} X^{m-1}_{x,j+v+y,k+w+z}\]同时我们知道， \(X^{m+1}_{uvw},X^{m-1}_{uvw}\) 之间的形式应该为：
\[X^{m+1}_{ijk} = \sum_{abc} W^{m+1,m-1}_{ijkabc} X^{m-1}_{a,j+b,k+c}\]和线性层不同，这里的 \(X^{m-1}_{a,j+b,k+c}\) 与 \(X^{m-1}_{x,j+v+y,k+w+z}\) 并不是一一对应的，而是多对一的关系，因此无法直接得到新参数，不过可以通过对应关系，得到
\[W^{m+1,m-1}_{ijkabc} = \sum \{ W^{m+1, m}_{ijkuvw} W^{m, m-1}_{uvwxyz} | \forall a=x, b=v+y, c=w+z \}\]在代码实现中，为了避免计算每组 $(a,b,c)$ 对应的 $(x,y,z,v,w)$ ，我们从 $(i,j,k,u,v,w)$ 的角度循环，把每次的结果累加到 $a=x,b=v+y,c=w+z$ 的位置。这种方法的缺点在于，循环 $v,w$ 具有依赖性，因此我们只批量化了循环 $i,j,k,u$ 。好在 $v,w$ 一般是很小的数，因此对于性能的影响并不大。
def backsubsitute_conv(self, lowerW: Tensor, lowerB: Tensor, upperW: Tensor, upperB: Tensor, stride, padding):
if self.preLayer is None:
return self.getBoundUnderWB(
lowerW, lowerB, upperW, upperB, stride, padding)
out_ch_1, in_ch_1, k_h_1, k_w_1 = self.lowerW.shape
out_ch_2, out_h_2, out_w_2, in_ch_2, k_h_2, k_w_2 = lowerW.shape
s1 = 2
p1 = 1
s2 = stride
p2 = padding
s = s1 * s2
p = p1 + s1*p2
k_h_merge = k_h_1 + s1*k_h_2 - s1
k_w_merge = k_w_1 + s1*k_w_2 - s1
assert(out_ch_1 == in_ch_2)
newLowerW = torch.zeros((out_ch_2, out_h_2, out_w_2,
in_ch_1, k_h_merge, k_w_merge))
newUpperW = torch.zeros((out_ch_2, out_h_2, out_w_2,
in_ch_1, k_h_merge, k_w_merge))
newLowerB = torch.zeros((out_ch_2, out_h_2, out_w_2))
newUpperB = torch.zeros((out_ch_2, out_h_2, out_w_2))
W1 = torch.stack([self.lowerW, self.upperW], dim=1)
B1 = torch.stack([self.lowerB, self.upperB], dim=1)
i, j, k, u = torch.meshgrid(
torch.arange(out_ch_2),
torch.arange(out_h_2),
torch.arange(out_w_2),
torch.arange(in_ch_2)
)
i, j, k = i.flatten(), j.flatten(), k.flatten()
u = u.flatten()
for v in range(k_h_2):
for w in range(k_w_2):
coef = lowerW[:,:,:,:,v,w] # [i,j,k,u]
idx_neg = (coef < 0).long().flatten()
lW = W1[u, idx_neg].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2, in_ch_1, k_h_1, k_w_1))
lB = B1[u, idx_neg].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2))
coef_ex = coef.unsqueeze(4).unsqueeze(5).unsqueeze(6)
newLowerW[:,:,:,:,s1*v:s1*v+k_h_1, s1*w:s1*w+k_w_1] += (coef_ex * lW).sum(3)
newLowerB += (coef * lB).sum(3)
coef = upperW[:,:,:,:,v,w]
idx_pos = (coef > 0).long().flatten()
uW = W1[u, idx_pos].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2, in_ch_1, k_h_1, k_w_1))
uB = B1[u, idx_pos].reshape((out_ch_2, out_h_2, out_w_2, in_ch_2))
coef_ex = coef.unsqueeze(4).unsqueeze(5).unsqueeze(6)
newUpperW[:,:,:,:,s1*v:s1*v+k_h_1, s1*w:s1*w+k_w_1] += (coef_ex * uW).sum(3)
newUpperB += (coef * uB).sum(3)
newLowerB += lowerB
newUpperB += upperB
return self.preLayer.backsubsitute(newLowerW, newLowerB, newUpperW, newUpperB, s, p)
具体代码中，我们还增加了步长和填充的影响。对于卷积ReLU层，和线性层类似，只需要对每个像素取 $\lambda$ 即可。