Defending against Reconstruction Attack in Vertical Federated Learning

Inspired by prior work (Li et al., 2019; Feutry et al., 2018; Goodfellow et al., 2014), AR simulates an adversarial attacker who aims to train a reconstructor $\mathcal{R(.)}$ that maps the embedding $\mathcal{F(X)}$ to the input $\mathcal{X}$ by minimizing the following reconstruction loss:

where $\mathcal{R(.)}$ can be any model (e.g. a MLP).

Ideally we can formulate the protection as a max-min problem that maximizes the minimized $\mathcal{L}_{r}$. However we empirically find that such optimization is unstable and hard to tune. Instead, we use Gradient Reversal Layer (GRL) from prior work (Ganin & Lempitsky, 2015; Feutry et al., 2018) that demonstrated promising results in stabilizing adversarial training. As shown in Figure 1, GRL is inserted between the feature extractor $\mathcal{F(.)}$ and the adversarial reconstructor $\mathcal{R(.)}$. In forward pass, GRL just performs identity transformation. In backward pass, it multiplies the corresponding gradient w.r.t to the cut layer by $-\lambda$ ($\lambda>0$, i.e. $\lambda=1$) and passes $-\lambda\frac{\partial{\mathcal{L}_{r}}}{\partial{\mathcal{F}}}$ to the preceding layer. Intuitively, GRL leads to the opposite of gradient descent that is performing gradient ascent on the feature extractor $\mathcal{F(.)}$ with respect to maximize the adversarial reconstruction loss (as an attacker). Therefore it roughly achieves the goal of maximizing the minimized reconstruction loss. After inserting GRL, we just need to minimize $\mathcal{L}_{r}$ as adversarial training. We update both $\mathcal{F(.)}$ and $\mathcal{R(.)}$ during training; after training is finished, we discard $\mathcal{R(.)}$ and save $\mathcal{F(.)}$ as the more robust feature extractor.

Noise Regularization also simulates an adversarial reconstructor $\mathcal{R}^{\prime}(.)$ but with a different goal. It is designed to reduce information about $\mathcal{X}$ in $\mathcal{F(X)}$ by misleading the reconstructed input $\mathcal{R^{\prime}(F(X))}$ toward a random direction (and therefore degrading reconstruction quality). During training, we generate random Gaussian noise $\mathcal{N}_{noise}$¹¹1Random noise from other distributions such as Uniform distribution are effective too. and minimize the following noise regularization loss:

Note that if only NR is used, we can train an independent reconstructor (without GRL) as $\mathcal{R(.)}^{\prime}$. If AR and NR are used together, we can reuse the AR reconstructor $\mathcal{R(.)}$ as $\mathcal{R(.)}^{\prime}$ ²²2Since the adversarial training is effective, $\mathcal{R(.)}$ and $\mathcal{R(.)}^{\prime}$ can achieve similar reconstruction performance. . In our experiments, we use the same reconstructor for both AR and NR. For notation simplicity, we will use $\mathcal{R(.)}$ to represent both AR and NR reconstructor.

As shown in Figure 1, $\mathcal{F(X)}$ are fed into the reconstructor $R$ directly (without GRL). NR calculates $\mathcal{L}_{n}$ and computes the gradients of the $\mathcal{L}_{n}$ w.r.t. $\mathcal{F(.)}$, i.e. $\nabla\mathcal{L}_{n}(\mathcal{F})$, and updates $\mathcal{F(.)}$ accordingly via backpropogation. Note that $\mathcal{R(.)}$ only works as a reconstructor and provides the reconstructed result $\mathcal{R(F(X))}$ as an input for $\mathcal{L}_{n}$ optimized by NR. NR only has effects on $\mathcal{F(.)}$ and does not update any parameters of $\mathcal{R}$ during backpropogation phase.

Distance correlation is designed to make input $\mathcal{X}$ and embedding $\mathcal{F(X)}$ less dependent, and therefore reduces the likelihood of gaining information of $\mathcal{X}$ from $\mathcal{F(X)}$. Distance correlation measures statistical dependence between two vectors³³3Two vectors can have different length.  (Vepakomma et al., 2018). The distance correlation loss is the following:

where we minimize the (log of) distance correlation loss ($\mathcal{L}_{d}$) during the model training. It can be interpreted as $\mathcal{X}$ being a good proxy dataset to construct $\mathcal{F(X)}$ but not as vice versa in terms of reconstructing $\mathcal{X}$ from $\mathcal{F(X)}$  (Vepakomma et al., 2018).

Note that dCor computates pairwise distance between samples and requires $O(n^{2})$ time complexity where $n$ is the batch size. In practice, there are some faster estimators of dCor (Chaudhuri & Hu, 2019; Huang & Huo, 2017). In addition, dCor is sensitive to the $n$ and a larger $n$ can give a more accurate estimation of the distance correlation.

We can unify all three modules into one framework. The overall loss function is a combination of four losses: adversarial reconstruction loss ($\mathcal{L}_{r}$), noise regularization loss ($\mathcal{L}_{n}$), distance correlation loss ($\mathcal{L}_{d}$), and normal label prediction loss ($\mathcal{L}_{c}$). In this paper, we focus on classification and use categorical cross entropy as label prediction loss. Optimizing $\mathcal{L}_{c}$ makes sure the model maintain a good utility; optimizing $\mathcal{L}_{r}$, $\mathcal{L}_{n}$, and $\mathcal{L}_{d}$ increases model privacy. Uniting them in one framework can help us defend against the reconstruction attack while maintaining the accuracy of the primary learning task. The overall loss function is:

where $\alpha_{d}\geq 0$, $\alpha_{n}\geq 0$ and $\alpha_{r}\geq 0$ are weights for distance correlation, noise regularization, and adversarial reconstructor module respectively. Note that during training, only the passive party optimizes these three modules while there is no change on the active party’s training optimization.

We evaluate the performance of using noise regularization module alone, i.e. minimizing $\mathcal{L}_{c}+\alpha_{n}\mathcal{L}_{n}$. First, we evaluate the impact of noise choice on model privacy and utility. In Figure 2 (a) and (b), we compare two different types of random noise: Gaussian noise and Uniform noise . With the same $\alpha_{r}$, both random noises achieve similar MSE and AUC. Therefore the NR module is not sensitive to the type of the random noise. In addition, we can see the tradeoff between privacy (MSE) and utility (AUC) by varying the value of $\alpha_{r}$. A larger $\alpha_{r}$ leads to a higher MSE (therefore better privacy) but a lower AUC score (therefore worse utility).

We evaluate the performance of minimizing distance correlation alone, i.e. minimizing $\mathcal{L}=\mathcal{L}_{c}+\alpha_{d}\mathcal{L}_{d}$. Figure 2 (c) and (d) show the dCor performance with different values of $\alpha_{d}$. First, vanilla vFL without any privacy protection, can naturally reduce the $dCor(\mathcal{X},\mathcal{F(X)})$ during the training (dropping from $0.6$ at the beginning of the training to $0.45$). Second, unsurprisingly optimizing dCor can reduce $dCor(\mathcal{X},\mathcal{F(X)})$ more than vanilla ($0.45$ for vanilla and $0.2$ for dCor with $\alpha_{d}=0.1$) while AUC drops less than $0.01$. It indicates that dCor module can achieve a reasonable privacy-utility tradeoff with an appropriate $\alpha_{d}$. Third, log(dCor) is more robust to $\alpha_{d}$ than dCor since the gap of log(dCor) between $\alpha_{d}=0.01$ and $\alpha_{d}=0.1$ is much smaller.

We now demonstrate the effectiveness of optimizing all losses together, i.e. DRAVL, by comparing it with individually optimizing each module. Figure 3 (a) shows that using NR module alone can reduce more $dCor(\mathcal{X},\mathcal{F(X)})$ than vanilla but its AUC is lower than dCor (Figure 3 (c)). Figure 3 (b) shows that dCor has lower MSE than NR, meaning NR preserves more privacy. This is unsurprising given NR is specifically designed to degrade the reconstruction quality. In addition, DRAVL helps gain the advantages of each module-it can reduce $dCor(\mathcal{X},\mathcal{F(X)})$ and increase MSE simultaneously. Unfortunately, DRAVL also hurts AUC more than optimizing any of modules alone. However, in terms of finding the best overall tradeoff, DRAVL shows more promising results among all competitors: compared to vanilla, on average DRAVL increases MSE by $9.5\%$ and decreases the dCor by $41.44\%$ with a cost of AUC drop by only $2.25\%$.

We also compare DRAVL with a straightforward yet effective protection baseline: adding random noise to the cut layer embedding. We generate a random noise from a zero-mean Gaussian and add it to the embedding. We only tune the standard deviation of the Gaussian noise to control the noise strength. As shown in Figure 3 (d), (e), and (f), with increasing the amount of noise added to the embedding, we can get a better privacy protection (lower dCor and higher MSE) but worse model utility (lower AUC). When noise strength is large enough (standard deviation $\geq 25$), the cut layer embedding is covered by the noise . As a result, the AUC drops to $0.5$, which is equivalent with a random guess. Overall, with a good control of the amount of the random noise added to the cut layer embedding, it might be an effective protection strategy.

We compare this noise perturbation with DRAVL in Figure 3 (d), (e), and (f). DRAVL and noise perturbation method with $std=12.5$ can achieve similar AUC and MSE, since both models can make the reconstructed input be similar to the mean of the raw input. However, DRAVL reduces dCor more than the noise perturbation. Another drawback of perturbation based is that compared to DRAVL, empirically it is much harder to tune in order to find a good trade-off between model utility and privacy.