中国科学院物理研究所
北京凝聚态物理国家研究中心
M02组供稿
第57期
2024年06月03日
基于自旋轨道力矩型磁性隧道结(SOT-MTJ)实现受限玻尔兹曼机

  ChatGPT等大语言模型的出现进一步促进了人工智能(AI)技术的发展,这些生成式模型本质上是以概率和统计原理为基础的大规模概率神经网络。尽管这些模型具有非凡的能力,但训练这些模型是一项计算密集型任务。随着传统CMOS技术发展放缓,经典的冯·诺依曼计算架构面临着严重的“内存墙”和“功耗墙”问题,而基于新兴的自旋轨道力矩驱动型磁性隧道结(SOT-MTJ)等非易失性存储技术,因其数据非易失性和高速高效特点,有望从硬件层面促进生成式AI技术的快速发展。

  受限玻尔兹曼机(RBM)是实现生成式人工神经网络的经典算法之一。在RBM算法中,除了矩阵乘法之外,被称为吉布斯采样(Gibbs Sampling)的随机采样操作也同样至关重要,RBM必须精准地根据网络状态预定义的概率分布函数(PDF)进行随机采样和状态更新,它才能最终收敛到所需的理想状态——在优化任务中找到网络能量的最低值或在学习任务中精准复现测试数据的分布规律。

  中国科学院物理研究所/北京凝聚态物理国家研究中心磁学室M02课题组开发了高性能自旋轨道转力矩(Spin-orbit Torque, SOT)驱动型磁性隧道结(MTJ)器件单元。受脉冲电压调控的SOT-MTJ可以在高低电阻态之间以0%到100%连续可调的概率产生二进制0/1随机数——即伯努利随机数发生器。并进一步将这种具有概率可调特性的SOT-MTJ器件,首次用于RBM的吉布斯随机概率采样操作,实现了多种应用场景的演示。

  SOT-MTJ的电压调控驱动磁矩翻转特性主要有两点,第一,它可以实现概率可调的真随机数发生器;第二,它的翻转概率和驱动电压满足sigmoid函数关系(图1)。而在受限玻尔兹曼机的核心操作—吉布斯概率采样中,人们正好需要sigmoid函数和概率可调的随机数。从这个角度讲,SOT-MTJ恰好可用于吉布斯采样操作的硬件加速。如果用传统CMOS电路实现相同的功能,则还需要sigmoid函数运算、伪随机数生成等操作。对于CMOS电路而言,这些操作无疑会显著耗费运算资源且效率低下。而采用SOT-MTJ器件只须输入一个合适大小的脉冲电压,随后测量MTJ的电阻值,便可完成上述吉布斯采样操作,简单易行、省时省力。

  以此为基础,该项工作首次在实验上验证了SOT-MTJ加速RBM吉布斯概率采样操作的可行性,实时进行了手写数字的图像识别和口述数字的语音识别,对于数字0 (或1),分别达到了100% (93%)和90% (95%)的识别准确率,这为SOT-MTJ下一步应用于AI硬件加速提供了可行的技术方案和器件原理验证(图2)。更进一步,利用语音和图像识别学习到的网络参数,将两个已知网络恰当地对接起来,组合出的新网络还能进行跨模态的识别和生成。例如给定0或1的语音,可以生成0或1的图像;反之亦然(图3)。这展示了该网络架构的联想能力和生成能力。该工作还探索了其他应用场景,包括失真图像恢复、数据加密解密、整数因式分解和可逆布尔逻辑等等(图4)。该研究表明利用SOT-MTJ可高效实现玻尔兹曼机吉布斯采样,并展示了利用自旋电子硬件加速生成式人工神经网络运算的可行性(图5)。

  该工作已在《Nano Letters》刊发。中国科学院物理研究所M02课题组韩秀峰研究员、万蔡华副研究员为论文共同通讯作者构思和指导了该项研发;中国科学院物理所博士生李枭翰是该论文第一作者。其他合作者参与了器件制备、测试、数据分析和论文写作。该工作得到了中国科学院、科技部重点研发计划和基金委重点基金等项目经费支持。[X. H. Li, C. H. Wan, X. F. Han, et al. Nano Lett. 24 (2024) 5420−5428]   https://doi.org/10.1021/acs.nanolett.3c04820.

图1.  SOT-MTJ的随机翻转性能。(a) SOT-MTJ的堆叠结构。(b) 低势垒和高势垒MTJ的能量分布示意图。(c) R-H测试,沿着MTJ的易磁化轴扫描平面内磁场。(d) 由50ns脉冲电流驱动的零场翻转。插图是MTJ的扫描电子显微镜(SEM)图像的顶视图。(e) 作为Vin的函数的切换概率(P)及其与匹配良好的S形曲线的拟合结果。(f)-(h) 在1.1 V (f)、1.2 V (g)和1.3 V (h)的特定电压下连续测试获得的磁矩翻转导致磁电阻高(1)低(0)变化的实验结果。 

图2. 基于SOT- MTJ实现的受限玻尔兹曼机(RBM)原理图以及图像和语音识别。(a),以典型的RBM结构为例,25 个可见节点(vi)和2个 隐藏节点(hj)。训练过程分为三个部分,RBM训练(无监督)、逻辑回归(有监督)和最终测试(计算准确率)。每个节点的吉布斯采样可以都通过SOT-MTJ来实现。具体而言,根据S形函数sigmoid函数的概率采样P(vi=1|h)和P(hj=1|v)可以直接用SOT-MTJ的简单磁矩翻转调控来代替。可见(隐藏)节点的条件概率P取决于来自其连接的隐藏(可见)节点的总体加权求和。这里的参数AB用于重整化。(b) 手写数字数据集(MNIST)的“0”和“1”示例以及平均池化后的最终训练数据。(d),自由发音数字集(Free-Spoken-Digit-Datasets)的“0”和“1”示例以及快速傅立叶变换(FFT)和平均池化后的最终训练数据。(c) &(e),图像识别(c)和语音识别(e)的实验和模拟识别准确率。

图3. 在合并的RBM上进行图像和语音跨模态学习的实验结果。(a)和(b)展示自由发音数字集中0(a)和1(b)音频。(c) 两个RBM通过一个独热编码(one-hot)层连接。左边的输入层包含音频信息,最右边的输出层包含图像信息,反之亦然。(d)和(e)则分别用(a)和(b)中的语音数据使用(c)中的网络来生成0(d)和1(e)的图像。

图4。基于SOT-MTJ的RBM展望。(a-c),RBM在生成(a)、图像修复(b)和数据加密(c)中的模拟结果。(d),可逆XOR门的实验结果,从给定的[C]中推断出可能的[AB]组合。(e) 由网络权重定义的能量(E)是可逆XOR门的输入(x,y)和输出(z)的空间坐标的函数。左图是具有所有零初始化权重的相位空间;右边是训练后的相空间。黄点对应于1的输出的正确XOR状态,而银点表示XOR门应避免的错误状态。正确状态的选择是通过正确塑造这里的能量空间来实现的。(f) 在实验中使用基于SOT-MTJs的RBM的整数分解过程。(g) 基于SOT-MTJ的吉布斯采样可以与同样基于MTJ阵列的矩阵乘法相结合,以构建全自旋电子硬件RBM加速器。


图5. 利用SOT-MTJ实现图像和语音跨模态学习的示意图。