Source-linked AI summary
Deep Kronecker neural networks: A general framework for neural networks with adaptive activation functions
Ameya D. Jagtap, Yeonjong Shin, Kenji Kawaguchi, George Em Karniadakis
TL;DR
Neural networks face training challenges such as gradient pathologies, plateaus, and difficulty learning high-frequency components. The paper introduces Kronecker neural networks and Rowdy activations to address these issues through wide, parameter-efficient architectures with adaptive nonlinearities. Theory establishes faster initial loss decay and conditional global convergence, while experiments report improved speed and accuracy across approximation, PDE, and benchmark tasks.
Problem
Neural-network training is affected by vanishing or exploding gradients, plateaus, limited training time, and spectral bias against high-frequency components.
Method
The paper constructs KNNs with Kronecker-product weight matrices and adaptive activation functions, including Rowdy activations with trainable sinusoidal components.
Results
Across function approximation, physics-informed PDE inference, and MNIST, CIFAR, and SVHN benchmarks, experiments report improved training speed and predictive accuracy, alongside theoretical faster initial loss decay and conditional global convergence.
Takeaways & Limitations
KNNs provide a parameter-efficient way to build wider networks while adapting activation functions across feed-forward, recurrent, and convolutional architectures.
Takeaways & Limitations
The convergence results require technical assumptions on initialization, activation regularity, and related conditions.
Abstract
from arXiv · showhide
We propose a new type of neural networks, Kronecker neural networks (KNNs), that form a general framework for neural networks with adaptive activation functions. KNNs employ the Kronecker product, which provides an efficient way of constructing a very wide network while keeping the number of parameters low. Our theoretical analysis reveals that under suitable conditions, KNNs induce a faster decay of the loss than that by the feed-forward networks. This is also empirically verified through a set of computational examples. Furthermore, under certain technical assumptions, we establish global convergence of gradient descent for KNNs. As a specific case, we propose the Rowdy activation function that is designed to get rid of any saturation region by injecting sinusoidal fluctuations, which include trainable parameters. The proposed Rowdy activation function can be employed in any neural network architecture like feed-forward neural networks, Recurrent neural networks, Convolutional neural networks etc. The effectiveness of KNNs with Rowdy activation is demonstrated through various computational experiments including function approximation using feed-forward neural networks, solution inference of partial differential equations using the physics-informed neural networks, and standard deep learning benchmark problems using convolutional and fully-connected neural networks.
1. Introduction
The paper introduces Kronecker neural networks as a general framework for adaptive activation functions, targeting training difficulties including gradient pathologies, plateaus, and spectral bias. It develops theoretical results and evaluates Rowdy activations across approximation, physics-informed, and benchmark tasks.
- Motivation: Training neural networks remains constrained by vanishing or exploding gradients, plateaus, limited training time, and spectral bias.Spectral bias specifically limits learning high-frequency components.
- Contribution: KNNs use the Kronecker product to construct neural networks with adaptive activation functions.The framework includes existing adaptive-activation networks as special instances.
- Results: Theoretical analysis proves that, initially, two-layer KNNs have strictly smaller loss than feed-forward networks.The result concerns gradient-flow dynamics.
- Results: Global convergence of gradient flow is established for two-layer KNNs under technical conditions.The convergence analysis uses assumptions on initialization and activation functions.
- Contribution: Rowdy activations inject bounded sinusoidal harmonic fluctuations to remove saturation regions and support exploration and faster learning.The base activation can be ReLU, tanh, ELU, sine, Swish, Softplus, or another standard choice.
2. Mathematical Setup and Kronecker Neural Networks
KNNs widen feed-forward networks through block-structured Kronecker weights while retaining a compact parameterization and allowing trainable activation behavior. The construction recovers standard and adaptive feed-forward architectures as special cases.
- Feed-forward setup: A feed-forward network composes affine transformations and element-wise nonlinear activations across hidden layers, with an identity output activation.Each layer is parameterized by a weight matrix and bias vector.
- Kronecker construction: Block activation functions apply potentially different element-wise activation functions to the K replicated blocks.Scaling parameters ωl and αl can be trainable or fixed, making the effective activation adaptive.
- Kronecker construction: KNNs replicate hidden-layer neurons K times while increasing the parameter count by only 2K(D −1).The Kronecker product creates the wider block structure without requiring explicit construction of the product.
- Special cases: KNNs recover standard feed-forward networks when K = 1 and suitable scaling parameters equal 1.Other choices recover Parametric ReLU, ELU, SELU, layer-wise locally adaptive, and self-learnable activation networks.
- Implementation: The KNN representation can be implemented without explicitly constructing block weight matrices and bias vectors.Its composition uses the Kronecker product and block-wise activation operations.
- Architecture: A three-hidden-layer KNN differs from a traditional architecture because each neuron’s output can pass to more than one activation function.The schematic depicts the replicated neurons and activation pathways.
3. Gradient Flow Analysis of the Kronecker networks
This section analyzes Kronecker and feed-forward networks under supervised square-loss training, establishing faster early loss decay for Kronecker networks and zero-training-loss convergence under technical assumptions.
- Setup: The analysis considers supervised learning on m training pairs with the square loss and compares Kronecker networks against standard feed-forward networks.The loss trajectories are denoted L_K(t) and L_FF(t), respectively.
- Setup: Gradient descent updates network parameters using the loss gradient, while gradient flow is obtained as the continuous-time limit when the learning rates approach zero.The resulting dynamics describe continuously evolving network parameters.
- Early loss decay: The comparison uses matched initialization, so the Kronecker and feed-forward networks begin with identical loss values.The Rowdy initialization augments the feed-forward parameters with α and ω while preserving the initial network output.
- Early loss decay: At least during the beginning of training, gradient-flow analysis proves that Kronecker networks have strictly smaller loss than feed-forward networks.The theorem holds with probability 1 over initialization under Assumptions 3.1 and 3.2.
- Early loss decay: The early faster loss decay does not imply that Kronecker-network training loss remains smaller throughout training.The stated result is limited to the beginning of training.
- Global convergence: Under the stated assumptions, two-layer Kronecker networks trained by gradient flow can achieve zero training loss when only the first-layer parameters are trained.The convergence analysis fixes ω and c and trains the parameters {w_i, b_i}^N_{i=1}.
- Global convergence: For bounded activation functions with B = 1, a sufficient condition includes K scaling as O(m), while the corresponding Kronecker-network parameter count is O(m) + N(d + 2).This result contrasts with prior global-convergence analyses of significantly over-parameterized two-layer feed-forward networks.
1004. Computational Examples
The computational examples evaluate Rowdy-Net across function approximation, high-frequency learning, PDE inference, and standard deep-learning benchmarks. Rowdy-Net generally converges faster or achieves lower losses than fixed and locally adaptive activations, while scaling and learning-rate choices affect sensitivity and computational cost.
- Rowdy activation functions: Rowdy activation functions combine a standard activation with sinusoidal or cosinusoidal fluctuating terms to inject bounded, non-monotonic effects and remove layer-output saturation regions.The scaling factor n controls these terms and has no universal rule for selection.
- Nonlinear discontinuous function approximation: Rowdy-Net approaches machine-zero loss quickly for a discontinuous nonlinear function using only five training points and a single hidden layer with 40 neurons.The experiment uses a cosine activation function and learning rate 8.0e-6.
- Learning-rate sensitivity: With learning rate 4.0e-3, Rowdy-Net loss decreases to 1.0e-11 before rising, whereas fixed and L-LAAF losses reach 1.0e-6; lower rates or annealing can avoid this sensitivity.The sensitivity is attributed to optimization changes in the trainable Rowdy activation parameters.
- Helmholtz equation: For the Helmholtz equation, Rowdy-Net5 converges faster than fixed and L-LAAF activations, but its computational cost is higher; switching to L-LAAF after 1000 iterations is presented as a remedy.The transfer-learning switch initially increases loss and error, which then decay.
- Standard deep learning benchmark problems: Across fully connected and convolutional benchmark networks, Rowdy-Nets achieve lower training and testing losses or errors than fixed activation functions and L-LAAF.The convolutional results show the same qualitative behavior on the larger CIFAR-100 dataset, and the authors report potential improvement without prior physics information.
5. Summary
The paper introduces KNNs as a low-parameter framework for adaptive activations and reports theoretical and empirical evidence of faster training and broad effectiveness.
- KNNs use the Kronecker product to construct wide networks while keeping the number of trainable parameters low.
- Theoretical analysis shows that KNN training loss is strictly smaller than feed-forward network loss during the beginning of gradient-descent training.
- Under an over-fluctuating case, the authors establish global convergence of KNN gradient descent.
- Rowdy activations inject sinusoidal fluctuations that remove saturation zones and support exploration and faster learning.
- Experiments report substantial improvements in training speed and predictive accuracy across function approximation, PDE inference, and deep-learning benchmarks.
Appendix A. Proof of Theorem 3.3
The appendix proves that a feature matrix associated with the Rowdy network is full rank under stated assumptions, supporting a strictly smaller initial training loss than standard feed-forward networks.
- The Rowdy network output includes adaptive activation terms, while standard feed-forward networks arise by dropping the corresponding α and ω terms.
- The proof defines parameter blocks for output weights, network weights, adaptive coefficients, and frequencies, then writes their gradient-flow dynamics using residuals.
- The proof concludes that the Rowdy network has smaller training loss during an initial interval because its loss initially matches the feed-forward loss and then decreases faster.
- Under the non-degeneracy and activation assumptions, the feature matrix Ψ is full rank with probability 1 over initialization.
- The full-rank conclusion follows by showing that a linear combination of the relevant feature rows can vanish only when all coefficients δ1 through δK are zero.
Appendix B. Proof of Theorem 3.4
The appendix analyzes conditions supporting global convergence by controlling the parameter-dependent feature map through Lipschitz bounds and initialization constraints.
- The proof uses the result of Theorem 3.3 as the starting point for the convergence argument.
- The proof is completed under the bounds maxj∥x̃j∥=1 and ∥α∥∞≤1.
- The argument assumes initialization satisfying K∥α∥∞≤1 and uses lemmas governing the relevant parameter norms.
- The feature map Φc is Lipschitz in V, with constant K∥α∥∞B∥c∥·∥X∥.
- Probability bounds in the proof take the form 1−e^−mδ2 for the controlled quantities.