Efficient and tight neural network verification in JAX

Efficient nonconvex reformulations of neural network convex relaxations [Hinder et al, NeurIPS 2020]: Verification algorithms for neural networks are often derived from a convex relaxation, that replaces the nonlinear relations between network activations by a weaker set of convex (often linear) constraints between inputs and outputs of neurons (Ehlers 2017). This enables incomplete verification via convex optimisation, with tightness governed by the gap between the weaker convex and the original nonlinear constraints. However, off-the-shelf convex optimisation solvers still don't scale efficiently to modern neural networks, and most attempts at developing scalable methods has required using weaker relaxations (Fast-Lin(Wong and Kolter 2017, Weng et al 2018), CROWN (Zhang et al 2018)). In this work, we develop a novel non-convex reformulation of convex relaxations of neural network verification. Despite the nonconvexity, we are able to derive algorithms that are guaranteed to converge quickly to the global optimum.