[Baselines] Add Entity Augmentation VFL baseline (arXiv:2406.17899)

Type

New Baseline

Description

New baseline proposal

Paper: Entity Augmentation for Efficient Classification of Vertically Partitioned Data with Limited Overlap Authors: Avi Amalanshu, Viswesh Nagaswamy, G.V.S.S. Prudhvi, Yash Sirvi, Debashish Chakravarty (IIT Kharagpur) Link: https://arxiv.org/abs/2406.17899

I am one of the authors of this paper and would like to contribute it as a Flower Baseline.

Summary of the method

Traditional VFL requires entity resolution, private set intersection and entity alignment, so guests can only train on the intersection of their entity sets and must discard the rest of their data. Entity Augmentation removes the alignment step entirely: guests forward whichever entities they hold, in whatever order, and the host synthesises a label for each concatenated activation by interpolating the labels of the contributing entities, weighted by each guest's activation dimension. This is functionally a form of CutMix in activation space, so it also acts as a regulariser. With only 5% entity overlap on CIFAR-10, adding the misaligned data via Entity Augmentation raises test accuracy from 48.1% to 69.48%.

Experiments to reproduce

Primary scope — CIFAR-10, ResNet-18, 2 guests, 60 epochs, Adam (β1=0.9, β2=0.999), lr=0.001, with each guest owning a CNN branch plus half of the first fully-connected layer and the host owning the remaining fully-connected layers:

1. Limited overlap (Table 4). Dataset split into x% entity-aligned data, with the remaining (100−x)% shuffled and divided evenly between the two guests as strictly non-overlapping samples. Train on the aligned subset only, versus aligned + misaligned via Entity Augmentation:

   Setting	Aligned only	Aligned + Misaligned
   CIFAR-10, x = 5%	48.1%	69.48%
   CIFAR-10, x = 10%	54.34%	70.4%

2. Full alignment vs full misalignment (Table 2). Train on fully entity-aligned data versus fully misaligned data with Entity Augmentation — CIFAR-10 / ResNet-18: 72.92% vs 74.34%.

3. Training curves (Figure 2). Test accuracy per epoch for the above settings, showing the faster convergence and smoother/more stable training that Entity Augmentation produces.

Planned extensions once the CIFAR-10 pipeline is reviewed:

• MNIST with ResNet-18: 97.1% → 98.37% (x = 5%), 97.92% → 98.5% (x = 10%), and 99.25% vs 98.20% for fully aligned vs misaligned.
• Deeper architectures on both datasets: ResNet-56 (77.09% vs 79.12% on CIFAR-10) and ResNeXt-29 8x64 (81.08% vs 82.06%).

Implementation notes

• Built on the Message API: the host ServerApp drives mini-batch VFL rounds through Grid — a query message collects guest embeddings, a train message delivers the corresponding embedding gradients.
• Guest model and optimizer state persist across messages via context.state, following the pattern established by the reworked examples/vertical-fl.
• Unlike examples/vertical-fl (full-batch, tabular, entity-aligned), this baseline performs mini-batch VFL training on images with partial or zero entity overlap. As far as I can tell it would be the first VFL entry in baselines/.
• Data loading and vertical partitioning through flwr-datasets; runs with a single flwr run . and supports --run-config overrides.

Planned Implementation

No response

Additional Context

No response