Source-linked AI summary
Position Matters: Feature Inversion Attacks in ViT Split Inference with Token Reduction and Shuffling
Stefano Leggio, Giulio Rossolini, Alessandro Biondi
TL;DR
The paper asks whether token reduction and shuffling truly protect privacy in ViT split inference against feature inversion attacks. It introduces SARA to realign and reconstruct manipulated tokens, finding that shuffling alone provides weak protection, while reduction leaves leakage when semantic and positional information remains; an edge-side defense reduces leakage without changing the cloud model.
Problem
Feature inversion privacy is insufficiently understood for token-reduced or shuffled ViT representations because conventional attacks assume fixed spatial organization.
Method
The paper introduces SARA, which predicts token positions, restores spatial arrangement, reconstructs missing embeddings, and recovers the input, alongside positional-embedding removal with progressive edge-side knowledge distillation.
Results
Token shuffling provides only apparent privacy, token reduction offers stronger but limited protection, and the proposed defense substantially reduces reconstruction leakage while preserving task performance.
Takeaways & Limitations
Token reduction and shuffling are limited as standalone privacy mechanisms, while attenuating positional information on the edge is a practical defense that leaves the cloud-side model unchanged.
Takeaways & Limitations
Utility is evaluated only through downstream classification accuracy, while other computer-vision tasks and domains remain future work.
Abstract
from arXiv · showhide
Vision Transformers (ViTs) are increasingly used in split-inference systems, where edge devices transmit intermediate token representations to a remote cloud. In this setting, token reduction lowers computation and communication costs, while token shuffling disrupts the spatial organization of the transmitted tokens, potentially limiting information leakage. However, their privacy benefits remain unclear against feature inversion attacks, which attempt to reconstruct the input from the transmitted embeddings. In this work, we show that, despite disrupting the spatial structure required by conventional reconstruction attacks, transmitted token embeddings retain substantial positional information. Based on this observation, we introduce the Spatially Aligned Reconstruction Attack (SARA), a unified pipeline that predicts token positions, restores their spatial layout, reconstructs missing embeddings using a feature-space masked autoencoder, and recovers the input image. Our results demonstrate that token shuffling provides only apparent privacy, as SARA largely reconstructs the original token organization. Token reduction offers stronger protection, but significant leakage persists when the retained tokens preserve sufficient semantic and positional information. Finally, we introduce a lightweight edge-side defense that removes positional embeddings and progressively adapts the edge-side transformer blocks through knowledge distillation. It substantially reduces attack performance against SARA, while preserving downstream task accuracy and requiring no changes to the cloud-side model.
1 Introduction
This section frames token reduction and shuffling as efficiency-oriented operations whose privacy benefits are uncertain under feature inversion attacks. It introduces SARA to test leakage and motivates an edge-side defense that attenuates positional information while preserving task behavior.
- ViTs encode local visual content and spatial position in patch tokens, making intermediate embeddings suitable for token-level manipulation.
- Token reduction lowers subsequent computation and communication by dropping or merging tokens, whereas shuffling permutes their sequence and disrupts explicit spatial correspondence.
- Conventional inversion attacks may fail after token manipulation because they assume fixed spatial organization, not because information leakage has genuinely decreased.
- SARA predicts transmitted-token positions, restores spatial arrangement, reconstructs missing embeddings, and recovers the input image for evaluating leakage under both manipulations.
- Experiments find shuffling alone offers only apparent privacy, while reduction provides stronger but utility-costly protection that remains limited when retained tokens preserve semantic and positional information.
- The proposed lightweight edge-side defense attenuates positional cues and reduces reconstruction leakage while maintaining a favorable task-performance trade-off.
2 Preliminaries and Threat Model
This section defines ViT token representations, edge-side token operations, and the split-inference threat model. It formalizes how reduction and shuffling alter transmitted representations and how a passive cloud attacker reconstructs private inputs.
- ViT tokenization and token operations: A ViT partitions an image into non-overlapping patches, projects each patch into a d-dimensional embedding, and combines it with a learnable positional embedding.
- ViT tokenization and token operations: Original ViT blocks preserve the patch-token sequence length, whereas token reduction progressively decreases it by dropping or merging selected tokens.
- ViT tokenization and token operations: With a fixed reduction amount r after each of the first ℓ blocks, the available patch-token count is N_ℓ = max{0, N_0 − ℓr}.Once N_ℓ = 0, only the class token remains; increasing r improves efficiency but may reduce downstream accuracy.
- ViT tokenization and token operations: Token shuffling preserves token count but permutes the sequence after a selected block, concealing the original spatial arrangement from the transmitted representation.
- Threat model: The honest-but-curious cloud observes the manipulated intermediate representation but not the original image, token provenance, or shuffle permutation, and remains passive during inference.
- Threat model: The attacker trains a reconstruction model on an auxiliary dataset and evaluates leakage by comparing reconstructed and private images with pixel-level and perceptual metrics.
3 Attack Pipeline
SARA addresses the failure of conventional inversion attacks on shuffled or reduced ViT tokens by recovering spatial organization before reconstructing missing features and the input image.
- Motivation: Conventional decoders depend on fixed token arrangement, with shuffling reducing ImageNet SSIM from 0.700 to 0.246.The comparison uses aligned and randomly shuffled ViT-B/16 tokens.
- Motivation: Token-position predictors show that intermediate ViT embeddings retain accessible spatial information across layers.Transformer-based prediction maintains high accuracy deeper in the network than a linear predictor.
- Token position prediction and placement: SARA predicts each transmitted token’s original patch position and places tokens into a spatially indexed full-length representation.Confident duplicate assignments retain the token with the highest confidence; unresolved positions receive a void value.
- Missing-token reconstruction: A feature-space masked autoencoder infers missing token embeddings from the spatially aligned representation.If no positions are missing, this stage is bypassed and the aligned representation is used directly.
- Image reconstruction: A convolutional decoder projects, reshapes, and upsamples the completed representation to reconstruct the input image.The overall pipeline comprises position prediction and alignment, missing-token recovery, and image decoding.
4 Proposed Defense
The proposed defense removes positional embeddings on the edge and progressively distills the original model’s behavior into the altered edge-side transformer blocks.
- Defense objectives: The defense attenuates positional information in smashed representations to reduce SARA’s token-position prediction effectiveness.It also preserves task-level behavior without modifying the cloud-side ViT.
- Defense design: A student ViT is created by removing positional embeddings from the edge-side model while retaining the original pretrained model as a frozen teacher.The student is initialized with the teacher’s pretrained parameters before adaptation.
- Progressive adaptation: Knowledge distillation uses logit KL divergence to adapt the edge-side transformer blocks while preserving the original model’s output behavior.The student’s initial representation uses patch embeddings without the teacher’s positional embeddings.
- Progressive adaptation: Progressive block-wise adaptation provides more stable task-objective optimization than jointly fine-tuning the entire edge-side model.The paper evaluates this design choice and a possible min–max optimization strategy in ablations.
5 Experimental Results
Experiments evaluate SARA and the defense on ImageNet-1K with ViT-B/16 and MAE-B/16, using reconstruction, classification, and combined privacy–utility metrics.
- Models and datasets: Experiments use ImageNet-1K with supervised ViT-B/16 and self-supervised MAE-B/16 models.The models share the ViT architecture but differ in pretraining strategy.
- Models and datasets: MAE-B is more vulnerable to reconstruction attacks than supervised ViT-B, including at deeper split points.The paper attributes this pattern possibly to MAE-B’s reconstruction-oriented pretraining.
- Token operations and metrics: The evaluation includes token shuffling, ToMe similarity-based token merging, and random token dropping.Reduction methods discard or merge a fixed or progressively selected number of tokens after ViT blocks.
- Token operations and metrics: Higher SSIM, PSNR, and FSIM indicate better reconstruction quality, greater information leakage, and weaker privacy protection.These metrics quantify image-reconstruction quality from the attacker’s perspective.
- Privacy–utility evaluation: PURI combines relative classification utility and privacy degradation through a weighted harmonic mean.The reported setting uses λ = 0.7, assigning greater importance to utility preservation.
- Experimental setup: SARA components are trained independently for each split point, and all components are retrained when evaluating a defended model.The attack uses separate training procedures for its position predictor, masked autoencoder, and decoder.
5.2 Attack results
SARA remains effective against shuffled and reduced token representations, showing that shuffling alone does not prevent reconstruction while token reduction provides model- and depth-dependent protection.
- Shuffling: SARA fully reconstructs inputs after token shuffling for both ViT and MAE models, with quality comparable to the unshuffled baseline across split points.The conventional convolutional decoder fails because shuffling disrupts token order, whereas SARA infers the original spatial positions.
- Shuffling: Token shuffling alone does not defend against SARA because the original token order remains recoverable even at deep ViT split points.This indicates that the apparent privacy benefit of shuffling results from breaking conventional decoder assumptions rather than removing positional information.
- Token reductions: Optimal reduction settings preserve downstream classification accuracy near baseline while balancing privacy against reconstruction attacks.The optimal reduction amount varies by split point, with lower values generally selected at deeper points.
- Token reductions: The reported metrics do not clearly distinguish ToMe from random token dropping, possibly because they ignore the image regions most relevant to semantics.Both approaches are evaluated using representative settings selected by maximizing PURI.
- Token reductions: Token reduction substantially degrades SARA reconstructions only at later ViT-B/16 split points, whereas MAE-B/16 maintains SSIM above 0.5 across all depths.For MAE-B/16, reconstruction quality remains remarkably stable as the split point moves deeper.
5.3 Defense results
The defense attenuates positional information to reduce SARA reconstruction leakage under shuffling and token reduction, while preserving a favorable task-accuracy trade-off. Progressive finetuning and adversarial training address compatibility and deep-layer limitations, respectively.
- Shuffling: The defense consistently reduces reconstruction quality under shuffling for ViT-B/16, while MAE-B/16 remains more vulnerable at deeper split points.MAE-B/16 retains stronger spatial structure because of its masked-patch reconstruction objective.
- Token reductions: The defense shifts the privacy–utility trade-off toward lower reduction amounts, allowing more tokens to be retained while maintaining favorable privacy scores.The optimal reduction values are consistently lower than for undefended models, and PURI curves approach their maxima at small reductions.
- Token reductions: SSIM remains below 0.4 across all split points for defended ViT-B/16 and MAE-B/16 under optimal token reduction.The defense substantially attenuates the increasing deep-layer reconstruction trend observed for MAE-B/16 under shuffling alone.
- Progressive finetuning: Progressive finetuning generally provides more stable results and a better privacy–utility trade-off than directly finetuning the entire edge-side model.Updating one transformer block at a time better preserves compatibility with the fixed cloud-side network and limits abrupt representation changes.
- Deep MAE layers: Adversarial training limits the shuffling-induced SSIM increase from 0.117 without adversarial training to 0.022 with adversarial training.Without adversarial training, SSIM rises from 0.384 at k=2 to 0.501 at k=8; with it, SSIM rises from 0.355 to 0.377.
- Deep MAE layers: The adversarial formulation can cause a non-negligible accuracy drop, particularly for ViT-B/16, and training at k=10 exhibited stability issues.The paper therefore retains simpler progressive finetuning as its primary defense.
6 Related work
Prior work studies feature inversion attacks and token-level privacy mechanisms in split inference. The paper distinguishes its client-only defense from approaches that require costly or unstable learning or alter the cloud-side model.
- Feature inversion attacks: Feature inversion attacks map transmitted intermediate features back toward the private input in split-inference settings.The paper places these attacks within the broader model inversion attack literature.
- Feature inversion defenses: Existing defenses include cryptographic, data-modification, and learned approaches with different computational, communication, and training costs.Cryptographic methods provide strong guarantees but generally incur substantial overhead, while learned approaches may require costly or unstable adversarial training.
- Proposed defense: The proposed defense operates only on the client side, avoids unstable learning trends, and leaves the cloud-side model unchanged.This design targets split-inference deployments with multiple edge devices sharing a fixed server-side model.
- Token-based privacy mechanisms: Token shuffling permutes transmitted token order, whereas token reduction discards or merges tokens as lightweight representation manipulations.These mechanisms have been investigated for privacy enhancement while preserving inference performance.
7 Conclusions and limitations
The paper finds that token reduction provides some obfuscation, but shuffling—especially for MAE-B/16—does not reliably prevent high-quality reconstruction. A progressive client-side defense that attenuates positional information reduces feature inversion effectiveness.
- Conclusions: SARA shows that token shuffling and token reduction have limited standalone privacy against feature inversion attacks.Token reduction provides some feature obfuscation, whereas shuffling can still permit high-quality reconstructions, particularly for MAE-B/16.
- Conclusions: The proposed defense removes client-side positional encoding through progressive fine-tuning to improve intermediate-feature obfuscation.Its reported effect is reduced effectiveness of feature inversion attacks.
- Limitations: The evaluation measures utility only through downstream classification accuracy and does not yet cover other computer-vision tasks.Semantic segmentation is identified as a future evaluation direction.
- Limitations: Adapting SARA and the defense to language models and testing more sophisticated token reordering against adaptive attacks remain open directions.These directions define the current scope boundary beyond the evaluated vision setting.
A.1 Optimal ratios with PURI
The defense changes the PURI-optimal reduction settings across split points and techniques, enabling favorable privacy scores at lower reduction amounts. This can retain more tokens and reduce accuracy loss while preserving privacy–utility balance.
- Optimal ratios: The defense shifts the PURI-optimal reduction amount toward lower values, rather than making smaller reduction automatically preferable.Lower reduction amounts retain more transmitted tokens and can improve accuracy while maintaining favorable privacy scores.
- Optimal ratios: Increasing the reduction amount lowers transmitted-token count and communication and cloud-side computation, but may degrade task accuracy.The reported optimal values therefore encode a privacy–utility trade-off at each split point.
A.2 Additional Illustrations.
Additional reconstructions are provided for the original and defended models across different split points, using PURI-optimal reduction amounts.
- Figures 14 and 13 show reconstructions at PURI-optimal reduction values for original and defended models across different split points.