Network and Compiler Optimizations for Efficient Linear Algebra Kernels in Private Transformer Inference

Large language model (LLM) based services are primarily structured as client-server interactions, with clients sending queries directly to cloud providers that host LLMs. This approach currently compromises data privacy as all queries must be processed in the cloud and in the clear. Fully Homomorphic Encryption (FHE) is a solution to this data privacy issue by enabling computations directly upon encrypted queries. However, running encrypted transformer inference is challenging as programmers must map standard kernels to the constrained instruction set provided by FHE. In this work, we explore implementations of linear algebra kernels needed for transformer inference in FHE and understand how network optimization can help mitigate FHE costs while remaining performant. We leverage the Orion PyTorch to FHE framework to benchmark several linear algebra kernels in order to profile two linear transformation methods, packed row and BSGS, and find that BSGS outperforms packed row methods by up to $13.7 \times$ at transformer-level scales. We also incorporate network-level pruning strategies that reduce FHE runtimes of feed forward layers by up to $11.46\times$. Furthermore, we extend Orion to include ciphertext-ciphertext matrix-matrix products, a key component in the self-attention blocks. Finally, we perform a roofline analysis of FHE primitives and encrypted linear transformations and find that (SIMD encoded) implementations are memory-bound with primitives having roughly $0.1$ integer operations per byte of DRAM traffic. These findings illustrate the need for exploring alternative encoding schemes and models of computation within CKKS to unlock scalable private transformer inference. We conduct all experiments using the Orion framework which can be found at: https://github.com/baahl-nyu/orion.

💡 Research Summary

This paper, titled “Network and Compiler Optimizations for Efficient Linear Algebra Kernels in Private Transformer Inference,” addresses the critical challenge of making privacy-preserving transformer inference practical using Fully Homomorphic Encryption (FHE). While FHE allows computation on encrypted data, solving privacy concerns in cloud-based LLM services, its massive computational overhead and programming complexity remain significant barriers. The work focuses on optimizing the core linear algebra kernels—matrix-vector and matrix-matrix products—that form the computational backbone of transformer models under FHE constraints.

The research is built upon the Orion framework, which translates PyTorch code for execution under FHE. A primary contribution is the detailed comparison and benchmarking of two methods for plaintext-ciphertext matrix-vector products, essential for transformer feed-forward layers. The first is a row-packing method, which packs matrix rows into separate plaintexts. The second is a diagonal-packing method utilizing the Baby-Step Giant-Step (BSGS) algorithm. Extensive experiments reveal that the BSGS method dramatically outperforms the (packed) row method, achieving speedups of up to 13.7x for transformer-scale matrices found in models like GPT-2, Phi-3-mini, and Llama 3 8B. This highlights the paramount importance of algorithmic choice in FHE performance.

Secondly, the paper explores network-level optimizations to mitigate FHE costs. Non-linear activation functions like GeLU in feed-forward networks are expensive to evaluate homomorphically and often trigger bootstrapping. The authors demonstrate that strategically pruning these activation functions can lead to substantial latency reductions—up to 11.46x for the feed-forward layer—presenting a tangible trade-off between model accuracy and the efficiency of private inference.

Third, the work extends the Orion compiler to support ciphertext-ciphertext matrix-matrix multiplication, a crucial operation within the self-attention mechanism (e.g., for calculating QK^T). This enables computations where both operands are encrypted, expanding the scope of possible private inference scenarios.

Finally, to understand the fundamental performance limits, the authors conduct a roofline analysis of FHE primitive operations and encrypted linear transformations. A key finding is that even highly optimized, SIMD-encoded implementations are severely memory-bound, with an arithmetic intensity of only about 0.1 integer operations per byte of DRAM traffic. This indicates that the memory subsystem, not compute throughput, is the primary bottleneck under current encoding schemes.

In conclusion, the paper makes a compelling case that scalable private transformer inference requires a multi-faceted approach: compiler-level optimizations like advanced linear algebra algorithms (BSGS), network-level architectural adjustments (pruning), and ultimately, the exploration of novel encoding schemes and computational models within the CKKS FHE scheme to overcome the inherent memory bottleneck. The integration of these optimizations into the Orion framework provides a valuable toolkit and benchmark for future research in privacy-preserving machine learning.