Time-aware large kernel convolutions

论文地址：

补充：这篇论文提供了cuda代码。

整体思路以及计算方式

整体思路是利用卷积的方式进行序列建模，看完之后感觉非常赞，这里详细理一下计算思路。

步骤一，simple conv1d，序列建模的方式。

以一维序列 $x_0,\ldots, x_{n-1}$ 为例，利用部分和（kernel值为1的卷积）的形式得到输出：

o_i=\sum_{j=\alpha_i^l}^{\alpha_i^r} x_j

其中 $\alpha_{i}^l , \alpha_i^r$ 为 $o_i$ 对应的边界值，注意上式的计算复杂度太高，但是可以构造前缀和降低计算复杂度：

\left\{\begin{array}{l}\mathcal{S}_0=0 \\ \mathcal{S}_i=\mathcal{S}_{i-1}+x_i, \quad 1 \leq i \leq n .\end{array}\right.

那么：

o_i=\mathcal{S}_{a_i^r}-\mathcal{S}_{a_i^l-1}

那么现在的问题就是如何计算 $\alpha_i^r, \alpha_i^r$ ，这在步骤二中可以解决。

步骤二，确定上界和下界。

首先构造可学习的参数：

\tilde{a}_i^{\{l, r\}}=\sigma\left(f^{\{l, r\}}\left(x_i\right)\right) \in[0,1]

然后利用下式计算边界：

\begin{aligned} a_i^l & =i-\tilde{a}_i^l \cdot l_{\max } \\ a_i^r & =i+\tilde{a}_i^r \cdot r_{\max }\end{aligned}

其中 $l_{\max}, r_{\max}$ 是超参。现在的一个问题是， $a_i^l, a_i^r$ 不一定是整数，但是我们只能计算整数下标的 $\mathcal S_k, k\in \mathbb Z$ ，这一点利用插值即可解决：

\begin{aligned} \mathcal{S}_{a_i^l-1} & =\gamma^l \cdot \mathcal{S}_{\left\lfloor a_i^l\right\rfloor-1}+\left(1-\gamma^l\right) \cdot \mathcal{S}_{\left\lceil a_i^l\right\rceil-1}\\ \mathcal{S}_{a_i^r} & =\left(1-\gamma^r\right) \cdot \mathcal{S}_{\left\lfloor a_i^r\right\rfloor}+\gamma^r \cdot \mathcal{S}_{\left\lceil a_i^r\right\rceil}\end{aligned}

步骤三：归一化和鲁棒性。

这里作者为了让算法work，增加了归一化和dropout：

\tilde{o}_i=o_i \cdot\left(\frac{1}{l_{\max }+r_{\max }+1}\right)

步骤四：

之前讨论的的都是一维的情形，然后作者将其推广到 $d$ 维度时候发现性能一般，这里感觉主要问题是映射 $\tilde{a}_i^{\{l, r\}}=\sigma\left(f^{\{l, r\}}\left(x_i\right)\right) \in[0,1]$ 不太稳定，为了缓解这点，作者将 $d$ 维拆成 $h\times (d/ h)$ ，每 $d/h$ 个特征共享一个 $\alpha_i^{\{l,r\}}$ ，并且由这 $d/h$ 共同确定。

时间复杂度

因为是查找表，所以时间复杂度是 $O(nd)$ 。

代码

简评

非常赞的一个思路，其作用机理优点类似于window attention，其中window范围由输入确定。

PreviousSimple Hardware-Efficient Long Convolutions for Sequence Modeling NextResurrecting Recurrent Neural Networks for Long Sequences

Last updated 2 years ago