Source-linked AI summary
MPCFormer: fast, performant and private Transformer inference with MPC
Dacheng Li, Rulin Shao, Hongyi Wang, Han Guo, Eric P. Xing, Hao Zhang
TL;DR
Private Transformer inference must protect user data and model weights without incurring prohibitive latency or sacrificing model quality. MPCFORMER combines MPC-friendly approximations with knowledge distillation to produce faster private-inference models. Evaluations report similar performance with substantial speedups across IMDb, GLUE, model sizes, and trained Transformer weights.
Problem
Private Transformer inference protects user data and model weights, but existing approaches can be much slower or reduce inference quality.
Method
MPCFORMER replaces Transformer bottlenecks with MPC-friendly approximations and uses knowledge distillation from the original model to train a performant student.
Results
Across evaluations, MPCFORMER substantially speeds up MPC inference while preserving similar or near-baseline performance, including 5.3× speedup on IMDb and 97% of BERTBASE performance with 2.2× speedup on GLUE.
Takeaways & Limitations
The framework remains effective with different trained Transformer weights and larger models, supporting practical private inference across varied model settings.
Takeaways & Limitations
Speedups and performance are evaluated on a single MPC system, and the teacher and student are constrained to the same model size.
Abstract
from arXiv · showhide
Enabling private inference is crucial for many cloud inference services that are based on Transformer models. However, existing private inference solutions can increase the inference latency by more than 60x or significantly compromise the inference quality. In this paper, we design the framework MPCFORMER as a practical solution, using Secure Multi-Party Computation (MPC) and Knowledge Distillation (KD). Through extensive evaluations, we show that MPCFORMER significantly speeds up Transformer inference in MPC settings while achieving similar ML performance to the input model. On the IMDb dataset, it achieves similar performance to BERTBASE, while being 5.3x faster. On the GLUE benchmark, it achieves 97% performance of BERTBASE with a 2.2x speedup. MPCFORMER remains effective with different trained Transformer weights such as ROBERTABASE and larger models including BERTLarge. Code is available at https://github.com/MccRee177/MPCFormer.
1 INTRODUCTION
MPCFORMER addresses the tension between private Transformer inference and practical latency by combining MPC-friendly approximations with knowledge distillation. It preserves strong task performance while substantially accelerating MPC inference across benchmarks and model sizes.
- Private inference matters because deployed Transformer services can require users to expose sensitive inputs or providers to release proprietary model weights.The paper motivates this trade-off with code-generating services such as GitHub Copilot.
- ∼60 seconds of MPC inference for BERTBASE contrasts with less than 1 second without MPC, motivating a faster private-inference design.Directly adapting approximation-and-retraining methods from CNNs substantially decreases Transformer performance.
- Knowledge distillation addresses training difficulty and limited downstream data by providing intermediate supervision and a data-efficient objective.These properties help the approximated Transformer retain performance despite aggressive function substitutions.
- MPCFORMER replaces bottleneck functions with MPC-friendly approximations, then distills the original Transformer into a faster model for private inference.The resulting model is deployed over an MPC engine such as CrypTen.
- 5.3× faster IMDb inference reaches similar performance to BERTBASE, while 2.2× faster GLUE inference reaches 97% of BERTBASE performance.The framework also reaches similar IMDb performance to BERTLARGE with a 5.9× speedup and works with RoBERTaBASE weights.
2 BACKGROUND
Transformer inference can be expressed as private two-party computation, but MPC communication makes nonlinear functions and overall runtime costly. Secret sharing preserves input privacy while adding communication overhead that dominates inference.
- In two-party computation, the user supplies data and the model provider supplies the Transformer model, and they jointly compute the inference result.The protocol guarantees that each party learns only its own input and the result.
- Secret sharing splits each input number into two shares so neither party alone learns the original value.The user retains one share while the model provider receives the other.
- 59.0 seconds with MPC versus under 1 second without MPC highlights the substantial runtime overhead of private BERTBASE inference.
- 79% of the whole inference process is communication, with overall communication taking 46.4 seconds.Softmax functions alone account for 50.3GB of communication and 34.1 seconds.
- MPC multiplication preserves privacy by masking shares but requires one extra communication round compared with ordinary multiplication.This added communication helps explain why vanilla Transformer models are slow under MPC.
- GeLU and Softmax are major Transformer bottlenecks because their evaluations require many multiplications, and Softmax also requires comparisons.GeLU uses a high-order Taylor expansion, while Softmax uses iterative squaring and a maximum operation for numerical stability.
3 RELATED WORK
Prior work establishes MPC for private computation, Transformer architectures for broad machine-learning applications, MPC-friendly approximations for CNNs, and knowledge distillation for transferring model representations. MPCFORMER builds on these directions for private Transformer inference.
- MPC research provides privacy-preserving computation systems suitable for Transformer inference without requiring MPCFORMER to implement a new MPC system.The paper instead targets an algorithmic solution for speeding Transformer models.
- Transformer models have become a dominant paradigm across language, vision, and other understanding tasks.
- Earlier MPC-friendly approximation work accelerated CNN computation but could cause accuracy loss when quadratic activations destabilized gradient descent.
- Knowledge distillation transfers teacher knowledge to students by matching hidden representations, with prior Transformer objectives including attention-matrix matching.Some prior methods train on pre-training data and pursue goals different from MPCFORMER.
4 METHOD
MPCFORMER converts a trained Transformer into a faster private-inference model through MPC-friendly approximations followed by knowledge distillation. The method uses aggressive quadratic replacements while distillation restores performance, including through layer-wise representation matching and teacher-based initialization.
- Workflow: MPCFORMER requires a trained Transformer T, downstream dataset D, and MPC-friendly approximations A, then performs private inference with S.Its workflow is Convert: S = MPCFORMER(T, D, A), followed by Inference: y = MPCS(X).
- MPCFORMER first constructs a fast approximated architecture S′, then distills it into S to preserve the input model’s performance.The framework is designed to make S faster than T in MPC while maintaining similar ML performance.
- Stage 1: Approximation: The first stage replaces GeLU and Softmax bottlenecks with MPC-friendly approximations to construct S′.
- Approximating GeLU: GeLU is approximated with the quadratic 0.125x^2 + 0.25x + 0.5, called Quad, because quadratic nonlinearities are fast in MPC.
- Approximating Softmax: 2Quad replaces ReLU inside a Softmax approximation with a quadratic function, enabling an aggressive speed-oriented approximation whose numerical values differ substantially from Softmax.Distillation bridges the resulting performance gap; Figure 4 compares approximation and original-function runtimes.
- Stage 2: Knowledge Distillation: Layer-wise distillation matches representations at the embedding, attention, hidden-state, and final prediction layers using mean squared error.
- Student Initialization: Initializing S′ with T’s weights outperforms random initialization, especially on smaller downstream datasets.
5 EXPERIMENTS
Experiments evaluate MPCFORMER across approximations, datasets, model sizes, initialization choices, and distillation-data budgets. Across these settings, MPCFORMER preserves performance more reliably than baselines while retaining substantial MPC speedups, though evaluation remains limited to one MPC system and same-size students.
- 5.1 COMPARISON WITH BASELINES ON DIFFERENT BENCHMARKS: 5.26× speedup with almost no accuracy drop on IMDb, and MPCFORMER consistently preserves accuracy across approximations from 1.24× to 2.65× speedup.Both baselines lose at least 6.8% accuracy at 5.26× speedup.
- 5.1 COMPARISON WITH BASELINES ON DIFFERENT BENCHMARKS: 2.2× speedup with 97% of BERTBASE performance on GLUE, while baselines incur severe performance drops at comparable speed.At 1.93× speedup, the two baselines lose 19.5 and 26.2 average-score points.
- 5.2 MORE COMPARISONS WITH DIFFERENT MODELS: 98% average score is preserved with RoBERTa-base, while the method also scales to BERT-large with 5.9× speedup and 94.5% IMDb accuracy.For BERT-large on selected GLUE tasks, MPCFORMER raises average score from 81.5 to 84.6 as the input model improves from BERT-base to BERT-large.
- 5.3 ABLATION STUDY: Initializing the student with teacher weights helps on smaller datasets, whereas the initialization gap is small on larger datasets such as QNLI and SST-2.The ablation compares teacher-weight and random initialization after extended training of the random-initialized student.
- 5.3 ABLATION STUDY: Roughly 5% of small GLUE datasets and 2% of medium datasets suffice to learn a good distilled student.The study covers RTE and MRPC as small datasets, and SST-2 and QNLI as medium datasets.
- 5.4 LIMITATION AND FUTURE DIRECTION: The evaluation uses one MPC system, and the teacher and student have the same model size because only functions differ.Extending evaluation to more MPC systems and smaller student models is left for future work.
6 CONCLUSION
The paper proposes MPCFORMER for fast, performant private Transformer inference and finds compatibility with varied MPC-friendly approximations and trained Transformer models.
- MPCFORMER achieves fast and performant private Transformer model inference with MPC.
- MPCFORMER is compatible with various MPC-friendly approximations and trained Transformer models.
A.1 A CONCRETE SYSTEM IMPLEMENTATION OF MPC: CRYPTEN
CrypTen implements private computation through secret sharing, arithmetic and binary routines, and communication-intensive conversions and comparisons.
- Secret sharing splits values across parties so either share alone reveals no information while both reconstruct the original value.
- Arithmetic-share addition is local, while multiplication uses Beaver triples and reveals masked intermediate values.
- Comparisons convert between arithmetic and binary shares before extracting the sign bit.
- With 64-bit tensors, each multiplication requires one communication round, comparisons require 7 rounds, and max over N elements requires O(log2(N)) rounds.
- Communication accounts for 79% of MPC runtime, making FLOP counts inaccurate estimators of Transformer inference time.
A.2 2QUAD IMPLEMENTATION DETAILS
The 2Quad implementation modifies Softmax masking to avoid explosion when exponentials are replaced by quadratic approximations.
- Directly replacing the exponential with a quadratic expression causes masked positions to explode during the forward pass.
- 2Quad instead implements masking by multiplying the exponential function by zero or one, preserving the masking purpose without exploding masked positions.
- The implementation empirically uses c = 5, and this constant appears flexible in the reported experiments.
A.3 ROBUSTNESS OF THE STUDENT MODEL
The robustness evaluation finds that approximations have a moderate empirical effect on the student model’s adversarial robustness.
- The student model has moderately higher attack success rates on three score-based attacks but lower attack success with the gradient-based HotFlip attack.
- The authors characterize the overall robustness effect of the approximations as empirically moderate.
A.4 HYPER-PARAMETER CHOICE
The section selects hyper-parameters through grid searches for baseline models and limited tuning for MPCFORMER. STS-B is used for baseline exploration because its regression performance varies widely, while MPCFORMER uses fixed choices guided by distillation loss and dataset size.
- Baseline hyper-parameter search: Baseline tuning searches learning rates, batch sizes, and epochs on STS-B, with BERTBASE evaluated extensively and smaller searches for BERTLarge and ROBERTABASE.The tested learning rates range from 1e-6 to 5e-4, batch sizes from 16 to 256, and epochs from 3 to 200.
- MPCFORMER choices: MPCFORMER minimizes hyper-parameter tuning so its performance is representative of future users who prefer not to tune settings.
- MPCFORMER choices: MPCFORMER selects epochs from embedding and Transformer-layer distillation losses, uses 5 epochs for prediction-layer distillation, and sets batch size by dataset size.Batch size is 8 for CoLA, MRPC, and RTE, and 32 for MNLI, QQP, SST2, and STS-B.
- Baseline hyper-parameter search: STS-B is chosen for baseline grid searches because it is a regression task whose performance varies across a large range.