Region Seeding via Pre-Activation Regularization: A Geometric View of Piecewise Affine Neural Networks

TL;DR

This paper introduces a regularizer that encourages neural networks to partition the input space early during training, enhancing their expressive capacity and improving performance on various datasets.

Contribution

It provides a geometric theory linking neuron switching surfaces to local partitioning, and proposes a practical regularizer for seeding data-relevant regions early in training.

Findings

Regularizer increases the number of realized affine regions.

Improves early-stage accuracy on toy datasets.

Achieves comparable or slightly better final accuracy on ImageNet-1k.

Abstract

Deep networks with continuous piecewise affine activations induce polyhedral partitions of the input space, making the number of realized affine regions a natural measure of expressive capacity and a key determinant of how well the model can approximate nonlinear target functions. In practice, standard training realizes far fewer region refinements in data-visited neighborhoods than the architecture could in principle support, while existing region-count theory is primarily architectural and offers little guidance on how optimization shapes the realized partition near the data. Our theory provides a sufficient condition under which bringing neuron switching surfaces sufficiently close to data points ensures their intersection with local neighborhoods, which in turn implies a strict increase in the local affine-region count, yielding a principled training-time handle for seeding data-relevant partitions early in optimization. Guided by these results, we propose a plug-and-play region-seeding regularizer that encourages early partitioning while allowing task-driven refinement to dominate later in training. Experiments show that the regularizer increases the number of realized affine regions via exact enumeration and improves overall performance on toy datasets, while also improving early-stage accuracy and achieving comparable (or slightly improved) final accuracy on ImageNet-1k for classical models.