Neural Network Verification with Branch-and-Bound for General Nonlinearities

TL;DR

The paper tackles neural network verification for models with general nonlinearities by introducing GenBaB, a general branch-and-bound framework that extends beyond ReLU activations. It combines linear bound propagation with a novel Bound Propagation with Shortcuts (BBPS) heuristic and a pre-optimization of branching points stored in a lookup table, enabling efficient and tighter subproblem relaxations. GenBaB applies to Sigmoid, Tanh, Sine, GeLU, as well as multi-input nonlinear operations in LSTMs and ViTs, and to nonlinear computation graphs such as ACOPF. Empirically, GenBaB verifiably outperforms baselines across MNIST, CIFAR-10, LSTMs, ViTs, and ML4ACOPF, with ablations confirming the benefits of BBPS and pre-optimized branching points, indicating a practical advance for verification of non-ReLU networks and complex architectures.

Abstract

Branch-and-bound (BaB) is among the most effective techniques for neural network (NN) verification. However, existing works on BaB for NN verification have mostly focused on NNs with piecewise linear activations, especially ReLU networks. In this paper, we develop a general framework, named GenBaB, to conduct BaB on general nonlinearities to verify NNs with general architectures, based on linear bound propagation for NN verification. To decide which neuron to branch, we design a new branching heuristic which leverages linear bounds as shortcuts to efficiently estimate the potential improvement after branching. To decide nontrivial branching points for general nonlinear functions, we propose to pre-optimize branching points, which can be efficiently leveraged during verification with a lookup table. We demonstrate the effectiveness of our GenBaB on verifying a wide range of NNs, including NNs with activation functions such as Sigmoid, Tanh, Sine and GeLU, as well as NNs involving multi-dimensional nonlinear operations such as multiplications in LSTMs and Vision Transformers. Our framework also allows the verification of general nonlinear computation graphs and enables verification applications beyond simple NNs, particularly for AC Optimal Power Flow (ACOPF). GenBaB is part of the latest $α$,$β$-CROWN, the winner of the 4th and the 5th International Verification of Neural Networks Competition (VNN-COMP 2023 and 2024). Code for reproducing the experiments is available at https://github.com/shizhouxing/GenBaB.

Figures (7)

• Figure 1: Illustration of the more complicated nature of branching for general nonlinearities (branching a ReLU activation v.s. branching for nonlinearities in an LSTM). Notations for part of an LSTM follows PyTorch's documentation (https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html). Nodes in orange are being branched. For general nonlinearities, branching points can be non-zero (0.86 in the LSTM example here), a nonlinearity can take multiple input nodes ($f_{t+1}\odot c_t$ here), and a node can also be followed by multiple nonlinearities ($c_t$ is followed by a multiplication and also Tanh, and branching $c_t$ affects both two nonlinearities).
• Figure 2: Illustration of branching the intermediate bounds of a neuron connected to the Sine activation sitzmann2020implicit.
• Figure 3: Total number of verified instances against running time threshold on feedforward networks for CIFAR-10 with various activation functions. "Base BaB" means that in the most basic BaB setting, we use a generalized BaBSR heuristic and always branch in the middle point of intermediate bounds. "Base + BBPS" uses our BBPS heuristic. Our full GenBaB uses both BBPS and pre-optimized branching points.
• Figure 4: Illustration of our new GenBaB framework summarized in \ref{['sec:framework']}.
• Figure 5: Linear relaxation for a Sin activation in an input range $[-1.5,1.5]$ where the function is S-shaped.