Integrated electro-optic attention nonlinearities for transformers
Authors: Luis Mickeler, Kai Lion, Alfonso Nardi, Jost Kellner, Pierre Didier, Bhavin J. Shastri, Niao He, Rachel Grange
Organizations: Optical Nanomaterial Group, Department of Physics, ETH Zurich, Zurich, Switzerland · Optimization & Decision Intelligence Group, Department of Computer Science, ETH Zurich, Zurich, Switzerland · Department of Physics, Politecnico di Milano, Milan, Italy · Centre for Nanophotonics, Department of Physics, Engineering Physics & Astronomy, Queen’s University, Kingston, Canada
Transformers achieve state-of-the-art performance in language processing and computer vision. Their computation increasingly depends on hardware that can keep pace not only with large matrix multiplications, but also with nonlinear operations like Softmax within the core attention mechanism. On digital accelerators, Softmax accounts for only a small fraction of total operations, yet it can disproportionately impact overall inference latency. Here, we use thin-film lithium niobate Mach-Zehnder modulators (MZM) as analog computing elements for low-latency Transformer attention. Our high-speed nonlinearities based on the intrinsic sinusoidal MZM response achieve performance close to digital baselines in Vision Transformers and GPT-2, even under 4-bit input and output quantization. We experimentally characterize the electro-optic response at symbol rates up to 10 GBaud and assess model robustness under corresponding noise conditions. Accounting for the full signal chain, our analysis indicates more than an order of magnitude lower latency relative to previous custom electronic and photonic Softmax implementations. These results establish analog electro-optic devices as nonlinear function units well suited for integration into emerging co-packaged electronic-photonic hardware, offering a path toward high-speed Transformer inference.
Figures & tables
Figure 1: Accelerating attention nonlinearities in transformer with co-packaged optics (a) Latency-critical applications using transformer models, where each Transformer block consists of multi-head attention and feed-forward layers [ 57 ] . The self-attention computation relies on a nonlinear reweighting using Softmax/Sigmoid nonlinearity (targeted in this work). Matmul refers to matrix multiplication. Mask refers to the application of a triangular mask to prevent attending to future tokens in the sequence. (b) Conventional digital accelerators rely on fast tensor cores for matrix multiplications, but scarce low-throughput special-function units (orange) set a latency bottleneck. (c) A hybrid accelerator leveraging co-packaged optics as active nonlinear computational elements to accelerate attention. (d) The Mach-Zehnder modulators in co-packaged optics exhibit a nonlinear sinusoidal optical response to the supplied voltage.
Figure 2: Architecture and experimental validation of the Optmax and Optmoid attention mechanisms . (a) The Optmax architecture approximating the Softmax function: the rising sinusoidal slope of a first MZM mimics the exponential numerator, while the falling slope of a second MZM approximates the reciprocal normalization. The time-multiplexed optical intensities are integrated on a low-bandwidth photodiode to modulate the second MZM before high-speed detection. (b) The Optmoid architecture approximates the Sigmoid function using the full min-to-max swing of a single MZM, thereby eliminating the reciprocal-normalization arm required by Optmax . (c) A comparison between the simulated (black) and measured (colored) responses for a sequence at 5-bit resolution and 10GBaud . (d) The transfer function mapping digital input values to sampled optical amplitudes at 10GBaud for Optmax and Optmoid ; black markers denote the target response. (e) The statistical distribution of the relative error for both nonlinearities across symbol rates of 100 MBaud, 1 GBaud, and 10 GBaud.
Figure 3: Vision Transformer (ViT) performance using electro-optic nonlinearities. (a) A vision transformer classifies images using stacked transformer blocks. We drop-in replace different attention nonlinearities. (b) Final test accuracy comparison for ViT models on MNIST, CIFAR-10, and SVHN datasets. The proposed Optmax and Optmoid demonstrate competitive performance relative to standard Softmax and Sigmoid benchmarks. (c) Training loss and validation accuracy curves for the CIFAR-10 dataset, illustrating stable convergence across all tested activation functions. (d) Quantization sensitivity analysis showing test accuracy as a function of bit-depth (full precision down to 4-bit). (e) Noise robustness showing final test accuracy under additive Gaussian noise N(0,σ2) (no noise during training).
Figure 4: GPT-2 FineWeb-Edu results using electro-optic nonlinearities . (a) A large-language model predicts the next word-fragment in a sequence using stacked transformer blocks. We drop-in replace different attention nonlinearities. (b) Final test loss comparison for GPT-2 on FineWeb-Edu using 5k training steps. The proposed Optmax and Optmoid demonstrate competitive performance relative to standard Softmax and Sigmoid benchmarks. (c) Training and validation loss curves for the Optmax and Optmoid , illustrating stable convergence. (d) Quantization sensitivity analysis showing test loss as a function of bit-depth (full precision down to 4-bit) using 1.5k training steps. (e) Noise robustness showing final test loss under additive Gaussian noise N(0,σ2) (no noise during training) using 2.5k training steps.
Architecture
Latency (s) per Sequence
Energy (J) per Sequence
Type
Technology
Ref
nMOS SMA
5.5e-04
1.9e-08
Electronic
Analog
[ 51 ]
Softermax
7.7e-04
1.3e-08
Electronic
Digital
[ 54 ]
Softonic
1.7e-05
4.5e-11
Optic
Analog
[ 13 ]
VEXP
2.2e-07
5.0e-08
Electronic
Digital
[ 58 ]
Optmax
1.3e-08
1.0e-08
Electro-Optic
Analog
This work
Optmoid
6.5e-09
4.7e-09
Electro-Optic
Analog
This work
Table 1: Performance Comparison: Latency and Energy consumption for a sequence length n=64 . Optmax and Optmoid values are calculated at 10GBaud .
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Figure S1: GPU runtime of GPT-2 for different sequence lengths broken down by operation. To enable a clear attribution of runtime within scaled dot product attention, we use the PyTorch MATH backend [ 48 ] .
Figure S2: Experimental setup and signal waveforms. (a) Schematic of the MZM characterization setup, including the RF driving chain and optical detection path. (b) Time-domain traces of a 10GBaud sequence at the DAC output, RF amplifier output, and photodiode output. Noise is maximized at the MZM quadrature point where voltage-to-power transconductance is highest. (c) A full time trace of a n=2048 Optmax response under 10GBaud .
Figure S3: Calibration of Optmax and Optmoid electro-optic units. (a–b) Least-squares fit of MZM transmission T(V) and the resulting exponential approximation fexp(x) mapped to the rising slope for x∈[0,4] . (c–d) Characterization of the normalization stage MZM and the optimized factor N(z) for z∈[6,14] . The fit minimizes error against 1/z by tuning gain α and extinction ratio β . (e) Optmoid transfer function fsig(x) using the full min-to-max MZM swing with bias b=−3.93 . Solid lines represent hardware-aware forward models (including quantization qin,qout ), while dashed lines indicate ideal mathematical targets. (g-h) Simulation frameworks of Optmax and Optmoid without noise injection.
Figure S4: Noise effects in Optmax and Optmoid for ViT training (a) Additive noise model. (b) Multiplicative noise model. (c) Test accuracy versus additive noise applied only during inference, showing strong degradation under 4-bit quantization. (d) Test accuracy with additive noise during both training and inference, demonstrating improved robustness and performance. (e) Test accuracy under multiplicative noise applied only during inference, showing minor degradation. (f) Test accuracy under multiplicative noise during both training and inference, indicating stable performance. (g) Additive noise in Optmoid without output quantization, producing out-of-range values. (h) Additive noise in Optmoid with output quantization, illustrating threshold lifting effects. (i) Multiplicative noise in Optmoid in full precision, preserving non-negativity. (j) Multiplicative noise in Optmoid with quantization, maintaining zeroed values.
Figure S5: Latency and compute efficiency scaling of Optmax and Optmoid (a) Estimated latency per Softmax operation as a function of the sequence length N for three baud rates ( fB=1 , 10 , and 100GBaud ). The analysis follows the pipelined latency model described in Supplementary Sec. S5 , including digital-to-analog conversion, optical propagation through the TFLN MZMs, photodetection with transimpedance amplification, and analog-to-digital conversion. For all operating points, Optmoid exhibits lower latency than Optmax because it requires only a single modulation stage. (b) Estimated energy per Softmax operation as a function of the sequence length N for three baud rates ( fB=1 , 10 , and 100GBaud ). The energy is computed from the total electrical power consumption of each architecture, including the laser source, TFLN MZMs, DAC, RF driver, photodetection stage, and ADC, multiplied by the corresponding operation latency. Optmoid yields a lower energy per operation than Optmax across the full parameter range owing to its single-stage architecture.
Parameter
Value
Embedding dimension
256
Hidden dimension
512
Number of heads
8
Number of layers
6
Patch size
4
Number of channels
3
Appendix
Table 2: ViT model configuration.
Learning rate
3e-4
5e-4
7e-4
9e-4
Validation accuracy (%)
78.10
79.02
78.10
75.02
Appendix
Table 3: Softmax ViT validation accuracy (%) on CIFAR-10 for different learning rates.
xmin
xmax
Val. acc. (%)
-14
14
77.72
-12
12
77.18
-10
10
77.46
-6
6
76.52
-4
4
76.16
-2
2
75.88
Appendix
Table 4: Optmax ViT validation accuracy (%) for different xmin and xmax settings using vanilla 1/z normalizations.
zmin
zmax
Val. acc. (%)
0.5
5.5
79.78
0.5
6.5
79.90
0.5
7.5
79.66
0.5
8.5
79.96
0.5
9.5
79.60
0.5
10.5
79.36
Appendix
Table 5: Optmax ViT validation accuracy (%) for different zmin and zmax settings using [xmin,xmax]=[0,10] .
Bias b
Sigmoid Val. acc. (%)
Optmoid Val. acc. (%)
-17.16
79.22
79.54
-16.16
79.24
79.54
-15.16
80.08
79.54
-14.16
80.00
79.06
-13.16
80.18
79.06
-12.16
80.04
79.28
Appendix
Table 6: Sigmoid and Optmoid ViT validation accuracy (%) for bias b .
xmin
xmax
Val. loss
-14
14
5.985
-12
12
5.933
-10
10
5.928
-6
6
5.860
-4
4
5.832
-2
2
5.800
Appendix
Table 7: Optmax GPT-2 validation loss for different xmin and xmax settings using vanilla 1/z normalizations.
zmin
zmax
Val. loss
11.0
13.0
5.553
10.0
14.0
5.555
9.0
15.0
5.559
8.0
16.0
5.536
7.0
17.0
5.531
6.0
18.0
5.518
Appendix
Table 8: Optmax GPT-2 validation loss for different zmin and zmax settings using [xmin,xmax]=[0,4] .
Bias b
Sigmoid Val. loss
Optmoid Val. loss
-8.93
5.645
5.967
-7.93
5.575
5.967
-6.93
5.580
5.967
-5.93
5.759
5.967
-4.93
5.921
5.759
-3.93
5.903
5.713
Appendix
Table 9: Sigmoid and Optmoid GPT-2 validation loss for bias b .
Transformer inference studies how trained Transformer models produce outputs efficiently, accurately, and securely across hardware and context lengths.
Softmax is a key operation in Transformer attention, but its exponentiation and normalization add significant overhead in compute-in-memory (CIM) accelerators, especially when analog attention scores must first be converted to the digital domain. This work presents a tunable-temperature analog softmax circuit in GlobalFoundries 22-nm fully depleted silicon-on-insulator (FDSOI) technology that operates directly on CIM-generated score voltages without intermediate analog-to-digital conversion. Each input score is converted into a time-domain event using a shared falling ramp. The corresponding comparator transition samples an RC-decaying reference to generate an exponential weight, which is then processed by an in-circuit normalization stage. In contrast to analog softmax circuits that rely on transistor weak-inversion behavior for exponentiation, the proposed architecture controls the softmax response through the ramp slope and RC time constant, enabling programmable effective temperature. The 128-element architecture is evaluated using transistor-level and post-layout extracted simulations, including multi-level input vectors, capacitance variation and mismatch, process and temperature variation, monte carlo analysis, and shared-interconnect parasitics. The complete 128-element implementation occupies 9453.42~μm2 including the shared global ramp circuitry, while each replicated softmax element occupies 70.2~μm2. The circuit achieves a 242.97-ns evaluation latency at 13.44mW total power, corresponding to 25.5pJ per output element. The simultaneous 128-element evaluation achieves an RMSE of 24.46~mV relative to the ideal softmax response. The extracted circuit characteristics are further incorporated into a MemTorch-based hardware-aware Transformer model, where the proposed softmax achieves a validation loss within 2.5% of the ideal-softmax baseline.
Ankur Singh, Ashish Gautam, Shruti R. Kulkarni +1
Department of Electrical and Computer Engineering, University of Wisconsin–Madison, Madison, WI, USA. · ORNL Oak Ridge, Tennessee, USA.
We introduce Lookup-table (LUT)-aware Efficient Attention on FPGAs (LEAF). LEAF is the first nanosecond-scale LUT-aware transformer for FPGAs. Deploying transformers that are large enough to perform well, small enough to fit on an FPGA, and fast enough to meet the sub-microsecond timescales required by real-time systems such as those at CERN remains difficult, since the same attention mechanism that gives these models their high expressivity also makes them costly to implement in hardware. LEAF converts the query, key, and value projections and the final output feed-forward layer into FPGA-native layers with LUT-aware training, creating the first LUT-aware transformer and linformer layers built for ultra-low-latency inference. Attention resists LUT-aware training because it multiplies two activations where the method needs a constant weight to tabulate; LEAF tabulates the attention block's constant matrix-vector multiplications and leaves the rest in arithmetic form. On a standard jet tagging benchmark this improves accuracy by 1.8 percentage points while using 2.9 times fewer LUTs than the previous state-of-the-art sub-microsecond transformer at the same input size. Our most efficient LEAF-based linformer with 64 input particles reaches 80.4% classification accuracy at 28k LUTs and 35ns.
Arianna Cox, Chang Sun, Lauri Laatu +7
Imperial College London, United Kingdom · California Institute of Technology, USA · University of Bristol, United Kingdom
Recently, photonic transformer accelerators (PTAs) have successfully achieved significant speedup and energy efficiency improvements over electronic accelerators for expediting Transformer inference. However, state-of-the-art rely on expensive multi-wavelength light generation and large dot-product units due to active phase-shifter components, thus making their approach inefficient and impractical. To address this, we propose MDTransformer, a novel hardware-software co-design of PTA based on mode-division optical dataflow and operations. Specifically, MDTransformer performs complex matrix operations using spatial-mode interference, that leverages the inverse-designed multi-mode couplers, crossings, and Mach-Zehnder IQ modulators into a compact mode-division photonic tensor core (MPTC), capable of executing matrix multiplications in the optical domain. Its each guided mode (i.e., TE0-TE3) acts as an independent computational lane, enabling four-fold parallelism-per-waveguide without spectral filtering or free-spectral-range limitations. Moreover, its coherent detection and IQ modulation jointly encode amplitude and phase, realizing complex-valued arithmetic for full-range operations in transformers. MDTransformer offers analog multiplication with sub-4-bit effective precision and inter-modal crosstalk below -30 dB. Its inverse-designed approach also offers scalable and full compatibility with single-laser continuous-wave operation at 1550 nm. Experimental results show that MDTransformer achieves 40.4% area reduction, 63.6% power saving, 40.6% energy saving, and comparable latency over the state-of-the-art PTA across different workloads (i.e., DeiT-Tiny/Small/Base and BERT-Base/Large). These results show that MDTransformer offers a practical solution for high-performance and energy-efficient transformer-based systems.
Solomon Micheal Serunjogi, Rachmad Vidya Wicaksana Putra, Ayat Taha +2
Photonic Research Lab (PRL), Division of Engineering, New York University (NYU) Abu Dhabi, United Arab Emirates · eBRAIN Lab, Division of Engineering, New York University (NYU) Abu Dhabi, United Arab Emirates