Lifted Geometry-Aligned GNNs for Min-$k$-Partition

Min-$k$-Partition is a fundamental NP-hard combinatorial optimization problem on graphs. Although learning methods have shown promise, generic architectures often fail to exploit the underlying problem geometry. This limitation can impair generalization and require costly adaptation for each new instance during inference. We introduce a KL family with node-specific weights for the probability simplex relaxation, accounting for node heterogeneity in the objective while respecting the simplex constraints. We select the node weights by minimizing a tractable spectral surrogate for relative smoothness, yielding the Geometry-Aligned KL (GA-KL) divergence. Mirror descent (MD) induced by GA-KL then provides an analytical solver aligned with the relaxed Min-$k$-Partition problem. We develop Lifted Geometry-Aligned Graph Neural Networks (LGA-GNNs) by lifting the GA-KL MD updates into latent spaces of higher dimension and adding trainable transformations to enhance representation ability. Recurrent application of the learned operator with shared and fixed parameters enables deeper refinement at inference without parameter updates or growth in model size. Experiments on random regular graphs with up to $10^6$ nodes, Gset instances, and diverse real-world networks demonstrate strong solution quality and practical scalability under matched runtime budgets. Moreover, LGA-GNN generalizes effectively without retraining to graphs that are $100\times$ larger and to heterogeneous real-world structures.