Skip to content

← Flax step 96 of 100

Medium Primitives

Partial Init - Warm-Start From Smaller Checkpoint

Jump to the task

Reuse the parameters of a smaller model when constructing a larger one. Initialize the complete larger model first, then replace its matching early layers with arrays from the smaller model. The new tail layers keep their fresh parameters.

This exercise demonstrates the copying operation with two freshly initialized models. It does not load a trained checkpoint or perform training. In a real warm start, the source arrays would come from a saved model.

The two model shapes

Use one MLP(nn.Module) class with num_layers and features fields. Its compact forward method repeatedly applies nn.Dense(self.features). There are no activations between layers in this exercise.

Let L = num_full_layers + num_extra_layers, F = features, and let x have shape (N, D). The small model has num_full_layers Dense layers; the big model has L. The first kernel in each has shape (D, F). Every later kernel has shape (F, F), and each bias has shape (F,).

In this construction, unnamed Dense children are named Dense_0, Dense_1 and so on. Corresponding early layers have both matching paths and matching shapes. For other architectures, check shapes and semantics as well as names before transferring arrays. Copying a dictionary entry does not validate compatibility; an incompatible array can fail when the model is applied.

Copy complete layer entries

new_inner = dict(big_variables["params"])
for i in range(num_full_layers):
    name = f"Dense_{i}"
    new_inner[name] = small_variables["params"][name]
new_variables = {"params": new_inner}

Each replacement transfers both the kernel and bias. The shallow dictionary copy is enough for replacing whole entries without changing the big model's original mapping. It does not deep-copy every nested dictionary. Avoid mutating shared nested containers later; JAX array updates themselves return new arrays.

The extra layers remain as initialized in the big model. Adding them does not generally preserve the smaller model's output, even when all reused arrays match exactly. This operation also does not transfer optimizer state.

Count arrays and scalar parameters separately

Each Dense layer contributes two array leaves, but many scalar parameters. Count scalars with sum(int(a.size) for a in jax.tree_util.tree_leaves(...)). For L >= 1, the expected total here is:

first layer:       D * F + F
remaining layers: (L - 1) * (F * F + F)
total:             D * F + F + (L - 1) * (F * F + F)

For D=4, F=2 and L=4, that is 10 + 3*6 = 28 scalar parameters in eight arrays. Batch size does not affect the parameter count.

The count is a shape check, not proof of copying. A freshly initialized big model has the same count. To check the transfer itself, compare the early arrays with the small model and the tail arrays with the original big model.

Problem

Implement partial_init_param_count(seed, x, num_full_layers, num_extra_layers, features):

  1. Cast seed and layer/feature counts to integers. Construct the small and big models described above, using default Dense settings, including bias.
  2. Set rng = jax.random.PRNGKey(int(seed)), then rng_small, rng_big = jax.random.split(rng) in that order.
  3. Initialize both models on x with their respective keys.
  4. Copy the small model's complete early layer entries into a copy of the big model's params mapping. Preserve the outer params collection.
  5. Count scalar elements in the resulting tree and return a JAX array [float(total_param_count), float(L)] of shape (2,).

x is a finite float32 JAX array of shape (N, D), with positive dimensions. num_full_layers and features are positive integer-valued scalars; num_extra_layers is a nonnegative integer-valued scalar. The seed is also nonnegative and integer-valued. All scalar inputs may arrive as floats. Leave x and both original variable mappings unchanged. The returned count depends on shapes and layer counts, not on the seed or input values.