Pinet: Optimizing hard-constrained neural networks with orthogonal projection layers

TL;DR

The approach, $\Pi$net, leverages operator splitting for rapid and reliable projections in the forward pass, and the implicit function theorem for backpropagation, to deploy as a feasible-by-design optimization proxy for parametric constrained optimization problems and obtain modest-accuracy solutions faster than traditional solvers when solving a single problem, and significantly faster for a batch of problems.

Abstract

We introduce an output layer for neural networks that ensures satisfaction of convex constraints. Our approach, $Π$net, leverages operator splitting for rapid and reliable projections in the forward pass, and the implicit function theorem for backpropagation. We deploy $Π$net as a feasible-by-design optimization proxy for parametric constrained optimization problems and obtain modest-accuracy solutions faster than traditional solvers when solving a single problem, and significantly faster for a batch of problems. We surpass state-of-the-art learning approaches by orders of magnitude in terms of training time, solution quality, and robustness to hyperparameter tuning, while maintaining similar inference times. Finally, we tackle multi-vehicle motion planning with non-convex trajectory preferences and provide $Π$net as a GPU-ready package implemented in JAX.

Figures (22)

• Figure 1: Illustration of the $\Pi$net architecture. The infeasible output of the backbone network is projected onto the feasible set through an operator splitting scheme. To train the backbone network, we use the implicit function theorem to backpropagate the loss through the projection layer.
• Figure 2: Scatter plots of RS and CV on the small and large non-convex problems on the test set. The red dashed lines show the thresholds to consider a candidate solution optimal.
• Figure 3: Comparison of the learning curves in terms of average RS and CV on the validation set, on the small and large non-convex problems. The solid lines denote the mean and the shaded area the standard deviation across 5 seeds. The learning curves for JAXopt on the large dataset are reported only in \ref{['sec:additional-results']} because of the orders of magnitude longer training times.
• Figure 4: (Top) The $\Pi$net approach to constrained multi-vehicle motion planning with arbitrary differentiable objective functions $\varphi$. (Bottom) From left to right, we show examples of the synthesized trajectories for 3 different objectives: $\varphi_\text{left} = \texttt{effort},\, \varphi_\text{mid} = \varphi_\text{left} + \texttt{preference}, \varphi_\text{right} = \varphi_\text{mid} + \texttt{coverage}$. We refer the reader to \ref{['appendix:multi-robot']} for formal definitions and additional plots.
• Figure 5: Scatter plots of RS and CV on the small and large convex problems on the test set. The red dashed lines show the thresholds to consider a candidate solution optimal.

Theorems & Definitions (3)

• Remark
• Remark
• Remark