Integrated electro-optic attention nonlinearities for transformers

📄 arXiv: 2604.09512v1 📥 PDF

作者: Luis Mickeler, Kai Lion, Alfonso Nardi, Jost Kellner, Pierre Didier, Bhavin J. Shastri, Niao He, Rachel Grange

分类: cs.LG, physics.optics

发布日期: 2026-04-10

💡 一句话要点

利用集成电光注意力非线性单元加速Transformer推理

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 电光计算 Transformer加速 非线性计算 薄膜铌酸锂 马赫-曾德尔调制器 模拟计算 低功耗 硬件加速

📋 核心要点

  1. Transformer模型中的Softmax操作虽占比小，却是推理延迟的重要瓶颈。
  2. 利用薄膜铌酸锂马赫-曾德尔调制器实现模拟非线性计算，替代传统数字Softmax和Sigmoid。
  3. 实验表明，该系统在保持精度的情况下，能有效降低非线性计算延迟，并具有良好的噪声鲁棒性。

📝 摘要（中文）

Transformer模型已成为主流神经网络架构，在语言处理和计算机视觉领域取得了领先性能。注意力机制是这些模型的核心，它需要使用Softmax函数进行非线性、非负映射。尽管Softmax运算在总运算量中占比不到1%，但它们可能会不成比例地成为整体推理延迟的瓶颈。本文利用薄膜铌酸锂(TFLN)马赫-曾德尔调制器(MZM)作为模拟非线性计算单元，大幅降低非线性计算的延迟。我们实现了数字Softmax和Sigmoid的电光替代方案，并在Vision Transformer和大型语言模型中评估了它们的性能。即使在模拟单元的激进的4位输入输出量化下，我们的系统也能保持极具竞争力的精度。我们进一步表征了高达10 GBaud编码速度下的系统噪声，并评估了模型在各种噪声条件下的鲁棒性。我们的研究结果表明，TFLN调制器可以作为混合共封装硬件中的非线性函数单元，从而实现高速和高能效的非线性计算。

🔬 方法详解

问题定义：Transformer模型在推理时，注意力机制中的Softmax等非线性函数计算成为性能瓶颈，尽管其计算量占比不高，但由于其复杂的计算过程，导致整体推理速度受限。现有方法主要依赖数字电路实现这些非线性函数，存在功耗高、延迟大的问题。

核心思路：本文的核心思路是利用薄膜铌酸锂(TFLN)马赫-曾德尔调制器(MZM)的电光特性，将非线性函数计算转移到模拟域进行，从而大幅降低计算延迟和功耗。通过控制MZMs的偏置电压，可以实现Softmax和Sigmoid等非线性函数的近似模拟。

技术框架：该方法构建了一个混合计算系统，其中线性计算仍然由数字电路完成，而Softmax和Sigmoid等非线性计算则由TFLN MZMs完成。输入数据首先被编码为电信号，然后通过MZMs进行非线性变换，最后将输出信号解码为数字数据。整个系统包括数据编码、MZMs非线性计算、数据解码三个主要阶段。

关键创新：最重要的技术创新点在于利用TFLN MZMs实现了模拟非线性计算，这与传统的数字电路实现方式有着本质区别。模拟计算具有天然的并行性和低功耗特性，可以显著加速非线性函数的计算过程。此外，该方法还探索了在低精度量化下的系统性能，证明了其在实际应用中的可行性。

关键设计：关键设计包括MZMs的偏置电压设置，用于控制其非线性响应曲线，使其近似于Softmax或Sigmoid函数。此外，还研究了输入输出数据的量化位数，发现即使在4位量化下，系统也能保持较高的精度。噪声分析也是关键设计的一部分，通过评估不同噪声条件下的模型鲁棒性，验证了系统的可靠性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，即使在激进的4位输入输出量化下，基于TFLN MZMs的非线性计算系统也能在Vision Transformer和大型语言模型中保持极具竞争力的精度。此外，系统在高达10 GBaud的编码速度下表现出良好的噪声鲁棒性，验证了其在实际应用中的可行性。

🎯 应用场景

该研究成果可应用于各种需要高速、低功耗Transformer推理的场景，如边缘计算设备、自动驾驶系统、移动设备等。通过将非线性计算卸载到电光硬件上，可以显著提高这些设备的推理速度和能效，从而实现更强大的AI应用。

📄 摘要（原文）

Transformers have emerged as the dominant neural-network architecture, achieving state-of-the-art performance in language processing and computer vision. At the core of these models lies the attention mechanism, which requires a nonlinear, non-negative mapping using the Softmax function. However, although Softmax operations account for less than 1% of the total operation count, they can disproportionately bottleneck overall inference latency. Here, we use thin-film lithium niobate (TFLN) Mach-Zehnder modulators (MZMs) as analog nonlinear computational elements to drastically reduce the latency of nonlinear computations. We implement electro-optic alternatives to digital Softmax and Sigmoid, and evaluate their performance in Vision Transformers and Large Language Models. Our system maintains highly competitive accuracy, even under aggressive 4-bit input-output quantization of the analog units. We further characterize system noise at encoding speeds up to 10 GBaud and assess model robustness under various noise conditions. Our findings suggest that TFLN modulators can serve as nonlinear function units within hybrid co-packaged hardware, enabling high-speed and energy-efficient nonlinear computation.