Keras
TensorFlow
kernel regularization
custom layer
machine learning

How to apply kernel regularization in a custom layer in Keras/TensorFlow?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

When you build a custom Keras layer, regularization does not happen automatically just because a layer has weights. You need to attach a regularizer to the weight created by add_weight, or add a custom loss term yourself, so that the penalty is included in the model loss during training.

Pass a Regularizer into add_weight

The standard pattern is to accept a kernel_regularizer argument in __init__, convert it with keras.regularizers.get, and pass it to the kernel weight when the layer is built.

python
1import tensorflow as tf
2from tensorflow import keras
3
4
5class RegularizedDense(keras.layers.Layer):
6    def __init__(self, units, kernel_regularizer=None, **kwargs):
7        super().__init__(**kwargs)
8        self.units = units
9        self.kernel_regularizer = keras.regularizers.get(kernel_regularizer)
10
11    def build(self, input_shape):
12        input_dim = int(input_shape[-1])
13
14        self.kernel = self.add_weight(
15            name="kernel",
16            shape=(input_dim, self.units),
17            initializer="glorot_uniform",
18            regularizer=self.kernel_regularizer,
19            trainable=True,
20        )
21        self.bias = self.add_weight(
22            name="bias",
23            shape=(self.units,),
24            initializer="zeros",
25            trainable=True,
26        )
27
28    def call(self, inputs):
29        return tf.matmul(inputs, self.kernel) + self.bias
30
31    def get_config(self):
32        config = super().get_config()
33        config.update(
34            {
35                "units": self.units,
36                "kernel_regularizer": keras.regularizers.serialize(
37                    self.kernel_regularizer
38                ),
39            }
40        )
41        return config

Now the layer can be used just like a built-in one:

python
1model = keras.Sequential(
2    [
3        keras.layers.Input(shape=(8,)),
4        RegularizedDense(16, kernel_regularizer=keras.regularizers.l2(1e-4)),
5        keras.layers.ReLU(),
6        keras.layers.Dense(1),
7    ]
8)
9
10model.compile(optimizer="adam", loss="mse")

Keras automatically adds the regularization penalty to model.losses, and the training loop includes it in the total loss.

This is the same mechanism used by built-in layers such as Dense and Conv2D. If you follow the same pattern in your custom layer, the rest of the framework continues to work normally with fit, evaluate, model saving, and custom training loops that sum model.losses.

Verify That the Penalty Is Being Collected

If you want to confirm that the regularizer is active, inspect the layer or model losses after the layer has been built.

python
1import tensorflow as tf
2
3x = tf.random.normal((4, 8))
4_ = model(x)
5
6print(model.losses)
7print(sum(model.losses).numpy())

You should see at least one scalar tensor representing the regularization penalty. That check is helpful when refactoring custom layers because forgetting one keyword argument can silently remove the regularizer.

Add a Custom Penalty When Needed

regularizer= on add_weight is the cleanest solution for kernel regularization, but add_loss is available for more specialized cases. For example, you might regularize only part of a tensor or add a nonstandard penalty.

python
1class CustomPenaltyDense(keras.layers.Layer):
2    def __init__(self, units, rate=1e-4, **kwargs):
3        super().__init__(**kwargs)
4        self.units = units
5        self.rate = rate
6
7    def build(self, input_shape):
8        self.kernel = self.add_weight(
9            name="kernel",
10            shape=(int(input_shape[-1]), self.units),
11            initializer="glorot_uniform",
12            trainable=True,
13        )
14
15    def call(self, inputs):
16        self.add_loss(self.rate * tf.reduce_sum(tf.square(self.kernel)))
17        return tf.matmul(inputs, self.kernel)

This works, but use it only when the built-in regularizer interface is not expressive enough. For ordinary L1 and L2 penalties, attaching the regularizer directly to the weight is simpler and easier to serialize.

You can apply the same pattern to other weights too. If your layer has a bias term, an embedding table, or another trainable matrix, each weight can have its own regularizer. The important point is that kernel regularization is not a separate training option; it is metadata attached to the weight creation step.

Common Pitfalls

  • Creating the kernel with add_weight but forgetting the regularizer= argument. In that case, no penalty is applied.
  • Passing a regularizer object in __init__ but not serializing it in get_config. That makes saved models harder to reload correctly.
  • Assuming regularization appears before the layer is built. model.losses is populated only after the relevant weights exist.
  • Using add_loss inside call for ordinary L2 regularization when regularizer= would be clearer and less error-prone.

Summary

  • In a custom Keras layer, kernel regularization is usually attached through add_weight.
  • Accept kernel_regularizer in __init__ and resolve it with keras.regularizers.get.
  • Keras collects those penalties automatically in model.losses.
  • Use add_loss only for custom penalties that do not fit the built-in regularizer interface.
  • Serialize the regularizer in get_config so the layer remains portable.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.