Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Traveling Salesman-Based Token Ordering Improves Stability in Homomorphically Encrypted Language Models

Donghwan Rho, Sieun Seo, Hyewon Sung, Chohong Min, Ernest K. Ryu

TL;DR

This work tackles privacy-preserving inference for large language models by using CKKS-based fully homomorphic encryption and introducing a Traveling Salesman Problem–driven token reordering to preserve semantic coherence under encryption noise. It couples a CKKS-friendly inverse transform sampling method with a post-processing step to bound and reduce embedding errors, and proves that TSP reordering further mitigates damage from imperfect sampling. The approach yields theoretical error bounds and improves text quality in experiments, achieving reduced corruption and enhanced coherence, especially with domain-specific fine-tuning. Altogether, the method advances practical and privacy-preserving LLM inference by enabling more stable encrypted text generation.

Abstract

As users increasingly interact with large language models (LLMs) using private information, secure and encrypted communication becomes essential. Homomorphic encryption (HE) provides a principled solution by enabling computation directly on encrypted data. Although prior work has explored aspects of running LLMs under HE, the challenge of text generation, particularly next-token prediction, has received limited attention and remains a key obstacle to practical encrypted interaction. In this work, we propose a TSP-based token reordering strategy to address the difficulties of encrypted text generation, together with a post-processing step that further reduces approximation error. Theoretical analysis and experimental results demonstrate that our method prevents collapse, improves coherence in generated text, and preserves data privacy throughout. Overall, our contributions advance the feasibility of practical and privacy-preserving LLM inference.

Traveling Salesman-Based Token Ordering Improves Stability in Homomorphically Encrypted Language Models

TL;DR

This work tackles privacy-preserving inference for large language models by using CKKS-based fully homomorphic encryption and introducing a Traveling Salesman Problem–driven token reordering to preserve semantic coherence under encryption noise. It couples a CKKS-friendly inverse transform sampling method with a post-processing step to bound and reduce embedding errors, and proves that TSP reordering further mitigates damage from imperfect sampling. The approach yields theoretical error bounds and improves text quality in experiments, achieving reduced corruption and enhanced coherence, especially with domain-specific fine-tuning. Altogether, the method advances practical and privacy-preserving LLM inference by enabling more stable encrypted text generation.

Abstract

As users increasingly interact with large language models (LLMs) using private information, secure and encrypted communication becomes essential. Homomorphic encryption (HE) provides a principled solution by enabling computation directly on encrypted data. Although prior work has explored aspects of running LLMs under HE, the challenge of text generation, particularly next-token prediction, has received limited attention and remains a key obstacle to practical encrypted interaction. In this work, we propose a TSP-based token reordering strategy to address the difficulties of encrypted text generation, together with a post-processing step that further reduces approximation error. Theoretical analysis and experimental results demonstrate that our method prevents collapse, improves coherence in generated text, and preserves data privacy throughout. Overall, our contributions advance the feasibility of practical and privacy-preserving LLM inference.
Paper Structure (31 sections, 3 theorems, 33 equations, 12 figures, 7 tables, 1 algorithm)

This paper contains 31 sections, 3 theorems, 33 equations, 12 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Prior work
  4. Language models under FHE.
  5. Next-token prediction methods of language models.
  6. Secure sampling methods under FHE.
  7. Preliminaries
  8. Cheon--Kim--Kim--Song (CKKS).
  9. Computation and threat model.
  10. Next token prediction of language models.
  11. Inverse transform sampling.
  12. Reordering tokens using TSP to prevent corrupted text generation
  13. With imperfect sampling, adjacent tokens affect each other
  14. TSP order mitigates damage from imperfect sampling
  15. CKKS-friendly and efficient SIMD-based inverse transform sampling
  16. ...and 16 more sections

Key Result

Theorem 1

Let $\mathbf{e}_{i(r)}$ denote the one-hot vector determined by inverse transform sampling from a random sample $r \sim U(0,1)$. For any homomorphic algorithm that computes an approximation $\tilde{I}$ of $\mathbf{e}_{i(r)}$, the following inequality holds: where the expectation is taken over the random sample $r$.

Figures (12)

  • Figure 1: The effect of TSP. (Left) Without TSP, we obtain the linear combination of semantically irrelevant token embeddings. (Right) With TSP, semantically similar tokens can be combined.
  • Figure 2: Visualization of adjacent token embeddings in a 2D plane using UMAP umap, with vectors normalized to the unit circle for cosine similarity comparison. In both figures, embedding vectors of tokens adjacent to the token 'connection' ($\blacktriangledown$) are plotted. Each green line connects two adjacent tokens, and longer lines indicate lower similarity. In the left figure, there are four long green lines, indicating the low cosine similarity between adjacent tokens. In contrast, with TSP (right), only one long line remains, showing that adjacent tokens are now semantically similar.
  • Figure 3: An illustrative example of Algorithm \ref{['alg:homomorphic_its']}. (Top) With the exact Heaviside function, we obtain the exact one-hot vector. (Bottom) With an approximate Heaviside function, the resulting vector is different from a one-hot vector. Note that the steps are labeled (i)--(vii) for the exact version, with an additional step (viii) in the approximate version to enhance the one-hotness of $\tilde{I}$.
  • Figure 4: The schematic illustration of our work. (Left) Standard next token prediction. (Right) Our HE-based text generation. First, a user encrypts the embedding vectors of the prompt and the model performs the inference. However, since the output is a ciphertext, the model cannot select the index of the next-token. Instead, (a) the model approximates the one-hot vector of the next token (red), and (b) computes the approximated embedding vector (purple) and concatenates it into the model input. After generation, the model sends the concatenation of (a) to the user, who decrypts and applies argmax to recover to token. See Figure \ref{['fig:tsp']} and Section \ref{['sec:TSP_reordering']} for details on how (a) and (b) affect text generation.
  • Figure 5: (Left) The train and evaluation loss and (Right) the ratios of zero-probability tokens during fine-tuning. After 60 steps, both the train and evaluation loss converge to zero, and the ratios of zero-probability tokens increase. In the right graph, we can see that the ratios of zero-probability tokens increase when TSP or post-processing (PP) is applied. Baseline means the case when the naive probabilistic sampling (not our algorithm) was used.
  • ...and 7 more figures

Theorems & Definitions (4)

  • Theorem 1
  • Definition 3
  • Theorem 5
  • Theorem 6