On the (In-)Security of the Shuffling Defense in the Transformer Secure Inference

Technical Summary: On the (In-)Security of the Shuffling Defense in the Transformer Secure Inference

Problem Statement
Secure inference for Transformer models using Multi-Party Computation (MPC) or Homomorphic Encryption (HE) faces a critical efficiency bottleneck: the computation of nonlinear layers (e.g., Softmax, GELU, LayerNorm). Full-model encryption (FME) requires numerous communication rounds and heavy data transmission for these layers, accounting for 75%–90% of total latency. To mitigate this, Linear-Only Encryption (LOE) inference was proposed, where intermediate activations are revealed to the client in plaintext to compute nonlinear layers locally. However, revealing activations allows adversaries to extract model weights by solving linear systems.

To counter this, recent works introduced a shuffling defense, where intermediate activations are randomly permuted before being revealed to the client. The prevailing security claim is that recovering the specific permutation is computationally infeasible (probability $1/d!$ for dimension $d$) and that shuffled activations destroy the correlation necessary for weight extraction attacks. This paper challenges that claim, demonstrating that the shuffling defense is not robust against weight extraction.

Methodology
The authors propose an attack that aligns differently shuffled activation vectors to a common permutation without knowing the specific permutation applied. The methodology proceeds in three stages:

1. Vector Alignment via Proximity: The adversary queries the Transformer with multiple inputs that are $\epsilon$-close to each other. Due to the Lipschitz continuity of neural networks, the resulting intermediate activation vectors are also close in value ($\|x_a - x_b\| < L\epsilon$). Although each vector is shuffled by a different, unknown permutation, the values within the vectors remain identical. The adversary aligns these vectors by solving an assignment problem (using the Hungarian algorithm) to minimize the distance between elements of the shuffled vectors, effectively recovering the correspondence between elements across different queries.
2. Constructing Close Inputs in GPT: A challenge arises because Transformer inputs are discrete token indices, making it difficult to generate $\epsilon$-close inputs directly. The authors exploit a property of secure inference protocols: the use of fixed-point arithmetic with truncation. Secure truncation introduces stochastic 1-bit errors (e.g., magnitude $2^{-18}$). By repeatedly feeding the same prompt, the adversary leverages these inherent probabilistic errors to generate the necessary minor perturbations in intermediate activations, creating the required $\epsilon$-close vectors.
3. Weight Extraction: Once the activation vectors are aligned to a common (but unknown) permutation $\pi$, the adversary constructs matrices $X^{(l)}\pi^{(l)}$ and $X^{(l+1)}\pi^{(l+1)}$. They solve for the weights using the linear relationship $W^{(l)'} = \pi^{(l)-1} X^{(l)-1} X^{(l+1)} \pi^{(l+1)}$. The resulting weights $W^{(l)'}$ differ from the original weights $W^{(l)}$ only by row and column permutations. Crucially, the paper notes that neural network functionality is invariant to such shuffling; thus, the extracted weights perform the same forward propagation as the original model.

Key Contributions

• Breaking the Shuffling Defense: The authors demonstrate that an adversary can successfully align activation vectors shuffled differently but containing close values, bypassing the assumption that permutation makes extraction impossible.
• Exploiting Protocol Noise: The work identifies and leverages the stochastic errors inherent in secure truncation protocols to generate close intermediate activations, overcoming the limitation of discrete token inputs in Transformers.
• Theoretical and Practical Extraction: The paper proves that aligned activations allow for the extraction of model weights that are functionally equivalent to the original, despite unknown row/column permutations.
• Empirical Validation: The attack is evaluated on Pythia-70m and GPT-2, demonstrating high alignment accuracy and successful weight recovery.

Experimental Results
The authors evaluated the attack on Pythia-70m and GPT-2 models using various fixed-point precisions (14, 16, and 18 bits).

• Alignment Accuracy: The proposed method aligns shuffled activations with Mean Squared Errors (MSE) ranging from $10^{-9}$ to $10^{-6}$. Less than 2% of elements are mismatched, with higher precision yielding better alignment.
• Weight Recovery: The extracted weights exhibit L1-norm differences from the oracle weights ranging from $10^{-4}$ to $10^{-2}$. Smaller models (Pythia-70m) showed lower error rates than larger ones (GPT-2).
• Cost Efficiency: The attack requires a query count proportional to the model dimension (e.g., ~32,768 for Pythia-70m). Using commercial API pricing, the total cost is approximately $1.
• Functional Equivalence: The stolen weights, when used as a surrogate model, achieve perplexity scores close to the original model (e.g., 31.81 vs. 44.46 for Pythia-70m). After a brief fine-tuning stage (at most 6 minutes), the recovered model matches the victim's perplexity (32.43 vs. 31.81), effectively transforming subsequent attacks from black-box to white-box scenarios.

Significance and Claims
The paper claims to reveal a significant information leakage in the shuffling defense, which was previously assumed to be secure against weight extraction. The authors argue that the defense fails because it does not account for the ability to align vectors based on numerical proximity rather than positional knowledge.

The significance of this work lies in exposing a vulnerability that allows adversaries to recover model weights with high fidelity at a low cost. The extracted weights are sufficient to analyze private models or serve as effective initialization for surrogate models, thereby strengthening existing black-box attacks. The authors conclude that the security of the shuffling defense in secure inference and related privacy-preserving settings requires reevaluation.

Limitations
The authors modestly note that while the attack is theoretically guaranteed and validated on small-scale models, the fidelity of extracted weights may degrade as model scale increases. However, they assert that successful validation on small-scale models is sufficient to demonstrate the fundamental vulnerability of the defense mechanism.