Integrated electro-optic attention nonlinearities for transformers

Luis Mickeler; Kai Lion; Alfonso Nardi; Jost Kellner; Pierre Didier; Bhavin J. Shastri; Niao He; Rachel Grange

cs.LGApr 10, 2026

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

Abstract

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 $10\text{\,}\mathrm{G}\mathrm{B}\mathrm{a}\mathrm{u}\mathrm{d}$ . (d) The transfer function mapping digital input values to sampled optical amplitudes at $10\text{\,}\mathrm{G}\mathrm{B}\mathrm{a}\mathrm{u}\mathrm{d}$ 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 $\mathcal{N}(0,\sigma^{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

\mathcal{N}(0,\sigma^{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

10\text{\,}\mathrm{G}\mathrm{B}\mathrm{a}\mathrm{u}\mathrm{d}

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 $10\text{\,}\mathrm{G}\mathrm{B}\mathrm{a}\mathrm{u}\mathrm{d}$ 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 $10\text{\,}\mathrm{G}\mathrm{B}\mathrm{a}\mathrm{u}\mathrm{d}$ .

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

f_{\exp}(x)

mapped to the rising slope for

x\in[0,4]

. (c–d) Characterization of the normalization stage MZM and the optimized factor

N(z)

for

z\in[6,14]

. The fit minimizes error against

1/z

by tuning gain

\alpha

and extinction ratio

\beta

. (e) Optmoid transfer function

f_{\text{sig}}(x)

using the full min-to-max MZM swing with bias

b=-3.93

. Solid lines represent hardware-aware forward models (including quantization

q_{\text{in}},q_{\text{out}}

), 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 (

f_{B}=1

10

, and

100\,\mathrm{GBaud}

). 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 (

f_{B}=1

10

, and

100\,\mathrm{GBaud}

). 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.

$x_{\min}$	$x_{\max}$	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

x_{\min}

and

x_{\max}

settings using vanilla

1/z

normalizations.

$z_{\min}$	$z_{\max}$	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

z_{\min}

and

z_{\max}

settings using

[x_{\min},x_{\max}]=[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

$x_{\min}$	$x_{\max}$	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

x_{\min}

and

x_{\max}

settings using vanilla

1/z

normalizations.

$z_{\min}$	$z_{\max}$	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

z_{\min}

and

z_{\max}

settings using

[x_{\min},x_{\max}]=[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

Transformer inference studies how trained Transformer models produce outputs efficiently, accurately, and securely across hardware and context lengths.

See all 57 papers

Compute-in-Memory Attention: A Time-Domain Analog Softmax Circuit with RC-Tunable Temperature
Sep 2, 2026Ankur Singh, Ashish Gautam, Shruti R. Kulkarni +1Softmax Attention Analog Computing
Department of Electrical and Computer Engineering, University of Wisconsin–Madison, Madison, WI, USA. · ORNL Oak Ridge, Tennessee, USA.
LEAF: LUT-aware Efficient Attention on FPGAs
Oct 19, 2026Arianna Cox, Chang Sun, Lauri Laatu +7Efficient Transformer Inference FPGA-Based Neural Network Acceleration
Imperial College London, United Kingdom · California Institute of Technology, USA · University of Bristol, United Kingdom
MDTransformer: A Hardware-Software Co-Design of Mode-Division Photonic Transformer Accelerator with Inverse-Designed Coherent Crossbar
Jul 28, 2026Solomon Micheal Serunjogi, Rachmad Vidya Wicaksana Putra, Ayat Taha +2Efficient Transformer Inference Photonic Computing
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

Sep 2, 2026J/K move · Enter open · S save

Compute-in-Memory Attention: A Time-Domain Analog Softmax Circuit with RC-Tunable Temperature

Ankur Singh, Ashish Gautam, Shruti R. Kulkarni, Guojing Cong

Department of Electrical and Computer Engineering, University of Wisconsin–Madison, Madison, WI, USA. · ORNL Oak Ridge, Tennessee, USA.

Softmax Attention Analog Computing Compute-in-Memory Transformer Attention

Open paper AskPDF arXiv

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~

μ\mathrm{m}^{2}

including the shared global ramp circuitry, while each replicated softmax element occupies 70.2~

μ\mathrm{m}^{2}

. The circuit achieves a 242.97-ns evaluation latency at 13.44~~mW total power, corresponding to 25.5~~pJ 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.

Abstract

Transformer Inference

Explore similar work

Compute-in-Memory Attention: A Time-Domain Analog Softmax Circuit with RC-Tunable Temperature

LEAF: LUT-aware Efficient Attention on FPGAs

MDTransformer: A Hardware-Software Co-Design of Mode-Division Photonic Transformer Accelerator with Inverse-Designed Coherent Crossbar