Input Convex Graph Neural Networks: An Application to Optimal Control and Design Optimization

Despite the success of modeling networked systems via graph neural networks (GNN), applying GNN for the model-based control is pessimistic since the non-convexity of GNN models hinders solving model-based control problems. In this regard, we propose the input convex graph neural networks (ICGNN) whose inputs and outputs are related via convex functions. When ICGNN is used to model the target objective function, the decision-making problem becomes a convex optimization problem due to the convexity of ICGNN and the corresponding solution can be obtained efficiently. We assess the prediction and control performance of ICGNN on several benchmarks and physical heat diffusion problems, respectively. On the physical heat diffusion, we further apply ICGNN to solve a design optimization problem, which seeks to find the optimal heater allocations while considering the optimal operation of the heaters, by using a gradient-based method. We cast the design optimization problem as a bi-level optimization problem. In there, the input convexity of ICGNN allows us to compute the gradient of the lower level problem (i.e., control problem with a given heater allocation) without bias. We confirm that ICGNN significantly outperforms non-input convex GNN to solve the design optimization problem.