MSP-GNN: Learning First-Stage Decisions in Multistage Stochastic Programming via Subtree Permutation Invariance

Scenario-based multistage stochastic programming (MSP) provides a principled framework for sequential decision-making under uncertainty, but the exponential growth of scenario trees with the planning horizon poses a major computational barrier to real-time optimization. Learning-based solution predictors offer a promising alternative. However, fully connected architectures that flatten scenario trees into fixed-length vectors overlook their hierarchical symmetries and require parameter counts that grow with tree size. We focus on first-stage, or here-and-now, decisions, which are operationally critical because they must be committed before future uncertainty unfolds while accounting for all subsequent recourse. Our central observation is that optimal first-stage decisions are invariant to recursive permutations of child subtrees at any node of the scenario tree. We exploit this symmetry to introduce MSP-GNN, an end-to-end graph neural network that enforces subtree permutation invariance by construction and has a parameter count independent of scenario-tree size. To characterize its expressive power, we develop a Tree-Weisfeiler-Lehman (TWL) framework tailored to MSP. We prove that TWL-indistinguishable scenario trees admit identical optimal first-stage decisions and that MSP-GNN can simulate the TWL test. Building on these results, we establish a universal approximation guarantee for the optimal first-stage decision mapping on compact domains. Experiments on multistage asset allocation and weighted stochastic linear-quadratic regulator problems demonstrate that, across balanced and asymmetric scenario trees, MSP-GNN achieves predictive accuracy comparable to or better than fully connected networks and DeepSets while using substantially fewer parameters. These findings establish subtree permutation invariance as an effective inductive bias for learning first-stage decisions in MSP.