Compilation pipeline

Compilation happens in two stages.

The first stage traces your model with torch.export, lowers it to MLIR, and attaches two annotations: which inputs are encrypted, and the value range of every operation that must be approximated. The MLIR contains the weights, so it is the only file uploaded.

The second stage chooses the CKKS parameters, inserts bootstrapping where the circuit runs out of levels, and generates the GPU and client code.

What compiles#

CKKS computes additions and multiplications on encrypted real numbers. How a neural network compiles to a sequence of additions and multiplications depends on the operations it contains.

The following layers do not compile as they are currently not supported:

• nn.MaxPool1d and nn.MaxPool2d. Use AvgPool instead.
• Grouped convolutions (groups > 1), except a depthwise nn.Conv1d or nn.Conv2d with groups == in_channels == out_channels, which compiles. Use groups=1.
• softmax and log_softmax. A classifier can return the logits and apply softmax after decryption.
• ConvTranspose and Upsample.
• Indexing with a tensor, such as x[:, idx] or torch.gather. Slices with fixed bounds, such as x[:, :8], compile.
• A sum or mean over several dimensions at once, such as x.mean(dim=(2, 3)) for global pooling. Use nn.AdaptiveAvgPool2d(1). A sum can also reduce one dimension at a time, x.sum(3).sum(2).

A check looks at operations on the encrypted input and on the values computed from it. An operation on the weights alone, such as a max pool of a parameter, runs before encryption and compiles.

To write the MLIR anyway, for example to try a compiler that supports more operations, pass allow_unsupported_ops=True to belfort_ml.compile. It logs the same list as a warning. The Belfort compiler still rejects the model.

Each multiplication consumes one level of the ciphertext. Once a ciphertext has used its levels, the compiler inserts a bootstrap, which refreshes it and is the most expensive GPU operation. A model shallow enough to avoid bootstrapping runs much faster:

• Prefer polynomial activations (x * x) where your accuracy allows it. The mlp and lola examples use squares; a model that keeps ReLU pays for its approximation.
• An approximated operation is accurate only inside its calibrated range; values outside it produce wrong results without an error.
• Every approximated operation except ReLU, reciprocal and division becomes a polynomial of degree 5 over its calibrated range. The degree cannot be changed yet. A function that a degree-5 polynomial fits poorly over that range, such as a steep sigmoid or tanh over a wide range, loses accuracy; compare the encrypted results with the plaintext model. Tighter calibration ranges, from representative calibration inputs, improve the fit.
• ReLU uses a composite approximation by default: accurate, but deep enough to force bootstrapping. use_composite_relu=False trades accuracy for a shallower model and smaller keys; see ReLU approximation.

A ciphertext holds at most 32,768 values: the largest ring Belfort uses has degree 2^16, and a CKKS ciphertext holds half that many. Every linear layer has to fit in one ciphertext: the longer side of an nn.Linear weight matrix, and for a convolution, channels times spatial size of both its input and its output. A convolution with many channels on a large input can exceed it. In this case, downsampling the input can sometimes help.

Neural network that are too deep can be problematic. If you encounter issues, you can try to reduce the input size and/or remove layers.

The example input fixes the shape#

The graph is specialized to the shape of the example input, including the batch dimension. The deployed model accepts exactly that shape. Use batch size 1; one inference carries one input. The whole input travels as one ciphertext, however many numbers it holds.

Calibration#

calibration_inputs is a list of tensors with the example input's shape. The compile step runs each of them through the model and records the full range of values reaching every approximated operation. It then doubles each range about its centre, except for operations whose domain is not all real numbers (sqrt, log, reciprocal, div and similar), which keep the observed range. Only these ranges go into the MLIR.

Use real samples from your training or validation data, up to 512 of them. A small set can miss the inputs your end clients send, and then the encrypted result is wrong. Random inputs give ranges that do not match what your end clients will send. With no calibration inputs the MLIR carries no ranges, and the compiler assumes a fixed range, [-1, 1] for most operations.

Encrypted and plaintext inputs#

By default every argument of forward is encrypted. If your model also takes inputs the server may see, mark only the secret ones with belfort_ml.Secret:

from typing import Annotated

import torch
from belfort_ml import Secret


class Model(torch.nn.Module):
    def forward(self, x: Annotated[torch.Tensor, Secret], scale: torch.Tensor):
        ...

The Belfort ML client currently sends one ciphertext per inference, so a model can have at most one encrypted argument.

Changing a model#

A model compiles once, from the MLIR uploaded with it. Model names are unique in your account. provider.models.get(name=, create_from=artifacts) returns the model of that name when it was built from the same MLIR, and raises ConflictError name_taken when the MLIR changed. To change the model, delete the old one with model.delete() first, or upload the new one under another name.