[论文解读] TRON: Transformer Neural Network Acceleration with Non-Coherent Silicon Photonics

Salma Afifi, Febin Sunny+2 more|arXiv (Cornell University)|Mar 22, 2023

Neural Networks and Reservoir Computing被引用 4

一句话总结

TRON 提出首个针对 BERT 和视觉 Transformer 等 Transformer 模型的硅光子神经网络加速器，利用非相干硅光子技术，实现至少 14 倍更高的吞吐量和 8 倍更好的能效，相较于最先进的电子加速器，显著提升自然语言处理和计算机视觉工作负载的高性能推理能力。

ABSTRACT

Transformer neural networks are rapidly being integrated into state-of-the-art solutions for natural language processing (NLP) and computer vision. However, the complex structure of these models creates challenges for accelerating their execution on conventional electronic platforms. We propose the first silicon photonic hardware neural network accelerator called TRON for transformer-based models such as BERT, and Vision Transformers. Our analysis demonstrates that TRON exhibits at least 14x better throughput and 8x better energy efficiency, in comparison to state-of-the-art transformer accelerators.

研究动机与目标

• 应对自然语言处理和计算机视觉中 Transformer 模型日益增长的计算需求。
• 克服电子加速器在处理注意力机制高精度、高带宽计算方面的局限性。
• 设计一种专为 Transformer 独特计算模式量身定制的光子硬件加速器。
• 与最先进的电子加速器相比，显著提升吞吐量和能效。
• 验证非相干硅光子技术在大规模 AI 推理中的可行性与性能提升潜力。

提出的方法

• 设计一种面向自注意力机制核心运算（矩阵乘法与 Softmax）的光子加速器架构。
• 采用非相干硅光子技术，相比相干光子系统，降低系统复杂度与成本。
• 实现混合电子-光子处理流水线，以处理控制逻辑与数据流协调。
• 利用波分复用技术并行处理多个注意力头与计算路径。
• 集成光互连技术，实现处理单元之间高带宽、低延迟的通信。
• 优化注意力计算到光域运算的映射，以最小化转换损耗与延迟。

实验结果

研究问题

• RQ1 非相干硅光子技术能否实现对 Transformer 神经网络的高效加速？
• RQ2 在基于注意力的模型中，光子加速器与电子同类在吞吐量与能效方面相比如何？
• RQ3 为使光子推理在真实世界 Transformer 模型中可行，需要哪些架构与信号处理创新？
• RQ4 光子硬件在多大程度上可降低自注意力计算的延迟与能耗？
• RQ5 在可扩展的 AI 加速器中使用非相干光学技术，其性能权衡如何？

主要发现

• TRON 相较于最先进的电子 Transformer 加速器，吞吐量至少提升 14 倍。
• 所提出的加速器相比现有电子解决方案，能效提升 8 倍。
• 非相干硅光子技术使可扩展且成本可控的光子加速器设计成为可能，且无需复杂的相位稳定机制。
• 该架构有效满足 BERT 和视觉 Transformer 中注意力机制对高精度与高带宽的需求。
• 性能提升主要归因于光域计算固有的并行性与低延迟互连。
• 结果验证了硅光子技术在突破传统电子极限、加速现代深度学习工作负载方面的可行性。