A Training Framework for Optimal and Stable Training of Polynomial Neural Networks

TL;DR

This work tackles the key barrier to deploying Polynomial Neural Networks (PNNs) in privacy-preserving inference with Homomorphic Encryption (HE) by addressing training instability and gradient explosion. It introduces two synergistic techniques: a Boundary Loss that exponentially penalizes activation inputs outside a predefined stable range and a Selective Gradient Clipping scheme that excludes BatchNorm parameters to preserve normalization dynamics, with the overall objective $L_{\text{total}} = L_{\text{CE}} + \lambda \sum_{i=1}^{N} L_i$ and $L_i = \frac{1}{|X_i|} \sum_{j=1}^{|X_i|} \left( e^{\max(|X_{i,j}| - \alpha B, 0)} - 1 \right)$. The framework enables stable training of PNNs up to high polynomial degrees (e.g., $d \le 22$) across seven diverse datasets and HE-friendly architectures, with accuracies closely approximating ReLU baselines and feasible encrypted inference (demonstrated on a CKKS/OpenFHE setup with $N=2^{15}$ and 128-bit security). Comprehensive ablations validate the complementary roles of Boundary Loss and Selective Gradient Clipping and provide practical hyperparameter guidance for $B$, $\alpha$, and $\lambda$. Overall, the approach advances practical privacy-preserving deep learning by delivering robust, HE-compatible PNNs with strong performance and stability traits.

Abstract

By replacing standard non-linearities with polynomial activations, Polynomial Neural Networks (PNNs) are pivotal for applications such as privacy-preserving inference via Homomorphic Encryption (HE). However, training PNNs effectively presents a significant challenge: low-degree polynomials can limit model expressivity, while higher-degree polynomials, crucial for capturing complex functions, often suffer from numerical instability and gradient explosion. We introduce a robust and versatile training framework featuring two synergistic innovations: 1) a novel Boundary Loss that exponentially penalizes activation inputs outside a predefined stable range, and 2) Selective Gradient Clipping that effectively tames gradient magnitudes while preserving essential Batch Normalization statistics. We demonstrate our framework's broad efficacy by training PNNs within deep architectures composed of HE-compatible layers (e.g., linear layers, average pooling, batch normalization, as used in ResNet variants) across diverse image, audio, and human activity recognition datasets. These models consistently achieve high accuracy with low-degree polynomial activations (such as degree 2) and, critically, exhibit stable training and strong performance with polynomial degrees up to 22, where standard methods typically fail or suffer severe degradation. Furthermore, the performance of these PNNs achieves a remarkable parity, closely approaching that of their original ReLU-based counterparts. Extensive ablation studies validate the contributions of our techniques and guide hyperparameter selection. We confirm the HE-compatibility of the trained models, advancing the practical deployment of accurate, stable, and secure deep learning inference.

Figures (5)

• Figure 1: Example of a neural network layer structure for our use cases. It only consists of layers which can be either dropped or converted into affine/polynomial functions
• Figure 2: Example of fitted polynomial activation functions for polynomial degrees of 2, 4, and 8. The activation function was fitted for the range $[-35, 35]$. The black dotted lines show the desired range of activation (if the value gets out of this bound we apply an exponential loss on the excess output)
• Figure 3: Evaluation of training stability and accuracy via ablation study on CIFAR-10 (top) and CIFAR-100 (bottom). The proposed training methodology demonstrates superior convergence rates and accuracy compared to baseline and partial implementations (boundary loss only, gradient clipping only) across varying polynomial degrees (2, 4, 8).
• Figure 4: Impact of polynomial fitting boundary $B$ on test accuracy for CIFAR-10 (left) and UCI-HAR (right) datasets for polynomial degrees 2, 4, 8, and 12. Each dot represents a training run with varied parameters (e.g., dropout, gradient clipping, $\lambda$, fit samples). Small $B$ leads to instability, while excessively large $B$ degrades performance. Higher degrees generally perform better but require larger $B$ for stability. The dashed purple lines indicates ReLU baseline accuracy.
• Figure 5: Effect of including (purple bars) versus excluding (red bars; our proposed Selective Gradient Clipping) Batch Normalization parameters from gradient clipping for the CIFAR-10 dataset. Results are shown for PNNs with polynomial degrees 2, 4, and 8, fitted with B-values of 13, 20, and 35, respectively.