Introduction

Custom weight initialization can change training speed, stability, and even whether a model trains at all. In TensorFlow, the old tf.layers.dense API and the modern Keras Dense API both let you control initialization, but current code should usually prefer tf.keras.layers.Dense because the legacy tf.layers path is deprecated.

Why initialization matters

Dense layers start with random weights. If those values are too small, signals can shrink as they pass through the network. If they are too large, activations and gradients can explode.

That is why common defaults exist:

• Glorot or Xavier for tanh-like or balanced activations
• He initialization for ReLU-family activations

Custom initialization is useful when:

• you need reproducible controlled experiments
• you are implementing a paper with a specific initializer
• you want to start from a domain-specific prior
• you are debugging unstable training

Legacy tf.layers.dense pattern

In TensorFlow 1 style code, custom initialization is typically passed through kernel_initializer:

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5x = tf.compat.v1.placeholder(tf.float32, shape=[None, 4])
6
7dense = tf.compat.v1.layers.dense(
8    inputs=x,
9    units=3,
10    activation=tf.nn.relu,
11    kernel_initializer=tf.compat.v1.keras.initializers.VarianceScaling(
12        scale=2.0, mode="fan_in", distribution="truncated_normal"
13    ),
14    bias_initializer=tf.zeros_initializer(),
15)
16
17with tf.compat.v1.Session() as sess:
18    sess.run(tf.compat.v1.global_variables_initializer())

That works for older graph-based code, but most new projects should use Keras layers directly.

Modern Keras equivalent

The same idea in TensorFlow 2 is cleaner:

1import tensorflow as tf
2
3layer = tf.keras.layers.Dense(
4    units=3,
5    activation="relu",
6    kernel_initializer=tf.keras.initializers.HeNormal(),
7    bias_initializer="zeros",
8)
9
10x = tf.ones((2, 4))
11y = layer(x)
12
13print(y.shape)
14print(layer.kernel.shape)

This is the preferred path because it integrates naturally with eager execution, Model, and GradientTape.

Writing a truly custom initializer

If built-in initializers are not enough, create a callable or subclass tf.keras.initializers.Initializer.

1import tensorflow as tf
2
3class SmallUniform(tf.keras.initializers.Initializer):
4    def __init__(self, low=-0.05, high=0.05):
5        self.low = low
6        self.high = high
7
8    def __call__(self, shape, dtype=None):
9        return tf.random.uniform(shape, minval=self.low, maxval=self.high, dtype=dtype or tf.float32)
10
11    def get_config(self):
12        return {"low": self.low, "high": self.high}
13
14layer = tf.keras.layers.Dense(
15    2,
16    kernel_initializer=SmallUniform(-0.1, 0.1),
17    bias_initializer="zeros",
18)
19
20out = layer(tf.ones((1, 5)))
21print(out.numpy())

Implementing get_config is important when the layer or model may be serialized.

Initializing from a fixed tensor

Sometimes you want a layer to start from known values instead of random ones.

1import tensorflow as tf
2import numpy as np
3
4initial_kernel = np.array([
5    [1.0, 0.0],
6    [0.0, 1.0],
7    [0.5, 0.5],
8], dtype=np.float32)
9
10layer = tf.keras.layers.Dense(
11    2,
12    use_bias=False,
13    kernel_initializer=tf.keras.initializers.Constant(initial_kernel),
14)
15
16x = tf.constant([[2.0, 3.0, 4.0]])
17print(layer(x).numpy())

This can be useful in transfer-style setups, structured models, or controlled debugging experiments.

Match initializer to activation and architecture

A custom initializer is not automatically better than the defaults. In many dense networks:

• ReLU often pairs well with He initialization
• tanh often pairs well with Glorot initialization
• linear output layers often follow the hidden-layer choice rather than requiring a special rule

If training is unstable, change one variable at a time. An initializer change mixed with optimizer and learning-rate changes makes debugging harder.

Common Pitfalls

The most common mistake is applying a custom initializer without checking whether the built-in initializer is already appropriate for the activation function. Another frequent issue is using a fixed constant initializer accidentally on all neurons, which can create symmetry problems and prevent useful learning. Developers also sometimes rely on legacy tf.layers.dense examples in modern TensorFlow code, then run into deprecation issues or serialization quirks. Forgetting get_config in a custom initializer is another practical problem because saved models may not reload cleanly. Finally, people often blame initialization when the real issue is an excessive learning rate or incorrect data scaling.

Summary

• Both legacy tf.layers.dense and modern Keras dense layers allow custom initialization.
• New TensorFlow code should usually use tf.keras.layers.Dense.
• Use kernel_initializer and bias_initializer to control how parameters start.
• Subclass Initializer when built-in options are not enough.
• Match the initializer to the activation and architecture instead of choosing arbitrarily.
• Test initializer changes in isolation so their effect is measurable.