AffineLens: Capturing the Continuous Piecewise Affine Functions of Neural Networks

TL;DR

AffineLens offers a practical method to analyze and visualize the geometric structure of piecewise affine neural networks, enhancing interpretability and understanding of their expressivity.

Contribution

It introduces AffineLens, a unified framework for computing and visualizing the hyperplane arrangements and regions of PANNs, supporting modern neural network components.

Findings

Enables enumeration of affine sub-regions within bounded input domains.

Supports visualization of decision boundaries and region partitioning.

Facilitates empirical analysis of how architecture affects network expressivity.

Abstract

Piecewise affine neural networks (PANNs) provide a principled geometric perspective on neural network expressivity by characterizing the input--output map as a continuous piecewise affine (CPA) function whose complexity is governed by the number, arrangement, and shapes of its affine regions. However, existing interpretability and expressivity analyses often rely on indirect proxies (e.g., activation statistics or theoretical upper bounds) and rarely offer practical, accurate tools for enumerating and visualizing the induced region partition under realistic architectures and bounded input domains. In this work, we present AffineLens, a unified framework for computing the hyperplane arrangements and polyhedral structures underlying PANNs. Given a calibrated (bounded) input polytope, AffineLens identifies the subset of neuron-induced hyperplanes that intersect the domain, enumerates the resulting affine sub-regions in a layer-wise manner, and returns provably non-empty maximal CPA regions together with interior representatives. The framework further provides visualizations of region partitioning and decision boundaries, enabling qualitative inspection alongside quantitative region counts. By exploiting the affine restriction property of CPA networks under fixed activation patterns, AffineLens supports a broad class of modern components, including batch normalization, pooling, residual connections, multilayer perceptrons, and convolutional layers. Finally, we use AffineLens to perform a systematic empirical study of architectural expressivity, comparing networks through region complexity metrics and revealing how design choices influence the geometry of learned functions.