Parameter Management

Managing Parameters

A neural network is a tree of parameters — the weight matrices and bias vectors gradient descent updates. Training is one thing you do with them; this deck covers the others.

  • Inspection — debug, sanity-check init, view features.
  • Iteration — optimizers, weight decay, checkpointing all walk every parameter.
  • Sharing (“tying”) — make two layers refer to the same tensor (tied embeddings, autoencoders).

The parameter tree

A nested module is just a tree. Each module is a node; each parameter is a leaf:

net  (Sequential)
├─ 0: Linear      ├─ weight  (8, 4)
│                 └─ bias    (8,)
├─ 1: ReLU         (no params)
└─ 2: Linear      ├─ weight  (1, 8)
                  └─ bias    (1,)

Two access patterns:

  • By path: net[2].weight — direct.
  • By traversal: walk the tree, yield every leaf.

Frameworks give you both, plus serialization built on the same traversal.

A toy model

from d2l import jax as d2l
from flax import linen as nn
import jax
from jax import numpy as jnp
net = nn.Sequential([nn.Dense(8), nn.relu, nn.Dense(1)])

X = jax.random.uniform(d2l.get_key(), (2, 4))
params = net.init(d2l.get_key(), X)
net.apply(params, X).shape
(2, 1)

Direct access

Index into a Sequential like a list; each layer exposes its parameters as attributes:

params['params']['layers_2']
{'kernel': Array([[ 0.39065108],
        [-0.47905394],
        [ 0.17323323],
        [ 0.22871736],
        [ 0.46032798],
        [-0.22263312],
        [ 0.19490093],
        [-0.21752292]], dtype=float32),
 'bias': Array([0.], dtype=float32)}

Two parameters per Linear layer — weight matrix and bias vector. The output object is a Parameter (PyTorch) or similar wrapper that carries the tensor + gradient + extra metadata.

Tensor inside the parameter

.data (PyTorch) unwraps the parameter to a plain tensor for inspection:

bias = params['params']['layers_2']['bias']
type(bias), bias
(jaxlib._jax.ArrayImpl, Array([0.], dtype=float32))

.grad is the gradient buffer — populated by backward(), otherwise None. Useful for custom optimizers or diagnosing dead neurons:

Recursive traversal

For everything-at-once, use named_parameters(). It walks the whole tree and yields (name, param) pairs at the leaves — names use dotted paths through the nesting:

jax.tree_util.tree_map(lambda x: x.shape, params)
{'params': {'layers_0': {'bias': (8,), 'kernel': (4, 8)},
  'layers_2': {'bias': (1,), 'kernel': (8, 1)}}}

This is the iterator optim.SGD(net.parameters(), …) consumes. It’s also what gets pickled when you save a checkpoint with state_dict(). Walk-tree-once, use many ways.

Parameter tying

Reuse the same module instance at multiple positions in your architecture, and the framework treats them as one parameter set — same memory, gradients accumulate across uses.

Common cases:

  • Tied embeddings: input embedding and output softmax projection in a language model share weights — saves |V| \cdot d parameters.
  • Autoencoders: decoder uses transposed encoder weights.
  • Recurrent layers: same kernel applied at every time step (the original tying mechanism).
# We need to give the shared layer a name so that we can refer to its
# parameters
shared = nn.Dense(8)
net = nn.Sequential([nn.Dense(8), nn.relu,
                     shared, nn.relu,
                     shared, nn.relu,
                     nn.Dense(1)])

params = net.init(d2l.get_key(), X)

# Check whether the parameters are different
print(len(params['params']) == 3)
True

Modify net[2].weight and net[4].weight reflects the same change — they are the same tensor, not just equal.

Recap

  • A module is a tree; parameters live at the leaves.
  • Direct access: net[i].weight, .bias, .grad.
  • Recursive traversal: named_parameters() / state_dict() walks the whole tree.
  • Same iterator powers optimizers, weight decay, checkpointing.
  • Tied parameters = reuse the same module instance — gradients accumulate; one buffer in memory.