Reducing the peak memory of secure transformer computation step by step

Background

Secure neural network inference lets a client run a model held by a server without either side seeing the other's data: the client's input and the server's weights stay hidden. The systems compute on secret shares, function secret sharing keys or homomorphic ciphertexts. These representations are much larger than the plaintext values, so memory can become the limiting resource. In the unmodified SHARK system, each of the two online parties holds its whole preprocessing file in memory for BERT-base, about 108 GB for both parties together, which does not fit on a 95 GB cluster node.

In my master's thesis, Optimizing the Memory Footprint of Secure Neural Network Inference for Transformer Architectures (University of Münster), I measured the peak memory of these research systems, decomposed it into the objects that hold it, and reduced it one change at a time. The companion repository is github.com/Jwrede/snni-memory-reduction (tag thesis-submission).

What was reduced

The repository holds 15 lines over nine systems and their dealer, GPU and ViT variants. They cover four paradigms: function secret sharing (SHARK, LLAMA, SIGMA), arithmetic secret sharing (SHAFT, PUMA), hybrid HE and MPC (BumbleBee, BOLT) and fully homomorphic encryption (MOAI, ARION). Every line runs BERT-base at sequence length 128, and two transfer lines run ViT-Base/16 on SHAFT. All runs used the PALMA cluster of the University of Münster.

Two examples from the published STEPS.md files:

The procedure

Each step changes one thing and is measured on a fixed machine. The rule for choosing the next step is fixed in the README:

1. Measure the peak, and decompose it at the peak instant into the objects that hold it.
2. Rank the objects by resident size and take the largest one that has a free fix, meaning a fix whose run-time change stays within the measured spread.
3. Re-measure and repeat. When no free fix remains, take paid fixes ranked by memory saved per unit of run time.

The fixes fall into four kinds: free a dead object, keep one copy of a duplicated object, reload an object when it is read, and create an object at its first use. A change that trades accuracy, security or portability is never a step; it is reported separately in each system's off-path.md.

The peak and the object table come from the same run. A native recorder logs every allocation of at least 64 KiB with its call site and reads its resident pages from /proc/<pid>/pagemap. The process is stopped with SIGSTOP for each capture, so objects and mappings describe one instant. Each step runs three times and the median replicate is published. A correctness gate compares each step's output with the baseline's output at a declared tier, from byte-identical to equal within a numerical tolerance. Steps that repartition work change the protocol randomness and therefore the output; they are marked separately and rest on an unchanged transcript and a plaintext test.

Which steps gave the largest reductions

In SHARK, the baseline reads the whole preprocessing file (the correlated key material from the offline dealer) into one buffer before the computation starts. The step s3_key_stream reads the file sequentially while it is consumed. This single step lowered the peak by 52.53 GiB, from 53.68 GiB to 1.14 GiB, at a measured run-time cost of 5.1%, which classifies it as paid. Two further steps brought the peak to 0.45 GiB, against a derived target of 217.4 MB for the key material and buffers of one truncation batch.

In BumbleBee, the largest free step was b1_dot_encode_chunk, which encodes matrices in slices instead of encoding a whole operand at once. It lowered the peak by 6.51 GiB at a run-time change of 1.0%. The later steps b5 to b8 reduce the number of parallel oblivious transfer instances from 16 to 1. Each instance holds its own buffers, so these steps lowered the peak from 6.44 GiB to 1.39 GiB, and the wall time grew from 388 s to 1342 s. steps.csv classifies all four as paid.

In both lines, the large reductions came from objects that were held whole although the computation used them piece by piece: the key file, encoded matrices and per-instance protocol state.

Regenerating the numbers

Every row of a steps.csv and every peak-objects.md is generated from the raw logs stored next to it. These commands check that the published files regenerate and that every step follows the procedure:

python3 tools/make_steps.py systems/<system> --check
python3 tools/make_impl.py systems/<system> --check
bash tools/test_conformance.sh
bash tools/check_policy.sh

Container images are not part of the repository. Each step's BUILD.md names the pinned upstream commit and the build commands, and each run's RUN_META records the image digest.

The code, data and documentation are at github.com/Jwrede/snni-memory-reduction.