Skip to content

← Flax step 3 of 100

Medium Primitives

Custom Parameter Initializer

Jump to the task

An initializer chooses a parameter's starting values. In this exercise you will set the kernel and bias initializers of a Flax Linen Dense layer explicitly, then compute one forward pass. There is no training step.

What the layer computes

For input x of shape (fan_in,) and features output units, Dense computes x @ kernel + bias. The kernel has shape (fan_in, features) and the bias has shape (features,). Here, fan_in means the number of input features.

We will use lecun_normal() for the kernel and zeros for the bias:

nn.Dense(
    features=self.features,
    kernel_init=nn.initializers.lecun_normal(),
    bias_init=nn.initializers.zeros,
)(x)

These are also Dense's default initializers. Specifying them explicitly shows where to change the initialization in another model.

Why scale weights by input width?

A wider layer adds more terms to each output. To see how this affects scale, suppose the input components and weights are independent, zero-mean random variables, with input variance v_x and weight variance v_w. With zero bias, each output has variance:

output variance = fan_in * v_w * v_x

Choosing v_w = 1 / fan_in preserves that variance under these assumptions. For 64 input features, the target weight standard deviation is sqrt(1 / 64) = 0.125. The measured standard deviation of a finite sampled kernel will vary; it is not guaranteed to equal 0.125 exactly.

JAX's lecun_normal() uses a truncated normal distribution, rescaled to target variance 1 / fan_in. It does not sample an unrestricted normal distribution. Flax exposes this initializer through nn.initializers.

Other choices use different target variances. Glorot initialization uses 2 / (fan_in + fan_out), the reciprocal of the arithmetic mean of the two widths. He initialization uses 2 / fan_in, motivated by preserving the mean squared value through ReLU under a symmetric input distribution. These are starting-point heuristics, not guarantees that training will converge. Activations, network structure and optimization also matter.

Factories versus ready-to-use functions

lecun_normal() is a factory call: it returns the initializer function. zeros is already an initializer function, so pass it without parentheses. Both resulting functions accept a random key, a shape and a dtype, and return an array. The zeros initializer ignores the key.

Flax supplies those arguments when it creates the parameters during init. A later apply uses the stored parameters rather than initializing them again. Zero bias is a common choice and adds no initial offset. Nonzero biases can also be useful; they do not inherently create a dead neuron.

Dense distinguishes param_dtype, which controls parameter initialization, from dtype, which controls its computation. Computing in bfloat16 does not require storing parameters in bfloat16. Keep the defaults for this exercise, which uses float32 inputs and parameters.

Problem

Implement custom_init_dense_forward(seed, x, features):

  1. Define CustomInitDense(nn.Module) with a features field and an @nn.compact method containing exactly the unnamed Dense layer above. Do not add an activation or another layer.
  2. Construct it with features=int(features).
  3. Initialize it with jax.random.PRNGKey(int(seed)) and x, without splitting the key first.
  4. Apply it with the full initialized variables mapping and x.

x is a finite float32 JAX array of shape (fan_in,), with fan_in >= 1. features is a positive integer-valued scalar; seed is a nonnegative integer-valued scalar. Either scalar may arrive as a float. Return a JAX array of shape (features,) and leave x unchanged.

Keep the Dense layer unnamed: Flax gives this child the path Dense_0, and parameter paths affect seeded initialization. The initialized kernel and bias are in variables["params"]["Dense_0"].