RAPTOR: Role-Aware Private Training for Mixture-of-Experts
Authors: Duc Dm, Khai Le-Duc, Nguyen Do, Minh Son Hoang, Florent Draye, Thai Hoang, Hoang Phuong Dam, Jiarui Liu, Chris Ngo, Terry Jingchen Zhang, Anh Le Duc Tran, Nhat Do Minh, Minh Ngoc Le, My T. Thai, Ran Xu, Silvio Savarese, Mona Diab, Bernhard Schölkopf, Zhijing Jin, Huy L. Nguyen, Daeyoung Kim
Organizations: KAIST · University of Toronto · Vector Institute · Knovel Engineering Lab · University of Florida · MPI for Intelligent Systems, Tübingen · Salesforce AI Research · Carnegie Mellon University · University of Oxford · Jinesis Lab, University of Toronto & Vector Institute · Hanoi University of Science and Technology · Vietnam National University, Hanoi · Stanford University · ELLIS Institute Tübingen · Northeastern University
Differentially private (DP) fine-tuning methods treat sparse Mixture-of-Experts (MoE) models as a single dense block, ignoring that shared layers see all data while experts only see routed records. We identify and formally characterize three resulting failure modes: global clipping suppresses expert gradients, batch-level normalization dilutes sparse expert updates, and fixed privacy noise degrades signal-to-noise ratio on low-load experts. We introduce RAPTOR - a Role-Aware Private Training framework, which alternates shared and expert optimization and targets each failure directly, using expert-specific clipping and noise together with a public expected-owner denominator and a count-independent update schedule that avoids conditioning on private, realized expert counts. We prove the resulting mechanism satisfies (ε,δ)-DP: because each record is assigned to exactly one owner expert, per-expert mechanisms within a layer compose in parallel, so updating all E experts costs no more, in privacy terms, than updating one, with shared and expert streams composing sequentially across training. We further derive a bias-variance decomposition of the public-denominator estimator showing its bias grows predictably with routing imbalance, yielding a privacy-free rule for selecting which layer to protect from routing entropy measured on a small public corpus. Experiments on Switch Transformer and OLMoE fine-tuning across GLUE tasks, and on DeepSeek-VL2-Tiny, show consistent gains over standard DP baselines across several privacy levels (ε), with the largest margins typically at the tightest budgets. Code and models are publicly available: https://github.com/leduckhai/RAPTOR
Figures & tables
Figure 1: Applying differential privacy to sparse MoE is not straightforward : treating the model as dense leads to three structural failures - clipping interference, update-scale dilution, and load-dependent signal degradation. Our role-aware method is designed to address these mismatches.
Figure 2: Overview of RAPT OR at the selected sparse layer l⋆ . (a) The frozen router performs the base MoE model’s standard token-level top- k routing. Our framework aggregates gate weights across all tokens of record xi and selects one deterministic owner expert al⋆(xi) . The forward pass is unchanged: all activated experts contribute, but only the owner receives the expert-parameter gradient; activated non-owner experts use stop-gradient. This assignment partitions Dpriv into disjoint owner groups Dl⋆,e={(xi,yi):al⋆(xi)=e} . (b) Each cycle draws independent Poisson subsamples for the shared and expert streams. Shared updates use clipping Cs , noise multiplier σs , and denominator B . Expert records are partitioned by owner and updated using Ce , σe , and the public denominator De=B/E , with noise-only updates for empty subsets. Expert mechanisms compose in parallel within l⋆ , while shared and expert releases compose sequentially, with εsh=ρε and εexp=(1−ρ)ε .
ε
Method
SST-2
MNLI
QNLI
QQP
∞
Non-priv. LoRA
94.69
86.55
91.20
90.96
Non-priv. FFN
94.50
86.36
91.09
91.28
1
Monolith. DP-Adam LoRA
90.94
76.91
80.38
81.19
Monolith. DP-Adam
87.04
72.72
81.44
81.84
Ours, first sparse
92.90
77.40
82.60
83.20
Ours, last sparse
92.66
77.87
83.30
83.80
Table 1: Switch-base-8 under (ε,δ) -DP with δ=1/∣Dpriv∣1.1 . Entropy-selected layers fixed from public routing statistics before training: block 1 (SST-2), block 7 (MNLI), block 11 (QNLI), block 9 (QQP). Bold : best per task–budget; underline : second-best distinct method.
ε
Method
SST-2
MNLI
QNLI
QQP
∞
Non-priv. LoRA
96.67
90.19
94.65
92.03
Non-priv. FFN
95.18
89.72
93.13
91.07
1
Monolith. DP-Adam LoRA
92.78
85.31
84.24
85.89
Ours, entropy-sel.
95.41
86.15
88.74
86.59
4
Monolith. DP-Adam LoRA
95.07
85.74
88.56
86.43
Ours, entropy-sel.
95.30
87.44
90.44
86.89
Table 2: OLMoE-1B-7B ( E=64 , top- 8 routing) under (ε,δ) -DP, δ=1/∣Dpriv∣1.1 . Entropy-selected layers fixed from public routing statistics: block 1 (SST-2, QNLI, QQP), block 15 (MNLI).
Table 3: MIA audit on QNLI ( Switch-base-8 ). 64 shadow models; neval=20,000 ( 10 k/ 10 k members/non-members). Single run per cell. Bold : highest attack success (AUC/TPR) among DP methods at that ε - the least favorable privacy result, not the best method; underline: second-highest/tied. Acc. bolding follows the opposite convention (highest accuracy), as in Tab. 1 .
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Search grid
Epochs
{10,20,30}
Batch size B
{500,1000,2000,3000}
Learning rate η
{10,5,3,1,0.5}×10−4
Clipping C
{0.1,1.0,10}
Budget split ρ
{0.7,0.8,0.9,0.95}
Appendix
Table 4: Hyperparameter search grid (SST-2). Selected values bolded; configuration frozen for all other tasks.
Figure 3: Validation accuracy on ScienceQA for DeepSeek-VL2-Tiny across ε∈{1,4,8} . Single run per ε .
ε≈8
ε=1
Method
SST-2
MNLI
QNLI
QQP
SST-2
MNLI
QNLI
QQP
Monolithic DP-Adam LoRA
93.92
79.36
82.48
82.68
90.94
76.91
80.38
81.19
RAPT OR (full)
94.15
81.32
85.26
85.92
92.90
78.42
83.30
84.08
w/ De=B
94.10
80.96
84.28
85.55
92.83
77.98
82.79
83.55
w/ full-model clip
94.09
81.25
84.64
85.84
92.82
78.34
83.22
83.99
w/o alternating
94.02
80.21
83.29
84.09
92.75
77.15
81.68
81.99
Appendix
Table 5: Component ablation on Switch-base-8 at ε≈8 and ε=1 . Bold: full RAPT OR . Each row removes or replaces one component; all variants use the entropy-selected sparse layer.
ε
Method
SST-2
MNLI
QNLI
QQP
1
Monolith. DP-Adam LoRA
90.94
76.91
80.38
81.19
Baseline A (global)
92.55
76.80
82.90
83.10
Baseline B (shared-only)
91.80
76.40
82.50
82.90
Baseline C (expert-only)
67.78
32.75
49.72
63.66
Ours, entropy-sel.
92.90
78.42
83.30
84.08
4
Monolith. DP-Adam LoRA
92.66
78.33
81.95
82.26
Appendix
Table 6: Matched-scope diagnostics on Switch-base-8 : Baselines A–C share Ours’s layer, scope, budget, and step count, varying only the update rule (Sec. 8.1 ). Monolithic DP-Adam LoRA and Ours (entropy-selected) shown for reference; full results including first-/last-sparse variants and non-private upper bounds are in Tab. 1 . Single run per cell.
Figure 4: SST-2 accuracy (%) of RAPT OR on Switch-base-8 across ε∈{1,2,4,8} and ρ∈{0.7,0.8,0.9,0.95} .
Task
Calibration source
n
Sparse layer
Encoder block
H(q)/logE↑
logE−H(q)↓
χ2(q∥u)↓
MNLI
ANLI R2
45,460
0
1
0.307
1.442
4.946
MNLI
ANLI R2
45,460
1
3
0.674
0.677
1.580
MNLI
ANLI R2
45,460
2
5
0.641
0.746
1.617
MNLI
ANLI R2
45,460
3
7
0.682
0.660
1.262
MNLI
ANLI R2
45,460
4
9
0.643
0.743
1.356
MNLI
ANLI R2
45,460
5
11
0.646
0.736
1.381
Appendix
Table 7: Routing statistics for all six sparse encoder layers of Switch-base-8 on external public corpora (labels ignored; disjoint from GLUE splits and Dpriv ). H(q)/logE ( ↑ ), entropy deficit logE−H(q) in nats ( ↓ ), and χ2(q∥u) ( ↓ ) for E=8 experts. Bold : entropy-selected layer l⋆ per task. All selections are unambiguous; no privacy budget is consumed.
Figure 5: Normalised routing entropy H(q(l))/logE as a function of calibration examples n (log scale) for all six sparse encoder layers of Switch-base-8 , across four tasks. Each line is one candidate layer; the entropy-selected layer at the full corpus size is marked in bold in Tab. 7 . Entropy estimates stabilise by n≈100 – 500 for most layers; the main exception is MNLI block 1 (blue), which is heavily collapsed and grows slowly even at large n . The frozen router is applied to unlabelled inputs only; no private data are accessed.
Figure 6: Public-calibration vs. GLUE-development normalised routing entropy H(q)/logE for all 24 layer-task combinations (six sparse encoder layers × four tasks). Each point is labelled with its encoder block index. The dashed line is the diagonal y=x . Transfer is closest for QNLI; MNLI shows a systematic gap, with all six layers above the diagonal. The largest deviations are MNLI block 1 and SST-2 blocks 7 and 9, none of which is selected by the entropy criterion; entropy-selected layers sit closer to the diagonal than collapsed, non-selected ones.
Figure 7: Public-calibration vs. GLUE-development normalised routing entropy H(q)/logE for all 24 layer–task combinations (six sparse encoder layers × four tasks). Each point is labelled with its encoder block index. The dashed line is the diagonal y=x . Points near the diagonal indicate that entropy estimated on the public corpus transfers to the private training distribution. The two largest off-diagonal deviations (SST-2 blocks 7 and 9, blue circles) are collapsed layers not selected by the entropy criterion; the entropy-selected layer for each task lies near or above the diagonal.
Task
Method
Mrole↓
SNR↑
SNR-CV↓
SST-2
Monolithic DP-Adam LoRA
0.320
1.08×10−4
2.35
Monolithic DP-Adam
0.591
6.38×10−6
1.20
RAPT OR
0.401
1.80×10−5
1.69
MNLI
Monolithic DP-Adam LoRA
0.494
9.65×10−5
3.54
Monolithic DP-Adam
0.560
3.04×10−5
1.24
RAPT OR
0.337
2.25×10−4
1.86
Appendix
Table 8: Mechanism diagnostics on Switch-base-8 at ε≈8 . Mrole : total variation distance between gradient-norm share and parameter-count share across roles (lower = more proportionate signal allocation). SNR : mean per-expert signal-to-noise ratio (higher = stronger expert update signal). SNR -CV: coefficient of variation of per-expert SNR (lower = more uniform update quality across experts).
Task
Public corpus
Layer (block)
CVl
Hl/logE
SST-2
IMDb unsup.
selected (1)
0.344
0.972
collapsed (9)
2.185
0.304
MNLI
ANLI R2
selected (7)
1.123
0.682
collapsed (1)
2.224
0.307
QNLI
SQuAD context
selected (11)
0.627
0.902
collapsed (1)
1.121
0.761
Appendix
Table 9: Corpus-level routing diagnostics for the entropy-selected and most-collapsed sparse layers of Switch-base-8 , computed on external public calibration corpora with labels ignored (see Appendix F for corpus details). Higher CVl and lower Hl indicate stronger routing imbalance; the two statistics are consistent across all tasks. Values are identical across private training methods because the router is frozen.
Role
Params
ΔLval
λ+
Shared
9,438,722
−2.175×10−4
2.305×10−5
Experts
7,864,320
−1.361×10−6
0
Appendix
Table 10: Noise-injection sensitivity on SST-2 at ε≈8 . λ+ floors non-positive ΔLval at zero (Eq. ( 40 )). A non-positive expert value indicates sensitivity below the measurement noise floor.
Method
Trainable params
% params
Peak GB
Sec./step
GPU-hours
Global DP-Adam LoRA
14.16 M
4.60%
130.74
4.07
1.47
Global DP-Adam
307.85 M
100.00%
165.51
7.23
2.61
RAPT OR
17.30 M
5.62%
112.16
4.51
1.65
Appendix
Table 11: Implementation footprint of fine-tuning Switch-base-8 on SST-2 at ε≈8 . Raw implementation-level costs, not a controlled same-scope benchmark—scope differs by design. Percentages relative to the full trainable Switch-base-8 model used by Global DP-Adam. Bold: lowest peak memory.
Mixture-of-Experts (MoE) transformers scale capacity by activating only a few experts per token, but this sparsity creates a hidden reliability problem: when routing is imperfect, load-balanced models may send tokens to experts that are insufficiently trained for the assigned inputs. We propose Distributionally Robust MoE Training (DRMoET), a drop-in objective that treats layer-wise experts as endogenous robustness groups and optimizes high-loss routing outcomes rather than merely equalizing traffic. DRMoET updates a per-layer expert distribution by an entropy-regularized softmax rule on EMA-smoothed, activation-weighted expert losses, strengthening plausible non-top routing paths while preserving standard MoE computation. Under the FLAME-MoE recipe at 746M-total and 10.3B-total scales, DRMoET improves downstream averages over both standard FLAME-MoE and auxiliary-loss-free balancing. At 10.3B total parameters and 67B training tokens, DRMoET improves the seven-task average from 0.6625 to 0.6767, while the auxiliary-loss-free baseline achieves 0.6431. Mechanistic analyses show lower expert-loss variance with nearly unchanged mean loss, 4.3% lower excess loss under forced mid-k misrouting, and improved domain-expert specialization. These results position routing robustness-not only utilization balance-as a practical objective for reliable sparse MoE scaling. Project page and code are available at: https://drmoet.github.io/.
Xin Teng, Muxiao Li, Hongyi Wen
New York University Center for Data Science, NYU Shanghai
Mixture-of-Experts (MoE) models scale capacity by combining specialized experts, but most existing approaches assume centralized access to training data. In practice, data are distributed across clients and cannot be shared due to privacy constraints, making unified MoE training challenging. We propose MetaMoE, a privacy-preserving framework that unifies independently trained, domain-specialized experts into a single MoE using public proxy data as surrogates for inaccessible private data. Central to MetaMoE is diversity-aware proxy selection, which selects client-domain-relevant and diverse samples from public data to effectively approximate private data distributions and supervise router learning. These proxies are further used to align expert training, improving expert coordination at unification time, while a context-aware router enhances expert selection across heterogeneous inputs. Experiments on computer vision and natural language processing benchmarks demonstrate that MetaMoE consistently outperforms recent privacy-preserving MoE unification methods. Code is available at https://github.com/ws-jiang/MetaMoE.
Weisen Jiang, Shuhao Chen, Sinno Jialin Pan
Department of Computer Science and Engineering, The Chinese University of Hong Kong · Department of Computer Science and Engineering, Hong Kong University of Science and Technology · Department of Computer Science and Engineering, Southern University of Science and Technology
Standard LoRA fine-tuning of Mixture-of-Experts (MoE) models applies adapters to every expert, yet our profiling shows that per-layer expert routing is highly skewed: a small subset of experts handles most tokens in each layer, while many others are rarely activated ("cold"). We propose MoE-Sieve, a simple routing-guided framework for LoRA fine-tuning, and pair it with a systematic profiling study of expert routing across architectures and tasks. The method is simple: profile routing counts on a small calibration set, select the top-k most-routed experts per layer, and apply LoRA only to those experts. Across two architecturally distinct MoE models and three diverse tasks, tuning only the top 25% routed experts per layer remains competitive with full LoRA, with mean differences within +/-1 percentage point across all conditions. This reduces LoRA trainable parameters by 70-73%, adapter checkpoint size by 71-73%, and wall-clock training time by up to 50%. We also observe a non-monotonic relationship between expert count and seed-to-seed variance, consistent with the hypothesis that adapting cold experts can introduce gradient noise without improving accuracy. Further ablations show that random expert selection at matched budget is about 2.5 percentage points worse, indicating that the routing signal matters, while greedy per-layer budget optimization does not improve over uniform top-k.