Marshaled Learning: Bridging Large Neural Networks with Memory-Constrained Trusted Execution Environments in Federated Learning

Despite the privacy-oriented design, federated learning (FL) remains vulnerable to privacy breaches due to the exposure of model update snapshots throughout training. Trusted Execution Environments (TEEs) offer hardware-based isolation to safeguard data and computations, providing a compelling foundation for privacy-preserving FL. However, the limited memory available in mainstream TEEs hinders the deployment of large-scale neural networks, such as GPT models, within these secure enclaves. To address this limitation, we propose Marshaled Learning, a novel FL framework that enables large neural network training across memory-constrained TEEs while ensuring strong privacy guarantees for both data and model owners. To achieve this, Marshaled Learning partitions a model into subnets and distributes them across clients according to their memory capacities, coordinating forward and backward passes across TEE-isolated environments. To mitigate the impact of heterogeneous data distributions and straggler clients, we introduce a dynamic knowledge propagation mechanism that facilitates cross-client learning and accelerates convergence. We present both theoretical convergence guarantees and empirical evaluations, demonstrating that Marshaled Learning outperforms existing FL methods by around 2% to 5% accuracy with much faster convergence rates. We also implement Marshaled Learning on commercial Azure Confidential VMs to prove its feasibility and show that it incurs only a 1 ~ 3× computational overhead compared to non-TEE settings, validating its practicality in real-world deployments.