Beyond NaN: Resiliency of Optimization Layers in The Face of Infeasibility

TL;DR

This paper investigates the vulnerability of optimization layers in neural networks to infeasibility issues caused by rank deficiency, proposing a defense mechanism to improve robustness in critical applications.

Contribution

It identifies a new failure mode in optimization layers due to input matrix rank deficiency and introduces a condition number control method as a defense.

Findings

The proposed defense prevents undefined outputs in optimization layers.

Adversarial rank deficiency can cause optimization failures.

Edge cases reveal bugs in popular solvers.

Abstract

Prior work has successfully incorporated optimization layers as the last layer in neural networks for various problems, thereby allowing joint learning and planning in one neural network forward pass. In this work, we identify a weakness in such a set-up where inputs to the optimization layer lead to undefined output of the neural network. Such undefined decision outputs can lead to possible catastrophic outcomes in critical real time applications. We show that an adversary can cause such failures by forcing rank deficiency on the matrix fed to the optimization layer which results in the optimization failing to produce a solution. We provide a defense for the failure cases by controlling the condition number of the input matrix. We study the problem in the settings of synthetic data, Jigsaw Sudoku, and in speed planning for autonomous driving, building on top of prior frameworks in end-to-end learning and optimization. We show that our proposed defense effectively prevents the framework from failing with undefined output. Finally, we surface a number of edge cases which lead to serious bugs in popular equation and optimization solvers which can be abused as well.