基于具有原位训练能力的片上衍射光学器件的多模深度学习
DOI: https://doi.org/10.1038/s41467-024-50677-3
![]()
随着人工智能生成内容（AIGC）的爆炸性增长，多模态深度学习技术因其能够处理和分析多种数据类型的能力，成为支持AIGC领域的重要技术之一，目前最新版本的GPT-4已成功实现了对图像、音频等多模态数据的处理。然而，传统的微电子处理器在处理复杂的多模态任务时，受到摩尔定律的限制，难以进一步提高吞吐量和能效，难以满足日益增长的AI计算需求。因此，研究人员开始探索基于光学的新型计算平台，特别是光子神经网络（ONN），以应对这些挑战。光子神经网络利用光作为信息载体，通过光信号在光子器件中的传播实现计算，其具有高并行性和高能效的优势。近年来，多种ONN架构已被提出并成功应用于AI任务，包括基于集成马赫-曾德尔干涉仪（MZI）网格的相干架构、基于微环调制器阵列的波分复用（WDM）处理架构，以及基于相变材料（PCM）的交叉波导路由网络等。这些架构在特定数据集任务上展现了优异的性能，如高达每秒11万亿次运算的吞吐量（光子器件|Nature 面向光学神经网络的11TOPS光子卷积加速器）、超过每秒每平方毫米1万亿次的计算密度等，已在某些经典数据集任务上与数字计算机相媲美。然而，这些架构仍面临一些限制。光子器件单元的尺寸难以进一步缩小，硬件集成度有限，计算规模通常局限于较小的矩阵向量乘法。此外，由于光子计算核心尺寸较小，单个推理任务需要频繁调用光子计算核心，导致频繁的光电转换和额外的能耗。为了实现大规模的片上ONN架构，已经有许多工作提出了片上衍射光学神经网络（DONN）的概念。DONN通过在硅基平台上集成一系列硅槽实现隐藏层的功能，硅槽中填充二氧化硅作为“神经元”。然而，传统DONN芯片训练过程是在数字计算机上实现且在制造完成后其功能和权重参数固定，缺乏可重构性，难以适应多模态数据的处理需求。基于此，华中科技大学Cheng等人开发一种具有原位训练能力的多模态衍射光学神经网络（TDONN）架构，旨在利用光的传播即计算特性实现多模态数据的处理和分类。TDONN芯片集成了大规模可调谐元件，能够在光域内实现原位训练和快速收敛。经过训练，TDONN芯片已经成功地对不同的模态数据进行了分类，并在多模态测试集上达到了85.7%的准确率。我们的 TDONN 芯片具有高计算密度 （447.7 TOPS/mm2）、高能效 （7.28 TOPS/W） 和低光延迟 （30.2 ps） 的潜在吞吐量，可实现每秒 217.6 万亿次运算 （TOPS） 的潜在吞吐量。神经形态光子学的目标是将复杂的物理模型映射到光学神经网络的抽象模型中。如图1a展示了该工作用于多模态分类任务的光学神经网络模型，它由三部分组成：输入层、五个隐藏层和一个输出层。在特征提取和特征融合后，从不同模态（如视觉、音频和触觉）的数据集中获得特征向量，并将其作为神经网络的输入。特征向量的维度与输入层中的神经元数量相匹配，每个向量元素通过强度调制编码到光信号中。在隐藏层中，神经元按照多层布局排列，并且在训练过程中调整每个神经元之间的连接权重以实现目标函数。因此，可训练的神经元是实现可重构TDONN的重要前提。流入输出层的数据被视为输出向量，每个输出向量的元素对应于输出通道。每个通道的输出功率映射到推理结果中的概率。通过比较每个端口的输出功率，可以获得具有最高概率的标签作为TDONN架构的推理结果。图1 当前基于OTTDL的波束成形网络架构，具有M延迟级的开关OTTDL方案
TDONN芯片是在具有220nm顶层硅和2μm厚埋氧基底的绝缘体上硅（SOI）晶片上制造的。图1b描绘了TDONN芯片的结构，包括对应于神经网络模型的输入层、隐藏层和输出层。输入层由16个强度调制单元组成，用于加载输入数据，并将特征向量的元素值编码到光信号的强度中。隐藏层由五层可调谐衍射单元组成，每层16个单元。这些衍射单元用于模拟隐藏层中的神经元。为了调整每个衍射单元的参数，使用金属微加热器阵列。通过调整施加到每个加热器的电压，可以精确控制每个独立衍射单元附近的局部温度，从而改变隐藏层中相应区域的局部有效折射率。输出层包含四个输出端口，输出光信号由四通道光电探测器（PD）阵列接收。通过光电探测器实时检测输出光功率，可以获得每个标签的概率分布。通过原位训练，可以根据具体的多模态任务配置TDONN芯片的权重参数和功能。训练完成后，仅需要一次光的前向传播即可实现视觉、音频和触觉等多模态信息的处理和计算。为了便于理解TDONN架构中衍射的物理特性，该工作建立了一个衍射神经网络的通用理论模型来描述光的前向传播和误差后向传播：其中，𝐄m 是第m个隐藏层的光场，𝐅 是离散傅里叶变换矩阵，𝐅† 是离散傅里叶变换矩阵的共轭转置。𝚽m 是一个对角矩阵，其对角元素是第m个隐藏层中各个神经元的透射系数。由于振幅透射系数近似为1，透射系数即为相位延迟因子。𝐏m 是第m个和第（m+1）个隐藏层之间的衍射矩阵，zm 是衍射距离，ξ 是元素的行（列）编号，νξ 是相应的空间频率。
在实际实验中，波导边界会反射一些光，这些光可能会到达下一个衍射层，因此光子芯片的训练并不严格遵循理论模型，不能直接获得显式梯度，因此该共作将TDONN视为一个具有许多可训练权重参数的“黑箱”，通过实时检测输出层的光学响应来更新衍射网络的权重参数。具体而言，在TDONN原型中，光子芯片的输出由PD阵列实时检测，实时成本函数值作为数字后端的评估指标计算得出。然后，通过基于FPGA的控制框架配置应用于芯片上衍射单元的电压，以改善成本函数值，最终实现目标函数。
强大的可重构性能是多模态深度学习计算硬件的一个重要能力。以前的片上衍射网络大多不具有可重构性，一旦制造完成，其功能就无法改变，因此只能处理特定模态状态的数据。而TDONN芯片支持在光域内的原位训练和计算，并实现了视觉、音频和触觉数据的四分类。TDONN芯片的训练分为两步：第一步是对不同模态状态的输入数据进行预处理以提取特征；第二步是训练芯片的可调谐衍射单元以实现目标函数。为了实现这一目标，开发了定制化的梯度下降算法和光子神经元的光学dropout机制，并在第二步中应用。该工作使用MNIST数据集作为视觉数据集，使用Spoken_numbers_pcm数据集作为音频数据集。同时使用商用体感手套进行数据采集，手套上设置了五个传感器以检测五个手指的弯曲程度，并将探测信号的强度作为触觉数据。每个模态数据集分为训练集和测试集（训练集:测试集=4:1），并设置相同的四个标签：1、2、3、4。第一步的预处理流程如图2所示，该工作针对三种模态数据设计了不同的预处理方法以提高计算效率。预处理完成后，每个多模态输入被转换为16个特征值，并可以通过输入层中的16个强度调制单元编码到光信号上。接下来是通过训练隐藏层中的片上衍射单元来实现多模态数据的四分类任务。
图3 TDONN芯片的训练：a 原位训练的流程图; b 图像分类任务中的成本函数与训练迭代的对比; c–e 片上衍射单元的归一化功率分布; f Drop out算法的概念图; g 在部署具有不同层数的 drop-out 算法之前和之后的迭代次数差异图3a展示了原位训练算法的流程图，训练过程中需要定义一个成本函数（CF）作为评估训练收敛度的品质因数（FOM）。在算法开始迭代之前定义CF，并训练TDONN直到其收敛。CF：
即相关性函数，该工作的CF是统计学中的皮尔逊相关系数的一种变式，其中Mi是目标向量，Mexp_i是由四个输出端口在迭代过程中的实时光功率组成的实验向量。该工作以图像分类任务的训练过程为例，CF随着迭代次数的增加而逐渐增加，并在1500个epochs后收敛（图3b）。在训练500个epochs后，归一化的CF达到0.8，表明预测该标签的可能性明显高于其他三个标签。随后的迭代旨在进一步增大对比度，从而获得令人满意的分类结果。尽管在80个加热器的训练过程中可能存在串扰，但整体参数是联合配置的，因此不会影响训练结果。一旦训练完成，TDONN就可以以光速执行任务，并且可以通过重新配置施加到加热器的电压来切换到另一个任务。为了降低训练的复杂性，该工作引入了dropout检查过程，具体而言，如果某个加热器在N次迭代后对CF的改进没有贡献，则在下一次迭代中将其排除。在训练电子神经网络时，dropout算法广泛用于增强模型的泛化能力。具体来说，当数据向前传播时，神经元会以一定的概率停止工作。受此启发，该工作为光子神经元的原位训练开发了光学dropout算法，以加速TDONN的收敛。图3f展示了dropout算法的示意图。首先，为TDONN隐藏层中的每个神经元设置一个迭代阈值Titer，实验中Titer=10。在迭代过程中，如果某个神经元在T次调整后仍不能增加CF，则将其设置为非激活状态，并在后续迭代中不再调整。随着训练的进行，非激活神经元的数量逐渐增加，并且只需要调整剩余的激活神经元，这显著减少了训练成本，即时间开销。该工作也测试了在TDONN中部署dropout算法对不同层数的影响，并展示了在部署dropout算法前后完成相同任务所需的迭代次数（图3g）。可以观察到，当TDONN的层数少于3层时，dropout算法的加速收敛效果不明显，但随着层数的增加，加速效果变得更加明显，当层数为5层时，加速效果可达到36.5%。TDONN芯片包括一个输入层（16个神经元）、五个隐藏层（80个神经元）和一个输出层（4个神经元），总共有100个片上神经元。经过对多模态输入数据的预处理后，可以获得图像、音频和触觉信息的16个特征，加载至输入层并编码为光信号。图4a-d、e-h、i-l分别展示了一组视觉、音频和触觉模态下的分类概率分布，其中每种情况下峰值概率都占据主导地位。实验结果表明，TDONN芯片能够通过光域内的原位训练实现多模态分类任务。为了进一步验证TDONN芯片推理性能的普遍性，我们使用TDONN芯片检查了每种模态下的100个测试数据。图5m-o描绘了三种模态测试数据的混淆矩阵，其中分类准确率分别可以达到86%、82%和89%，平均准确率为85.7%。结果表明，TDONN芯片成功地对每种模态下的四个类别数据进行了分类，性能与数字计算机相当。通过扩展输入维度和网络深度，可以进一步提高分类准确率。为了更好地展示TDONN芯片的能力，该工作进一步使用TDONN芯片来分类手势5-8。由于手势5-8的含义不仅仅取决于手指的数量，这个分类任务更具挑战性。实验结果表明，经过训练的TDONN芯片能够成功识别手势5-8，并且分类结果如图5所示。
TDONN芯片的估计延时，及完成整个任务所需的训练和推理时间如下表所示。原型系统的延迟主要由四部分组成：光传播、TiN加热器的响应时间、DAC和ADC。光传播的延迟仅为每次迭代30.2ps，而原型的操作频率受限于TiN加热器的响应时间。为了匹配热光调制的响应时间，原型系统的操作频率在多模态任务演示中设置为10kHz。完成整个任务的训练过程大约涉及1000次迭代，训练完成后，只需一次前向传播即可实现推理。更重要的是，在TDONN芯片的衍射区域中的加热器状态在训练完成后无需改变，因此可以使用高速电光调制器和PD进行多模态推理任务，从而显著降低后续推理延迟。为了评估光学训练的相对效率，我们使用数字计算机（Intel® Core™ i9-12900K CPU，32 GB RAM）训练相同的多模态分类模型，并记录所需的训练时间。在数字计算机上，完成模型训练需要1226.47秒。在10千赫兹的操作频率下，TDONN仅需0.21秒即可完成模型训练，实现了超过5800倍的加速。总的来说，TDONN芯片在处理视觉、音频和触觉等多模态数据时，表现出了优异的分类性能和计算效率且芯片吞吐量可高达217.6 TOPS，计算密度为447.7 TOPS/mm2，系统级能效为7.28 TOPS/W，与其他ONN架构相比具有明显的优势，为实现低功耗、大规模光子AI模型提供了新的途径。
[1]Cheng, J., Huang, C., Zhang, J. et al. Multimodal deep learning using on-chip diffractive optics with in situ training capability. Nat Commun 15, 6189 (2024).
