Federated Distributionally Robust Neural Combinatorial Optimization with Convergence Guarantees

Neural combinatorial optimization (NCO) solvers generalize poorly under distribution shift, yet robustness across heterogeneous problem distributions is essential for real-world deployment where no single client should be left behind. Group distributionally robust optimization (Group DRO) offers a principled remedy, but requires pooling data that often cannot be shared due to confidentiality or regulatory constraints. Federated learning removes data-sharing requirements; however, existing federated DRO methods cannot accommodate the policy-gradient estimators used in modern NCO. Specifically, prior analyses rely on Lipschitz losses with bounded variance. In contrast, the REINFORCE policy-gradient estimator multiplies stochastic costs by score functions, causing variance to scale with cost magnitude and violating these standard assumptions. We close this gap by replacing Lipschitz-loss conditions with a bounded-cost assumption. This allows us to control cost-dependent variance and establish the first federated Group DRO framework with $O(1/\sqrt{T})$ convergence guarantees. Our analysis applies to any REINFORCE-trained model under bounded costs, including NCO solvers as a special case. Experiments on capacitated vehicle routing—one of many applicable domains including scheduling and packing—demonstrate the necessity of federated training for minority-distribution clients. We show that while na\"ive Group DRO suffers from catastrophic weight collapse, our KL-regularized formulation avoids this collapse. Furthermore, it provides a formal worst-group robustness certificate with no empirical performance penalty compared to standard federated averaging.