Martins P., Marinho Z. and Martins A. \(\infty\)-former: Infinite Memory Transformer. arXiv preprint arXiv:2109.00301, 2021.

在transformer中引入一种长期记忆机制.

主要内容

假设\(X \in \mathbb{R}^{L \times d}\), 即每一行\(x_i\)代表一个token对应的特征.

Attention需要进行如下的步骤:

\[Q = XW^Q, K = X W^K, V = XW^V, \\
Z = \mathrm{softmax}(\frac{QK^T}{\sqrt{d}})V.
\]

为了符号简易起见, 我们不考虑multi-head的情形, 下面的思想可以直接应用之.

我们知道, 可以通过径向基函数来逼近任意的连续函数:

\[\sum_{k} b_k \psi_k (t) \rightarrow f(t).
\]

现在, 我们令\(t_i = \frac{i}{L}\), 即对\(L\)个tokens冠以时序, \(X\)的每一列都可以看成一个特殊的\(f_j(t)\)的位于\(t_i, i=0,1,\cdots, L-1\)处的值.

给定\(N\)个基函数\(\psi_k (t), k=0,1,\cdots, N-1\), 我们要通过求解系数\(\bm{b}_j = [b_{j0}, b_{j1},\cdots b_{j,N-1}]^T\)来逼近\(f_j\)(\(X\)的第\(j\)列).

设\(\Psi \in \mathbb{R}^{N \times L}, \Psi_{ki}=\psi_{k}(t_i)\), \(B \in \mathbb{R}^{d \times N}, B_{jk} = b_{jk}\).

作者通过岭回归来求解系数\(b\):

\[B = \arg \min_{B} \|B \Psi - X^T\|_F^2 + \lambda \|B\|_F^2,
\]

其显示表达式为:

\[B = X^T\Psi^T(\Psi\Psi^T + \lambda I)^{-1}.
\]

\[X^T \approx B\Psi \rightarrow x_i \approx B \psi (t_i).
\]

现在我们用\(\tilde{X} := \Psi^T B^T\)来代替\(X\), 则

\[K = \tilde{X} W^K = \Psi^TB^TW^K, \tilde{V} = \tilde{X}W^V = \Psi^TB^TW^V.
\]

注意, 我们并不对\(Q\)进行替换, 因为这个只是用作长期的记录用, Q每次重新计算.

对于每个\(q_i\), 我们构建一个其关于\(t\)的密度函数\(p_i(t)\), 文中假设其满足高斯分布:

\[\mathcal{N}(t; \mu_i; \sigma_i^2).
\]

\(\mu_i, \sigma_i^2\)分别通过如下估计:

\[\mu_i = \mathrm{sigmoid} (w_{\mu}^T K q_i)
=\mathrm{sigmoid} (w_{\mu}^T B^TW^K q_i), \\
\sigma^2_i = \mathrm{softplus} (w_{\sigma}^T K q_i)
=\mathrm{softplus} (w_{\sigma}^T B^TW^K q_i). \\
\]

注意最后令\(w^T\Psi^T = w^T\)既然\(\Psi\)是事先确定的.

我们知道

\[\mathrm{softmax}(\frac{Kq_i}{\sqrt{d}})
\]

实际上求解的是一个离散化的\(p_i(t)\), 即\(q_i\)和\(k_j\)的相合程度, 而

\[\mathrm{softmax}(\frac{Kq_i}{\sqrt{d}})^TV
\]

实际上就是求解期望

\[\mathbb{E}_{p_i}[v(t)].
\]

现在我们近似了一个连续的\(p_i(t)\), 也可以通过这种方式得到最后的\(z_i\):

\[\mathbb{E}_{p_i}[v(t)]
=\mathbb{E}_{p_i}[\psi^T(t)B^TW^V]
=\mathbb{E}_{p_i}[\psi^T(t)]B^TW^V.
\]

当我们取\(\psi\)为高斯径向基函数的时候, 上述是由显示解的.

现在来剖析一下, 好在哪里?

原本的\(K\)是\(L\times d\)的, 现在由于我们只需要计算\(B^TW\), 故实际上只有\(N \times d\), 我们可以选取很大的\(L\)但是选择较小的\(N\)来避免较高的复杂度.

如何扩展?

难不成每一次都要重新计算\(B\)? 倘若真的是这样就谈不上是长期记忆了.

作者采取了一种比较巧的方法, 实际上, 现在的\(B\psi(t)\)可以看成是一个\(d\)维的向量函数.

我们首先将其进行压缩至\([0, \tau], \tau \in (0, 1)\):

\[B\psi(t /\tau),
\]

如此一来, 整个函数的能量集中在\([0, \tau]\)中, 我们可以用剩下的\((\tau, 1]\)来放置新的\(X\).

我们首先从\([0, \tau]\)中采样\(M\)个点\(t_0, \cdots, t_{M-1}\), 并得到:

\[X_{past} = [x_0, \cdots, x_{M-1}]^T \in \mathbb{R}^{M \times d}, x_m=\psi^T(t_m/\tau)B^T.
\]

加上新的\(X_{new}\), 我们有

\[X = [X_{past}^T, X_{new}^T]^T \in \mathbb{R}^{(M + L) \times d},
\]

对\(X\)按照上面的逻辑重新估计\(B\)即可更新记忆.

关于如何采样这\(M\)个点, 作者提了一种sticky memories的方法, 将其与密度函数联系在一起, 便不细讲了.

实验细节

在看这篇论文的时候, 困扰我的就是这个径向基函数是怎么选的?

举一个作者在Language Modeling中的例子便可:

选取150个高斯径向基函数\(\mathcal{N}(t;\mu, \sigma^2)\), 其中

\(\mu\)从\([0, 1]\)中均匀采样, \(\sigma \in \{0.01, 0.05\}\).

还有用KL散度防止一般化就不讲了. 感觉本文有趣的点就是压缩这个地方, 还有对\(\Psi\)的处理.

随机推荐

  1. A Child's History of England.16

    CHAPTER 5 ENGLAND UNDER CANUTE THE DANE Canute reigned eighteen years. He was a merciless King at fi ...

  2. Spark基础:(六)Spark SQL

    1.相关介绍 Datasets:一个 Dataset 是一个分布式的数据集合 Dataset 是在 Spark 1.6 中被添加的新接口, 它提供了 RDD 的优点(强类型化, 能够使用强大的 lam ...

  3. Hive(八)【行转列、列转行】

    目录 一.行转列 相关函数 concat concat_ws collect_set collect_list 需求 需求分析 数据准备 写SQL 二.列转行 相关函数 split explode l ...

  4. jQuery无限载入瀑布流 【转载】

    转载至 http://wuyuans.com/2013/08/jquery-masonry-infinite-scroll/ jQuery无限载入瀑布流 好久没更新日志了,一来我比较懒,二来最近也比较 ...

  5. 【vector的输出问题】 洛谷 P1996 约瑟夫问题

    题目:P1996 约瑟夫问题 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn) 可恶啊,本来是一道不算难的题,硬是因为cin,cout同步流卡了我一天qwq 关闭cin,cout同步流 ...

  6. Linux系统根目录下各文件夹介绍

    参考自:[1]Linux 系统根目录下各个文件夹的作用 https://www.cnblogs.com/jiangfeilong/p/10538795.html[2]了解Linux根目录"/ ...

  7. cordova配置与开发

    1.环境配置 1.1.安装ant 从 apache官网 下载ant,安装并配置,将ant.bat所在目录加到path环境变量,如c:\apache-ant\bin\.在cmd中运行以下语句如不报错即可 ...

  8. JDK的简介,卸载和安装过程

    JDK的简介,卸载和安装过程 JDK JRE JVM JDK:Java Development Kit JRE:Java Runtime Environment JVM:Java Virtual Ma ...

  9. 我的第一篇博客blog,笑哭

    我的第一篇博客blog Markdown学习 一级标题:#加一个空格 加 文字, 二级标题:加2个##以此类推 字体 粗体:hello world!字体前有二个星号,字体后有二个星号 斜体:hello ...

  10. 『学了就忘』Linux服务管理 — 77、RPM包安装基于xinetd的服务的管理

    目录 1.基于xinetd服务的启动管理 (1)telnet服务安装 (2)telnet服务启动 2.基于xientd服务的自启动管理 现在Linux系统中基于xinetd的服务越来越少了,但Linu ...