A novel method for input privacy from LLMs

One of the most interesting privacy technologies that I have come across is called Stained Glass Transform (SGT). This was invented by folks at Protopia AI (their team includes my talented friend and collaborator Sid Roy) and in this blog I am looking into their technical paper [1]. The problem it addresses is one that anyone building on cloud LLM APIs encounters: you want the model’s intelligence, but you don’t want the LLM provider to see your prompt/data.

The problem

When you call any hosted LLM endpoint (think ChatGPT, Claude.ai, OpenRouter, HuggingFace), you hand your prompt in the clear to a third-party server, which stores it in their database logs. This is a major concern given the increasingly personal nature of prompts and the mechanics of the data economy. A mechanism to let users benefit from LLMs while preserving the privacy of their input is therefore critical.

Existing solutions

There are a few different ways to resolve this prompt privacy challenge.

• $\textbf{Local hosting.}$ Host the model yourself so the prompt never leaves your environment. Ollama makes this straightforward, letting you run Llama, Mistral, Gemma, and other open-weight models on consumer hardware with a single command. The obvious limitation is compute: a capable model needs a GPU with sufficient VRAM. Beyond that, you forfeit all the infrastructure that comes for free with hosted endpoints: load balancing, auto-scaling, automatic retries, hardware maintenance, and the operational overhead of keeping a model server healthy in production.
• $\textbf{Fully Homomorphic Encryption (FHE).}$ FHE allows computations directly on encrypted data so your prompt is encrypted on-device and the server processes it without ever decrypting it. This Belfort Labs demo is a live in-browser experience that gives a feel for what FHE-based inference looks like in practice. On the open-source side, Zama’s Concrete ML is the leading library tackling the underlying hard cryptographic engineering. The downsides are steep: FHE inference is slower than plaintext, LM endpoints need significant re-engineering to operate over encrypted arithmetic (plaintext-ciphertext), and key management at scale is a non-trivial operational challenge.
• $\textbf{Trusted Execution Environments (TEEs).}$ TEEs (e.g. Intel SGX/TDX, AMD SEV, Confidential Containers) create hardware-isolated enclaves where code and data are hidden even from the host OS and cloud provider. This can be used to perform two-sided privacy where the server cannot see the user’s prompt and the model provider’s weights can simultaneously remain confidential. In practice, the user must still trust the hardware vendor’s attestation, GPU TEE support (needed for performant inference) is relatively new (NVIDIA Hopper is the first generation with production-ready confidential computing), and trust questions around the TEE hosting entity can undermine the privacy guarantees entirely.

Stained Glass Transform (SGT)

The Stained Glass Transform is a novel solution to the same problem with a well-studied and rigorous notion of privacy. The solution involves sending obfuscated embeddings instead of raw text to the LLM provider and letting the provider’s endpoint do the rest.

In other words, it moves the initial preparatory stages used by all LLMs (tokenization and embedding) to the user’s side. Using a trained machine learning model (their secret sauce), the embedding (and thus the prompt) is obfuscated. The key insight, however, is that this obfuscated prompt provides two empirically validated guarantees:

• $\textbf{(Utility preservation)}$ The LLM output on the obfuscated prompt is close to the LLM output on the raw text.
• $\textbf{(Privacy guarantee)}$ The raw text prompt is hard to reverse-engineer from the obfuscated embeddings.

Conventional Flow

Private Flow

User Prompt "How does a black hole form?"

Embedder

d = 2,048

Embedder

d = 2,048

Stained Glass Transform

obfuscate embeddings
using a trained model

LLM

LLaMA 3.2 1B · frozen

LLM

LLaMA 3.2 1B · frozen

Response (conventional)

Response (prompt private)

Token Embeddings

d = 2,048 dims · 7 tokens

LLaMA 3.2 1B

Training the SGT

The SGT paper is well-written and in this post I have simply followed their approach. While I describe my implementation choices such as architecture (which may not be fully detailed in the paper for IP reasons), I encourage the reader to refer to the paper for further details. The high-level idea is that you run a small local network (called the SGT) that takes the embedding sequence and replaces it with a perturbed version. Thus, the server never sees tokens or raw embeddings; it only ever processes the scrambled version. The crux of the work is showing how to efficiently train the SGT to preserve embedding privacy while retaining utility — that is, LLM output quality should not degrade.

Embedding transformation

The transform is stochastic:

The SGT predicts a deterministic shift $\mu_\theta$ and a per-dimension noise scale $\sigma_\theta$. Adding Gaussian noise with a learned variance means no two passes produce the same obfuscated embeddings — which is important for resisting repeated-query attacks.

Architecture

SGT is trained per model and in this blog, I use the model from the paper — Llama 3.2 1B. I chose the following small post-norm transformer encoder as the SGT module placed in front of the frozen LLM:

class SGT(nn.Module):
    def __init__(self, embed_dim=2048, num_layers=2, nhead=8):
        super().__init__()
        enc_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=nhead,
            dim_feedforward=embed_dim * 2, dropout=0.0,
            batch_first=True, norm_first=False,  # post-norm keeps output bounded
        )
        self.encoder = nn.TransformerEncoder(
            enc_layer, num_layers=num_layers, norm=nn.LayerNorm(embed_dim)
        )
        self.mu_head        = nn.Linear(embed_dim, embed_dim)
        self.log_sigma_head = nn.Linear(embed_dim, embed_dim)
        # initialize as identity: mu=0, small sigma
        nn.init.zeros_(self.mu_head.weight);   nn.init.zeros_(self.mu_head.bias)
        nn.init.zeros_(self.log_sigma_head.weight)
        nn.init.constant_(self.log_sigma_head.bias, -2.0)

    def forward(self, x, padding_mask=None):
        h         = self.encoder(x.float(), src_key_padding_mask=padding_mask)
        mu        = self.mu_head(h)
        log_sigma = self.log_sigma_head(h).clamp(-6.0, 3.0)
        eps       = torch.randn_like(h)
        x_tilde   = x.float() + mu + log_sigma.exp() * eps
        return x_tilde.to(x.dtype), mu.to(x.dtype), log_sigma.to(x.dtype)

For Llama 3.2 1B (embed_dim=2048), this SGT has 75.5 M parameters — about 6% the size of the LLM it protects. It runs locally in milliseconds per token; the LLM never needs to move.

Loss functions

Training balances four objectives simultaneously: one utility loss and three obfuscation loss components (refer to the paper for more details — the authors explain the challenges and their choices well). I trained over 40K OpenOrca examples for 5000 steps on a Google Colab T4 GPU (best checkpoint was at step 4500).

Utility — the obfuscated sequence should produce the same distribution of next tokens as the clean sequence. I use KL divergence instead of hard-label cross-entropy, because a 128 K-vocab LM’s probability mass is spread across many tokens. Hard-label gradients are too sparse to compete with the obfuscation losses through 16 frozen transformer layers.

def loss_utility(logits_obf, logits_clean):
    log_p_obf = F.log_softmax(logits_obf.reshape(-1, V).float(), dim=-1)
    p_clean   = F.softmax(logits_clean.detach().reshape(-1, V).float(), dim=-1)
    return F.kl_div(log_p_obf, p_clean, reduction="batchmean")

AbsCosine — push the obfuscated embedding orthogonal to the original. If $\lvert\cos(\tilde{x}, x)\rvert$ is near zero, the nearest-neighbour attack can’t find the original token:

def loss_abscosine(x, x_tilde):
    cos = F.cosine_similarity(x.reshape(-1, D), x_tilde.reshape(-1, D), dim=-1)
    return cos.abs().mean()

Norm penalty — keep obfuscated norms close to clean norms per token, so the LLM’s internal normalizations behave as expected:

def loss_norm_penalty(x, mu):
    clean_norms  = x.float().norm(dim=-1).detach()
    shifted_norms = (x.float() + mu.float()).norm(dim=-1)
    return (shifted_norms - clean_norms).abs().mean()

Mutual information — a minibatch Monte Carlo estimate of \(I(\tilde{x}; x)\) in nats per dimension, computed in float64 to avoid cancellation. This directly minimizes how much information \(\tilde{x}\) retains about \(x\) across the learned distribution, not just pointwise:

def loss_mi(x_tilde_A, mu_A, log_sigma_A, x_clean_B, mu_B, log_sigma_B):
    # H(x̃ | x) from diagonal Gaussian component entropy
    H_comp = (0.5 * LOG_2PIE + log_sigma_A_64).sum(dim=(-1, -2)).mean()
    # H(x̃) ≈ -E[log p_mix(x̃)] via minibatch GMM
    log_prob = -0.5 * (diff.pow(2) / var_B + log_const).sum(dim=(-1, -2))
    H_mix    = -torch.logsumexp(log_prob, dim=1).mean() + math.log(B_B)
    return ((H_mix - H_comp) / (T * d)).float()

The final combined loss uses weights \((\alpha_u, \alpha_\text{acs}, \alpha_\text{norm}, \alpha_\text{mi}) = (2.0, 0.3, 0.05, 0.15)\). Getting these weights right took three iterations — the main failure mode is \(\alpha_u\) so large that the utility loss keeps \(\sigma\) tiny, leaving mutual information high throughout training. The loss curves are below:

Note that the training is probably a reasonable local optimum given that the privacy metrics and utility are worse than those reported in the paper.

Does the LLM output give away the input?

The paper covers simple attack baselines and the same authors also construct a better reconstruction attack called BeamClean [2]. Given that their attack only considers the embedding vector, I was curious to see if a stronger attacker — one that can also see the text output produced by the model — could improve on BeamClean. BeamClean [2] finds the top vocabulary candidates at each token position by cosine similarity to the obfuscated embedding, scores them with a language-model prior, and runs beam search.

I implemented two extensions that use the observed LLM output as an additional signal with regularization to prevent the language model from exploiting quirks in the garbled output (e.g., preferring Does over does for superficial reasons). The results, however, have been mixed and not significant enough to generalize broadly. While the current evidence suggests that BeamClean+output is no stronger than BeamClean alone, I leave it as an open question to rigorously verify.

Takeaways

SGT is a genuinely clever idea. The key insight is that if the embedding layer can be made public, you can separate it from the inference pipeline to achieve strong privacy. This can be a great middle ground where the model owner retains ownership of the model while the user gets prompt privacy.

• It is genuinely surprising to me that a model can be trained to achieve two contrary objectives well: (1) obfuscation and (2) utility preservation. In this regard, SGT feels just as innovative as fully homomorphic encryption.
• I was able to train the model from scratch with limited resources. This is largely a credit to the paper being well-written and speaks to the academic community’s culture of knowledge sharing.
• Fully reproducing the paper’s reported NN-FR of 0.93 likely requires significantly more than 5,000 training steps. My checkpoint is a useful proof of concept that the approach works directionally, though not yet at the privacy levels claimed for production use.

References

1. J. Roberts, K. Mylonakis, S. Roy, and K. Kale. “Learning Obfuscations Of LLM Embedding Sequences: Stained Glass Transform.” arXiv:2506.09452, 2025. To appear at IEEE S&P 2026. arxiv.org/abs/2506.09452
2. K. Kale, K. Mylonakis, J. Roberts, and S. Roy. “BeamClean: Language Aware Embedding Reconstruction.” arXiv:2505.13758, 2025. arxiv.org/abs/2505.13758