Underperforming SELUs - How to correctly constrain layer weights in TF/Keras?

Question

The promise of SELUs and SNNs

I first read up about the 'power' of SELUs on a machine learning blog post. The promise of a Self-normalizing Neural Network (SNN) sounds too good to be true; its ability to self normalise and self regularize internally promises simpler deep network design and theoretically solves both dying and exploding gradients.

However in my experience it cannot outperform ReLU DNNs, and in the case of deep CNNs it learns nothing at all.

The network architecture

I have a deep CNN for classification of Places 365 images (which are greyscale transformed) involving 12 Conv2D layers and 3 Dense FC layers and an additional Dense FC output layer.

When using ReLUs it can achieve 40% top-1-accuracy and 70% top-5-accuracy on the validation set... but when using SELUs the accuracy (on both training and validation) converges to $\frac{1}{365}$ which suggests the output is meaningless. Additionally when using SELUs in this network the cross-entropy loss also converges to a constant (around 6).

• When using SELUs I used the correct LeCun Normal weight initialisation and zeros bias initialisation.
• I tried training the network with various regularization enabled and disabled, including: alpha dropout, weight decay, and batch normalization.
• I also tried the Adam optimizer with various learning rates and epsilon values, and also the Adadelta optimizer with learning rates varying from $1.0$ (the value corresponding to behaviour of the original Adadelta paper) to $0.001$ (the default in Keras).
• I am also certain that my input pipeline is not broken in any way; I verified that the when transforming from RGB to Lab the network receives a Lightness channel with values in the range $[0,1]$ and data input can be correctly displayed using pyplot.
• I also verified that the behaviour is the same regardless of using float16 or float32, and the behaviour is also the same with or without Softmax at the final layer (and the corresponding from_logits parameter being correctly set for the loss function).

My Hypothesis

According to the paper which introduces SNNs and SELUs, their Theorem 1 states that "...SELU networks are self-normalizing under mild conditions on the weights."

But prior to the theorem and proof they state that "For the weight initialization, we propose ω = 0 and τ = 1 for all units in the higher layer. ... Of course, during learning these assumptions on the weight vector will be violated. However, we can prove the self-normalizing property even for weight vectors that are not normalized, therefore, the self-normalizing property can be kept during learning and weight changes." Which sounds contradictory to me as there is no mention of such conditions needing to be maintained post initialization.

My hypothesis therefore is that these 'mild conditions' are more important than alluded to in the paper, and the lack of constraints on the layer weights (and perhaps biases) is resulting in the underperforming of SELUs in practice.

My question is then how do I impose constraints on layer weights such that ω = 0 and τ = 1 during training? I know that Keras Dense and Conv2D layers both have a kernel_constraint parameter but I'm unsure which class of keras.constraints, if any, can fulfil this property. Is this something which cannot be done using just the constraints parameter and requires some other method?

Edit:

I have implemented a weight centralization constraint as suggested by Redefining The Self-Normalization Property which adjusts the weight matrix for each layer with $\hat{W} = W - mean(W)$. The SELU net is no longer converging to an accuracy of $\frac{1}{365}$ which is an improvement, however I do not know how well it will perform over time. I'm running the training overnight and will report back in ~12 hours with code for the constraint if successful.

Answer 1

After 12 hours using the weight constraint method I suggested in my question edit the model has continued to learn using SELUs. Additionally the validation loss has decreased much more stably than when ReLUs were used, although the convergence is slightly slower.

Here is the constraint code in tensorflow and the corresponding way to use it in layers.

class WeightCentralization( keras.constraints.Constraint ):
    def __call__( self, w ):
        return w - tf.math.reduce_mean( w )

...

Conv2D(
    64,
    3,
    activation = 'selu',
    kernel_initializer = 'lecun_normal',
    kernel_constraint = WeightCentralization()
)

I hope this answers any questions others might have had with issues getting SELU to perform well.