SARA cracks ViT split inference privacy tricks

Problem

The paper studies privacy leakage in Vision Transformer split inference, where an edge device sends intermediate token embeddings to a cloud. Two lightweight edge-side operations are common: token reduction to lower computation and communication, and token shuffling to hide spatial layout. The question is whether these operations meaningfully prevent feature inversion attacks that reconstruct the original image from transmitted embeddings. This matters because an honest-but-curious cloud can observe smashed representations and attempt input reconstruction without altering the inference protocol.

Approach

The authors first show that intermediate ViT token embeddings retain substantial positional information even after several transformer blocks. They introduce the Spatially Aligned Reconstruction Attack, SARA, a three-stage pipeline: a token position predictor that estimates each transmitted token’s original patch index and restores a spatially aligned full-length representation; a feature-space masked autoencoder that imputes missing token embeddings; and a convolutional decoder that reconstructs the image from the completed token grid. The attacker is passive and trains SARA components on an auxiliary dataset. Experiments use ImageNet-1K with ViT-B/16 and MAE-B/16, and consider two token-reduction strategies, random dropping and ToMe merging, plus token shuffling. Reconstruction quality is measured by SSIM, PSNR and FSIM, and privacy–utility trade-offs are summarised by a Privacy–Utility Reconstruction Index, PURI.

Key Findings

• Positional information survives ViT processing: simple predictors can recover patch indices with high accuracy in early layers, and transformer-based predictors remain effective deeper, enabling spatial realignment of shuffled tokens.
• Token shuffling gives only apparent privacy: a conventional convolutional decoder fails on shuffled tokens, but SARA largely recovers original token organisation and restores reconstruction quality to near the aligned baseline; shuffling alone is insufficient.
• Token reduction improves privacy but incompletely: removing or merging tokens reduces leakage, yet significant reconstruction persists when retained tokens keep enough semantic and positional cues. Optimal reduction balances privacy and accuracy as quantified by PURI.
• MAE-pretrained models are more vulnerable: MAE-B/16 shows consistently high reconstruction quality across split points, reflecting its reconstruction-oriented pretraining that preserves spatial relations.
• Edge-only defence is effective and practical: removing positional embeddings on the edge and progressively fine-tuning edge blocks via knowledge distillation substantially reduces SARA performance while preserving downstream accuracy and requiring no cloud changes. This defence reduces SSIM of reconstructions under shuffling to below 0.4 for ViT-B/16, though MAE remains more challenging.
• Adversarial extension of the defence, which trains an auxiliary token-position predictor while adapting the edge to frustrate it, further stabilises protection across layers but can introduce accuracy loss and optimisation instability.

Limitations

Evaluation focuses on image classification utility; other tasks such as segmentation are not studied. The defence may be less effective or less stable on models pretrained for reconstruction, and the adversarial variant can reduce accuracy. Results are specific to the tested ViT variants and ImageNet data; additional token-reordering strategies or domains remain to be explored.

Implications

An attacker in an honest-but-curious cloud can recover spatial layout and reconstruct private inputs even when tokens are shuffled or partially reduced, by exploiting residual positional cues and imputing missing tokens. Lightweight edge-only defences that remove positional embeddings and fine-tune edge blocks can substantially reduce this leakage without changing cloud models, but adaptive attackers and reconstruction-oriented pretrained models remain a practical risk.