StitchFL: Heterogeneous Submodel Training in Federated Learning for Membership Inference Attacks Mitigation

Federated Learning (FL) enables organizations to collaboratively train models without sharing raw data, yet exchanged models can remain vulnerable to membership inference attacks (MIAs). This vulnerability often increases with model size, creating a privacy–utility trade-off, since bigger models tend to perform better. Submodel training in FL offers a natural way to limit model exposure by allowing each client to train only a fraction of the global architecture. However, many existing approaches impose structural constraints such as nested submodels, shared bases, full-model clients, coverage of every layer, or auxiliary data at the server. We present StitchFL, a heterogeneous submodel training method that allows clients to train arbitrary subsets of neurons across layers (including skipping layers entirely) and reconstructs the global model without requiring private or auxiliary data at the aggregator. StitchFL combines count-normalised averaging of overlapping slices, pseudo-inverse stitching of skipped layers, and anchor-based output recalibration using synthetic inputs. We instantiate the method for dense, convolutional, and recurrent architectures. Using a dense neural network, StitchFL reaches 0.89 accuracy compared to 0.83 achieved by the second best submodel training baseline. It also outperforms competing submodel methods on the convolutional case and forecasting task. Across six MIAs, restricting clients to 40% of each hidden layer reduces the strongest MIA effect by nearly 68.3%. These results show that flexible submodel allocation can substantially reduce client-side model exposure while preserving most of the utility of full-model federated training.