Activity Regularizer Not Normalized by Batch Size

[fwesel]

The documentation in keras/src/regularizers/regularizers.py lines 29-53 states:

"The value returned by the activity_regularizer is divided by the input batch size"

This is not true. The example in the docs claims ops.sum(layer.losses) equals 5.25, but it actually returns 25.25.

Minimal reproduction following the example in the docs in keras/src/regularizers/regularizers.py:

import os
os.environ["KERAS_BACKEND"] = "jax"
from keras.layers import Dense
from keras.regularizers import L1, L2
from keras import ops

layer = Dense(5, input_dim=5, kernel_initializer='ones',
              kernel_regularizer=L1(0.01), activity_regularizer=L2(0.01))
tensor = ops.ones(shape=(5, 5)) * 2.0
out = layer(tensor)

print(f"Total: {float(ops.sum(layer.losses))}")
# Expected: 5.25 (with normalization)
# Actual: 25.25 (without normalization)

The activity regularizer returns 25.0 instead of 5.0 (not divided by batch_size=5). This means activity regularization strength scales with batch size, making hyperparameter tuning inconsistent across different batch sizes.

Environment: Keras 3.x, JAX backend

[Abhinavexists]

Hi @fwesel , i was able to reproduce this error and it was a calculation miss in the layer.py where division was not happening via batch size , fix was relatively small.

[fwesel]

Hi @Abhinavexists, great work with the PR! I have one doubt: I think batch scaling should affect all types of regularization, not just activity regularization. From a maximum likelihood point of view, when training any deep learning model we maximize the likelihood. Activity regularization is then equivalent to placing a prior on the outputs of a layer, and kernel regularization is equivalent to placing a prior on the weights of a layer. These priors multiply the likelihood and are jointly optimized (so we maximize the posterior). In practice we maximize the log-likehood, or log-posterior if there are any priors, and when doing so, we normalize the log-likehood or log-posterior by the batch size. This means that all terms in the loss (which is the log-posterior) should be normalized by batch size, not just the likelihood and activity regularizer, but also kernel or any other type of layer regularization.

[Abhinavexists]

Hi @Abhinavexists, great work with the PR! I have one doubt: I think batch scaling should affect all types of regularization, not just activity regularization. From a maximum likelihood point of view, when training any deep learning model we maximize the likelihood. Activity regularization is then equivalent to placing a prior on the outputs of a layer, and kernel regularization is equivalent to placing a prior on the weights of a layer. These priors multiply the likelihood and are jointly optimized (so we maximize the posterior). In practice we maximize the log-likehood, or log-posterior if there are any priors, and when doing so, we normalize the log-likehood or log-posterior by the batch size. This means that all terms in the loss (which is the log-posterior) should be normalized by batch size, not just the likelihood and activity regularizer, but also kernel or any other type of layer regularization.

@fwesel its a great point indeed and it honeslty needed me to investigate a bit so what i understood is that intutively the thought is correct on maximising the log-likelihhod but Activity regularization is a per-sample prior placed on layer outputs and it grows with batch size so it needs to be normalised but in case of Kernel Regularization is based on weights as being model-only so it dosen't depend on data, that keeps it constant wrt to batch size.

if i miss something, do tell

[fwesel]

Hi @Abhinavexists, you are correct!