Secure LLM Inference with MP-SPDZ

LLMs are now the standard benchmark in encrypted AI (PPML) research. Producing real next-token predictions using research frameworks, however, remains challenging. In this post, we show how to evaluate a real HuggingFace LLM end-to-end inside an MPC framework without rewriting it from scratch.

In this post, we evaluate BERT’s [DCLT19] masked language modeling task end-to-end in MP-SPDZ. Parties jointly evaluate a pretrained HuggingFace model on secret-shared input sentences and obtain token predictions without revealing intermediate activations or model parameters.

Concretely, given a sentence such as

"Paris is the [MASK] of France"

we want to predict the masked token using secure multi-party computation.

Our goal is not to introduce a new protocol, but to examine how far an existing research framework can be pushed toward realistic transformer inference.

We will use an off-the-shelf BERT model from the HuggingFace Transformers library, widely adopted in both research and production systems. However, HuggingFace models are designed for flexible PyTorch execution, with dynamic control flow and various optimizations. This makes them difficult to directly translate into MP-SPDZ’s tracing-based compilation model.

MP-SPDZ

MP-SPDZ is a framework for secure multi-party computation designed for implementing and benchmarking MPC protocols [Kel20]. It supports a wide range of protocols across different security models and is widely used in research and industry prototyping [LI25] [JPMC25].

MP-SPDZ provides a high-level, Python-like language for writing MPC programs that provides secure types such as sint, sfix, and sfloat. These types have operations that are traced into a static computation graph during compilation. The graph is then compiled to bytecode, which the virtual machine executes at runtime using the selected MPC protocol.

Conceptually, the flow looks as follows:

A key constraint is that the computation graph must be fully determined at compile time. Control flow that depends on secret values, such as data-dependent branching or loops with secret bounds, cannot be expressed directly. Instead, such computations must be rewritten as arithmetic operations or implemented using techniques such as oblivious selection. In practice, this requires structuring models in a static, traceable way.

Machine Learning Support

MP-SPDZ includes an ml module with support for common machine learning layers and building blocks. It has been used for models such as LeNet and ResNet, and was recently extended to support transformer-style architectures.

One particularly powerful feature is the ability to convert PyTorch models directly into MP-SPDZ layers using layers_from_torch(). This function leverages PyTorch’s tracing mechanism to capture a model’s computation graph and translate each operation into its MP-SPDZ equivalent.

As a result, pre-trained PyTorch models can be executed in MPC without manually rewriting the model in MP-SPDZ's DSL. This is the key ingredient that makes running Huggingface's transformer library realistic in practice.

BERT and Masked Language Modeling

BERT, short for Bidirectional Encoder Representations from Transformers, consists of a stack of transformer encoder layers, each built around multi-head self-attention and position-wise feed-forward networks [VSP+17].

Masked Language Modeling

To perform next-token style prediction, we use BERT with a masked language modeling head. During pretraining, BERT optimizes the masked language modeling objective. Given an input sequence where some tokens are replaced by a special [MASK] symbol, the model must predict the original tokens at those positions.

The model processes the entire sequence and produces contextual embeddings for each position. The MLM head first projects the encoder output back into the hidden dimension, followed by GELU activation and layer normalization, and finally a large dense layer that projects to the vocabulary size (typically 30,000+ tokens).

Formally, for a hidden representation $h$, the MLM head computes:

$$\text{MLM}(h) = \text{LayerNorm}(\text{GELU}(hW_1 + b_1))W_{\text{vocab}} + b_{\text{vocab}}$$

The result is a vector of logits over the vocabulary for each position.

Why Transformers Are Expensive in MPC

MPC protocols handle additions and multiplication with public scalars efficiently. However, transformer models rely heavily on non-linear operations that are significantly more expensive in secure computation.

Matrix multiplications, even large ones, map well to MPC because they consist of structured additions and multiplications. In contrast, operations such as exponentiation, division, comparison, and square roots are significantly more expensive, and must be implemented via bit decompositions, polynomial approximations or interactive subprotocols.

In each transformer layer, the main computational bottlenecks are:

• Softmax in self-attention: $O(n^2)$ exponentiations and divisions per attention head
• LayerNorm: Two instances per layer, each requiring a square root followed by a division
• GELU: Typically implemented via polynomial or rational approximations in MPC.

These non-linear components dominate the cost of secure transformer inference.

MP-SPDZ incorporates a wealth of optimizations for these operations from research. When converting a PyTorch model using layers_from_torch(), PyTorch operations are automatically mapped to optimized MPC implementations. Below we highlight a few of the most commonly used optimizations for transformer inference.

Softmax is defined as

$$\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}$$

which requires both exponentiation and division.

MP-SPDZ uses the numerically stable formulation common in modern ML frameworks:

1. Compute $m = \max(x)$ for numerical stability
2. Evaluate $e^{x_i - m}$ using polynomial approximations of base-two exponentiation
3. Divide by the sum using secure fixed-point division

The exponentiation protocol is described in detail by Keller and Sun [KS22] as part of their secure quantized training work. The fixed-point division uses the Goldschmidt iteration of Catrina and de Hoogh [CdH10].

Layer Normalization computes

$$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta.$$

The costly part is the square root followed by a division. Rather than computing them sequentially, MP-SPDZ combines these into a single inverse square root operation:

$$\frac{1}{\sqrt{\sigma^2 + \epsilon}}$$

following the approach of Lu et al. [LFHH20]. Their method decomposes the input into mantissa and exponent components and applies polynomial approximations separately, avoiding iterative methods such as Newton–Raphson.

GELU The GELU activation is defined as

$$\text{GELU}(x) = x \cdot \Phi(x)$$

where $\Phi$ is the Gaussian cumulative distribution function. Directly computing the Gaussian CDF is expensive in MPC. Instead, MP-SPDZ uses the piecewise polynomial approximation proposed by Dong et al. [PUMA23]:

$$ \text{GELU}(x) \approx \begin{cases} 0 & x < -4 \ p_0(x) & -4 \leq x < -1.95 \ p_1(x) & -1.95 \leq x \leq 3 \ x & x > 3 \end{cases} $$

where $p_0$ is a degree-3 polynomial and $p_1$ is a degree-6 polynomial. This avoids computing the CDF directly by replacing it with low-degree polynomials and a few comparison operations.

Implementing Next-Token Prediction in MP-SPDZ

Let’s shift to what this actually looks like in practice. At a high level, it largely resembles a standard PyTorch implementation of BERT’s MLM task.

Handling Control Flow

We start from a pretrained BERT model from HuggingFace, just as we would in any PyTorch project. We use the masked language modeling variant, which adds the MLM head described above on top of the encoder.

from transformers import BertForMaskedLM

MODEL_NAME = 'prajjwal1/bert-small'  # Small BERT (4 layers, 512 hidden dimension size)
model = BertForMaskedLM.from_pretrained(MODEL_NAME)

This is the only place where we fix the model architecture and load pretrained weights. In our case we choose the small 29M parameter variant to keep compilation time manageable. Nothing in the pipeline fundamentally depends on this choice, so larger models work in the same way, just with higher cost.

From our description above, we would hope we could simply pass this model to ml.layers_from_torch() and be done. Unfortunately, it is not that straightforward.

The HuggingFace implementation contains many conditionals in the forward pass. These handle optional arguments such as attention masks, head masks, and flags that control whether attention weights are returned. These conditionals are not secret-dependent, but they still break tracing because they change the graph structure.

A simplified example from the encoder looks like this:

# Example from HuggingFace BertModel
if attention_mask is not None:
    extended_attention_mask = ...
if head_mask is not None:
    ...
if output_attentions:
    ...

This is perfectly fine in PyTorch. However, MP-SPDZ requires a static computation graph. If the framework encounters such conditionals during tracing, the compiler throws an error.

To make the model compatible, we wrap the relevant components in a forward pass that fixes all optional arguments and removes the argument-dependent branching. The wrapper enforces a single, fixed execution path with a fixed output structure.

class BertMLM(nn.Module):
    """BERT encoder wrapper for MP-SPDZ."""

    def __init__(self, encoder, mlm, config):
        super().__init__()
        self.encoder = encoder
        self.mlm = mlm
        self.config = config

    def forward(self, hidden_states, attention_mask=None):
        encoder_outputs = self.encoder(hidden_states, attention_mask=attention_mask)
        cls_outputs = self.mlm(encoder_outputs[0])
        return cls_outputs

The encoder returns a tuple (hidden_states, attentions). We only use the first element and pass it to the MLM head. Note that the wrapper stores the model's config object, which layers_from_torch uses to determine the hidden size, intermediate size, and number of attention heads when constructing the MPC layers.

With this wrapper in place, we can convert the model into MP-SPDZ layers. The ml.layers_from_torch function additionally requires us to specify the shape of the input tensor, the party that will provide the model weights at runtime (input_via), and the batch size for the computation, which is used to allocate the appropriate amount of memory for each layer.

Here, embeddings is the secret-shared input tensor that we construct in the next section.

bert_wrapped = BertMLM(model.bert.encoder, model.cls, model.config)

mpc_layers = ml.layers_from_torch(
    bert_wrapped,
    embeddings.shape,
    input_via=0,
    batch_size=BATCH_SIZE
)

The resulting mpc_layers contain:

• one BertLayer object per transformer block
• dense layers for the MLM transform and vocabulary projection
• activation and normalization layers

Feeding Secret Inputs

One important point to note is that we pass embeddings rather than token IDs.

The embeddings are computed in plaintext by the client and then secret shared. This assumes that the embedding matrix is public, which is common in MPC based ML systems, although not always aligned with production setups.

To implement this, we can simply run the HuggingFace tokenizer and embedding layer at compile time:

# Tokenize input
encoded = tokenizer(text, padding='max_length', max_length=MAX_LENGTH,
                    return_tensors='pt')

# Compute embeddings (in plaintext)
embeddings_np = model.bert.embeddings(
    encoded['input_ids'],
    token_type_ids=encoded['token_type_ids']
).detach()

# Secret-share for MPC
embeddings = sfix.input_tensor_via(0, embeddings_np.numpy())

The call input_tensor_via(player_id, numpy_array) instructs the VM that:

• player 0 will provide an input tensor of this shape at runtime
• the tensor is serialized and stored so that player 0’s VM can load it
• the resulting object is an sfix tensor with the same shape

From this point on, embeddings behaves like a regular tensor in Python, except that operators such as + and * are overloaded. Instead of performing local arithmetic, they emit bytecode instructions that will later be executed securely by the MPC virtual machine.

Inference and Prediction

We are now ready to execute the forward pass.

In MP-SPDZ, layers are executed through an Optimizer abstraction. The name is slightly misleading in our case, since we are not training and will not update any weights. Internally, this component simply orchestrates the forward pass through the layer stack.

optimizer = ml.SGD(mpc_layers)

# Forward pass through the MPC model (we ignore the returned accuracy and loss)
optimizer.reveal_correctness(embeddings, reference_output, batch_size=BATCH_SIZE)

Calling reveal_correctness emits the bytecode instructions for the forward pass through all transformer layers and the MLM head. The function is designed for testing classifiers and additionally compares the output against a reference, revealing the accuracy and loss. Since this is not a classification task, we pass the plaintext PyTorch encoder output (secret-shared like the embeddings) as reference_output and ignore the returned values.

If we want the predicted token itself, we can extract the logits at the masked position and apply the argmax function.

# MLM logits; the last layer is the output layer added by layers_from_torch
mpc_logits = mpc_layers[-2].Y
predicted_idx = []

for sample_idx in range(N_SAMPLES):
    mask_pos = mask_positions[sample_idx]
    sec_argmax = ml.argmax(mpc_logits[sample_idx][mask_pos])
    predicted_idx.append(sec_argmax.reveal())

This gives us the predicted token ID for each masked position. Note that argmax is implemented via secure comparison, which adds noticeable overhead, especially over large vocabularies.

While MP-SPDZ does not allow us to decode the token ID into a string within the MPC computation, we can compare the token ID with the one we would have gotten from plaintext PyTorch inference to verify correctness.

# pt_encoder_output: plaintext PyTorch encoder output for all samples
with torch.no_grad():
    for i in range(N_SAMPLES):
        # Use the PyTorch decoder to get token text
        mlm_logits = model.cls(pt_encoder_output)
        mask_pos = int(mask_positions[i])
        mask_logits = mlm_logits[i, mask_pos, :]
        predicted_id = torch.argmax(mask_logits).item()
        predicted_token = tokenizer.decode([predicted_id])
        matches_spdz_output = predicted_idx[i] == predicted_id

        sentence_with_pred = EXAMPLES[i].replace('[MASK]', f'**{predicted_token}**')
        print_ln("  %s (%s, match=%s): %s", i, predicted_id, matches_spdz_output, sentence_with_pred)

Running the code

We set up MP-SPDZ, install dependencies and generate keys for each party.

git clone -b blog/bert-next-token-prediction https://github.com/pps-lab/MP-SPDZ.git
cd MP-SPDZ
make -j8 replicated-ring-party.x
pip install torch transformers datasets torchinfo
./Scripts/setup-ssl.sh 3  # Generate certificates for 3-party setting

Running the Program

./Scripts/compile-run.py -E ring bert_next_token

This compiles and runs the program in a 3-party setting with replicated-secret sharing.

=== BERT Next Token Prediction (MPC) ===
Model: prajjwal1/bert-small
Samples: 3

Running MPC forward pass through BERT encoder...
Input embeddings are SECRET-SHARED across parties

=== Decoded Predictions ===
  0 (3007, match=1): Paris is the **capital** of France.
  1 (3191, match=1): I like to **read** books.
  2 (18360, match=1): Birds **fluttered** in the sky.

=== Complete ===
The following benchmarks are including preprocessing (offline phase).
Time = 55.317 seconds
Data sent = 904.104 MB in ~262182 rounds (party 0 only)
Global data sent = 2255.03 MB (all parties)

We benchmark on a single machine (localhost) using three-party replicated secret sharing over a 64-bit ring (ring protocol). The model is prajjwal1/bert-small with 4 encoder layers, 512 hidden dimension, and 8 attention heads (29M parameters).

Compilation takes ~3.5 minutes, since the computation graph is large. MP-SPDZ is good at optimizing the computation graph and re-using bytecode between attention blocks using features such as @_layer_method_call_tape. Nevertheless, the large number of operations, especially in the MLM head and argmax leads to a long compilation time.

The match=1 indicates the MPC computation produced the same prediction as plaintext PyTorch inference.

Runtime

The numbers below include both the online and preprocessing phases. All three parties run on the same machine (Intel Macbook Pro 2019), so network latency is negligible. In a real deployment with network round-trips, the ~262k rounds would significantly impact total runtime.

Conclusion

In this post, we showed how to use MP-SPDZ to run a BERT MLM task using a pretrained HuggingFace model. By introducing a thin wrapper that enforces a static forward pass, we were able to convert the transformer encoder and MLM head into MP-SPDZ layers and run masked language modeling under multi-party computation. The resulting predictions matched plaintext PyTorch inference in our experiments.

At the same time, this experiment shows that modern transformer architectures can be integrated into an MPC framework without rewriting them from scratch. As long as the forward pass can be expressed without data-dependent branching, layers_from_torch provides a viable path from PyTorch models to secure execution.

Closing the gap between feasibility and practicality will require further improvements in compiler reuse, nonlinear approximations, and round complexity. But structurally, there is no fundamental barrier to running real transformer models under MPC.

The full implementation of this demo is available at Programs/Source/bert_next_token.mpc.

References

• [Kel20] M. Keller. MP-SPDZ: A Versatile Framework for Multi-Party Computation. In CCS 2020, pages 1575–1590. ACM, 2020.
• [KS22] M. Keller and K. Sun. Secure Quantized Training for Deep Learning. In ICML 2022, volume 162 of Proceedings of Machine Learning Research, pages 10912–10938. PMLR, 2022.
• [DCLT19] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACL-HLT 2019, pages 4171–4186. ACL, 2019.
• [VSP+17] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention Is All You Need. In NeurIPS 2017, pages 5998–6008, 2017.
• [PUMA23] Y. Dong, W.-j. Lu, Y. Zheng, H. Wu, D. Zhao, J. Tan, Z. Huang, C. Hong, T. Wei, and W. Chen. PUMA: Secure Inference of LLaMA-7B in Five Minutes. arXiv preprint arXiv:2307.12533, 2023.
• [LFHH20] W.-j. Lu, Y. Fang, Z. Huang, C. Hong, C. Chen, H. Qu, Y. Zhou, and K. Ren. Faster Secure Multiparty Computation of Adaptive Gradient Descent. In PPMLP 2020 (CCS Workshop), pages 47–49. ACM, 2020.
• [CdH10] O. Catrina and S. de Hoogh. Improved Primitives for Secure Multiparty Integer Computation. In SCN 2010, pages 182–199. Springer, 2010.
• [LI25] LinkedIn Job Posting: Cryptography Research Engineer (MPC). LinkedIn, 2025.
• [JPMC25] JP Morgan Chase. Applied Research Scientist – Privacy-Enhancing Technologies. Oracle HCM, 2025.