Deep Kronecker neural networks: A general framework for neural networks with adaptive activation functions

In plain words

Each unit in a neural network passes a weighted sum through a fixed bend, the activation function. This paper lets every unit use several bends at once, mixed by weights the network learns, so the shape of the bend is itself trained. The construction uses a Kronecker product. It makes the layer wider while the number of parameters stays close to that of an ordinary network.

The authors prove that, for a small two-layer network, this version lowers the loss faster than an ordinary network started from the same weights, at least at the beginning of training. They also prove that gradient descent reaches zero training loss in that setting, under stated conditions. Their concrete choice, called Rowdy, adds sine and cosine terms with trainable frequencies to a standard activation, so that units do not stay in flat, saturated regions. They test it on curve fitting, on physics-informed networks for a Helmholtz equation, and on standard image benchmarks. They report faster training and better accuracy in their own runs.

TL;DR

Kronecker neural networks (KNNs) are feed-forward nets whose units each apply a weighted sum of K activation terms. Rowdy is one choice: a standard activation plus sinusoids that carry trainable frequencies. The paper proves a faster initial loss decrease and convergence for two-layer nets. Experiments are small and author-reported, and the only released code covers one high-frequency fit. Read the Pitfalls before implementing.

Problem

Standard fixed activations can slow training. Networks trained by gradient descent show spectral bias, so they learn high-frequency components of a target slowly §1. Theory on convergence of gradient descent usually needs heavily over-parameterised networks §1. The authors also observe that a unit can sit in a saturated region of its activation, where little signal flows §4. Earlier adaptive-activation work (layer-wise locally adaptive activation, L-LAAF) changes only the slope of the activation fig. 8.

Method

Kronecker neural network

• A feed-forward net of depth D has weights W^l and biases b^l per layer §2 eq. 3.
• The KNN replaces each layer's matrices by block versions built with all-ones blocks and the Kronecker product: \bm{1}_{K\times K}\otimes W^l and \bm{1}_{K\times 1}\otimes b^l §2.
• Each hidden layer carries two extra K-vectors, \omega^l (a column) and \alpha^l (a row). The scaled block weights are \tilde W^l = (\omega^l\otimes\alpha^l)\otimes W^l and \tilde b^l = \omega^l\otimes b^l for hidden layers, with \tilde W^D = \bm{1}_{1\times K}\otimes W^D and \tilde b^D = b^D at the output §2.
• The input is repeated K times, \bm{1}_{K\times 1}\otimes z, and the composition is u^{\mathcal K} §2 eq. 4.
• The authors show the same network equals an ordinary net with a pointwise adaptive activation. Each layer applies \tilde\phi^l(\mathcal L_l(z)) = \sum_{k=1}^K \alpha_k^l\,\phi_k(\omega_k^l\,\mathcal L_l(z)), where \mathcal L_l is the affine map §2 eq. 2 eq. 4. The Kronecker product need not be computed §2.
• Equation (1) is the general form \phi_{\alpha,\omega}(x)=\sum_k \alpha_k\phi_k(\omega_k x) with fixed \phi_k and trainable or fixed \alpha, \omega §1 eq. 1.
• Width: each hidden layer has K times as many neurons as the feed-forward net. The parameter count differs from the feed-forward net by 2K(D-1) §2.
• Special cases §2: K=1 with \omega=\alpha=1 gives the feed-forward net. K=1 with trainable \omega gives layer-wise locally adaptive activations. K=2 with ReLU and \max\{-x,0\} gives parametric ReLU. K=2 with ReLU and (e^x-1)\mathbb{I}_{x\le 0} gives ELU or, with a rescaled \omega, SELU. \omega=1 and \phi_k(x)=x^{k-1} gives self-learnable activations.

# One KNN hidden layer, in adaptive-activation form  [§2] [eq. 4]
# Input:  z (B, N_prev); W (N, N_prev); b (N,)
#         alpha (K,), omega (K,): one pair per layer  [§2]
#         phis: list of K activation functions  [§2]
# Output: h (B, N); the K-times wider layer is never built  [§2]

def kronecker_layer(z, W, b, alpha, omega, phis):
    # 1. Affine map, shared by all K branches  [§2] [eq. 2]
    pre = z @ W.T + b                          # (B, N)
    # 2. Sum the K activation branches on the same pre-activation
    #    [§2] [eq. 4]
    h = 0
    for k in range(len(phis)):
        h = h + alpha[k] * phis[k](omega[k] * pre)
    return h

Rowdy activation

• The first term \phi_1 is any standard activation, for example ReLU, tanh, ELU, sine, Swish or Softplus §4.
• For 2\le k\le K, the other terms are \phi_k(x)=n\sin((k-1)nx) or n\cos((k-1)nx), with a fixed scale n\ge 1 §4 eq. 11.
• The purpose is to add bounded but strongly non-monotone fluctuations, so that each layer's output avoids saturated regions §4.
• The authors state that there is no rule of thumb for n. In their regression runs, n\ge 1 speeds convergence, but large n makes optimisation sensitive §4.
• Initialisation §4 (the table gives f = fixed, t = trainable): fixed activation uses \alpha_1=1 (f), \alpha_k=0 for k\ge 2 (f), \omega_1=1 (f), \omega_k=0 for k\ge 2 (f). L-LAAF uses \omega_1=1 (t). Rowdy uses \alpha_1=1 (f), \alpha_k=0 for k\ge 2 (t), \omega_1=1 (t), \omega_k=1 for k\ge 2 (t).
• The text says the scale is absorbed so that n\,\omega=1 at initialisation §4. This conflicts with the table's \omega_k=1 for k\ge 2. See Pitfalls.

# Rowdy activation on one layer's pre-activations  [§4] [eq. 11]
# Input:  pre (B, N): pre-activations  [eq. 4]
#         alpha (K,): alpha[0] fixed to 1  [§4]
#         omega (K,): omega[0] trainable (cosine base)  [§4]
#         n >= 1: fixed scale  [§4]
# Output: out (B, N)

def rowdy_activation(pre, alpha, omega, n, K):
    # 1. Base term phi_1, here cosine  [§4.2]
    out = alpha[0] * cos(omega[0] * pre)
    # 2. Sinusoidal terms k = 2..K: phi_k(omega_k pre) with
    #    phi_k(x) = n sin((k-1) n x)  [eq. 11] [eq. 4]
    for k in range(2, K + 1):
        arg = (k - 1) * n * omega[k - 1] * pre
        out = out + alpha[k - 1] * n * sin(arg)
    return out

Gradient-flow analysis of two-layer networks §3

• Two-layer networks §3: u^{\text{FF}}(x)=\sum_{i=1}^N c_i\,\phi_1(w_i^\top x+b_i) and u^{\mathcal K}(x)=\sum_{i=1}^N c_i\sum_{k=1}^K\alpha_k\,\phi_k(\omega_k(w_i^\top x+b_i)).
• Loss, with m training pairs §3 eq. 5:

L(\Theta)=\frac{1}{2}\sum_{i=1}^m\bigl(u(x_i)-y_i\bigr)^2 \qquad [\text{eq. 5}]

• Training is gradient descent, analysed as gradient flow \dot\Theta(t)=-\nabla_\Theta L(\Theta(t)) §3 eq. 7.
• Assumption 1 §3: the c_i are drawn independently from a continuous distribution symmetric about 0. The vectors v_i=[w_i;b_i] are drawn from N(0,I_{d+1}), and \omega=\bm 1_{K\times 1}.
• Assumption 2 §3: each \phi_k is C^1, K\ge m, and for any m distinct scalars z_j the K\times m matrix [\Phi]_{kj}=\phi_k(z_j) has full rank.
• Matching start §3 eq. 8: set \omega(0)=\bm 1_{K\times 1} and \alpha(0)=(1,0,\dots,0). The two networks then coincide, so L^{\mathcal K}(0)=L^{\text{FF}}(0).
• Theorem 3.1 §3: under Assumptions 1 and 2 and the start in eq. 8, with probability 1 over initialisation there is a T>0 with L^{\mathcal K}(t)<L^{\text{FF}}(t) for all t\in(0,T). Proof sketch app. A: the time derivative of the loss is -\|\bm M\,\text{Res}\|^2, where \bm M stacks the gradient blocks \bm C,\bm B,\bm A,\bm\Omega. The \alpha-block gives the extra term \|\bm A\,\text{Res}\|^2. This term is positive because \bm A is nonsingular with probability 1 app. A.
• Theorem 3.2 §3: hold c and \omega fixed and train only \{w_i,b_i\} and \alpha. Inputs are augmented as \tilde x_i=[x_i;1/\sqrt 2], and \|\tilde x_i\|=1 is assumed §3 eq. 9.
• Theorem 3.2 setting §3: c_i=\frac{\|y\|}{Kn\sqrt m}\,\xi_i with \xi_i Bernoulli with p=0.5. Each \phi_k is bounded by B. \alpha satisfies K\|\alpha\|_\infty\le 1 and \sum_k\alpha_k\,\mathbb E_{z\sim N(0,1)}[\phi_k(z)]=0. The width K satisfies eq. 10, with \delta\in(0,1) and \lambda_0 from app. A:

\left(\lambda_0\sqrt K-2\|y\|^2B\right)K\ \ge\ 2(1+\delta)\,\|y\|^2B^2 \qquad [\text{eq. 10}]

• Conclusion §3: with probability at least 1-e^{-m\delta^2/(2\|X\|^2)}, L^{\mathcal K}(t)\le L^{\mathcal K}(0)\,e^{-\lambda_0 t/2} for all t\ge 0.
• Remark §3: for B=1 the text gives the sufficient condition K\ge(1+\sqrt{1+4\lambda_0})\,\|y\|^2/\lambda_0, and K=\mathcal O(m). The parameter count is 2K+N(d+2).
• \lambda_0 is the m-th smallest singular value of the K\times m matrix \Psi built from the random initialisation app. A.

Experiments §4

• Function approximation with a shallow net §4.1: a discontinuous target on [-3,3] with 5 fixed training points, one hidden layer of 40 units, cosine base, learning rate 8.0e-6. Compares fixed, L-LAAF and Rowdy with 3, 6 and 9 terms fig. 3 tab. 1.
• Rowdy variants with other \phi_k for k\ge 2 §4.1: eq. 12 uses tanh and eq. 13 uses ReLU. A third variant uses a list of mixed activations, namely tanh, sigmoid, ELU, ReLU, tanh, tanh, softmax and Swish §4.1 fig. 4.
• High-frequency fits §4.2: \sin(m\pi x) with m\in\{1,100,200\} on [0,2\pi], 100 points, three hidden layers of 50 units, cosine base, first nine Rowdy terms, n=10 fig. 5 tab. 2.
• Higher learning rate §4.2.1: rate 4.0e-3 on \sin(\pi x). Learning rate annealing, from 4e-3 to 1e-4 after 500 iterations, is compared fig. 6.
• Helmholtz equation with PINNs §4.3: u_{xx}+u_{yy}+k^2u=g on [-1,1]^2 with k=1. The exact solution is u=\sin(\pi x)\sin(4\pi y), and g is given in the text §4.3 eq. 14. Rowdy uses 5 terms fig. 7 fig. 8. A variant switches from Rowdy to L-LAAF after 1000 iterations fig. 9. A higher-frequency case uses u=\sin(5\pi x)\sin(10\pi y) fig. 10 fig. 11.
• Fully connected nets on two-moons and two-circles §4.4: 1000 training and 1000 test points from scikit-learn. Three layers of 400 units; mini-batch SGD with batch 64, momentum 0.8, learning rate 0.001, weight decay 1e-4, binary cross-entropy. Rowdy-Net4 and Rowdy-Net8 are compared with n\in\{1,2,4\} fig. 12 fig. 13.
• Convolutional nets §4.4: a five-layer LeNet variant on Semeion, Fashion-MNIST, Kuzushiji-MNIST, CIFAR-10, CIFAR-100 and SVHN fig. 14. A pre-activation ResNet-18 is used for SVHN and CIFAR-100 fig. 15 fig. 16 tab. 3 tab. 4. Rowdy-Net2 and Rowdy-Net4 are used, with n=1.

Hyperparameters

Results

Author-reported. Setups are described under Method and Hyperparameters. The runs are small (few points, few layers, few trials), and no independent replication is cited.

Timing (fixed activation = 1), the paper's normalised cost §4.1 §4.2:

Accuracy on image benchmarks, mean test error in % with standard deviation in brackets §4.4:

Other results, stated in the text or in figures without a table:

• Discontinuous fit: the Rowdy loss approaches machine precision quickly, and more terms and larger n speed training §4.1 fig. 3.
• High-frequency fits: Rowdy-Net9 converges faster than fixed and L-LAAF for m=1, 100 and 200 §4.2 fig. 5.
• High learning rate: fixed and L-LAAF reach a loss near 10^{-6}. Rowdy's loss falls to about 10^{-11} and then rises suddenly. Annealing from 4e-3 to 1e-4 after 500 iterations removes the spike and reduces oscillation §4.2.1 fig. 6.
• Helmholtz, low frequency: Rowdy-5 has lower loss and relative L_2 error than fixed and L-LAAF over 30k iterations. Rowdy's trained activation curves become strongly oscillatory §4.3 fig. 7 fig. 8.
• Helmholtz, switch to L-LAAF: loss and error jump at the switch, then decay §4.3 fig. 9.
• Helmholtz, high frequency: fixed and L-LAAF do not converge within 20k iterations, and Rowdy-5 converges faster. The pointwise error of Rowdy is smaller §4.3 fig. 10 fig. 11.
• Moons and circles: training and test losses of Rowdy nets are lower than those of fixed and L-LAAF §4.4 fig. 12 fig. 13.
• CNN training losses are lower for Rowdy than for fixed and L-LAAF on the listed datasets §4.4 fig. 14 fig. 15 fig. 16.
• Theory: Theorem 3.1 gives a strict initial advantage, but only for an unquantified interval; Theorem 3.2 gives exponential decay with a rate tied to \lambda_0 §3. No experiment checks these rates.

Paper vs code

Searched: the first author's GitHub account (AmeyaJagtap) and the paper's own links. The repository AmeyaJagtap/Rowdy_Activation_Functions is the only match, and the owner is an author. Its LICENSE is MIT, copyright 2021 Ameya D. Jagtap LICENSE. The code is recorded at commit 785d887. Its only script covers the high-frequency one-dimensional example README.md.

• Scope. The repository has no code for the Kronecker block construction, the theorems, the discontinuous fit §4.1, the Helmholtz PINNs §4.3 or the benchmarks §4.4 Rowdy_Activation_Functions.py.
• Target. The code fits only \sin(200\pi x), the m=200 case, on [-\pi,\pi] with N+1=101 points Rowdy_Activation_Functions.py. The paper uses m\in\{1,100,200\} on [0,2\pi] with 100 points §4.2.
• Architecture. Layers are [1]+3\times[50]+[1], and the activation sums branches on a shared pre-activation Rowdy_Activation_Functions.py Rowdy_Activation_Functions.py. This matches the adaptive-activation identity §2 at width 50 per layer. It does not build the K-times wider layer.
• Terms. The code has one cosine base term and nine sinusoids, so ten terms in total Rowdy_Activation_Functions.py. The paper says "first nine terms" for Rowdy-Net9 §4.2. Under the paper's convention that K counts \phi_1, this means eight sinusoids. The paper's wording is ambiguous, so the code and text differ in this respect.
• Frequency initialisation. Frequencies are initialised at 0.1 and multiplied by 10=n in the argument Rowdy_Activation_Functions.py Rowdy_Activation_Functions.py. So n\omega=1 at start, which matches the text §4. The table in §4 gives \omega_k=1 instead. The code follows the text.
• Base term. The cosine base uses a trainable argument scale, initialised at 0.1 and multiplied by 10 Rowdy_Activation_Functions.py Rowdy_Activation_Functions.py. This gives \omega_1=1 at start, which matches §4. The base has no separate amplitude, which matches \alpha_1=1 fixed §4.
• Sinusoid amplitudes. Amplitudes start at 0 and are multiplied by n=10 Rowdy_Activation_Functions.py Rowdy_Activation_Functions.py. This matches §4.
• Loss. The code uses the mean of squared errors Rowdy_Activation_Functions.py. The theory uses one half of the sum of squared errors eq. 5. The paper does not state the loss for §4.2.
• Optimiser and steps. Adam with learning rate 4e-6 Rowdy_Activation_Functions.py, matching §4.2 for the rate. Adam is not named in §4.2. The loop runs 20001 updates Rowdy_Activation_Functions.py. The paper gives no iteration count for §4.2.
• Initialisation of weights. Weights and biases are drawn from N(0, 0.1^2) Rowdy_Activation_Functions.py Rowdy_Activation_Functions.py. The paper does not state weight initialisation for §4.2.
• Learning-rate annealing. Not in the code. It appears only in §4.2.1.
• Framework. TensorFlow 1 graph code, with tf.placeholder and tf.Session Rowdy_Activation_Functions.py. It will not run unchanged under TensorFlow 2.

Pitfalls

• Local theorem. Theorem 3.1 holds only on an interval (0,T) whose length is not given. The proof uses existence of the gradient flow near t=0 app. A. The authors note the advantage need not persist §3.
• Input normalisation is a restriction. Theorem 3.2 augments x with 1/\sqrt2 so that \|\tilde x\|=1. This needs \|x\|=1/\sqrt2. The paper calls it "without loss of generality", but it is a rescaling of the data §3 eq. 9. The appendix uses \tilde x=[x;1] instead app. A app. B.
• Constant in eq. 10. Eq. (10) uses the factor 2 on the \|y\|^2B term. The chain of inequalities in app. B yields \sqrt2 in the corresponding place. The card cannot tell which is intended.
• Sufficient condition for B=1. A quick check with \lambda_0=1, \|y\|^2=1 and K\approx3.24 gives a negative left side of eq. (10), so the text's sufficient condition does not satisfy eq. (10) under this reading. Treat the B=1 remark as unverified §3 eq. 10.
• \lambda_0 is data- and initialisation-dependent. It is the m-th smallest singular value of a random matrix app. A. Eq. (10) therefore cannot be checked before training.
• Notation. The theorem uses n for the width and \sqrt m in c_i, but the appendix uses \sqrt{mn} and N for width §3 app. B.
• Bernoulli values. The text says \xi_i is Bernoulli with p=0.5 and does not say whether it takes values in \{0,1\} or \{-1,1\}. Assumption 1 suggests symmetric signs, but this is not stated §3.
• Frequency initialisation conflict. The table sets \omega_k=1 for k\ge 2 §4, while the text sets n\omega=1 §4. These give different starting frequencies, k-1 versus (k-1)n. The code follows n\omega=1 Rowdy_Activation_Functions.py.
• Scale n. No rule of thumb is given. Large n makes optimisation sensitive, and the authors observed this at high learning rates §4 §4.2.1.
• Cost. Rowdy nets cost 1.26 to 1.75 times the fixed net in the shallow fit tab. 1 and 1.73 times in the high-frequency fit tab. 2. PINN costs rise too §4.3.
• High learning rate. Rowdy can blow up after reaching very low loss. Annealing or a lower rate is the authors' remedy §4.2.1.
• Helmholtz high frequency. The forcing term and the scale n are not stated for this case §4.3.
• Benchmarks. Epochs, schedule and seeds are not stated for the convolutional and residual nets. The text says all other settings follow the fully connected runs, which implies SGD with momentum 0.8, learning rate 0.001, weight decay 1e-4 and batch 64 §4.4. The Semeion split seed is not given §4.4. Benchmark code is absent from the repository.
• Figures. Figures are described only in words here. Figure and table numbers were counted from the LaTeX source §4. Trust the labels if the published PDF numbers differ.
• Equation numbers were counted by hand from the source: eq. 1 to eq. 14 are the numbered displays in the main text, with the two rows of the align environment counted separately.

Lineage

• Builds on: layer-wise locally adaptive activation functions (L-LAAF), by the same group, which are the K=1 case with trainable slope §2 §4.
• Special cases named by the paper: parametric ReLU, ELU, SELU, self-learnable activations and smooth adaptive activations §2.
• Related designs named in §4: noisy activations and probabilistic activations (Gulcehre and Lee). The paper names them but does not compare with them.
• Compared against: fixed activations and L-LAAF, in the authors' own runs §4.
• Superseded by: not known from content already in this wiki.