[论文解读] L2ight: Enabling On-Chip Learning for Optical Neural Networks via Efficient in-situ Subspace Optimization

Jiaqi Gu, Hanqing Zhu+4 more|arXiv (Cornell University)|Oct 27, 2021

Neural Networks and Reservoir Computing被引用 7

一句话总结

L$^2$ight 提出了一种可扩展、高效的片上学习框架，用于基于硅光子学的光学神经网络（ONNs），采用三阶段流程：原位校准、基于解析解与零阶优化的并行光子核心映射，以及用于实时梯度评估的多级稀疏子空间学习。该方法相比先前方法实现了千倍量级的可扩展性提升和30倍以上的效率提升，首次实现了片上百万参数ONN的训练。

ABSTRACT

Silicon-photonics-based optical neural network (ONN) is a promising hardware platform that could represent a paradigm shift in efficient AI with its CMOS-compatibility, flexibility, ultra-low execution latency, and high energy efficiency. In-situ training on the online programmable photonic chips is appealing but still encounters challenging issues in on-chip implementability, scalability, and efficiency. In this work, we propose a closed-loop ONN on-chip learning framework L2ight to enable scalable ONN mapping and efficient in-situ learning. L2ight adopts a three-stage learning flow that first calibrates the complicated photonic circuit states under challenging physical constraints, then performs photonic core mapping via combined analytical solving and zeroth-order optimization. A subspace learning procedure with multi-level sparsity is integrated into L2ight to enable in-situ gradient evaluation and fast adaptation, unleashing the power of optics for real on-chip intelligence. Extensive experiments demonstrate our proposed L2ight outperforms prior ONN training protocols with 3-order-of-magnitude higher scalability and over 30X better efficiency, when benchmarked on various models and learning tasks. This synergistic framework is the first scalable on-chip learning solution that pushes this emerging field from intractable to scalable and further to efficient for next-generation self-learnable photonic neural chips. From a co-design perspective, L2ight also provides essential insights for hardware-restricted unitary subspace optimization and efficient sparse training. We open-source our framework at https://github.com/JeremieMelo/L2ight.

研究动机与目标

• 解决由于硬件限制和非理想光子元件导致的光学神经网络（ONNs）中片上训练可扩展性与效率的关键挑战。
• 通过克服梯度评估与电路校准的局限性，实现在可编程光子芯片上的实时、原位学习。
• 开发一种协同设计框架，支持大规模ONN训练（最高达百万参数模型），并具备对器件差异与噪声的鲁棒性。
• 集成多级稀疏性——反馈、特征与数据采样——以减少片上训练能耗与运行时间，同时不损失准确性。
• 展示所学子空间在任务迁移与自学习中的可迁移性与适应性。

提出的方法

• 提出三阶段学习流程：(1) 与变化无关的身份校准，以稳定非理想的光子状态；(2) 基于交替投影的并行映射，结合解析求解与零阶优化；(3) 用于原位梯度评估的多级稀疏子空间学习。
• 采用受限酉子空间优化，仅训练奇异值（Σ），而酉矩阵（U, V*）在映射后固定，从而减少参数空间，实现高效的片上自适应。
• 引入平衡反馈采样，以减少梯度估计期间所需的光场测量次数，最小化片上能耗。
• 应用保持信息的列特征采样，仅选择性地更新最相关的光子通路，提升训练效率。
• 通过运行时减少的数据采样，对训练数据进行均匀子采样，以最小化精度损失的方式减少计算量，并起到正则化作用。
• 结合三种稀疏机制——反馈、特征与数据——在保持模型准确性的同时，实现高训练效率。

实验结果

研究问题

• RQ1 在真实硬件约束下，光学神经网络中的片上学习能否扩展至百万参数模型？
• RQ2 在可观测性有限的光子硬件中，如何实现高效且节能的原位梯度评估？
• RQ3 在电路存在非理想性的情况下，受限子空间优化在多大程度上能保持模型的可学习性与可迁移性？
• RQ4 在多级反馈、特征与数据采样中，稀疏性与准确性的最优权衡是什么，以实现高效的ONN训练？
• RQ5 所提出的框架能否实现鲁棒、自适应的ONNs，实现实时补偿器件差异与噪声？

主要发现

• L$^2$ight 相比先前的片上训练协议，实现了超过千倍量级的可扩展性提升，可扩展至约1000万参数，而先前工作仅支持约100至2500参数。
• 该框架相比先前的稀疏训练方法，实现了超过30倍的训练效率提升，其中集成数据采样带来3倍速度提升，并减少梯度计算开销。
• 通过优化酉基（U, V*）的子空间学习，将准确率提升5–6%，同时能耗与训练步数分别降低9.9倍。
• 即使映射准确率仅为60%，子空间学习仍能将准确率恢复至约90%，表明对次优校准与映射具有强鲁棒性。
• 利用继承的酉基进行原位迁移学习，相比从零开始训练，准确率提高1–2%，训练步数减少3–5倍，证实了强大的可迁移性。
• 在权重梯度计算中使用约1%的稀疏性时，模型准确率仅下降约1%，证明效率增益不会损害模型性能。