# LaMemo Language Modeling with Look-Ahead Memory

论文地址：

* <https://arxiv.org/abs/2204.07341>

## 整体思路以及计算方式

之前Transformer中使用memory的方式都是当前token和memory中token交互，但是memory中token无法和当前token交互，本文就是对这点进行改进。

符号：

* 当前token：$$\mathbf {X}*{\tau}=\left\[\mathbf {x}*{\tau+1}, \cdots, \mathbf {x}\_{\tau+N}\right] \in \mathbb{R}^{N \times d}$$
* memory：$$\mathbf {X}*{\tau-1}=\left\[\mathbf {x}*{\tau-M+1}, \cdots, \mathbf {x}\_{\tau}\right] \in \mathbb{R}^{M \times d}$$
* $$\tilde{\mathbf {X}}*{\tau-1}=\left\[\mathbf {x}*{\tau-N+2}, \cdots, \mathbf {x}\_{\tau+1}\right] \in \mathbb{R}^{N \times d}$$

计算：

* $$\mathbf{Q}*{\tau}=\mathbf{X}*{\tau} \mathbf{W}*{q}, \mathbf{K}*{\tau}=\mathbf{X}*{\tau} \mathbf{W}*{k}, \mathbf{V}*{\tau}=\mathbf{X}*{\tau} \mathbf{W}\_{v}$$
* $$\tilde{\mathbf{K}}*{\tau-1}=\tilde{\mathbf{X}}*{\tau-1} \mathbf{W}*{k}, \tilde{\mathbf{V}}*{\tau-1}=\tilde{\mathbf{X}}*{\tau-1} \mathbf{W}*{v}$$
* $$\mathbf{C}*{\tau}^{\leftarrow}=\operatorname{Softmax}*{\text{lower-triangle}} \left(\frac{\mathbf{Q}*{\tau} \tilde{\mathbf{K}}*{\tau}^{\top}}{\sqrt{d}}\right) \tilde{\mathbf{V}}\_{\tau}$$
* $$\mathbf{C}*{\tau}^{\rightarrow}=\operatorname{Softmax}*{\text{upper-triangle}} \left(\frac{\mathbf{Q}*{\tau} \tilde{\mathbf{K}}*{\tau}^{\top}}{\sqrt{d}}\right) \tilde{\mathbf{V}}\_{\tau}$$
* $$\mathbf{C}*{\tau-1}^{\leftrightarrow}=\mathbf{\alpha}*{\tau} \operatorname{sg}\left(\mathbf{C}*{\tau}^{\rightarrow}\right)+\left(1-\mathbf{\alpha}*{\tau}\right) \mathbf{C}\_{\tau}^{\leftarrow}$$
  * $$\mathrm{sg}$$表示不计算梯度；
  * $$\mathbf{\alpha}*{\tau}=\frac{\operatorname{sg}\left(\mathbf{s}*{\tau}^{\rightarrow}\right)}{\operatorname{sg}\left(\mathbf{s}*{\tau}^{\rightarrow}\right)+\mathbf{s}*{\tau}^{{\leftarrow}}+\varepsilon}$$；
  * 其中$$\mathbf{s}*{\tau}^{\rightarrow}$$表示$$\mathbf{C}*{\tau}^{\rightarrow}$$中Softmax矩阵归一化之前的元素和；

## 时间复杂度

时间复杂度为$$O(N(N+M)d)$$。

## 训练以及loss

不变。

## 代码

* <https://github.com/thu-coai/LaMemo>

## 实验以及适用场景

适用于encoder和decoder；论文只测试了lm(decoder)场景，获得了一定的提升。

## 细节

暂无。

## 简评

提供了一种新的memory交互方式。