NeurIPS2026 (marked accepted in the author list) Medical Imaging medical federated learning gradient privacy generative defense synthetic data privacy–utility trade-off

Aegis: Generative Gradient Masking for Privacy-Preserving Medical Federated Learning¶

Conference: NeurIPS2026 (marked accepted in the author list)
arXiv: 2609.38339
Area: Medical Imaging
Keywords: medical federated learning, gradient privacy, generative defense, synthetic data, privacy–utility trade-off

TL;DR¶

Aegis adds a task-relevant synthetic-image gradient update after each client-local epoch to increase sample mixing in linear leakage; it reduces the measured reconstruction rate to 9.38%–11.63% across three MedMNIST modalities, without providing zero-leakage or differential privacy guarantees.

Background & Motivation¶

Sharing raw chest X-rays, CT images, or pathology images across hospitals is often constrained by governance and authorization requirements, motivating federated learning (FL): data remain within institutions while model updates are aggregated. Those updates are not information-free summaries, however, and a server may infer training images from them. Secure aggregation protects the visibility of individual client updates during transmission and aggregation; it does not automatically eliminate additional leakage from aggregate results when a malicious server manipulates the broadcast model. The paper starts from this active-server threat, rather than treating “data never left the hospital” as a privacy proof.

Existing defenses mainly alter updates or protect communication. Differential privacy (DP) limits information leakage through calibrated random mechanisms, but requires explicit privacy accounting and may reduce task utility; pruning, Soteria, and adaptive-noise methods also change the learning signal. Aegis takes a different route: preserve normal training on real data and supply additional task-related synthetic updates, making the observed signal harder to associate with individual patient images. The medical setting creates a specific tension: synthetic samples must be numerous and diverse enough to provide mixing without steering the classifier away from its task through distribution mismatch.

The paper uses finite leakage capacity to motivate its defense: under the discussed linear-leakage mechanism, individual information becomes harder to separate when multiple samples contribute to the same leakage bin. A bin is a grouping of contributions in the leakage analysis, not a new client-side network module; having more samples than bins directly establishes, at most, that collisions exist, not that every real sample is protected. Core idea: after real-data training in each local epoch, add a large logical-batch synthetic gradient from a locally cached, task-relevant dataset, using structured learning signals for empirical privacy protection rather than substituting unconditional security claims for evaluation.

Method¶

Overall Architecture¶

Aegis takes the server-broadcast model, private client images and locally visible labels, and a conditional image generator available to the client. The client first constructs a synthetic defense cache offline. In each round, it performs normal local training, computes a large logical-batch gradient over that cache at the end of every local epoch, applies the update, and finally sends the accumulated parameter delta. The output remains a model update accepted by the underlying FL protocol; the cache is not uploaded, and no additional server aggregation step is required.

The defender trusts its execution environment and generator workflow, but not the server. The server may control the broadcast model and parameters, possess auxiliary distribution knowledge, and attempt to obtain patient images from early training updates; the client cannot reliably determine whether the model contains leakage structures. The experiments directly evaluate three linear-leakage methods—Robbing-the-Fed, LOKI, and Scale-MIA—not universal protection against every possible server behavior, membership inference, or client poisoning.

The diagram shows only the client-side defensive training flow. The generator synthesizes images here; it is not invoked during diagnostic inference, which uses only the trained classifier. Real labels and synthetic conditions provide training supervision, while communication still contains only accumulated updates.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Local images and labels<br/>trusted generator"] --> B["Task-Relevant Synthetic Cache"]
    B --> C["Per-Epoch Large-Batch Masking"]
    R["Broadcast model and real labels"] -->|real training in each epoch| C
    C --> D["Protocol-Compatible Accumulated Update"]
    D --> U["Server aggregation<br/>parameter updates only"]
    U -.->|after training| P["Diagnostic inference<br/>classifier on new images"]

The three defensive designs appear in order: Task-Relevant Synthetic Cache, Per-Epoch Large-Batch Masking, and Protocol-Compatible Accumulated Update. Server aggregation and final diagnostic inference are upstream/downstream processes, not new defense modules. Images can be reused across rounds, but masking gradients must be recomputed at the current model parameters; reusing images does not mean reusing stale gradients.

Key Designs¶

1. Task-Relevant Synthetic Cache: provide relevant additional contributions without unnecessarily disrupting classification

The client generates candidate images from local labels or text conditions, aligns their size, cropping, and channel layout with the task, and accepts samples according to an in-distribution budget. The default generator is Stable Diffusion 2.0 with OpenCLIP ViT-H/14. Appendix F.1 states that the default public checkpoint receives no dataset-specific fine-tuning, while other generators may be adapted on a sanctioned client-side training split. The defense does not replace real images with synthetic ones: synthetic images provide additional gradients, while all real data still participate in normal training.

Task relevance serves two purposes. Distributionally close samples may contribute to leakage groups relevant to real data, making the mixing more targeted; matched synthetic labels and images may also support the original task instead of directing the classifier toward unrelated objectives. Nevertheless, semantic similarity, small feature-space distance, and sufficient leakage-bin coverage are different properties. Good generation quality alone does not establish reliable protection.

The paper describes acceptance through an expected squared-distance budget between private and synthetic data:

\[ \mathbb{E}\lVert D_i-D_i'\rVert^2\leq H. \]

Here \(D_i\) denotes private data, \(D_i'\) the synthetic cache, and \(H\) the allowed discrepancy budget. This is the paper's constraint, not a complete implementation definition permitting direct subtraction of arbitrary datasets of different sizes. The main text mentions matched samples, while Appendix A.2 gives feature-space MSE in a frozen feature extractor and a small validation split as an example. It does not specify a definitive matching rule, feature extractor, modality-specific values of \(H\), or complete distance-estimation workflow. The budget therefore expresses a principle, not an unambiguously reproducible filter.

When rejection rates are high, the algorithm allows generator fine-tuning on a sanitized private subset. This expands the trust boundary: the generator, adaptation process, and cache storage must all be trusted. Keeping synthetic images local does not prove that fine-tuning avoids memorizing patient information. If the generator copies private samples, their gradients still enter shared updates and may reintroduce leakage. The paper provides no dedicated privacy accounting for generator memorization, membership inference, or this indirect disclosure pathway.

2. Per-Epoch Large-Batch Masking: increase mixed contributions through an additional task gradient

The default real mini-batch size is \(B=64\), the synthetic defense cache size is \(M_i=2048\), and the evaluated leakage-layer capacity is \(k=1024\). The intended design makes the number of real and synthetic contributions exceed that capacity; the synthetic logical batch is 32 times larger than the ordinary mini-batch. The client does not place 2048 synthetic samples into every real mini-batch. Instead, it finishes the real mini-batches in an epoch and then computes one gradient over the entire synthetic cache.

For limited memory, the paper recommends splitting the cache into micro-batches, accumulating one logical-batch gradient, and then updating parameters. Matching the full-batch mean gradient requires correct sample-count normalization and fixed parameters throughout accumulation; updating after every micro-batch is a different operation. Equivalence also requires a sample-decomposable loss. Batch-dependent operations or uncontrolled randomness prevent an unconditional claim of exact equivalence.

The paper presents \(B+M_i>k\) as the structural basis for protection, but this is not sufficient to establish mixing in every bin. The pigeonhole principle guarantees a collision in at least one bin; it does not exclude bins containing only real samples or ensure that synthetic samples cover every real sample's group. Distribution matching, bin occupancy, and gradient weights all influence protection. The nonzero reconstruction rates reported in the paper reflect this boundary.

The more precise interpretation is therefore that additional task-relevant contributions make sample separation harder in specific linear-leakage evaluations. The design provides no DP parameters, indistinguishability proof, or bound covering all attacks. If the server changes leakage capacity, exploits correlations from repeatedly reused caches across rounds, or uses an unevaluated inference mechanism, new defensive evaluation is needed; the default batch size does not establish security.

3. Protocol-Compatible Accumulated Update: place the defense in client training rather than change the aggregation interface

Algorithm 1 actually follows the sequence “multiple real mini-batch parameter updates → one synthetic parameter update.” This repeats for \(E\) local epochs before transmitting the final parameters relative to the global parameters at the start of the round. With the default \(E=3\), each participating client performs three synthetic logical-batch forward/backward computations per round, not just one. The server still receives and aggregates parameter deltas through FedAvg; when enabled, secure aggregation wraps that delta using the existing protocol.

Preserving the communication interface does not mean preserving the learning trajectory. Real gradients are not directly pruned or perturbed with random noise, but synthetic parameter updates change the starting point of the next epoch and consequently its real gradients. Thus, learning from all real data is a valid procedural description, not a statement that the final update is identical to unprotected training with an independent mask added only at upload time.

The main text uses two mathematical descriptions: one adds real and synthetic gradients with the same learning rate, while another, also used in Appendix G, defines a sample-count-weighted effective gradient:

\[ \tilde g_i=(1-\lambda)g_i+\lambda g_i',\qquad \lambda=\frac{M_i}{B+M_i}. \]

This is the effective direction used in the paper's analysis, not a stepwise identity for Algorithm 1's sequential updates. Default parameters give \(\lambda=32/33\). If both gradients come from mean losses, adding them with the same learning rate differs from mixing them with these weights. A real epoch may also contain multiple mini-batches, and the two gradient types are evaluated at different parameter points. Appendix B uses Adam experimentally, whereas the algorithm shows SGD-style updates; optimizer state further affects equivalence. The paper does not adequately reconcile normalization, rescaling, and implementation across these descriptions.

For deployment, compatibility means no new message type is needed and secure aggregation can be retained. Aegis does not replace secure aggregation or cover every risk in transport, client execution environments, or the generator supply chain. Evaluation mainly uses first-round aggregate results, so an unchanged protocol does not establish multi-round privacy composition.

Loss & Training¶

Aegis computes real and synthetic gradients using the same loss as the diagnostic classification task; it adds no separate attack loss. Synthetic conditions come from locally visible labels, and the images and conditions provide additional training supervision. The cache is constructed offline and reused across rounds, with reconstruction possible when the modality changes. The paper does not define an inference-time defense or require clinical inference inputs to pass through the generator.

The experiments simulate 100 non-overlapping clients, select 10% randomly per round, and restrict each client to at most five classes in the original non-IID configuration. Local training uses Adam with learning rate \(10^{-3}\), three epochs per round, and real mini-batches of 64. The default classifier contains three convolutional layers and a 512-unit fully connected layer; this classifier width must not be confused with the evaluated \(k=1024\) leakage-layer capacity.

Appendix G claims that the synthetic discrepancy budget enlarges constants but retains an \(\mathcal{O}(1/T)\) convergence rate under assumptions including smoothness, strong convexity, bounded gradient variance and norms, balanced client weights, and uniform sampling. This is neither a convergence theorem for deep CNNs or Transformers nor a privacy theorem. The fixed-learning-rate Adam experiments also differ from the theorem's diminishing-learning-rate setting.

The derivation itself needs clarification. A.1 splits the squared norm of a sum of two error terms into two expectations without explicitly handling the cross term. Smoothness in parameters is then used to bound gradient changes from data discrepancy, without a separate data-direction regularity assumption. The appendix also interprets zero squared paired distance as identical distributions, which are not generally equivalent. This note therefore reports the authors' theoretical claim and assumptions without treating the proof as fully verified or supplying missing constants.

Key Experimental Results¶

Main Results¶

Reconstruction rate (RR) is the fraction of evaluated images with PSNR above 18 dB. Lower PSNR and SSIM indicate worse reconstruction according to these image-similarity metrics, not comprehensive protection of patient identity or sensitive clinical details. The following table combines Robbing-the-Fed privacy results from main-text Table 2 with final task accuracy from Appendix Table A.4, using default \(M_i=2048\). PSNR and SSIM are averages over all reconstructions.

Dataset	Defense	ACC (%)	RR (%)	PSNR (dB)	SSIM
ChestMNIST	No defense	55.8	89.06	62.40	0.94
ChestMNIST	Aegis	55.5	9.50	18.70	0.46
OrganAMNIST	No defense	58.6	82.81	50.60	0.89
OrganAMNIST	Aegis	58.8	11.63	20.10	0.51
PathMNIST	No defense	57.3	78.13	45.20	0.86
PathMNIST	Aegis	57.2	9.38	17.30	0.38

Accuracy changes relative to no defense are −0.3, +0.2, and −0.1 percentage points, respectively, but these do not establish clinical diagnostic validity. Surviving reconstructions are particularly important: Table 2 reports Aegis successful-subset SSIM values of 0.91, 0.88, and 0.84, and corresponding PSNR values of 29.80, 31.40, and 27.90 dB. Lower average quality and unrecognizability of every patient image are different conclusions.

Ablation Study¶

Appendix Table A.5 sweeps the synthetic batch under the same Robbing-the-Fed configuration, with \(B=64\), \(E=3\), and \(k=1024\). The table retains accuracy and RR for all three modalities, showing that privacy improvement is not confined to a discrete threshold crossing.

Dataset	Synthetic batch \(M_i\)	ACC (%)	RR (%)
ChestMNIST	512	56.0	49.95
ChestMNIST	1024	55.9	27.03
ChestMNIST	1536	55.7	16.24
ChestMNIST	2048	55.5	9.50
OrganAMNIST	512	58.8	47.82
OrganAMNIST	1024	59.0	27.31
OrganAMNIST	1536	58.9	17.66
OrganAMNIST	2048	58.8	11.63
PathMNIST	512	57.5	44.34
PathMNIST	1024	57.6	24.53
PathMNIST	1536	57.4	15.21
PathMNIST	2048	57.2	9.38

Main-text Table 1 reports single-workstation wall-clock training time: 100 rounds for MNIST and 150 for CIFAR-10. These budgets differ, so the two rows do not directly compare task difficulty.

Dataset	No defense (s)	Aegis (s)	GC (s)	GD (s)	Soteria (s)
MNIST	290.5	1890.1	334.5	4622.4	19816.2
CIFAR-10	433.1	2339.3	512.1	8356.8	32482.0

Aegis is faster than GD and Soteria, but its training time is approximately 6.51 and 5.40 times the no-defense baseline; a blanket low-overhead description would be misleading. Appendix B reports approximately 5 seconds per accepted candidate. Generating 2048 candidates implies about 10240 seconds, or 2.84 hours of one-time client generation cost: a linear estimate from the stated speed, not a measured total initialization time. Table 1 does not clearly state whether this cost is included, and its RTX 3080 workstation measurements do not represent heterogeneous hospital devices.

Key Findings¶

Larger synthetic batches progressively reduce RR across modalities with small accuracy changes. Even when \(64+1024>1024\), RR remains 24.53%–27.31%, contradicting an interpretation of the capacity inequality as complete protection.
Appendix Table A.3 reduces CIFAR-10 RR from 87.50% to 12.50%, but Aegis successful-subset SSIM is 0.99. Main-text Figure 4 shows declining trends for all three evaluated methods, but the cache does not expose all curve values reliably; an exact three-method ranking cannot be invented.
Appendix F additionally varies generators, non-IID partitions, and backbones. Backbone comparisons share the same leakage front-end and do not establish resistance to every architecture-specific leakage mechanism. Tables omit variances and confidence intervals, so small utility differences should not be treated as significant advantages.

Highlights & Insights¶

Replacing purely random perturbation with task-relevant synthetic gradients connects the defensive signal to the learning objective. Its practical value depends jointly on protection, distribution shift, and cost, not merely on retaining real data.
Separating a large logical batch from micro-batch accumulation decouples memory use from contribution scale. This system design is reusable, provided gradient normalization and batch-dependent operation assumptions are explicit.

Limitations & Future Work¶

Protection is an empirical improvement against evaluated mechanisms, not DP or zero RR. Successful-subset quality, sensitive clinical-attribute leakage, and cross-round evaluation remain important.
Generator trust, memorization during private adaptation, and cache-reuse risks lack dedicated validation. The filtering budget and precise real/synthetic update normalization also need specification for rigorous reproduction.
Low-resolution MedMNIST classification is not evidence of clinical deployment validity. Generation quality, compute requirements, and diagnostic metrics for high-resolution CT volumes and whole-slide pathology remain unvalidated.
Theory covers only an idealized convex setting, and the identified derivation steps need clarification. Actual sequential Adam training and long-term utility require independent examination; convergence rates cannot substitute for security proofs.
The cache contains extraction artifacts: some thresholds appear as PSNR>>18, while the main text and Appendix B specify PSNR>18. Table A.2 contains a caption but no readable values, and image-based curves cannot be recovered from text. Extremely high PSNR values are retained as reported rather than silently corrected.

vs DP, Outpost, GC, and Soteria: Aegis adds task gradients rather than relying only on noise or pruning. Advantages at these parameter settings do not establish superiority to DP under an equivalent formal privacy budget.
vs secure aggregation: Secure aggregation limits visibility of individual client updates; Aegis changes their contents. They can be combined, but their protection targets and guarantee types differ.
vs Eloul et al.: The same-label batching and MSE baseline in Appendix E preserves utility but retains modality RR values of 85.94%, 79.69%, and 75.00%. Aegis uses an additional large synthetic batch for stronger empirical protection in the evaluated setting.

Rating¶

Novelty: 4/5 — Synthetic data specifically targets client-side leakage mixing, but the capacity argument is overstated.
Experimental Thoroughness: 3/5 — Multiple modalities and sensitivity studies are covered; complete statistics, multi-round privacy, and generator leakage audits are missing.
Writing Quality: 3/5 — The workflow is clear, but gradient mixing, sequential parameter updates, and the theoretical connection remain insufficiently reconciled.
Value: 4/5 — A useful defensive research direction, not a standalone privacy guarantee for clinical deployment.