Optimized Piecewise Affine Abstractions of Neural Networks with Learnable Activation Functions

Abstract

We present a generalized framework for the range verification of neural networks featuring non-linear activation functions. Our approach first constructs an ``optimized piecewise affine abstraction" of the network that replaces each non-linear activation function by a piecewise affine (PWA) function plus a bounded error. Such PWA functions are readily amenable to existing neural network verification techniques using specializations of linear arithmetic SMT solvers and mixed-integer optimization approaches. However, there are infinitely many ways to abstract each node, with a natural tradeoff between the number of pieces used, the global error bound, and the complexity of the resulting verification problem. We propose a dynamic programming (DP) algorithm to systematically compute the optimized PWA abstraction for general activation functions, guaranteeing tighter output bounds. The algorithm combines a local DP approximation at each node with a global error bound, yielding a variant of the knapsack problem for deciding how to allocate a fixed budget on the total number of pieces across units so as to minimize the worst-case error bound between the network and its approximation. Although the knapsack problem is itself NP-hard, we can use pseudo-polynomial DP algorithms as well as approximation schemes to solve it efficiently. Crucially, our approach is broadly applicable to diverse networks consisting of non-linear activations, including standard Multi-Layer Perceptrons (MLPs) and recently proposed architectures such as Kolmogorov-Arnold Networks (KANs). Over a series of KAN benchmarks spanning 20 to 22,000 parameters, our approach yields output bounds that are consistently of smaller width than uniform PWA allocation. The overall time taken is roughly comparable while the overhead for computing the optimized abstraction is subsumed by the time taken to compute output bounds.