🤖 AI

VMask: Tunable Label Privacy Protection for Vertical Federated Learning via Layer Masking

VMask is a novel, tunable label privacy protection framework for Vertical Federated Learning that effectively thwarts powerful model completion attacks by selectively applying secret sharing to mask critical layer parameters, achieving an optimal balance between privacy, utility, and computational efficiency.

Original authors: Juntao Tan, Lan Zhang, Zhonghao Hu, Kai Yang, Peng Ran, Bo Li

Published 2026-01-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Juntao Tan, Lan Zhang, Zhonghao Hu, Kai Yang, Peng Ran, Bo Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you and a friend are trying to solve a mystery together, but you each hold different pieces of the puzzle. You have the clues (data), and your friend has the final answer key (labels). This is how Vertical Federated Learning (VFL) works: you collaborate to build a smart AI without ever showing your raw data to each other.

However, there's a sneaky problem. Even though you don't share the raw clues, your friend can sometimes look at the "intermediate notes" you send them and figure out the answer key anyway. This is called a Model Completion (MC) attack. It's like your friend looking at your rough sketches and realizing, "Oh, if you drew a cat this way, you must have been trying to guess if it's a pet or a wild animal."

The paper introduces a new tool called VMask to stop this. Here is how it works, using simple analogies:

1. The Problem: The "Too-Perfect" Sketch

In a normal VFL setup, your friend (the attacker) trains their part of the AI so well that their internal notes become a perfect map of your data. If they add a small "decoder" to their end, they can read your private secrets (like credit scores or medical diagnoses) with scary accuracy.

Existing defenses were like trying to stop a thief by either:

Blurring the whole picture: Making the AI so fuzzy that it stops working well (sacrificing accuracy).
Building a fortress: Using heavy cryptography that takes forever to compute (sacrificing speed).

2. The Solution: VMask (The "Layer Masking" Trick)

VMask is a clever way to protect your secrets without ruining the picture or slowing things down. It uses a technique called Secret Sharing, which is like splitting a secret recipe into two halves. Neither half makes sense on its own, but when combined, they work perfectly.

Here is the step-by-step magic:

The "Randomizer" (Layer Masking): Instead of letting your friend see the entire map of your data, VMask randomly scrambles specific parts of their notes (the "layers" of the AI). Imagine you are drawing a picture, but every few seconds, you swap your pen for a random, scribbly marker for a split second. Your friend can't trace the perfect line anymore because the path is broken.
The "Smart Selector" (Layer Selection): You don't need to scramble every part of the drawing; that would be too much work. VMask is smart enough to figure out exactly which parts of the notes are the most dangerous to keep. It targets the specific "critical layers" that reveal the most about the answer key and scrambles only those. This keeps the system fast.
The "Shadow Coach" (Tunable Privacy): This is the coolest part. The person holding the answer key (the defender) has a "Shadow Model"—a practice version of the AI trained on a tiny bit of extra data. Before the real game starts, the defender runs a simulation: "If I scramble these layers, can my friend still guess the answer?"
- If the friend still guesses too well, the defender scrambles more layers.
- If the friend is guessing randomly, the defender stops scrambling.
- This allows the defender to set a Privacy Budget: "I'm okay with a little bit of risk, but not too much." They can adjust this dial based on how much they care about privacy versus speed.

3. The Results: Fast, Safe, and Accurate

The authors tested VMask on 13 different types of data (images, text, numbers) and 5 different AI models. Here is what they found:

It stops the thief: VMask reduced the attacker's ability to guess the answer down to the level of random guessing (like flipping a coin).
It keeps the AI smart: The main AI model still worked almost perfectly. In some cases, the accuracy dropped by less than 0.1% (which is barely noticeable).
It's incredibly fast: Compared to the heavy "fortress" cryptography methods, VMask was up to 60,000 times faster. Even compared to the standard, unprotected version, it was only about twice as slow, which the authors say is a fair trade-off for the safety it provides.

In a Nutshell

VMask is like a security guard who doesn't lock the whole building (which stops business) or leave the doors wide open (which invites thieves). Instead, the guard strategically places "fog machines" in the most critical hallways. The thieves can't see the path to the treasure, but the employees can still walk through easily. Plus, the guard can turn the fog on or off depending on how much risk they are willing to take that day.

Technical Summary: VMask – Tunable Label Privacy Protection for Vertical Federated Learning via Layer Masking

1. Problem Statement

Vertical Federated Learning (VFL) enables multiple parties to collaboratively train a machine learning model while keeping their local datasets private. In a typical VFL setup, an active party holds the labels, while passive parties hold feature subsets. Although VFL is generally considered privacy-preserving because only intermediate feature embeddings and gradients are exchanged, recent studies have demonstrated its vulnerability to label inference attacks.

Among various attack surfaces, the Model Completion (MC) attack is identified as the most powerful and severe threat. In this scenario, a semi-honest (honest-but-curious) passive party, who strictly follows the training protocol, can fine-tune a locally trained bottom model with an added inference head to accurately predict the active party's private labels. This is possible because the bottom model learns to generate feature embeddings that are strongly correlated with the private labels.

Existing defense mechanisms against MC attacks face three critical challenges:

Suboptimal Trade-off: Methods often sacrifice significant model utility to achieve privacy, or vice versa.
Impractical Overhead: Cryptography-based solutions (e.g., Homomorphic Encryption, Secure Multi-Party Computation) offer rigorous privacy but incur prohibitive computational costs, making them unsuitable for large-scale deployment.
Lack of Tunability: Current methods do not allow the defender to flexibly adjust the level of label privacy based on specific resource constraints or accuracy requirements.

2. Methodology: The VMask Framework

The authors propose VMask, a novel framework that defends against MC attacks through layer masking while preserving model utility and offering tunable privacy.

Core Insight

The fundamental insight is that the strong correlation between input data and the resulting feature embeddings allows attackers to infer labels. VMask disrupts this correlation by randomizing specific layer parameters within the attacker's bottom model. By masking these layers, the model's ability to generate representative feature embeddings is degraded, thereby reducing the attacker's inference accuracy.

Key Components

VMask integrates two modules and two procedures into a standard VFL system:

Secret Sharing (SS) for Layer Masking:
- Instead of simply dropping layers, VMask uses Additive Secret Sharing to mask linear layer parameters.
- The active party and passive parties share the parameters such that the passive party (attacker) holds a random share, while the active party holds the complementary share (plus random noise to prevent reconstruction).
- This preserves the integrity of linear computations (Add and Mul operations) during training, ensuring the main task accuracy remains high.
Critical Layer Selection Strategy:
- Naively masking all layers is computationally expensive. VMask employs a strategy to select only critical layers for masking.
- Criterion: Layers are selected based on their accumulated gradient norms over the training history. Layers with larger gradients are deemed more critical for label inference.
- Observation: Empirical results show that masking non-adjacent critical layers is more effective at reducing attack accuracy than masking consecutive layers.
Shadow Model for Estimation:
- Since the active party cannot directly access the attacker's true gradients, it maintains a shadow model (identical architecture to the passive party's bottom model) trained synchronously on a small auxiliary dataset.
- The active party uses this shadow model to locally simulate the MC attack and estimate the extent of label leakage.
Tunable Privacy Budget:
- The framework supports a tunable privacy budget ( $\epsilon$ ). The active party simulates the attack on the shadow model; if the estimated attack accuracy exceeds the budget, the system greedily selects additional critical layers to mask until the leakage falls within the defined limit.

Workflow

Initialization: Passive parties and the active party initialize models. The first layer is masked by default.
Layer Masking: Based on the selection strategy, specific layers are masked using SS.
Secure Model Update: The VFL model trains. Masked layers are computed via secret sharing (secure forward/backward), while unmasked layers use standard plaintext computation.
Shadow Model Update: The active party trains the shadow model on the auxiliary dataset.
Layer Selection: The active party simulates an MC attack on the shadow model. If leakage > budget, more layers are added to the masking list for the next epoch.
Iteration: This process repeats until the global VFL model converges.

3. Key Contributions

Novel Framework: The introduction of VMask, the first framework to utilize secret sharing for layer masking to defend against MC attacks while maintaining high model utility.
Tunable Privacy: The design of a tunable privacy budget mechanism, allowing defenders to flexibly control the trade-off between privacy and utility based on actual requirements, a feature absent in prior works.
Efficient Selection Strategy: A layer selection strategy based on accumulated gradient norms that targets critical, non-adjacent layers, significantly reducing the overhead of applying secret sharing to the entire model.
Comprehensive Evaluation: Extensive experiments across 5 model architectures (MLP3, LeNet5, VGG13, ResNet18, DistilBERT) and 13 datasets (tabular, image, text), comparing VMask against 12 existing defense methods.

4. Experimental Results

The evaluation demonstrates that VMask achieves the best privacy-utility trade-off among all compared methods:

Privacy Protection: VMask successfully thwarts MC attacks, reducing the label inference accuracy to a random guessing level (comparable to a randomized attacker's model). In all tested cases, the attack accuracy was reduced to near-random levels (e.g., ~10-20% depending on the dataset, compared to 40-90% for unprotected models).
Model Utility: The framework preserves model performance with minimal degradation.
- In Transformer-based models (DistilBERT), the average drop in VFL model accuracy was only 0.09%.
- Across all models, the average accuracy drop was at most 0.34%.
Efficiency:
- VMask is up to 60,846 times faster than cryptography-based methods (e.g., Homomorphic Encryption).
- Compared to standard VFL, the runtime overhead is marginal. For large models like DistilBERT (67M parameters), VMask's runtime is only 1.8 times that of standard VFL, which the authors deem generally acceptable.
Tunability: The framework consistently controls label leakage within the defined privacy budget across different privacy levels.

5. Significance and Claims

The paper claims that VMask addresses the urgent need for a practical defense against the Model Completion attack in VFL. Its significance lies in:

Breaking the Trade-off: It achieves high model utility and strong privacy simultaneously, overcoming the limitations of previous ML-based defenses that often sacrifice one for the other.
Practical Deployment: By avoiding heavy cryptographic operations, it offers a solution that is computationally feasible for real-world scenarios, unlike existing cryptographic approaches.
Flexibility: It is the first to offer a tunable privacy budget, empowering defenders to make informed decisions based on their specific resource constraints and privacy requirements without needing to retrain or redesign the system.

The authors conclude that VMask provides a robust, efficient, and flexible solution for protecting label privacy in Vertical Federated Learning, making it a viable candidate for deployment in sensitive domains such as finance and healthcare where label leakage is a critical concern.