How can I make non-trainable variable using nnx.Module?

[SangminLee0828]

Hi,

I am trying to create a normalization layer. This normalization layer has 'mean' and 'variance' inside, so when the values come in, the output values will be normalized value using the stored mean and variance.

https://www.tensorflow.org/api_docs/python/tf/keras/layers/Normalization

class Normalization(nnx.Module):
    """Normalization layer in JAX/Flax (nnx)."""
    def __init__(self, mean: jnp.ndarray, variance: jnp.ndarray, invert: bool =False):
        self.mean = mean
        self.variance = variance
        self.invert = invert  # Whether to denormalize instead of normalize.

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """Applies normalization or denormalization."""
        if self.invert:
            # De-normalize: x * sqrt(variance) + mean
            return x * jnp.sqrt(self.variance) + self.mean
        else:
            # Normalize: (x - mean) / sqrt(variance)
            return (x - self.mean) / jnp.sqrt(self.variance)

    @staticmethod
    def adapt(data: jnp.ndarray, axis: int = 0):
        """
        Computes the mean and variance of the dataset for normalization.
        
        Args:
            data: Input data to compute mean and variance.
            axis: Axis along which to compute mean and variance.
                  Use None for global statistics across all dimensions.
        
        Returns:
            A tuple of (mean, variance).
        """
        mean = jnp.mean(data, axis=axis, keepdims=True)
        variance = jnp.var(data, axis=axis, keepdims=True)
        return mean, variance
    

adapt_data = jnp.array([[0.],
                        [2.],
                        [0.],
                        [2.]], dtype=jnp.float32)

input_mean, input_variance = Normalization.adapt(adapt_data, axis = 0)
print(f'input mean: {input_mean}, input variance: {input_variance}')
input_norm_layer = Normalization(mean=input_mean, variance=input_variance)

input_data = jnp.array([[0., 7., 4.]], dtype=jnp.float32)

input_normalized = input_norm_layer(input_data)

How can I make 'self.mean' and 'self.variance' not trainable?

[DiagRisker]

My quick take is : by assigning those parameter to self, they are found by jax.grad as parameters to differentiate with (altough I am suprised because nnx.param exists for a good reason).

however, if you want something purely static, maybe the usage of nnx.Module or a class overall is not needed ? (if you want to make variance and mean on the fly)

[cgarciae]

The solution is to create a filter for the train trainable Variable and pass it to both Optimizer and nnx.grad so you effectively only train a subset of the weights.

class Classifier(nnx.Module):
  def __init__(self, embed_dim, num_classes, backbone, rngs):
    self.backbone = backbone
    self.head = nnx.Linear(embed_dim, num_classes, rngs=rngs)

  def __call__(self, x):
    x = self.backbone(x)
    x = self.head(x)
    return x

def load_model():
  return nnx.Linear(784, 1024, rngs=nnx.Rngs(0))

backbone = load_model()
classifier = Classifier(1024, 10, backbone, rngs=nnx.Rngs(1))

# filter to select only Params on head path
head_params = nnx.All(nnx.Param, nnx.PathContains('head'))

optimizer = nnx.Optimizer(
  classifier,
  tx=optax.adamw(3e-4),
  wrt=head_params,  # filter head params
)

# simple train step
@nnx.jit
def train_step(model, optimizer, x, y):
  def loss_fn(model):
    logits = model(x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()

  diff_state = nnx.DiffState(0, head_params) # filter head params of the first argument
  grads = nnx.grad(loss_fn, argnums=diff_state)(model)
  optimizer.update(grads)

x = jnp.ones((1, 784))
y = jnp.ones((1,), jnp.int32)
train_step(classifier, optimizer, x, y)

[DiagRisker]

@cgarciae a documentation of the filtering method with the corresponding function, would be of great help !

head_params = nnx.All(nnx.Param, nnx.PathContains('head'))

This is of interest to my post in #4514

[cgarciae]

@DiagRisker have you checked the Using Filters guide?

[DiagRisker]

my apologies, I don't know how I missed this. Many thanks
I believe, I've checked it before, however the answer I was looking for, was not in documentation, but in your answer.