Abstract
As large-scale foundation models become the backbone of modern AI, Split Inference (SI) has emerged as a standard for deploying these models on resource-constrained edge devices. While SI is often touted for its inherent privacy benefits—sending intermediate representations (IR) rather than raw data—this paper introduces DRAG (Data Reconstruction Attack using Guided Diffusion), a powerful method that demonstrates these representations are far from secure. By leveraging the generative power of Latent Diffusion Models (LDMs), DRAG can reconstruct high-fidelity original images from the deep layers of Vision Transformers (ViTs), effectively bypassing existing privacy defenses.
Introduction
Split Inference (SI) partitions a neural network into a lightweight "client" portion and a heavy "server" portion. The client processes the raw input into an intermediate representation (IR), or "smashed data," which is then transmitted to the cloud. The prevailing assumption has been that these abstract features protect user privacy by discarding pixel-level detail.
However, Data Reconstruction Attacks (DRAs) aim to invert this process. While previous attacks were limited to shallow layers or specific architectures like small CNNs, the rise of Vision Transformers (ViTs) and foundation models necessitates a re-evaluation of these risks. DRAG addresses this by treating reconstruction as a conditional generation task, utilizing the rich image priors found in pre-trained diffusion models to "un-smash" data even from the deepest layers of a model.
Methodology
DRAG introduces a novel framework for reconstruction that moves beyond the limitations of traditional GAN-based or optimization-based attacks.
1. Diffusion as a Prior
Instead of attempting to reconstruct an image from a blank slate, DRAG utilizes a pre-trained Latent Diffusion Model (LDM), such as Stable Diffusion. This model acts as a powerful "image prior," providing the system with an inherent understanding of what natural images look like, which helps fill in the gaps left by compressed or pruned representations.
2. Guided Reconstruction and Tweedie’s Formula
The attack begins with random noise and iteratively refines it. At each step, DRAG calculates a "guidance gradient" by comparing the IR of the current estimate to the target IR from the victim. A significant challenge is that client models are trained on clean images, while the diffusion process involves noisy intermediates. To bridge this gap, DRAG employs Tweedie’s formula to estimate a denoised version of the image at every iteration before feeding it into the client model for comparison.
3. DRAG++
To further enhance performance, the authors propose DRAG++, which incorporates an Inverse Network. This network provides a coarse initial guess of the original image, which the diffusion process then polishes into a high-quality, realistic reconstruction. This hybrid approach combines the speed of direct inversion with the generative fidelity of diffusion.
Experimental Setup
The researchers evaluated DRAG across a diverse set of datasets to ensure its general-purpose applicability:
- Natural Images: MSCOCO and ImageNet-1K.
- Human Faces: FFHQ.
- Specialized Domains: UCMerced LandUse (aerial imagery) and LSUN-bedroom.
The attack was primarily tested against Vision Transformers (ViTs), focusing on the vulnerability of representations as deep as the 12th layer.
Key Insights and Results
Foundation Models are Not Private
The study reveals that even highly compressed representations from the 12th layer of a ViT contain enough semantic information to allow for recognizable reconstruction. While older attacks (like rMLE or GLASS) produced unrecognizable blobs at these depths, DRAG consistently recovered images that closely resembled the originals.
Domain-Agnostic Capabilities
Unlike previous state-of-the-art attacks like GLASS, which require training a specific GAN for every target category (e.g., one for faces, one for cars), DRAG is remarkably robust. Using a single general-purpose LDM (Stable Diffusion), it can attack various data types. Even when the diffusion prior was trained on a different distribution (e.g., bedrooms) than the target (e.g., ImageNet objects), DRAG outperformed competing methods.
Bypassing Modern Defenses
DRAG was tested against several popular privacy-preserving techniques:
- Channel Pruning (DISCO/NoPeek): These defenses remove "sensitive" channels. The authors found that an adaptive version of DRAG can ignore these pruned channels and still recover the data.
- Token Shuffling: Some researchers suggested that shuffling the order of patches (tokens) in a ViT would protect privacy. DRAG researchers proved this is a weak defense by training a simple MLP to reorder shuffled tokens with high accuracy, making reconstruction possible again.
Efficiency and Fidelity
Quantitatively, DRAG significantly outperformed previous methods across metrics such as MS-SSIM (structural similarity) and LPIPS (perceptual loss). Furthermore, DRAG proved more efficient; at 4,000 iterations, it often achieved better results than GAN-based attacks requiring 20,000 iterations.
Conclusion and Future Work
DRAG demonstrates that current Split Inference defenses are insufficient for the era of foundation models and Vision Transformers. The ability of guided diffusion to extract high-fidelity private data from deep, abstract layers suggests that privacy must be baked into the architecture of the models themselves rather than added as an afterthought via pruning or shuffling. Future research is required to develop "transformer-native" privacy protections that can withstand the generative power of diffusion-based attacks.