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. To achieve robust privacy-preserving FL that protects both data and model privacy, Trusted Execution Environments (TEEs) offer a promising solution by isolating code and data within a secure memory enclave. However, the memory capacity of commonly used TEEs still constrains the training of large-scale neural networks, such as GPT, creating significant challenges within these secure enclaves and thereby limiting the full potential of TEEs in federated learning. To address this limitation, we propose Marshaled Learning, a solution to protect FL privacy for both data and model owners while enabling large neural network training within memory-constrained TEEs. To overcome memory constraints, we first partition the neural network and distribute subnetworks to clients in alignment with their TEE memory capacities. We also facilitate end-to-end training for global optimization by enabling knowledge propagation across client-side TEEs. Given the distributed nature of subnetworks, we introduce a dynamic knowledge propagation mechanism to enhance knowledge transfer in FL. This diversified propagation accelerates FL on heterogeneous data and mitigates critical straggler effects common in distributed training. We also analyze the convergence of Marshaled Learning under conditions of data heterogeneity. Our theoretical and empirical results demonstrate the effectiveness and efficiency of Marshaled Learning over existing FL algorithms in the constrained memory scenarios. Marshaled Learning outperforms the baseline methods by around 2% to 5% accuracy with much faster convergence. Furthermore, we implement Marshaled Learning in the real-world TEE and show that Marshaled Learning only incurs around 1 ∼ 3× computational overhead compared with non-TEE environments but ensures strong privacy preservation.