FedJigsaw: Multi-Agent Collaborative Model Reassembly for Decentralized Heterogeneous Federated Learning

Abstract

Model Heterogeneous Federated Learning (MHFL) addresses client-level resource heterogeneity by allowing each participant to train a personalized model architecture under a shared training objective. A prevalent paradigm, Partial Training (PT), achieves this by allowing each client to train a subnetwork of the global model. However, existing PT methods typically rely on predefined architectural templates or over-parameterized supernets, limiting fine-grained personalization and imposing substantial computational and memory overhead. We propose FedJigsaw, a novel framework that reshapes model personalization as a dynamic and decentralized model assembly problem. Instead of selecting subnetworks from a predefined supernetwork, each client constructs its model by assembling reusable modules learned from neighboring clients. At the client level, we introduce AttenAssemble to enable each participant to adaptively construct a tailored model based on local observations. To support efficient knowledge sharing under communication and privacy constraints, we design SymbioArchitect, a mechanism that allows clients to exchange granular model modules with their topological neighbors. To mitigate training instability introduced by decentralized module exchange, we design CoRe-Tune, an attention-enhanced centralized training with a decentralized execution strategy, which guides local policies to foster implicit collaboration and stabilize training dynamics, without compromising data privacy. Extensive evaluations demonstrate that FedJigsaw outperforms state-of-the-art MHFL baselines by up to 13.8% in relative accuracy while significantly shrinking cross-client performance variance, but also slashes decision-making latency and peak memory footprint compared to existing policy-driven methods.