Hyena Hierarchy: Towards Larger Convolutional Language Models

论文地址：

• https://arxiv.org/abs/2302.10866

整体思路以及计算方式

整体还是Toeplitz的思路，这里理一下如何计算。

首先是基本操作$f$：

$$f(\mathbf X_1, \mathbf X_2, \mathbf T)= \mathbf X_1 \odot (\mathbf T \mathbf X_2)$$

其中$\mathbf T$是Toeplitz matrix。

然后是Hyena的计算方式：

• 假设Hyena的循环次数为$m$，输入为$\mathbf X \in \mathbb R^{n\times d}$；

• 每次循环使用基本操作$f$；

• 第一步得到初始输入和每次循环中使用的$\mathbf X_1$，即$f$的输入1；

  • $\mathbf U = \mathbf X \mathbf W \in \mathbb R^{n\times (m+1)\times d}$；

  • $\mathbf V=\mathrm{Conv1d}(\mathbf U) \in \mathbf X \mathbf W \in \mathbb R^{n\times (m+1)\times d}$；

    • kernel size为3，这一步的作用是进行local token mixing；

  • 将$\mathbf V$按照第二个维度拆分为$m+1$个向量：

    • $\mathbf V^1,\ldots, \mathbf V^m, \mathbf X^0 \in \mathbb R^{n\times d}$；

• 第二步利用一个网络计算每一步的Toeplitz matrix：

  • $\mathbf T^1,\ldots, \mathbf T^m\in \mathbb R^{n\times n}$；

  • 使用了类似Tnn中的Rpe加上指数衰减；

• for i in $1,\ldots, m$:

  • $\mathbf X^{i}= f(\mathbf V^i, \mathbf X^{i-1}, \mathbf T^i)$

• return $\mathbf X^{m}$；

时间复杂度

$O(nmd\log n + nd^2)$。

代码

• https://github.com/HazyResearch/safari

简评

非常有意思的工作：

• local token mixing看起来比较关键；

• 如何去掉那个循环，是一个值得研究的问题；