Deep Learning Systems Lecture 22 Implicit Layers

这里回顾dlsys第二十二讲，本讲内容是Implicit Layers。

课程主页：

大纲

复杂函数作为深度网络中的层（运算符）；
implicit layer的微分；
implicit layer的例子；

复杂函数作为深度网络中的层（运算符）

深层网络中的“层”

术语“层”在深度学习中是一个重载。我们需要区分两件事：

Operators在计算图中定义“原子”操作，以及它的显式梯度传递：即矩阵乘法、卷积、逐元素的非线性等。
Modules定义了运算符的集合，其中反向传递是通过计算图的构造隐式创建的。

将函数定义为运算符而不是模块有什么好处？

例子：卷积

回想一下我们之前的讲座，卷积可以根据许多矩阵乘法运算来指定：

def conv_matrix_mult(Z, weight):
    N,H,W,C_in = 2.shape
    K,_,_,C_out = weight.shape
    out = np.zeros((N,H-K+1,W-K+1,C_out))

    for i in range(K):
        for j in range(K):
            out += z[:,i:i+H-K+l,j:j+w-K+l,:] @ weight[i,j]

    return out

使卷积成为原子运算符有什么好处？

可以以模块化方式优化正向传播；
无需将所有中间计算项（例如 im2col 矩阵）保存在计算图中的内存中；

更复杂的层？

如果我们想在深度网络中定义一个执行一些更复杂操作的“层”怎么办：解决优化问题、计算不动点、求解ODE 等？
我们可以简单地在我们的自动微分工具本身（即作为一个模块）中实现“求解器”，允许我们将它嵌入任何更大的计算图中。
但是，正如我们将看到的，将这些函数实现为原子操作（即作为运算符）也有一些显着的优势。

Explicit vs Implicit层

到目前为止，我们使用的几乎所有层（运算符）都是显式的，因为它们本身涉及输入和输出之间的“固定”转换。
相比之下，隐式层（运算符）只是定义一个更复杂的操作，以满足输入和输出的某些联合条件。

微分方程、不动点迭代、优化解等例子，都可以被包含在隐式层中；

例子：Deep Equilibrium Models

考虑传统的MLP：

我们现在以两种方式修改这个网络：通过在每一步重新注入输入，以及通过在每次迭代中应用相同的权重矩阵（权重绑定）：

深度权重绑定模型的迭代

使用这种形式的权重绑定模型，我们将相同的函数重复应用于隐藏单元：

$z_{i+1}=\sigma\left(W z_i+x\right)$

在许多情况下，我们可以设计网络，使该迭代收敛到某个不动点或平衡点：

$z^{\star}=\sigma\left(W z^{\star}+x\right)$

可以定义一个直接找到这个平衡点的层（通过不动点迭代或其他求解方法）。

在实际中，我们想找到一个更复杂的“单元”的平衡点，并将其用作我们的整个模型（加上一个额外的线性层）

如果我们将此操作实现为needle中的运算符，那么我们如何计算不动点𝑧⋆（例如，使用高级非线性求解器）并不重要，我们也不需要在计算图中存储中间状态。

但是为了让它成为一个Operator，我们需要能够可导（即，为这个操作实现compute_gradient）。

implicit layer的微分

对implicit layer微分

我们定义了一个隐式层来找到非线性方程的解：

$z^{\star}=f\left(z^{\star}, x\right)$

我们如何通过这一层进行微分？利用隐式微分法：

$\begin{aligned} z^{\star}(x)& =f\left(z^{\star}(x), x\right) \\ \frac{\partial z^{\star}(x)}{\partial x}& =\frac{\partial f\left(z^{\star}(x), x\right)}{\partial x} \\ \frac{\partial z^{\star}(x)}{\partial x}&=\frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)} \frac{\partial z^{\star}(x)}{\partial x}+\frac{\partial f\left(z^{\star}(x), x\right)}{\partial x} \\ \frac{\partial z^{\star}(x)}{\partial x}& =\left(I-\frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)}\right)^{-1} \frac{\partial f\left(z^{\star}(x), x\right)}{\partial x} \end{aligned}$

集成到自动微分中

回想一下，在我们的自动微分框架中，我们真的只是希望能够计算伴随和层梯度的乘积：

$\bar{v} \frac{\partial z^{\star}(x)}{\partial x}=\bar{v}\left(I-\frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)}\right)^{-1} \frac{\partial f\left(z^{\star}(x), x\right)}{\partial x}=\bar{v}^{\prime} \frac{\partial f\left(z^{\star}(x), x\right)}{\partial x}$

换句话说，要为隐式层实现反向传播，我们只需要将这种“特殊”的伴随计算添加到反向传播中即可。

详细说明：计算反向传递

让我们更仔细地看看这个反向传递的方案：

$\begin{aligned} & \bar{v}^{\prime}=\bar{v}\left(I-\frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)}\right)^{-1} \\ & \bar{v}^{\prime}=\bar{v}+\bar{v}^{\prime} \frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)} \end{aligned}$

这看起来很像不动点方程的解：

$z_{i+1}=\bar{v}+z_i \frac{\partial f\left(z^{\star}(x), x\right)}{\partial z^{\star}(x)}$

不动点算子的反向传递可以实现为AD伴随计算的不动点算子。

小结

最终的实现确实有点绕，所以如果你没有完全按照上面的做，也不用担心；
主要要点是，不动点算子的反向传递涉及解决另一个不动点操作，即各种伴随算子；
所有这些都可以在needle的Operator类中实现：即，在其compute_gradient()函数中创建FixedPointOp调用自身；

implicit layer的例子

多尺度深度均衡模型

关键思想：在DEQ模型的隐藏单元内维护多个空间尺度，并同时为所有这些尺度找到平衡点：

神经常微分方程

如果向量$z$遵循动力学$f$：

$\frac{d z}{d t}=f(z(t), t, \theta)$

可以通过从$z(t_0)$开始积分直到$t_1$得到$z(t_1)$：

$z\left(t_1\right)=z\left(t_0\right)+\int_{t_0}^{t_1} f(z(t), t, \theta) d t$

隐式层：

$y=\operatorname{odeint}\left(f, x, t_0, t_1, \theta\right)$

连续时间时间序列模型

神经ODE可以处理以不规则间隔收集的数据：

可微优化

DEQ和神经ODE都对层的性质施加实体结构，以获得实质性的表征能力。另一种强加不同（但相关）结构的常见策略是可微优化。层的形式为：

$z^{\star}=\underset{z \in \mathcal{C}(x)}{\operatorname{argmin}} f(z, x)$