TensorFlow
weight decay
machine learning
neural networks
deep learning

How to define weight decay for individual layers in 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

If you want different weight-decay strength on different layers in TensorFlow, the cleanest approach is usually to assign regularizers per layer instead of trying to force one global optimizer setting to do everything. In Keras, that often means setting kernel_regularizer differently on each layer, or adding manual parameter-specific penalties in a custom training step.

Per-layer regularization with Keras layers

For many models, layer-specific weight decay is easiest with built-in L2 regularizers.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(
5        128,
6        activation="relu",
7        kernel_regularizer=tf.keras.regularizers.L2(1e-4)
8    ),
9    tf.keras.layers.Dense(
10        64,
11        activation="relu",
12        kernel_regularizer=tf.keras.regularizers.L2(1e-5)
13    ),
14    tf.keras.layers.Dense(
15        10,
16        kernel_regularizer=tf.keras.regularizers.L2(0.0)
17    ),
18])
19
20model.compile(
21    optimizer="adam",
22    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
23    metrics=["accuracy"],
24)

Each layer contributes its own penalty term to model.losses, and Keras automatically adds those terms to the training loss during fit().

This is the most direct answer when you need different decay coefficients per layer.

What gets regularized

In dense and convolution layers, kernel_regularizer applies to the weight matrix or convolution kernel. There are also related options:

  • 'bias_regularizer'
  • 'activity_regularizer'

Most weight-decay use cases focus on kernels only, not biases. That is why kernel_regularizer is usually the relevant setting.

For example:

python
1layer = tf.keras.layers.Dense(
2    32,
3    activation="relu",
4    kernel_regularizer=tf.keras.regularizers.L2(1e-4),
5    bias_regularizer=None,
6)

This matches the common practice of decaying weights but leaving biases alone.

Inspect the added regularization losses

Keras stores regularization contributions in model.losses.

python
1import tensorflow as tf
2
3inputs = tf.random.normal([4, 20])
4_ = model(inputs)
5
6print(model.losses)
7print(tf.add_n(model.losses).numpy())

That is useful when you want to verify that only the intended layers are contributing regularization terms.

Manual layer-specific decay in a custom training step

If you need more control than kernel_regularizer provides, you can add penalties manually in a custom training loop. This is useful when:

  • different variables inside the same layer need different decay
  • some variables should be excluded entirely
  • the penalty depends on names or patterns
python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(128, activation="relu"),
5    tf.keras.layers.Dense(64, activation="relu"),
6    tf.keras.layers.Dense(10),
7])
8
9optimizer = tf.keras.optimizers.Adam()
10loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
11
12@tf.function
13def train_step(x, y):
14    with tf.GradientTape() as tape:
15        logits = model(x, training=True)
16        base_loss = loss_fn(y, logits)
17
18        decay = (
19            1e-4 * tf.nn.l2_loss(model.layers[0].kernel) +
20            1e-5 * tf.nn.l2_loss(model.layers[1].kernel)
21        )
22
23        loss = base_loss + decay
24
25    grads = tape.gradient(loss, model.trainable_variables)
26    optimizer.apply_gradients(zip(grads, model.trainable_variables))
27    return loss

This makes the penalty policy explicit and easy to audit.

Weight decay versus optimizer weight decay

Some optimizers, such as AdamW, expose a weight_decay argument. That is convenient, but it is usually global across the variables the optimizer updates. If the goal is per-layer differences, built-in per-layer regularizers or a custom training step are usually clearer.

If you do use an optimizer with decoupled weight decay, pay attention to whether the behavior matches classical L2 regularization for your use case. The names are related, but the mechanics are not always identical.

Common Pitfalls

The biggest mistake is assuming one optimizer-level weight_decay value can express different decay rates for different layers. It usually cannot without extra logic.

Another issue is regularizing biases and normalization parameters accidentally. In many architectures, people decay kernels but exclude biases and parameters such as batch-normalization scale and offset.

Developers also forget that Keras regularizers are added through model.losses. If training code bypasses normal Keras loss handling, those terms may need explicit inclusion.

Finally, mixing manual penalties and layer regularizers without checking the final loss can double-count regularization on some variables.

Summary

  • Use kernel_regularizer on each layer when you want different weight-decay strengths.
  • This is the simplest per-layer solution in TensorFlow Keras.
  • Use a custom training step when you need variable-level control beyond what layer regularizers provide.
  • Be deliberate about which parameters should and should not be decayed.
  • Check model.losses or the final loss computation so the regularization policy is doing exactly what you intend.

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.