Partition of Unity Neural Networks for Interpretable Classification with Explicit Class Regions

TL;DR

This paper introduces Partition of Unity Neural Networks (PUNN), an interpretable classification architecture that explicitly models class regions with learned partitions, matching the accuracy of traditional models while enhancing interpretability.

Contribution

PUNN provides a novel neural network architecture that directly models class probabilities through learned partitions, improving interpretability without sacrificing accuracy.

Findings

PUNN achieves comparable accuracy to standard neural networks on benchmarks.

Shape-informed gates reduce parameters significantly when data structure matches priors.

PUNN is dense in the space of continuous probability maps, ensuring expressive power.

Abstract

Despite their empirical success, neural network classifiers remain difficult to interpret. In softmax-based models, class regions are defined implicitly as solutions to systems of inequalities among logits, making them difficult to extract and visualize. We introduce Partition of Unity Neural Networks (PUNN), an architecture in which class probabilities arise directly from a learned partition of unity, without requiring a softmax layer. PUNN constructs $k$ nonnegative functions $h_1, \ldots, h_k$ satisfying $\sum_i h_i(x) = 1$, where each $h_i(x)$ directly represents $P(\text{class } i \mid x)$. Unlike softmax, where class regions are defined implicitly through coupled inequalities among logits, each PUNN partition function $h_i$ directly defines the probability of class $i$ as a standalone function of $x$. We prove that PUNN is dense in the space of continuous probability maps on compact domains. The gate functions $g_i$ that define the partition can use various activation functions (sigmoid, Gaussian, bump) and parameterizations ranging from flexible MLPs to parameter-efficient shape-informed designs (spherical shells, ellipsoids, spherical harmonics). Experiments on synthetic data, UCI benchmarks, and MNIST show that PUNN with MLP-based gates achieves accuracy within 0.3--0.6\% of standard multilayer perceptrons. When geometric priors match the data structure, shape-informed gates achieve comparable accuracy with up to 300$\times$ fewer parameters. These results demonstrate that interpretable-by-design architectures can be competitive with black-box models while providing transparent class probability assignments.