Skip to content

← Flax NNX step 92 of 100

Medium Primitives

NNX Surgery Zero Layer

Jump to the task

Why this matters

You can change a layer's parameters without removing the layer or changing its dimensions. Here you will set the final Linear's kernel to zero while leaving its bias and all earlier layers untouched.

For a Linear receiving a vector h, the output is h @ kernel + bias. With a zero kernel, finite h and zero bias, that output is zero. With a nonzero bias, it is the bias instead. A trained layer's bias need not be zero.

A complete example

import jax
import jax.numpy as jnp
from flax import nnx

class MLP(nnx.Module):
    def __init__(self, in_features, num_layers, features, rngs):
        layers = []
        width = in_features
        for _ in range(num_layers):
            layers.append(nnx.Linear(width, features, rngs=rngs))
            width = features
        self.layers = nnx.List(layers)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = MLP(3, 2, 4, rngs=nnx.Rngs(0))
last = model.layers[-1]
last.kernel[...] = jnp.zeros_like(last.kernel[...])
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
assert jnp.allclose(model(x), jnp.zeros(4))
assert len(jax.tree.leaves(nnx.state(model, nnx.Param))) == 4

# The bias still contributes when the kernel is zero.
last.bias[...] = jnp.array([1.0, -2.0, 0.5, 3.0], dtype=jnp.float32)
assert jnp.allclose(model(x), last.bias[...])

The example changes the bias only to demonstrate its role. In the exercise, leave the default zero bias untouched. Keep the Param wrapper by updating last.kernel[...], rather than replacing last.kernel with a raw array.

nnx.List holds the submodules and supports indexing and iteration. Use it here to make the sequence explicit. A temporary Python list during construction is fine; there is no reason to claim that every Python list silently loses its parameters. Container handling depends on the NNX version and context.

Interpreting a zeroing experiment

Zeroing a final kernel is different from bypassing the layer. Bypassing it would pass its input onward; zeroing it removes its input-dependent output. In this sequential model there is no skip connection around the final layer. There are also no activation functions between the Linear layers.

The zero result follows from this construction, not from a measurement of the layer's importance. To study a trained model, compare an appropriate metric before and after a controlled change across representative inputs. An unchanged result on one example does not establish that a layer is redundant.

The finite-value assumption matters: multiplying infinity or NaN by zero does not reliably produce zero. Use inputs and model sizes whose intermediate activations stay finite, as in the fixtures.

Problem

Implement zero_last_layer_kernel(seed, x, num_layers, features):

  1. Create int(num_layers) Linear layers in order, sharing one nnx.Rngs(int(seed)). The first maps x.shape[-1] to int(features); later layers map int(features) to itself. Keep default initializers and biases, and store the layers in nnx.List.
  2. Set only the final layer's kernel array to jnp.zeros_like of itself. Leave its bias and all earlier parameters unchanged.
  3. Apply the layers in order, without activations, and return the result.

All fixture outputs are zero. The numerical tests cannot distinguish the requested model edit from simply returning zeros, so follow the construction and mutation requirements even though the expected numbers are simple.

Inputs: seed is a nonnegative integer-valued scalar; num_layers and features are positive integer-valued scalars. They arrive as floats and must be cast to int. x is a nonempty float32 vector with finite values and finite intermediate activations. A positive layer count ensures there is a final layer to edit.

Output: vector of shape (int(features),), containing zeros for this newly initialized model and the stated finite-value domain.