Skip to content

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

Description

@PrudhviGudla

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions