CGF-Softmax: A Cumulant-Based Softmax Reformulation for Efficient Inference under Homomorphic Encryption
Status. Submitted to NeurIPS 2026 in May 2026 — currently under review.
Summary. Homomorphic encryption (HE) supports only addition and multiplication, so non-polynomial operations are handled with polynomial approximations — and higher-degree ones consume more multiplicative depth, driving up latency. This is especially costly for softmax, a core component of transformer attention, whose exponential and division operations are expensive under HE.
This work proposes CGF-softmax, which reformulates the softmax denominator statistically via the cumulant generating function (CGF). By absorbing normalization into a single exponential, it eliminates both homomorphic division and the comparison-heavy maximum subtraction, substantially reducing multiplicative depth. On Vision Transformers and large language models, CGF-softmax stays within 1% of exact-softmax accuracy while requiring 4–6× less depth.
My role. When I joined the lab, the CGF-softmax formulation existed but was difficult to use directly without additional training. My task was to establish its advantages over other training-based softmax methods. I designed and ran extensive experiments across a range of models and benchmarks to verify these advantages, and built the narrative that positions our method on the cost–accuracy Pareto frontier relative to both training-based and training-free approaches. I wrote about 80% of the manuscript draft. Through this work, I strengthened my ability to analyze, modify, and train AI models, as well as to implement non-linear operations under HE.