FedProphet: Memory-Efficient Federated Adversarial Training via Robust and Consistent Cascade Learning

FedProphet: Memory-Efficient Federated Adversarial Training via Robust and Consistent Cascade Learning

Speakers: Minxue Tang, Yitu Wang, Jingyang Zhang, Yiran Chen, Hai Helen Li

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into FedProphet, an innovative framework designed to enable memory-efficient federated adversarial training while maintaining high model robustness and utility. Presented at MLSys 2025 by Minxue Tang and co-authors, this work addresses a critical challenge in the confluence of federated learning (FL) and adversarial training (AT): the prohibitive memory demands of training robust, large-scale models on resource-constrained edge devices. FedProphet introduces a novel approach combining local adversarial cascade learning with a sophisticated central training coordinator, effectively mitigating the "objective inconsistency" that plagues prior memory-efficient FL methods.

The core problem FedProphet tackles stems from the fact that achieving strong adversarial robustness often necessitates larger models, which clash with the limited memory and computational power prevalent in edge devices participating in FL. Existing solutions either sacrifice accuracy/robustness or incur significant performance penalties due to memory swapping. FedProphet offers a comprehensive solution that not only achieves comparable accuracy and robustness to full joint training but also delivers substantial memory reduction and training speedups, making robust FL a practical reality for a wider array of real-world applications.

The significance of FedProphet lies in its ability to bridge the gap between theoretical advancements in robust machine learning and their practical deployment in privacy-preserving, distributed environments. By enabling the training of robust models on devices that would otherwise be excluded due to resource limitations, FedProphet expands the applicability of FL to safety-critical domains where adversarial attacks are a major concern, such as autonomous systems, healthcare, and sensitive data processing.

Background

▶ Watch: Introduction to FedProphet and federated learning background (0:00)

Federated learning (FL) has emerged as a prominent distributed machine learning paradigm, primarily lauded for its ability to protect training data privacy. In a typical FL setup, a central server orchestrates the training of a global model across numerous edge devices. Each device holds its own local dataset, which is never shared directly with the server or other devices. Instead, clients perform local model updates (e.g., using Stochastic Gradient Descent (SGD)) based on the global model received from the server, and then transmit only these updated model parameters back to the server for aggregation. This iterative process allows for collaborative model training without compromising individual data privacy.

Despite its privacy-preserving nature, standard FL does not inherently protect models from adversarial examples. Adversarial examples are inputs carefully crafted with imperceptible perturbations that can cause a machine learning model to misclassify with high confidence. To counter this vulnerability, adversarial training (AT) has become the de facto standard for enhancing model robustness. AT involves generating adversarial perturbations on training data to maximize the model's loss, and then minimizing this maximized loss through standard model training. This process forces the model to learn robust features that are less susceptible to adversarial manipulation.

However, integrating AT into FL introduces significant challenges. Adversarial training typically leads to a degradation in the model's clean accuracy (performance on unperturbed data). To compensate for this and achieve both high clean accuracy and strong adversarial robustness, larger and more complex models are often required. For instance, empirical evidence suggests that a larger model can achieve significantly higher clean and adversarial accuracy compared to a smaller counterpart. While desirable for performance, this requirement for larger models directly conflicts with the resource constraints of edge devices in an FL setting.

Training large models on devices with limited memory often leads to frequent memory swapping, where parts of the model or data are constantly moved between RAM and slower storage (e.g., disk). This swapping introduces severe data access latency, drastically increasing training time and rendering the process impractical for real-time or time-sensitive applications. Previous efforts to develop memory-efficient federated learning methods have attempted to address this by allowing clients to train only a subset or a smaller version of the global model. While these methods successfully avoid memory swapping and enable training on constrained devices, they often suffer from a critical issue known as objective inconsistency. Objective inconsistency arises because different clients, training different parts or sizes of the model, produce gradients that may not align with the optimization objective of the full global model. This discrepancy can lead to suboptimal aggregation, poor convergence, and ultimately, significantly lower clean accuracy and adversarial robustness compared to training the full model without memory constraints. FedProphet specifically aims to mitigate this objective inconsistency while simultaneously ensuring robustness and memory efficiency.

Key Findings

▶ Watch: FedProphet framework overview: local trainer and server coordinator (3:30)

FedProphet introduces a paradigm shift in federated adversarial training, demonstrating that it is possible to achieve high model utility and robustness on memory-constrained edge devices without incurring the penalties of previous memory-efficient FL approaches. The core findings are:

• Comparable Accuracy and Robustness: FedProphet successfully trains models that achieve clean accuracy and adversarial accuracy on par with Joint Federated Training (JFT), which represents the ideal scenario of training the full model without memory limitations. This is a significant breakthrough, as prior memory-efficient FL methods consistently exhibited substantial drops in both metrics. This finding underscores FedProphet's effectiveness in resolving the objective inconsistency problem.
• Significant Memory Reduction: The framework achieves up to 80% memory reduction during local training on client devices. This drastic reduction is critical for enabling large model training on devices with limited RAM, directly addressing the root cause of memory swapping.
• Substantial Training Speedup: By eliminating the need for frequent memory swapping, FedProphet delivers a remarkable 2x to 11x speedup in overall training time compared to traditional federated adversarial training methods that would trigger swapping on constrained devices. This performance gain makes robust FL practical for real-world deployments.
• Robustness via Strong Convexity: The novel application of strong convexity regularization within a cascade learning framework at the local client level is proven to be an effective and memory-efficient strategy for guaranteeing the joint robustness of cascaded modules without doubling batch sizes or increasing memory footprint.
• Adaptive Server-Side Coordination: The introduction of a central training coordinator with adaptive perturbation adjustment and differentiating module assignment ensures an optimal balance between robustness and utility across heterogeneous clients, further contributing to the framework's superior performance.

Technical Deep Dive

▶ Watch: Local Trainer: Adversarial cascade learning for memory efficiency (4:15)

FedProphet's architecture is bifurcated into two primary components: a local trainer on each client device, employing adversarial cascade learning, and a central training coordinator on the server side. This synergistic design addresses memory constraints, robustness, and objective inconsistency simultaneously.

Local Trainer: Adversarial Cascade Learning

The local trainer on each client adopts adversarial cascade learning to reduce memory requirements while preserving model robustness. In traditional cascade learning, a large model is broken down into smaller, sequential modules. Each module is trained independently, one by one, with an appended auxiliary output model (typically a small, single linear layer). This auxiliary model calculates the loss for the current module, allowing its gradients to be computed without needing to load the entire subsequent network into memory. The output of a trained module (ZM) then serves as the input for the next module in the cascade. This sequential training drastically reduces the peak memory footprint.

To guarantee the joint robustness of the entire cascaded model, a crucial condition must be met: the perturbation of the output of a previous module should not exceed the robustness budget of the subsequent module. Formally, this can be expressed as ensuring that the output perturbation of module M-1 is less than or equal to the robustness budget for module M. A naive approach to enforce this condition during adversarial training would involve adding this output perturbation directly into the loss function and performing AT on this modified loss. However, this method effectively doubles the batch size, thereby increasing memory consumption—a direct contradiction to the goal of memory efficiency.

FedProphet ingeniously overcomes this memory hurdle by proposing an alternative: strong convexity regularization. A key lemma demonstrates that if the loss function is strongly convex with respect to the output feature, an upper bound on the perturbation can be established. To achieve this strong convex loss, two specific design choices are made:

1. Single Linear Layer Auxiliary Output: Utilizing only a single linear layer as the auxiliary output model for each module naturally leads to a logistic-like loss function, which is amenable to strong convexity.
2. L2 Regularization: An L2 regularization term is explicitly added to this loss function. This regularization directly enforces strong convexity.

By conducting adversarial training on this L2-regularized, strongly convex loss, FedProphet effectively reduces the sensitivity (CM) of the loss. This not only guarantees the required robustness for each module (and thus the joint robustness of the whole model) but also minimizes memory consumption during training. An additional, significant benefit of achieving adversarial robustness through this method is a concomitant reduction in objective inconsistency. The theoretical results indicate that robustness at both the module level and the whole model level leads to smaller upper bounds for objective inconsistency metrics (beta M prime, C upper M, beta M, C little M), thereby contributing to better overall performance, accuracy, and robustness.

Server-Side: Training Coordinator

The central server hosts a sophisticated training coordinator with two main components: adaptive perturbation adjustment and differentiating module assignment.

Model Partitioning

Before training commences, the global model is partitioned into a set of modules. FedProphet employs a greedy method for this partitioning. Given a predefined minimum reserved memory (Imin) that all clients must be able to accommodate, layers or blocks are progressively added to a module until its memory footprint reaches Imin. This greedy strategy is proven to yield the minimum number of modules for a given Imin, optimizing the number of communication rounds and reducing overhead. In the experiments, this method successfully partitioned models like VGG16 and ResNet34 into seven modules, achieving 80% memory reduction.

Adaptive Perturbation Adjustment

A critical challenge in cascade learning for robustness is setting the perturbation budget (epsilon M) for each module. A naive approach—setting epsilon M as the maximum of all previous module perturbations—often results in an excessively large perturbation, causing training to diverge. Moreover, feature representations are often unnormalized or unbounded, making it difficult to predefine a suitable perturbation budget.

To address this, FedProphet introduces an adaptive perturbation adjustment mechanism using a scaling factor, alpha M. This alpha M automatically scales down the perturbation for each module. The adjustment of alpha M is dynamic, based on the ratio between the clean accuracy (CM) and the adversarial accuracy (AM) of the current module.

• If the ratio CM/AM is observed to be too high (indicating a lack of robustness relative to clean performance), alpha M is increased, thereby increasing the perturbation budget to enhance robustness.
• Conversely, if the ratio is too low, alpha M is decreased.

This adaptive scheme ensures that the ratio between clean and adversarial accuracy remains within a constant, optimal range, thereby achieving a dynamic balance between robustness and utility throughout the training process.

Differentiating Module Assignment

To leverage the heterogeneity of client resources, the training coordinator implements differentiating module assignment. This component allows clients with greater computational resources (referred to as "profits" in the talk) to train more modules simultaneously in a single round. Two crucial conditions govern this assignment:

1. Memory Constraint: The combined memory footprint of the cascaded modules assigned to a client must not exceed the client's available memory. This condition is paramount to prevent memory swapping.
2. Time Constraint: The total training time for the cascaded modules on a given client must not exceed the time it takes for the slowest client to train a single module. This condition is vital for minimizing synchronization delays in each communication round, a common bottleneck in federated learning.

By adhering to these conditions, differentiating module assignment not only avoids memory swapping but also significantly minimizes the overall training latency in the federated network.

Aggregation Rule: Partial Average

Given that different clients may train varying numbers of modules, the standard federated averaging (FedAvg) rule is not directly applicable. FedProphet adopts partial average as its aggregation rule. In this scheme, each specific module is averaged only over the subset of clients that trained that particular module in the current communication round. This method is a common practice in memory-efficient federated learning approaches where clients might only process parts of the global model.

Experimental Setup & Results

▶ Watch: Achieving robustness reduces objective inconsistency (6:30)

To validate FedProphet, experiments were conducted to simulate realistic systematic heterogeneity across client devices. Devices were sampled from two distinct pools, reflecting variations in computational power and memory capacity. The models evaluated were VGG16 and ResNet34, trained on standard image classification datasets (implied by typical benchmarks, though specific dataset names like CIFAR-10/100 were not explicitly stated in the transcript, the context suggests common CNN benchmarks).

The global model was partitioned into seven modules using FedProphet's greedy method, leading to an impressive 80% memory reduction during local training on client devices.

The key performance metrics were clean accuracy, adversarial accuracy, and total training time. FedProphet was compared against several baselines, including:

• Joint Federated Training (JFT): This serves as the upper bound, representing training the full model without memory constraints, typically requiring memory swapping on constrained devices.
• Other memory-efficient federated learning baselines: These methods aim to reduce memory but often compromise on accuracy or robustness due to objective inconsistency.

Headline Results:

1. Accuracy and Robustness: The results clearly demonstrated FedProphet's superiority. While JFT achieved high clean and adversarial accuracy, it did so at the cost of significant memory swapping and latency. All other memory-efficient baselines failed to achieve both high clean accuracy and high adversarial accuracy simultaneously, exhibiting a clear trade-off or significant degradation in one or both metrics. In contrast, FedProphet achieved comparable clean accuracy and adversarial accuracy to JFT, proving its ability to effectively minimize objective inconsistency while maintaining strong robustness. For instance, in the VGG16/ResNet34 settings, FedProphet consistently matched or closely approached JFT's performance across both metrics.

1. Training Time and Speedup: The impact of memory swapping was starkly evident in the training time comparisons. JFT, when forced to operate on memory-constrained devices, incurred very high data access latency due to frequent swapping, leading to prolonged training times. FedProphet, by entirely avoiding memory swapping, achieved a remarkable 2x to 11x speedup in training time compared to JFT under memory-constrained conditions. This substantial acceleration makes robust FL practical for deployment.

In summary, the experimental evidence unequivocally supports FedProphet's claims: it provides a robust, memory-efficient, and fast solution for federated adversarial training, effectively addressing the long-standing challenges of objective inconsistency and resource limitations on edge devices.

Practical Implications

▶ Watch: Server-side training coordinator: adaptive perturbation adjustment (7:50)

FedProphet carries significant practical implications for a wide range of stakeholders involved in the deployment and development of machine learning systems, particularly in distributed and privacy-sensitive contexts.

For practitioners and model builders, FedProphet liberates them from the dilemma of choosing between model size, robustness, and deployability on edge devices. Traditionally, building a robust model often meant increasing its complexity, making it unsuitable for resource-limited platforms like smartphones, IoT devices, or embedded systems. FedProphet enables the deployment of large, highly robust models in these environments, expanding the applicability of federated learning to critical domains such as:

• Autonomous Driving: Robust perception models are crucial for safety against adversarial attacks on sensors.
• Healthcare: Privacy-preserving training of robust diagnostic models on patient data at the source.
• Industrial IoT: Deploying robust predictive maintenance models on factory floor devices with limited compute.
• Personalized AI: Training robust personalization models on user devices without compromising privacy or performance.

For infrastructure teams and deployers, FedProphet offers substantial operational benefits. By eliminating memory swapping, it drastically reduces training latency (up to 11x speedup) and minimizes the computational overhead on client devices. This translates to more efficient resource utilization, lower power consumption, and faster model updates in production. The differentiated module assignment mechanism further optimizes resource allocation across a heterogeneous client base, ensuring that available compute power is leveraged effectively without penalizing slower devices or causing synchronization bottlenecks. This flexibility is crucial for real-world federated deployments, which inherently involve diverse hardware capabilities.

However, it's also important to consider tradeoffs and limitations. While FedProphet significantly reduces local memory footprint, the server-side coordinator introduces its own layer of complexity. Managing adaptive perturbation adjustments and differentiated module assignments requires careful design and monitoring. The framework's scalability to extremely large models (e.g., beyond the tested VGG16/ResNet34 architectures) or to federated networks with an exceptionally wide variance in client capabilities might still present challenges. The talk primarily focuses on systematic heterogeneity (device capabilities) and not explicitly on data heterogeneity (non-IID data), which is another common and significant challenge in federated learning. While the reduction of objective inconsistency should generally help, further research might be needed to assess its performance under extreme non-IID data distributions. Furthermore, while adversarial training enhances robustness against specific types of attacks, it does not guarantee certified robustness or protection against all possible adversarial threats. The talk also did not delve into other privacy-enhancing techniques like differential privacy, which, while complementary, can introduce their own trade-offs with accuracy and robustness.

Despite these considerations, FedProphet represents a major step forward, empowering practitioners to build and deploy robust, privacy-preserving AI systems on the edge, pushing the boundaries of what's achievable in distributed machine learning.

Key Takeaways

• Memory-Efficient Robust FL: FedProphet enables training of large, robust models on memory-constrained edge devices within a federated learning framework.
• Adversarial Cascade Learning: It utilizes a novel local training approach that sequentially trains model modules with strong convexity regularization to ensure robustness and memory efficiency.
• Server-Side Coordination: A central training coordinator adaptively adjusts perturbation budgets and assigns modules based on client resources, optimizing the balance between robustness and utility.
• Objective Inconsistency Mitigation: FedProphet effectively addresses the objective inconsistency problem inherent in previous memory-efficient FL methods, leading to high clean and adversarial accuracy.
• Significant Performance Gains: The framework achieves up to 80% memory reduction and 2x to 11x training speedup by eliminating memory swapping.
• Practical Deployment: It makes robust federated learning feasible for real-world applications in resource-limited, privacy-sensitive environments.

About the Speaker(s)

The work on FedProphet was presented by Minxue Tang, alongside co-authors Yitu Wang, Jingyang Zhang, Yiran Chen, and Hai Helen Li. While specific affiliations or detailed bios were not provided within the transcript, their collective contribution to this technical article at MLSys 2025 highlights their expertise in the intersection of machine learning systems, federated learning, and adversarial robustness. Their research focuses on practical solutions to deploy advanced AI models efficiently on heterogeneous computing infrastructures.