Regularization of over-excited activation values in Tensorflow 2.0 Keras

For a specific problem I want to regularize over-excited activations of a layer in my Keras model. Therefore I created a custom-loss that is based on the amount of activation values that are higher than a specific threshold.

The loss shows up but does not get reduced by the optimizer. Only the classification loss is minimized.

Using the latest version of tensorflow 2.0 with Keras to build the model. My current attempt adds a second output with a custom loss dependent on a specific layer. I already tried to provide the loss as a second input and add it to the dense loss, which gave me the same results, custom loss was not minimized.

I suppose the gradient of the second loss cannot be calculated, but I have a working example in PyTorch where this seems to work.

import tensorflow as tf
import numpy as np

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Conv2D, Flatten, MaxPooling2D, SpatialDropout2D, Input
from tensorflow.keras.optimizers import SGD

import tensorflow.keras.backend as K
from tensorflow.python.keras.layers import Layer

import get_prepare_dataset

# config
REGULARIZATION_HYPERP = 0.01
(train_images, train_labels), (test_images, test_labels) = get_prepare_dataset.mnist10()


class Taboo(Layer):
    def compute_output_signature(self, input_signature):
        pass

    def __init__(self, thresholds, **kwargs):
        super(Taboo, self).__init__(**kwargs)
        self.thresholds = thresholds

    def call(self, x, mask=None):
        return x


def taboo_loss(y_true, y_pred):
    above_threshold = tf.math.greater(y_pred, y_true[0])
    return REGULARIZATION_HYPERP * K.mean(tf.reduce_sum(tf.cast(above_threshold, tf.float32)))


if __name__ == "__main__":
    y_input = Input(shape=(28, 28, 1))
    layer1 = Conv2D(filters=10, kernel_size=(5, 5), strides=(1, 1), activation='relu', padding="same")(y_input)
    layer2 = MaxPooling2D(pool_size=(2, 2), strides=(1, 1), padding='valid')(layer1)
    layer3 = Conv2D(filters=20, kernel_size=(5, 5), strides=(1, 1), activation='relu', padding="valid")(layer2)
    layer4 = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid')(layer3)
    layer5 = SpatialDropout2D(0.5)(layer4)
    layer6 = Flatten()(layer5)
    output_layer = Dense(units=10, activation='softmax')(layer6)

    thresholds = [0.5]

    cr = Taboo(thresholds)(layer3)
    model = Model(inputs=y_input, outputs=[output_layer, cr])
    model.compile(optimizer=SGD(lr=0.1), loss=[K.binary_crossentropy, taboo_loss])

    model.fit(train_images, [train_labels, np.asarray([thresholds] * len(train_images))],
              validation_data=[test_images, [test_labels, np.asarray([thresholds] * len(test_images))]], epochs=40)

Classification loss is reduced, whereas the second loss get's bigger. I expect the second loss to be minimized.