A Unified Masked Jigsaw Puzzle Framework for Vision and Language Models

2025-10-14

IEEE Transactions on Pattern Analysis and Machine Intelligence

AI Summary

Transformer models in federated learning are vulnerable to gradient-based attacks, as gradients from positional encodings can be exploited to reconstruct original inputs and leak sensitive information. To address this, this work proposes a unified Masked Jigsaw Puzzle (MJP) framework that enhances privacy and robustness by randomly shuffling input tokens and replacing original positional encodings with learnable embeddings at unknown positions. This disrupts local spatial structure and encourages the model to learn more robust feature representations. MJP is the first approach to unify the jigsaw puzzle mechanism across both vision and language Transformers. Experiments demonstrate consistent improvements in downstream performance on ImageNet-1K image classification and sentiment analysis tasks on Yelp and Amazon text datasets, while simultaneously providing strong defense against gradient inversion attacks.

Technology Category

Computer Vision: Adversarial Attacks & RobustnessMachine Learning: Adversarial Learning & RobustnessNatural Language Processing: Safety and Robustness

Application Category

User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systemsSecurity and Privacy: Security and privacy of machine learning and AI applicationsSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 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 (e.g., ImageNet-1 K) and sentiment analysis for text (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.

Problem

• federated learning
• gradient attack
• Transformer
• position embedding
• vision and language models

Innovation

• Masked Jigsaw Puzzle
• Position Embedding Privacy
• Federated Learning
• Transformer Robustness
• Unified Vision-Language Framework