Large Memory Layers with Product Keys

发布 : 2021-06-29 分类 : 深度学习 浏览 :
image.png

3 Learnable product key memories


我们考虑函数 $m:\mathbb{R}^{d} \rightarrow \mathbb{R}^{n}$ 它将作为神经网络中的一个层。m的目的是在神经网络中提供大容量。

3.1 Memory design


高层结构: 我们memory整体结构如图 1 和图 2 所示。memory由三个组件组成:查询网络、包含两组子键的键选择模块和值查找表。它首先计算与积键集进行比较的查询。对于每个积键,它计算一个分数并挑选分数最高的k个积键。然后,通过与所选键相关联的值的加权总和,使用m(x)来产生输出。memory所有参数都是可训练的,但每个输入只更新k个memory slot。稀疏选择和参数更新使得训练和推理都非常有效。

查询生成:预处理网络: 函数 $q: x \mapsto q(x) \in \mathbb{R}^{d_{\mathrm{q}}}$ 简称查询网络,将d维输入映射到$d_q$维度的隐式空间。通常q是一个线性映射或者多层感知机,将维数从d减小到512.由于keys是随机初始化的,因此它们相对均匀地占据空间。在查询网络的顶部添加BN层,有助于训练时key的收敛。我们在4.5节中的消融实验证实了这一见解。

标准键分配与加权: 设q(x)表示查询网络,$\mathcal{T}_{k}$表示top-k运算符。给定键集 $\mathcal{K}=\left\{k_{1}, \ldots, k_{|\mathcal{K}|}\right\}$ 由$|\mathcal{K}|$个$d_q$维向量和一个输入x组成,我们使用查询$q(x)$选择前k个键来最大化内积:
image.png$\mathcal{I}$表示k个最相似的键(其中相似度量采用的是内积),w是表示与所选键关联的归一化分数的向量。所有这些操作都可以使用自动微分机制来实现,使我们的层可以在神经网络中的任何位置插入。

操作(2)和(3)仅取决于前k个索引,因此在计算上是有效的。相比之下,等式(1)的详尽比较对于大memories来说效率不高,因为它设计计算$|K|$个内积。为了规避这个问题,我们采用一组结构化的keys集合,称为product keys。

product key set 被定义为两个向量codebooks $C$和$C^\prime$的外积,与向量连接操作符有关。
image.png
这个笛卡尔积结构所引起的keys总数是$|\mathcal{K}|=|\mathcal{C}| \times\left|\mathcal{C}^{\prime}\right|$。集合 $C$ 和 $C^\prime$都包含一组维度为 $d_q/2$ 的子keys。我们利用这种结构来有效地计算最接近的keys $\mathcal{I} \in(1, \ldots, K)$。首先,我们将查询$q(x)$拆分为两个子查询 $q1$ 和 $q2$。然后,我们计算 $C$ 中最接近子查询 $q1$(相应 $q2$)的那 k 个子键(相应的 $C^\prime$):
image.png
可保证 K 中 k 个最相似的键的形式为$\left\{\left(c_{i}, c_{j}^{\prime}\right) \mid i \in \mathcal{I}_{\mathcal{C}}, j \in \mathcal{I}_{\mathcal{C}^{\prime}}\right\}$。图 2 显示了具有key选择过程的product keys示例。
image.png

本文作者 : HeoLis
原文链接 : http://ishero.net/Large%20Memory%20Layers%20with%20Product%20Keys.html
版权声明 : 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议。转载请注明出处!

学习、记录、分享、获得

微信扫一扫, 向我投食

微信扫一扫, 向我投食