CAFE: Catastrophic Data Leakage in Vertical Federated Learning

In plain words

Some machine-learning projects split the work between several parties that each hold different information about the same people. A bank might know a customer's payments, a hospital their test results, and a shop their purchases. Each party keeps its own data and only passes along intermediate numbers, so that together they can train one model. This setting is called vertical federated learning.

The usual belief is that these shared numbers are safe, and that training with bigger batches of records makes them even safer, because the numbers then mix many records together. This paper shows that the belief is wrong for this setting. The coordinating server knows which records are in each batch, because it has to tell every party which rows to use. With that knowledge, the server can work backwards from the mixed numbers. It first recovers how each party's first layer responded to each record, then it guesses the original inputs and keeps adjusting the guesses until the numbers they would produce match what the parties sent.

The authors report that this works on batches of 40 records and on the largest batch they tried, 100, and that it recovers images well. They also propose a defense: each party sends made-up numbers that look like its real ones, so the server's guesses converge to the wrong inputs. The defense has no formal privacy guarantee.

TL;DR

• Attacks vertical FL, where parties hold different features of the same records. The server knows the batch indices and the model weights, and sees the aggregated gradients §3.1 §3.3.
• Three optimization steps recover, in turn, the gradient at each worker's first layer outputs, the first-layer inputs, and then the training data and labels §3.3 alg. 1 alg. 2 alg. 4.
• Author-reported PSNR at batch size 40 (4 workers, 800 training images): 31.83 on CIFAR-10, 43.15 on MNIST, 33.22 on Linnaeus 5, against 3.86 to 18.94 for five baselines §4.1 tab. 2.
• Relevant for anyone who shares gradients in VFL and relies on batch size as a defense. The released code covers MNIST only and has several differences from the paper (see Paper vs code).

Problem

In federated learning, parties share gradients rather than raw data, and a common belief is that larger batches make gradient-based reconstruction harder §1. The earlier attacks in the paper's comparison table report maximal batch sizes of 5 to 48, and one reports 100 with mostly unrecognizable results tab. 1. For the DLG-style objective, the aggregated gradient has a fixed size while the number of unknown inputs grows with the batch size K, so the problem becomes underdetermined as K grows §3.2 eq. 3. Some earlier analyses rely on assumptions that the authors call impractical, for example that the number of classes greatly exceeds the batch size §1.

The paper's setting is an honest-but-curious server in VFL. It follows the protocol but tries to recover clients' data from the aggregated gradients. It assumes the server controls batch selection through data-index alignment: for every batch it sends the list of row ids to all parties, so the server knows which rows each batch contains §3.1 §3.3.

Method

Notation §3.1: M workers, N training rows, batch size K. Worker m holds a feature partition x_{n,m} of each row, and the labels sit at the server or at one worker. The first fully connected layer has weights Θ_1 (d1 x d2) and bias b_1 (d2). Θ_c are the parameters before that layer, and h(Θ_c, x) gives the inputs to it.

The server draws a batch index vector s uniformly from the set S of all K-subsets of rows §3.1. The batch loss and gradient are the 1/K-normalized sums over the selected rows eq. 1 eq. 2.

Why the batch is hard. Recovering the rows directly from the aggregated gradient gives the objective in eq. 3. Its unknown set grows with K, so many solutions fit. CAFE avoids this by using the known batch indices, which turns the system into one that can be determined §3.2 §3.3.

Step I: gradient at the first-layer outputs §3.3 alg. 1. Let u_n be the output of the first layer for row n eq. 4. The gradient with respect to all u_n is the matrix in eq. 5. Each batch's bias gradient is the sum of the rows selected by s eq. 6. The attacker does not see eq. 5 directly, so it estimates a matrix V (N x d2) by minimizing the expected squared error over random batches eq. 7. The paper proves this objective is strongly convex when K < N (Theorem 1, proof in app. B) and that SGD converges to the truth almost surely for fixed Θ §3.3. Algorithm 1 initializes V with uniform random values and updates it with the gradient of the Step I objective each iteration.

Step II: inputs to the first layer §3.3 alg. 2. Using the chain rule, the gradient of the first-layer weights equals H^T times the output gradients, where H (N x d1) stacks the layer inputs h_n §3.3. The attacker estimates H as Ĥ by minimizing the expected Frobenius-norm error between the batch's weight gradient and the sum of ĥ_n v_n^T over selected rows eq. 8. Theorem 2 states that, if N < d2 and Rank(V*) = N, the Step II objective is strongly convex and SGD converges to H for fixed Θ. Its proof is in app. C. The paper notes that for models whose only learned layers are fully connected, Steps I and II alone may already recover the data §3.3.

Step III: recover data and labels §3.3 alg. 3 alg. 4. The attacker initializes fake rows x̂_n and fake labels ŷ_n with uniform random values. It then minimizes the expectation over batches of the objective in eq. 9. This has three terms:

1. A gradient-matching term: the squared distance between the real batch gradient with respect to all parameters Θ and the gradient the fake batch produces, weighted by α.
2. A truncated total-variation term on the fake inputs of the batch, weighted by β. It is zero when the TV norm is below a threshold ξ.
3. An internal-representation term: the squared distance between the Step II estimate Ĥ* and the first-layer inputs that the fake rows produce, weighted by γ. The paper says this term contributes most to recovery §4.2.

Two update schedules are described. The nested-loops schedule runs Steps I and II to a fixed number of iterations, then runs Step III alg. 3. The single-loop schedule runs one update of each step per iteration alg. 4. The paper reports that single-loop needs fewer iterations §4.2 tab. 5.

# One CAFE iteration in single-loop form  [alg. 4]
# Input:  batch indices s (K of N rows); bottom models; server model
# State:  V (N, d2), H_hat (N, d1), fake inputs x_hat, fake labels y_hat
# Output: updated V, H_hat, x_hat, y_hat

def cafe_iteration(s, V, H_hat, x_hat, y_hat, models, server):
    # Regular VFL round: the real batch gradients, which the server sees  [§3.3]
    grad_theta = vfl_batch_gradients(models, server, s)
    # Step I: fit V so that its batch sums match the first-layer bias
    # gradient  [alg. 1] [eq. 6] [eq. 7]
    V = sgd_step_on_V(V, s, grad_theta.b1)
    # Step II: fit H_hat so that H_hat^T V matches the first-layer weight
    # gradient  [alg. 2] [eq. 8]
    H_hat = adam_step_on_H(H_hat, V, s, grad_theta.theta1)
    # Step III: fit fake inputs and labels to the real batch gradients, to
    # the recovered inputs, and to the TV term  [alg. 4] [eq. 9]
    x_hat, y_hat = adam_step_on_data(x_hat, y_hat, s, grad_theta, H_hat)
    return V, H_hat, x_hat, y_hat

The paper's own descriptions of the optimizer for Steps II and III are the generic update of alg. 2 and alg. 4. Their concrete form is not stated in the paper text; see Paper vs code.

Defense: fake gradients §3.4 alg. 5, app. E. Each worker replaces its true gradient with a look-alike before uploading it. The worker draws ν random gradient vectors with i.i.d. entries from N(0, σ²) and sorts each one's entries in descending order. It keeps the candidate nearest, in L2 distance, to its own true gradient sorted the same way. If that distance is above a threshold τ, it draws new candidates. Then each true entry is clipped to the range set by the candidate value of the same rank, and the result is uploaded. The authors argue that this keeps large and small entries in the same places, so training still converges §3.4 alg. 5. The defense is deterministic in essence and does not give differential privacy app. F.3.

Hyperparameters

Default experimental setting: M = 4 workers, K = 40, batch ratio 0.05, N = 800 training rows §4 tab. 2 tab. 9.

Results

Author-reported. Setup: VFL with M = 4 workers, K = 40, 800 training rows, compared with five baselines implemented by the authors. The paper does not state whether the model parameters were frozen for these tables. The hardware was an RTX 2080 and a TITAN V §4. The paper gives PSNR values only; the definition of PSNR is not in the compiled text (see Pitfalls).

Comparison with the state of the art §4.1 tab. 2:

The baselines replace the objective in eq. 9 with the named alternatives §4.1. The paper does not report the learning rates or iteration counts used for the baselines beyond the statement that all methods ran on the same model and iterations §4.1.

PSNR against batch size, 800 training rows §4.2 tab. 3:

The paper states that PSNR stays above 30 on CIFAR-10, above 32 on MNIST and above 28 on Linnaeus 5 across these batch sizes, so batch size is not an effective defense in this setting §4.2.

Effect of auxiliary regularizers, K = 40 §4.2 tab. 4:

The paper concludes that γ (the internal-representation term) contributes most, and that MNIST recovers well even without the TV term §4.2.

Nested-loops against single-loop, iterations to reach the stop PSNR (maximum in parentheses) §4.2 tab. 5. Column order follows the table header:

Effect of the number of workers, K = 40 §4.2 tab. 6. Column order follows the table header:

Attacking while training §4.3 tab. 7. The table header gives PSNR with the learning rate in parentheses for three settings. The paper does not define what each setting is beyond that learning rate:

Training while attacking on MNIST §4.3 tab. 8:

Fake-gradient defense, CIFAR-10 with unchanged Θ §4.4: PSNR of the recovered data is 28.68 with ordinary gradients and 7.67 with fake gradients. The paper reports that the training loss with fake gradients follows a similar trajectory to that with true gradients §4.4 (figure in §4.4).

Comparison with DP-based defense app. F.3 tab. 11. The paper reports training loss, not PSNR, for this comparison. Under DP at epsilon 10, 5, 1 and 0.1 the training loss stays between 2.69 and 2.92 across the listed iterations. Fake gradients and true gradients both fall; at 8000 iterations the values are 2.69 for every DP setting, 1.15 for fake gradients and 0.46 for true gradients. The paper says the noise DP needs to block leakage also keeps the loss from falling app. F.3.

Effect of the number of epochs and the batch size on Linnaeus 5 app. F.2 tab. 10: the paper says that increasing K and the number of iterations both contribute to attack performance.

Recovery of 32 x 32 faces (Yale): PSNR above 42 §4.5 (figure in §4.5).

Headline scale claim: CAFE leaks up to 800 images with 4 workers at batch size 100 on the authors' hardware §4.

Table 9: hyperparameters per dataset app. F.1 tab. 9:

Table 10: PSNR on Linnaeus 5 by epoch and batch size app. F.2 tab. 10. Columns are batch sizes K:

Table 11: training loss under DP, fake gradients and true gradients app. F.3 tab. 11:

Paper vs code

Official code: https://github.com/DeRafael/CAFE at code_commit, MIT license. The paper's abstract links it, and the repository's license header names the copyright holder DeRafael; the cafe.py header names Xiao Jin, an author. Code paths below are relative to the repository root. I did not run the code: it needs TensorFlow and the MNIST download, so it was read only.

• Dataset. The paper evaluates MNIST, CIFAR-10 and Linnaeus 5 §4. The loader reads MNIST only data_preprocess.py and its docstring says CIFAR-100 data_preprocess.py. No CIFAR, Linnaeus or Yale loader is present.
• Feature split. The paper does not describe how features are divided among workers. The code gives each worker one 14 x 14 quadrant of each 28 x 28 image, and keeps the labels with the server data_preprocess.py utils.py.
• Width of the first layer. The paper's worked example uses d2 = 1024 with N = 800 §3.3. The code's first fully connected layer per worker has 256 outputs model.py. Each worker's bottom model is two 5 x 5 convolutions with 64 and 128 filters, each followed by 2 x 2 max pooling, then dense layers of 256, 64 and 10 units model.py. The server is one 10-way softmax layer over the concatenated four 10-unit outputs model.py utils.py.
• Learning rates. Table 9 gives MNIST lr 1e-2 for all three steps. The code sets Step I to SGD with 5e-3, Step II to 1e-2 and Step III to 1e-2 config.py cafe.py. The 5e-3 Step I value is the paper's CIFAR-10 and Linnaeus 5 value, not its MNIST value.
• Optimizers. The paper states the step updates as objectives and does not name an optimizer for Steps II and III. The code uses Adam-style updates (beta1 0.9, beta2 0.999, epsilon 1e-7) with per-row moment buffers that are updated only when a row is in the batch, but bias-corrected with the global iteration count utils.py utils.py config.py.
• Gradient weights. The code multiplies the gradient-matching term by 100 and then by 1e-4, which gives the paper's alpha = 1e-2 utils.py utils.py. The internal-representation weight is 1e-3, matching gamma utils.py. The TV weight is 1e-4, matching beta utils.py.
• TV term. The paper uses a truncated TV norm with threshold xi §3.3. The code adds the TV gradient only when the batch-mean TV norm exceeds a fixed 25, which is the paper's MNIST xi utils.py. The TV value is not included in the loss utils.py.
• Labels. The code applies a softmax to the fake labels before the loss utils.py. The paper does not say whether labels are logits or probabilities.
• Initial shapes. Algorithms 1, 3 and 4 initialize V as N x d1, and Algorithms 2, 3 and 4 initialize Ĥ as N x d2 alg. 1 alg. 2 alg. 3 alg. 4. The text defines V in R^{N x d2} and Ĥ in R^{N x d1} §3.3. The code follows the text: V has 256 columns and Ĥ has 2048 cafe.py cafe.py.
• Initial ranges. Fake inputs are uniform on [0, 1) utils.py. V is uniform on [-8e-4, 8e-4] utils.py. Ĥ is uniform on [0, 0.08) utils.py. The paper gives no range for any of these.
• Training during the attack. The released entry point, cafe.py, updates every victim model with Adam at learning rate 1e-6 each iteration cafe.py config.py. The paper calls the unchanged-parameter case the strongest attack §4.4, and the repository has no switch for it. The paper does not say which setting produced Tables 2 and 3.
• Switch thresholds. The single-loop switching rules (1e-9 and 5e-9) do not appear in the code. cafe.py runs all three steps every iteration cafe.py.
• Nested-loops file. triple_shot.py calls the first-step and second-step functions with argument lists that do not match their definitions triple_shot.py triple_shot.py first_shot.py double_shot.py. As written it does not implement the nested schedule.
• Iterations and decay. The code runs 20000 iterations config.py. Its learning-rate decay is set to multiply by 1 config.py, so the rate stays constant. The paper gives no decay schedule.
• Defense. No code for Algorithm 5 is in the repository. Searching for sorting, normal sampling or sigma found no defense implementation; the only sigma reference is commented out utils.py.
• PSNR. The code clips dummy inputs to [0, 1], computes PSNR with peak 1.0 on each 14 x 14 worker slice, and averages over the four workers utils.py. The compiled paper has no PSNR formula; its formula is in commented source text.

No discrepancies found: the Step I and Step II objectives, the gradient-matching term over all parameters, and the internal-representation term match eq. 7 eq. 8 eq. 9 first_shot.py double_shot.py utils.py.

Pitfalls

• Theorem 2 does not cover the released configuration. The guarantee needs N < d2 and Rank(V*) = N §3.3. The paper's example, N = 800 and d2 = 1024, satisfies this. The code's d2 = 256 model.py does not, because 800 > 256 and Rank(V*) = N cannot then hold. Treat Step II results on the released setup as empirical.
• Step I and II convergence is stated only for fixed Θ §3.3 app. B app. C. Table 7 shows PSNR falling as the model's learning rate rises, and the paper calls that an effective defense §4.3. No theory covers that regime.
• Batch-size claims are limited to one ratio. The tested batch sizes go up to K = 100 out of N = 800, a ratio of 0.125 tab. 3. The paper does not test larger ratios or other N. Its statement that batch size is not a defense holds only over this range.
• Non-monotonic batch-size results. On CIFAR-10, K = 40 scores lower than K = 20 tab. 3. The paper does not explain this, and it reports one value per setting with no variance or seeds.
• Label recovery is not measured. Step III recovers labels too §3.3. The paper reports PSNR and training loss only. No label accuracy is given.
• Dependence on batch indices. The attack depends on the server knowing each batch's rows §3.3. The paper states that data-index alignment is needed in VFL, so the attack applies to protocols that send indices to every party.
• Server knowledge. Step III needs each worker's bottom model and its forward function h(Θ_c, ·) §3.1 §3.3. The paper says the server has the model parameters and their gradients §3.3. The attack therefore assumes more than gradient access alone.
• Table 7 settings. The header gives three settings only by their learning rate. The paper does not define what changes between them beyond that.
• Table 11 gap. The paper says fake gradients do not lose training efficacy §4.4. Table 11 shows a training loss of 1.15 for fake gradients against 0.46 for true gradients at 8000 iterations tab. 11. The gap is real, and the claim should be read as "similar trajectory", not "same".
• Defense parameters. The threshold tau, the number of retries and the learning rate that the defense experiments used are not fully stated. The SGD learning rate of 0.1 is given in §4.4 without saying which optimization it applies to.
• Baselines are the authors' implementations. They are reimplementations that swap the objective term §4.1. Their tuning is not reported, so the comparison in tab. 2 is not a comparison against each method's best published setup.
• PSNR is computed per worker slice. On the released code's MNIST split, each PSNR value averages four 14 x 14 slices rather than full 28 x 28 images utils.py. Compare numbers only with the same definition.
• Two different "dims". The Algorithm boxes and the text disagree on the shapes of V and Ĥ (see Paper vs code). Use the text's shapes and the code's values.
• Version. This card follows the arXiv v4 LaTeX source fetched by wikillm fetch. The published venue and any later version may differ.

Lineage

• Builds on: DLG, deep leakage from gradients [arXiv:1906.08935], the baseline whose objective Step III generalizes; GradInversion [arXiv:2104.07586], whose BN and group-consistency regularizers are compared. Neither has a card here, so their methods are not described.
• Compared against: DLG [arXiv:1906.08935]; iDLG; Inverting Gradients; SAPAG; R-GAP; the theory-oriented attack [arXiv:2010.13356]; GradInversion [arXiv:2104.07586]; and a framework for evaluating gradient leakage [arXiv:2004.10397]. Table 1 lists their reported maximal batch sizes §2 tab. 1.
• Concurrent work: the related-work section mentions a concurrent study of privacy leakage across multiple rounds of communication, which focuses on model inversion and its defense §2.
• Superseded by: not stated. No later work is covered in this wiki.