Skip to content

← Flax step 94 of 100

Medium Primitives

Param Surgery - Zero Last Layer

Jump to the task

A Dense layer computes x @ kernel + bias. Setting its kernel to zero makes the output equal to its bias. This exercise applies that edit to the last layer of a Dense stack, while keeping every bias and earlier kernel unchanged.

Build the stack

Define a Linen MLP with num_layers unnamed Dense layers inside a compact __call__. Every layer outputs features values. There are no activations or residual connections between layers in this exercise.

Earlier layers use the default zero bias. The last layer uses bias_init=nn.initializers.ones. Default kernel initialization is unchanged. The automatically assigned names are Dense_0 through Dense_{num_layers - 1}.

For num_layers=1, the first layer is also the last: it receives the ones bias and is the layer whose kernel you replace. For any depth, the first kernel has shape (in_dim, features) and later kernels have shape (features, features).

Edit only the last kernel

After initialization, select variables["params"][last_name], where last_name = f"Dense_{num_layers - 1}". Build a new mapping that replaces only that layer's kernel with jnp.zeros_like(old_kernel).

Keep the layer's bias, all earlier layer entries, and the original variables mapping intact. Copy the dictionary branches you change; the unchanged JAX arrays can be shared. Pass the full new variables mapping to model.apply.

Zeroing a kernel changes the model's parameters. It is different from freezing the layer or dropping a residual branch during training. This stack has no residual connection, so it does not become an identity function.

Why the output is ones

The final output is hidden @ 0 + ones, so every row is an all-ones vector of length features. The bias initializer makes preserving the bias visible: zeroing it too would return zeros instead.

The output alone does not establish that you edited the right kernel. For example, zeroing an earlier layer in this activation-free stack can also produce the same final output. When debugging, inspect the parameter trees: only the final kernel should change, and its entire array should be zero. Do not use output equality as proof that all other parameters were preserved.

Problem

Implement surgery_zero_layer(seed, x, num_layers, features) using the stack above. Cast the integer-valued scalars to int, initialize with jax.random.PRNGKey(seed) and x, replace the last kernel, then apply the model with the modified variables. Do not run an optimizer or update gradients.

x is a finite float32 array of shape (N, in_dim) with positive dimensions. num_layers >= 1 and features >= 1 are integer-valued scalars; seed is nonnegative and integer-valued. These scalars may arrive as floats. Return a JAX array of shape (N, features) and leave x unchanged. Carry out the requested tree edit and model application even though the expected output values are known in advance.