Tyche: Probabilistic Selection over Encrypted Data for Generative Language Models

Generative AI, a significant technological disruptor in recent years, has impacted domains such as augmented reality, coding assistance, and text generation. However, the use of these models requires users to trust the model owners with sensitive data given to prompt the model. Fully Homomorphic Encryption (FHE) offers a promising solution, and many earlier works have investigated the use of this technology for machine learning as a service (MLaaS) inference applications. Still, these efforts do not cater to generative models that operate probabilistically, allowing for diverse and creative outputs. In this work, we introduce three novel probabilistic selection algorithms for encrypted autoregressive generative AI using the TFHE cryptosystem at 128 bits of security: multiplication-scaled cumulative sum, heuristic cumulative sum, and the random-multiplication argmax. Each of these approaches presents distinctive challenges in optimizing the tradeoff between precision and timing performance, a balance intricately tied to the specific characteristics of the data under consideration. Our results show that the loss function of the argmax-based random multiplication method scales better than the cumulative sum methods, highlighting the importance of developing efficient encrypted argmax functions.