Regularisation by Hand: L2, Dropout and Batch Norm in TensorFlow

The Story — Training With a Handicap

In the last lesson a big network memorised 20 noisy points and got worse on new data. Early stopping was one cure. This lesson gives you three more, all written by hand: L2 weight decay, dropout and batch normalisation.

The Test Bench

Same problem as before: 20 noisy sine points to learn from, 100 for validation, and a network that is far too big.

import tensorflow as tf
import numpy as np, math

rng = np.random.default_rng(1)
def make(n):
    x = rng.uniform(-3, 3, (n, 1)).astype("float32")
    return x, (np.sin(x) + rng.normal(0, 0.3, x.shape)).astype("float32")
x_train, y_train = make(20)
x_val, y_val = make(100)

class Dense(tf.Module):
    def __init__(self, n_in, n_out, act=None):
        super().__init__()
        self.w = tf.Variable(tf.random.normal([n_in, n_out], stddev=math.sqrt(2 / n_in)))
        self.b = tf.Variable(tf.zeros([n_out]))
        self.act = act
    def __call__(self, x):
        z = x @ self.w + self.b
        return self.act(z) if self.act else z

class Net(tf.Module):
    def __init__(self, drop=0.0):
        super().__init__()
        tf.random.set_seed(0)
        self.l1, self.l2, self.l3 = Dense(1, 256, tf.nn.relu), Dense(256, 256, tf.nn.relu), Dense(256, 1)
        self.drop = drop
    def __call__(self, x, training=False):
        h = self.l1(x)
        if training and self.drop: h = tf.nn.dropout(h, rate=self.drop)
        h = self.l2(h)
        if training and self.drop: h = tf.nn.dropout(h, rate=self.drop)
        return self.l3(h)
    def weight_norm(self):
        return float(tf.sqrt(sum(tf.reduce_sum(tf.square(l.w)) for l in (self.l1, self.l2, self.l3))))

def train(l2=0.0, drop=0.0, epochs=1500):
    net, opt = Net(drop), tf.keras.optimizers.Adam(0.003)
    @tf.function
    def step():
        with tf.GradientTape() as tape:
            loss = tf.reduce_mean((net(x_train, training=True) - y_train) ** 2)
            if l2:
                loss += l2 * tf.add_n([tf.reduce_sum(tf.square(l.w)) for l in (net.l1, net.l2, net.l3)])
        grads = tape.gradient(loss, net.trainable_variables)
        opt.apply_gradients(zip(grads, net.trainable_variables))
    for _ in range(epochs):
        step()
    mse = lambda x, y: float(tf.reduce_mean((net(x) - y) ** 2))
    return mse(x_train, y_train), mse(x_val, y_val), net.weight_norm()

print("bench ready")

L2 Regularisation (Weight Decay)

We only penalise the weights (w), not the biases. Big weights let the network draw wild, twisting curves through every noisy point; small weights keep the curve smooth.

Dropout with tf.nn.dropout

During training, tf.nn.dropout(x, rate) sets a random fraction rate of values to zero, and multiplies the rest by 1 / (1 − rate) so the average stays the same. At inference time, you skip it.

tf.random.set_seed(3)
x = tf.ones([10])
print("rate 0.5 :", tf.nn.dropout(x, rate=0.5).numpy())      # kept values become 2.0

big = tf.ones([100_000])
for r in (0.2, 0.5, 0.8):
    y = tf.nn.dropout(big, rate=r)
    print(f"rate {r}: kept {float(tf.reduce_mean(tf.cast(y > 0, tf.float32))):.3f}, mean stays {float(tf.reduce_mean(y)):.3f}")

Batch Normalisation by Hand

Batch norm makes every feature of a layer's output have mean 0 and variance 1 within the batch, then lets the network rescale it with two learned vectors, γ (gamma) and β (beta).

class BatchNorm(tf.Module):
    def __init__(self, n, momentum=0.99, eps=1e-3):
        super().__init__()
        self.gamma = tf.Variable(tf.ones([n]))
        self.beta = tf.Variable(tf.zeros([n]))
        self.moving_mean = tf.Variable(tf.zeros([n]), trainable=False)
        self.moving_var = tf.Variable(tf.ones([n]), trainable=False)
        self.momentum, self.eps = momentum, eps

    def __call__(self, x, training=False):
        if training:
            mean, var = tf.nn.moments(x, axes=[0])
            m = self.momentum
            self.moving_mean.assign(m * self.moving_mean + (1 - m) * mean)
            self.moving_var.assign(m * self.moving_var + (1 - m) * var)
        else:
            mean, var = self.moving_mean, self.moving_var
        return self.gamma * (x - mean) / tf.sqrt(var + self.eps) + self.beta

tf.random.set_seed(5)
ours = BatchNorm(4)
keras_bn = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=1e-3)
for _ in range(5):
    batch = tf.random.normal([32, 4], mean=10.0, stddev=3.0)
    a = ours(batch, training=True)
    b = keras_bn(batch, training=True)
print("training output : mean", abs(a.numpy().mean(0)).round(3), " std", a.numpy().std(0).round(3))
print("max diff vs Keras (training) :", float(tf.reduce_max(tf.abs(a - b))))
test = tf.random.normal([8, 4], mean=10.0, stddev=3.0)
print("max diff vs Keras (inference):", float(tf.reduce_max(tf.abs(ours(test) - keras_bn(test, training=False)))))
print("trainable:", len(ours.trainable_variables), " non-trainable:", len(ours.variables) - len(ours.trainable_variables))

The Results

Now we train the over-sized network three ways on the same data and compare.

for name, kw in [("no regularisation", {}), ("L2, lambda=0.01  ", {"l2": 1e-2}), ("dropout 0.2      ", {"drop": 0.2})]:
    tr, va, norm = train(**kw)
    print(f"{name}: train {tr:.4f}  val {va:.4f}  weight norm {norm:6.2f}")

Golden Rules