Write a Dense Layer From Scratch: Weights, Bias and Initialisation

The Story — A Panel of Judges

The Maths and the Shapes

import tensorflow as tf

class Dense(tf.Module):
    def __init__(self, in_features, out_features, name=None):
        super().__init__(name=name)
        with self.name_scope:
            self.w = tf.Variable(tf.random.normal([in_features, out_features], stddev=0.1), name="w")
            self.b = tf.Variable(tf.zeros([out_features]), name="b")

    def __call__(self, x):
        return x @ self.w + self.b

tf.random.set_seed(0)
layer = Dense(4, 3, name="judges")
x = tf.constant([[8.0, 6.0, 9.0, 7.0],       # dish 1: taste, smell, look, texture
                 [3.0, 5.0, 2.0, 4.0]])      # dish 2
print("scores shape:", layer(x).shape)
print(layer(x).numpy().round(3))

Try It — Dense Layer Calculator

Why Initialisation Matters

Problem 1: All Zeros — Every Neuron Becomes a Twin

If all weights start equal, every output neuron computes the same thing and gets the same gradient. They stay twins forever, so 64 neurons act like 1.

w = tf.Variable(tf.zeros([3, 4]))                   # 4 neurons, all zero
x = tf.constant([[1.0, 2.0, 3.0]])
with tf.GradientTape() as tape:
    loss = tf.reduce_sum(tf.nn.sigmoid(x @ w) * [1.0, 2.0, 3.0, 4.0])
grad = tape.gradient(loss, w)
print("gradient columns (one per neuron):")
print(grad.numpy().round(3))
print("every column the same shape of numbers -> the neurons can never become different")

Problem 2: Wrong Scale — Signals Vanish or Explode

Each layer multiplies the signal by its weights. If the weights are a little too small, the signal shrinks layer by layer. A little too big, and it explodes.

def signal_through(depth, stddev_fn, act=tf.nn.relu, width=256):
    tf.random.set_seed(1)
    h = tf.random.normal([512, width])
    for _ in range(depth):
        W = tf.random.normal([width, width], stddev=stddev_fn(width))
        h = act(h @ W)
    return float(tf.math.reduce_std(h))

schemes = {
    "stddev 0.01  ": lambda n: 0.01,
    "stddev 1.0   ": lambda n: 1.0,
    "He sqrt(2/n) ": lambda n: (2.0 / n) ** 0.5,
}
for name, fn in schemes.items():
    print(name, "std after 10 relu layers:", f"{signal_through(10, fn):.3e}")

Try It — Initialisation Explorer

The Final Dense Layer

Now put it all together: He or Glorot initialisation, zero bias, an optional activation, and a helper for the number of parameters. We will use this class for the rest of the module.

import math

class Dense(tf.Module):
    def __init__(self, in_features, out_features, activation=None, init="glorot", name=None):
        super().__init__(name=name)
        if init == "he":
            w0 = tf.random.normal([in_features, out_features], stddev=math.sqrt(2.0 / in_features))
        else:  # glorot uniform
            limit = math.sqrt(6.0 / (in_features + out_features))
            w0 = tf.random.uniform([in_features, out_features], -limit, limit)
        with self.name_scope:
            self.w = tf.Variable(w0, name="w")
            self.b = tf.Variable(tf.zeros([out_features]), name="b")
        self.activation = activation

    def __call__(self, x):
        z = x @ self.w + self.b
        return self.activation(z) if self.activation else z

    @property
    def num_params(self):
        return int(tf.size(self.w) + tf.size(self.b))

tf.random.set_seed(42)
hidden = Dense(784, 128, activation=tf.nn.relu, init="he", name="hidden")
output = Dense(128, 10, name="logits")
batch = tf.random.normal([32, 784])
print("output shape:", output(hidden(batch)).shape)
print("parameters  :", hidden.num_params + output.num_params)
print("W std (He)  :", round(float(tf.math.reduce_std(hidden.w)), 4), "~ sqrt(2/784) =", round(math.sqrt(2 / 784), 4))

Golden Rules