清华大学李黄龙团队AEM：基于随机阻变阵列的稀疏自注意力机制实现方法

本文来源于Advanced Electronic Materials，欢迎浏览！

文献速读

针对大语言模型Transformer中注意力机制算法复杂度高的问题，本文提出利用具有电导分散性的忆阻器阵列实现稀疏自注意力机制的方法。将随机阻变阵列作为随机哈希矩阵，在硬件上一步完成特征向量的哈希分类，从而支持注意力分数的计算仅在相似向量间执行，降低了Transformer的计算复杂度。

正文导读

“注意力”是大语言模型Transformer的主要工作机制，可以处理输入序列元素之间的远程依赖关系。与传统的循环神经网络和卷积神经网络不同，Transformer架构只使用注意力机制，能够实现并行处理。随着实际应用需求的增加，Transformer模型被用于处理越来越长的序列。然而，原始的Transformer模型并不能很好地扩展到长度较长的序列，因为它的注意力分数计算复杂度与序列长度L成二次方关系，导致其内存消耗过高。为此，研究者提出了在注意力层中引入稀疏性来解决这一问题。在这些解决方案中，加州伯克利大学和谷歌的Kitaev等人提出的解决方案使用了基于随机投影（Random Projection, RP）的局部敏感哈希（Locality-Sensitive Hashing, LSH），使得查询向量和键向量被哈希分类到多个“桶”中，相似的向量以较大概率落在同一个桶中。因此，只需要对每个桶内的向量计算注意力分数，就可以得到与完全自注意力相近的结果，从而将复杂度降低到O（L logL）。然而，在实践中，虽然使用基于RP-LSH的稀疏自注意力模型可以提高内存效率，但其计算速度却相当慢，一个主要的瓶颈恰恰在于RP-LSH这一矩阵运算。

图1 基于稀疏注意力机制的Transformer模型原理图：（a）标准Transformer模型的流程图；（b）传统的完全注意力机制；（c）基于RP-LSH的稀疏注意力机制

本文作者提出RP-LSH稀疏自注意力模型实际应用的障碍可以通过新兴的非易失性存储器技术来解决。阻变器件在突破冯·诺依曼瓶颈方面的适用性已在大量文献中被报道，即在神经网络处理的线性加权求和步骤中，阻变器件被集成到阵列中执行一步点积计算。这种阻变阵列结构允许在数据存储的地方利用器件物理和其他电路定律进行原位计算。阵列中器件与器件（device-to-device, D2D）分散性是阻变阵列作为神经网络加速器执行线性加权求和时面临的主要挑战之一，传统观点认为必须减轻D2D分散性。本文作者则提出利用D2D分散性来实现RP-LSH过程。

图2 16 × 2 阻变阵列的随机电导分布：在（a1）RESET一次后的高阻态电导分布图；（a2）五次SET-RESET阻变操作后的高阻态电导分布图；（b1）-（b2）：（a）中器件的电导统计分布；（c1）-（c2）：（a）中相邻两列器件电导差值的统计分布

基于阻变阵列，本文以一种软硬件结合的方式，使用基于稀疏自注意力的Transformer模型执行序列预测任务，以更低的计算复杂度实现了70%以上的测试准确率。还可以通过利用阻变阵列中每次循环操作（cycle-to-cycle, C2C）之间的随机性，执行多轮RP-LSH，可以进一步提高准确率。

期刊简介

期刊Advanced Electronic Materials重点发表物理：应用、材料科学：综合、纳米科技相关方向的文章。

该期刊是一个跨学科论坛，在材料科学，物理学，电子和磁性材料工程领域进行同行评审，高质量，高影响力的研究。除了基础研究外，它还包括电子和磁性材料、自旋电子学、电子学、器件物理学和工程学、微纳机电系统和有机电子学的物理和物理性质的研究。期刊最新引文指标为0.9，最新影响因子为5.3（2023）。

点击投稿