A Unified Masked Jigsaw Puzzle Framework for Vision and Language Models

TL;DR

The paper tackles privacy risks in federated learning posed by gradient inversion attacks on Transformer-based vision and language models, focusing on the泄 spatial information embedded in Position Embeddings. It introduces Masked Jigsaw Puzzle (MJP), a unified framework that shuffles tokens and masks their position embeddings with a learnable unk embedding, complemented by a Dense Absolute Localization (DAL) loss for vision and an n-gram window with auxiliary losses for NLP. Empirical results show MJP preserves or even boosts performance on ImageNet-1K and sentiment/QA tasks while significantly reducing gradient leakage under both analytic and optimization-based attacks, across ViT/DeiT/Swin and BERT-BASE. The approach offers a practical privacy-preserving augmentation for federated learning with broad applicability to vision and language transformers, and the authors provide public code to facilitate adoption and further exploration.

Abstract

In federated learning, Transformer, as a popular architecture, faces critical challenges in defending against gradient attacks and improving model performance in both Computer Vision (CV) and Natural Language Processing (NLP) tasks. It has been revealed that the gradient of Position Embeddings (PEs) in Transformer contains sufficient information, which can be used to reconstruct the input data. To mitigate this issue, we introduce a Masked Jigsaw Puzzle (MJP) framework. MJP starts with random token shuffling to break the token order, and then a learnable \textit{unknown (unk)} position embedding is used to mask out the PEs of the shuffled tokens. In this manner, the local spatial information which is encoded in the position embeddings is disrupted, and the models are forced to learn feature representations that are less reliant on the local spatial information. Notably, with the careful use of MJP, we can not only improve models' robustness against gradient attacks, but also boost their performance in both vision and text application scenarios, such as classification for images (\textit{e.g.,} ImageNet-1K) and sentiment analysis for text (\textit{e.g.,} Yelp and Amazon). Experimental results suggest that MJP is a unified framework for different Transformer-based models in both vision and language tasks. Code is publicly available via https://github.com/ywxsuperstar/transformerattack

Figures (9)

• Figure 1: Illustration of privacy leakage through Transformer model gradients. An attacker reconstructs input data by minimizing the distance between the real gradients $\nabla{\mathbf{W}}$ and the dummy gradients $\nabla{\mathbf{W}'}$. This leakage is mitigated by shuffling the input data and employing the MJP method to disrupt the spatial information of the PEs while preserving model performance, with $x$ and $\tilde{x}$ denoting the original and shuffled data, respectively.
• Figure 2: Low-dimensional projection of PEs. (a) The 2D PCA projection (DeiT-S touvron2021training), it shows that the PEs from neighboring positions tend to cluster together, which reveals a consistent pattern that reflects the spatial relationships of the input patch positions. (b) The 3D PCA projection (DeiT-S touvron2021training), it also shows that the position information is well captured with PEs.(c) The 1D PCA projection (BERT$_\text{BASE}$2018BERT), a similar sinusoidal pattern shows the same order as the input token positions. (d) The 3D PCA projection (BERT$_\text{BASE}$2018BERT), position information can be captured by PEs. Note that the embedding of index 0 (highlighted in red) corresponds to the first token embedding.
• Figure 3: (a) A unified MJP framework: We feed the randomly shuffled tokens (vision for block-wise shuffling and text for n-gram shuffling) into the model, and the PEs of these selected tokens are masked correspondingly. Note that segment embedding is only used in NLP tasks. Dark yellow indicates selected PEs are masked with unknown (unk), and the rest of PEs remain unchanged. (b) Training process: After MJP operation, the embeddings are fed to the model directly. Note that the unk embeddings are defined as parameters which are learned in the training process.
• Figure 4: Ablation on the mask ratio $\gamma$ during inference: (a) comparisons among ViT-S, DeiT-S, and our method (trained with $\gamma=0.03$); (b) comparisons between Swin-T and our method (trained with $\gamma=0.03$); (c) comparisons of our method on DeiT-S trained with different $\gamma$.
• Figure 5: Ablation studies over the masking ratio $\gamma$ on the Yelp: (a) comparisons between the shuffle and MJP method on BERT$_\text{BASE}$ (under different training ratios $\gamma$); (b) comparisons between the shuffle and our method on BERT$_\text{BASE}$ (trained with fixed ratio $\gamma=0.5$) when evaluated with different inference masking ratios. (c) comparisons of our method on BERT$_\text{BASE}$ trained with different ratios.