Parameters, transforms, and regularization

Parameters, transforms, and regularization#

This tutorial is the guided tour: it builds three models of increasing complexity and uses them to connect parameters, transforms, regularization, and the model-level parameter API. Four related documents cover different ground:

Document

What it gives you

The parameter model

The why — why a change of variables beats clipping, why a penalty is a prior, why parameters form a tree

Choose parameter transforms

A catalog you look things up in — every transform by constrained domain, and how to pick one

Constrain and regularize parameters

Short task recipes for combining a constraint with a penalty

Parameter containers API

Exact signatures

 
import jax
import jax.numpy as jnp

import braintools
import brainstate
import brainstate.nn as nn

brainstate.random.seed(0)
brainstate.__version__

'0.5.2'

1. From ParamState to Param#

A brainstate.ParamState is a bare trainable array: an optimizer may move it anywhere in \(\mathbb{R}^n\). Real models mean more than that. A time constant must be positive, a mixing weight must lie in \([0, 1]\), and a weight matrix may want a penalty discouraging large values.

brainstate.nn.Param layers two orthogonal facilities on top of ParamState:

  • a transform (t=), a bijection from the unconstrained array the optimizer updates to the constrained value the model uses;

  • a regularization (reg=), a penalty contributed to the loss.

Two attributes matter. param.value() is the constrained value — use it in the forward pass. It is a method, not an attribute, because it computes the forward transform on each call rather than storing a constrained copy that could drift out of sync. param.val is the underlying ParamState the optimizer updates.

w = nn.Param(jnp.array([0.5, 1.0, 2.0]))
print('value()  :', w.value())
print('val      :', w.val)
print('trainable:', w.fit)

value()  : [0.5 1.  2. ]
val      : ParamState(
  value=ShapedArray(float32[3])
)
trainable: True

1.1 Const for values that are never trained#

Const is a Param with fit=False. It is excluded from ParamState collection, so optimizers and grad never see it — the right tool for a fixed scale or a lookup table.

Note that Const is a Param subclass, so parameter traversal still finds it. When trainability is what you mean, check .fit.

c = nn.Const(jnp.array(10.0))
print('Const value():', c.value(), '| fit:', c.fit)

Const value(): 10.0 | fit: False

2. Two orthogonal concerns#

A transform answers where is this parameter allowed to live. A regularization answers which values inside that domain do we prefer. They are independent: any transform composes with any penalty, and a parameter may carry one, both, or neither.

Section 3 builds a model using both. Sections 4 and 5 then show what changes once a model has enough parameters that managing them by hand stops being practical.

3. Model one: a constrained RNN cell#

A recurrent cell whose time constant must stay above 1.0 and whose gain must stay inside \((0, 5)\), with weight decay on the recurrent matrix. SoftplusT and SigmoidT enforce the domains; L2Reg supplies the penalty.

The value you pass to Param is interpreted in constrained space — tau starts at 5.0, and the inverse transform derives the unconstrained number stored in val.

class ConstrainedRNNCell(nn.Module):
    def __init__(self, n_in, n_hidden):
        super().__init__()
        self.n_hidden = n_hidden
        self.W_in = nn.Param(brainstate.random.randn(n_in, n_hidden) * 0.1)
        self.W_rec = nn.Param(brainstate.random.randn(n_hidden, n_hidden) * 0.1,
                              reg=nn.L2Reg(1e-3))
        self.b = nn.Param(jnp.zeros(n_hidden))
        self.tau = nn.Param(jnp.full((n_hidden,), 5.0), t=nn.SoftplusT(lower=1.0))
        self.gain = nn.Param(jnp.full((n_hidden,), 1.0), t=nn.SigmoidT(lower=0.0, upper=5.0))
        self.dt = nn.Const(jnp.array(1.0))

    def init_state(self, batch_size=1):
        self.h = brainstate.HiddenState(jnp.zeros((batch_size, self.n_hidden)))

    def update(self, x):
        drive = jnp.tanh(x @ self.W_in.value() + self.h.value @ self.W_rec.value()
                         + self.b.value())
        dh = (-self.h.value + self.gain.value() * drive) / self.tau.value()
        self.h.value = self.h.value + self.dt.value() * dh
        return self.h.value


cell = ConstrainedRNNCell(3, 8)
nn.init_all_states(cell, batch_size=4)
out = cell(brainstate.random.randn(4, 3))
print('output shape :', out.shape)
print('tau > 1      :', bool(jnp.all(cell.tau.value() > 1.0)))
print('0 < gain < 5 :', bool(jnp.all((cell.gain.value() > 0) & (cell.gain.value() < 5))))
print('W_rec penalty:', float(cell.W_rec.reg_loss()))

output shape : (4, 8)
tau > 1      : True
0 < gain < 5 : True
W_rec penalty: 0.0006242521340027452

3.1 Why a transform and not a clip#

Clipping enforces the same domain, but destroys the gradient at the boundary: below the bound the clipped output is constant, so the optimizer gets no signal to bring the parameter back. A bijection has a non-zero gradient everywhere.

theta = jnp.array(-3.0)
soft = nn.SoftplusT(lower=0.0)
clip_grad = jax.grad(lambda t: jnp.clip(t, 0.0, None).sum())(theta)
soft_grad = jax.grad(lambda t: soft.forward(t).sum())(theta)
print(f'at theta = -3.0   clip gradient = {float(clip_grad):.6f}'
      f'   Softplus gradient = {float(soft_grad):.6f}')

at theta = -3.0   clip gradient = 0.000000   Softplus gradient = 0.047426

The clipped parameter is stuck at zero gradient. The transformed one still moves.

For the full catalog and how to pick a transform, see Choose parameter transforms.

4. Model two: an HH population, and param_precompute()#

A Hodgkin–Huxley population has several parameters that must stay positive — the maximal conductances and the membrane capacitance — and it is simulated by stepping forward in time. That combination is what makes Module.param_precompute() worth having.

To keep the focus on parameters, this model is written dimensionless: voltages in mV and times in ms as plain floats, rather than brainunit quantities. For the same model with full units, see Hodgkin-Huxley neuron.

class HHPopulation(nn.Module):
    """Dimensionless Hodgkin-Huxley population: V in mV, t in ms."""

    def __init__(self, n):
        super().__init__()
        self.n = n
        pos = nn.SoftplusT(lower=0.0)
        self.gNa = nn.Param(jnp.full((n,), 120.0), t=pos, reg=nn.L2Reg(1e-5))
        self.gK = nn.Param(jnp.full((n,), 36.0), t=pos, reg=nn.L2Reg(1e-5))
        self.gL = nn.Param(jnp.full((n,), 0.03), t=pos)
        self.C = nn.Param(jnp.full((n,), 1.0), t=nn.SoftplusT(lower=0.1))
        self.ENa = nn.Const(jnp.array(50.0))
        self.EK = nn.Const(jnp.array(-77.0))
        self.EL = nn.Const(jnp.array(-54.387))

    def init_state(self, batch_size=None):
        shape = (self.n,) if batch_size is None else (batch_size, self.n)
        self.V = brainstate.HiddenState(jnp.full(shape, -65.0))
        self.m = brainstate.HiddenState(jnp.full(shape, 0.05))
        self.h = brainstate.HiddenState(jnp.full(shape, 0.6))
        self.n_gate = brainstate.HiddenState(jnp.full(shape, 0.32))

    def update(self, I_ext, dt=0.01):
        V, m, h, ng = self.V.value, self.m.value, self.h.value, self.n_gate.value
        m_a = 0.1 * (V + 40.0) / (1.0 - jnp.exp(-(V + 40.0) / 10.0))
        m_b = 4.0 * jnp.exp(-(V + 65.0) / 18.0)
        h_a = 0.07 * jnp.exp(-(V + 65.0) / 20.0)
        h_b = 1.0 / (1.0 + jnp.exp(-(V + 35.0) / 10.0))
        n_a = 0.01 * (V + 55.0) / (1.0 - jnp.exp(-(V + 55.0) / 10.0))
        n_b = 0.125 * jnp.exp(-(V + 65.0) / 80.0)
        I_Na = self.gNa.value() * (m ** 3) * h * (V - self.ENa.value())
        I_K = self.gK.value() * (ng ** 4) * (V - self.EK.value())
        I_L = self.gL.value() * (V - self.EL.value())
        dV = (I_ext - I_Na - I_K - I_L) / self.C.value()
        self.V.value = V + dt * dV
        self.m.value = m + dt * (m_a * (1 - m) - m_b * m)
        self.h.value = h + dt * (h_a * (1 - h) - h_b * h)
        self.n_gate.value = ng + dt * (n_a * (1 - ng) - n_b * ng)
        return self.V.value


hh = HHPopulation(200)
nn.init_all_states(hh)
print('parameter containers:', len(list(hh.param_modules())))

parameter containers: 7

4.1 The problem: the transform runs on every step#

Each call to update() reads gNa.value(), gK.value(), gL.value() and C.value(), and each of those applies a SoftplusT. Roll the model forward 100 steps and the same four transforms are recomputed 100 times over parameters that never changed.

Param.value() does not cache on its own. That is deliberate: caching automatically would capture a traced value inside jit and leak it into later calls. The cache is opt-in.

STEPS = 100
I0 = jnp.full((200,), 10.0)


def rollout_plain(I):
    for _ in range(STEPS):
        v = hh(I)
    return jnp.sum(v ** 2)


def rollout_precomputed(I):
    with hh.param_precompute():
        for _ in range(STEPS):
            v = hh(I)
        return jnp.sum(v ** 2)

4.2 Measuring the difference#

param_precompute() warms every parameter’s cache on entry and clears it on exit. Inside the block the transforms are computed once and reused. Counting equations in the traced jaxpr makes the saving concrete.

Two mechanics matter here. Use brainstate.transform.make_jaxpr, not jax.make_jaxpr — the plain JAX version cannot trace writes to a HiddenState and raises TraceContextError. And reset the hidden states before each trace, so the two counts are comparable.

nn.init_all_states(hh)
n_plain = len(brainstate.transform.make_jaxpr(rollout_plain)(I0)[0].eqns)
nn.init_all_states(hh)
n_pre = len(brainstate.transform.make_jaxpr(rollout_precomputed)(I0)[0].eqns)

print(f'jaxpr equations without param_precompute(): {n_plain}')
print(f'jaxpr equations with    param_precompute(): {n_pre}')
print(f'saved                                     : {n_plain - n_pre}')

jaxpr equations without param_precompute(): 8121
jaxpr equations with    param_precompute(): 7329
saved                                     : 792

4.3 Same numbers, same gradients#

The saving is pure redundancy removal, not an approximation. Both the forward value and every gradient are unchanged.

nn.init_all_states(hh)
r_plain = brainstate.transform.jit(rollout_plain)(I0)
nn.init_all_states(hh)
r_pre = brainstate.transform.jit(rollout_precomputed)(I0)
print('forward values agree:', bool(jnp.allclose(r_plain, r_pre)))

hh_params = hh.states(brainstate.ParamState)
nn.init_all_states(hh)
g_plain = brainstate.transform.grad(lambda: rollout_plain(I0), hh_params)()
nn.init_all_states(hh)
g_pre = brainstate.transform.grad(lambda: rollout_precomputed(I0), hh_params)()
print('gradients agree     :', all(bool(jnp.allclose(g_plain[k], g_pre[k])) for k in g_plain))

forward values agree: True

gradients agree     : True

4.4 The cache has a scope#

The cache is live only inside the block. On exit it is cleared — including when the block exits through an exception — so a stale constrained value can never escape into later code.

This is why the check below happens inside the with block. Calling value() outside it would leave valid at False, because value() does not populate the cache.

print('before:', hh.gNa.cache_stats)
with hh.param_precompute():
    print('inside:', hh.gNa.cache_stats)
print('after :', hh.gNa.cache_stats)

try:
    with hh.param_precompute():
        raise RuntimeError('something went wrong mid-simulation')
except RuntimeError:
    pass
print('after an exception:', hh.gNa.cache_stats)

before: {'valid': False, 'has_cached_value': False}
inside: {'valid': True, 'has_cached_value': True}
after : {'valid': False, 'has_cached_value': False}
after an exception: {'valid': False, 'has_cached_value': False}

5. Model three: an E-I network, and managing parameters across a tree#

The HH population was flat: every parameter sat directly on one module. Real networks nest. An excitatory–inhibitory network has two populations and four projections between them, each owning its own parameters, so paths look like exc.tau and syn_ee.w.

At this depth, managing parameters by hand starts to hurt — and the remaining three methods each answer one of those pains.

class LIFPopulation(nn.Module):
    def __init__(self, n, tau_init):
        super().__init__()
        self.n = n
        self.tau = nn.Param(jnp.full((n,), tau_init), t=nn.SoftplusT(lower=1.0))
        self.V_th = nn.Param(jnp.full((n,), 1.0), t=nn.SoftplusT(lower=0.1))
        self.V_rest = nn.Const(jnp.zeros(n))
        # A hard threshold has zero derivative everywhere, so no gradient would reach
        # tau or V_th. A surrogate gradient is the standard fix.
        self.spike_fn = braintools.surrogate.ReluGrad()

    def init_state(self, batch_size=None):
        shape = (self.n,) if batch_size is None else (batch_size, self.n)
        self.V = brainstate.HiddenState(jnp.zeros(shape))

    def update(self, I, dt=0.1):
        V = self.V.value + dt * (-(self.V.value - self.V_rest.value()) + I) / self.tau.value()
        spike = self.spike_fn(V - self.V_th.value())
        self.V.value = V * (1.0 - spike)
        return spike


class Projection(nn.Module):
    def __init__(self, n_pre, n_post, scale, sign, reg):
        super().__init__()
        self.sign = sign
        self.w = nn.Param(jnp.abs(brainstate.random.randn(n_pre, n_post)) * scale,
                          t=nn.SoftplusT(lower=0.0), reg=reg)

    def update(self, spike):
        return self.sign * (spike @ self.w.value())


class EINetwork(nn.Module):
    def __init__(self, n_exc=80, n_inh=20):
        super().__init__()
        self.exc = LIFPopulation(n_exc, 20.0)
        self.inh = LIFPopulation(n_inh, 10.0)
        self.syn_ee = Projection(n_exc, n_exc, 0.02, +1.0, nn.L2Reg(1e-4))
        self.syn_ei = Projection(n_exc, n_inh, 0.02, +1.0, nn.L2Reg(1e-4))
        self.syn_ie = Projection(n_inh, n_exc, 0.08, -1.0, nn.L1Reg(1e-4))
        self.syn_ii = Projection(n_inh, n_inh, 0.08, -1.0, nn.L1Reg(1e-4))
        self.readout = nn.Param(brainstate.random.randn(n_exc, 2) * 0.1, reg=nn.L1Reg(1e-3))

    def init_state(self, batch_size=None):
        self.exc.init_state(batch_size)
        self.inh.init_state(batch_size)
        e_shape = self.exc.n if batch_size is None else (batch_size, self.exc.n)
        i_shape = self.inh.n if batch_size is None else (batch_size, self.inh.n)
        self.s_exc = brainstate.HiddenState(jnp.zeros(e_shape))
        self.s_inh = brainstate.HiddenState(jnp.zeros(i_shape))

    def update(self, I_ext):
        se, si = self.s_exc.value, self.s_inh.value
        I_e = I_ext + self.syn_ee(se) + self.syn_ie(si)
        I_i = I_ext[..., :self.inh.n] + self.syn_ei(se) + self.syn_ii(si)
        new_e = self.exc(I_e)
        new_i = self.inh(I_i)
        self.s_exc.value, self.s_inh.value = new_e, new_i
        return new_e @ self.readout.value()


net = EINetwork()
nn.init_all_states(net)
print('E-I forward output shape:', net(jnp.ones(80) * 50.0).shape)

E-I forward output shape: (2,)

5.1 named_param_modules() — discovery without a registry#

Without it, you keep a hand-written list:

self._params = [self.exc.tau, self.exc.V_th, self.syn_ee.w, self.syn_ei.w, ...]

That list has to be updated every time a projection is added, and the names are invented separately from the structure, so the two drift apart.

named_param_modules() walks the graph and yields (dotted path, param). Identity comes from structure, so there is nothing to keep in sync — and the path doubles as a checkpoint key.

print(f'{"name":<16}{"shape":<12}{"transform":<14}{"regularizer":<14}trainable')
for name, p in net.named_param_modules():
    print(f'{name:<16}{str(tuple(jnp.shape(p.value()))):<12}'
          f'{type(p.t).__name__:<14}{type(p.reg).__name__ if p.reg else "-":<14}{p.fit}')

name            shape       transform     regularizer   trainable
exc.V_rest      (80,)       IdentityT     -             False
exc.V_th        (80,)       SoftplusT     -             True
exc.tau         (80,)       SoftplusT     -             True
inh.V_rest      (20,)       IdentityT     -             False
inh.V_th        (20,)       SoftplusT     -             True
inh.tau         (20,)       SoftplusT     -             True
readout         (80, 2)     IdentityT     L1Reg         True
syn_ee.w        (80, 80)    SoftplusT     L2Reg         True
syn_ei.w        (80, 20)    SoftplusT     L2Reg         True
syn_ie.w        (20, 80)    SoftplusT     L1Reg         True
syn_ii.w        (20, 20)    SoftplusT     L1Reg         True

Note exc.V_rest and inh.V_rest: they are Const, so they appear in the traversal but are not trainable. Traversal finds every parameter container; .fit tells you which ones an optimizer will touch.

Filtering is a list comprehension. These methods accept only allowed_hierarchy — they take no filter function.

synaptic = [(n, p) for n, p in net.named_param_modules() if n.startswith('syn_')]
print('synaptic weights:', [n for n, _ in synaptic])

regularized = [(n, p) for n, p in net.named_param_modules() if p.reg is not None]
print('regularized     :', [n for n, _ in regularized])

trainable = [n for n, p in net.named_param_modules() if p.fit]
print('trainable       :', len(trainable), 'of', len(list(net.param_modules())))

synaptic weights: ['syn_ee.w', 'syn_ei.w', 'syn_ie.w', 'syn_ii.w']
regularized     : ['readout', 'syn_ee.w', 'syn_ei.w', 'syn_ie.w', 'syn_ii.w']
trainable       : 9 of 11

allowed_hierarchy=(min_depth, max_depth) restricts the walk. It is a closed interval, and a Param attached directly to the root counts as depth 1 — here, only readout.

print('depth 1 only:', [n for n, _ in net.named_param_modules(allowed_hierarchy=(1, 1))])

depth 1 only: ['readout']

5.2 param_modules() — the same walk without the names#

When the path is not needed, param_modules() yields the containers alone. It is what you want for a bulk operation over every parameter.

print('parameter containers:', len(list(net.param_modules())))
print('scalar entries      :', sum(int(jnp.size(p.value())) for p in net.param_modules()))
print('with a transform    :',
      len([p for p in net.param_modules() if not isinstance(p.t, nn.IdentityT)]))

parameter containers: 11
scalar entries      : 10460
with a transform    : 8

5.3 reg_loss() — one call for the whole tree#

The penalties in this network sit at different depths: L2Reg on the excitatory projections, L1Reg on the inhibitory ones and on the readout. Summing them by hand means walking the tree yourself.

The manual version below is correct — the traversal recurses by default — but it is only correct if you know that. model.reg_loss() removes the assumption.

manual = sum(p.reg_loss() for p in net.param_modules())
builtin = net.reg_loss()
print(f'manual sum     : {float(manual):.8f}')
print(f'net.reg_loss() : {float(builtin):.8f}')
print('agree          :', bool(jnp.allclose(manual, builtin)))

manual sum     : 0.02653572
net.reg_loss() : 0.02653572
agree          : True

The failure mode is a wrong assumption about depth. Restricting the walk to depth 1 — a plausible-looking mistake — silently returns less than half the true penalty, and nothing raises an error:

shallow = sum(p.reg_loss() for p in net.param_modules(allowed_hierarchy=(1, 1)))
print(f'depth-1 only  : {float(shallow):.8f}   <-- misses every nested regularizer')
print(f'correct total : {float(net.reg_loss()):.8f}')

depth-1 only  : 0.01322584   <-- misses every nested regularizer
correct total : 0.02653572

Training would still run. The penalty would just be wrong, and nothing would say so.

6. Training the network with all four#

The task is homeostatic: under a fixed drive the untrained network fires at rate 0.20, and we ask it to settle at 0.10. The training step uses every piece — constrained parameters keep the time constants and weights inside their domains, reg_loss() supplies the whole-tree penalty, and param_precompute() wraps the 20-step rollout so each transform is computed once rather than once per timestep.

The surrogate gradient matters here. With a hard V > V_th threshold the derivative is zero everywhere, so tau and V_th would receive no gradient at all and would sit motionless — their constraints would then hold only because nothing ever moved them. With the surrogate every trainable parameter gets a gradient, so the constraint check below means something.

I_drive = jnp.ones(80) * 50.0
TARGET_RATE = 0.10
T = 20


def firing_rate():
    """Mean excitatory firing rate over a fresh T-step rollout."""
    nn.init_all_states(net)
    total = 0.0
    for _ in range(T):
        net(I_drive)
        total = total + jnp.mean(net.s_exc.value)
    return total / T


brainstate.random.seed(0)
net = EINetwork()
nn.init_all_states(net)
params = net.states(brainstate.ParamState)
opt = braintools.optim.Adam(lr=5e-2)
opt.register_trainable_weights(params)

rate_before = float(firing_rate())
tau_before = float(jnp.mean(net.exc.tau.value()))

def loss_fn():
    with net.param_precompute():
        total = 0.0
        for _ in range(T):
            net(I_drive)
            total = total + jnp.mean(net.s_exc.value)
        rate = total / T
    return (rate - TARGET_RATE) ** 2 + net.reg_loss()


@brainstate.transform.jit
def train_step():
    grads, loss = brainstate.transform.grad(loss_fn, params, return_value=True)()
    opt.update(grads)
    return loss


losses = []
for _ in range(60):
    nn.init_all_states(net)
    losses.append(float(train_step()))

print(f'loss        : {losses[0]:.6f} -> {losses[-1]:.6f}')
print(f'firing rate : {rate_before:.4f} -> {float(firing_rate()):.4f}   (target {TARGET_RATE})')
print(f'exc.tau mean: {tau_before:.4f} -> {float(jnp.mean(net.exc.tau.value())):.4f}')
print('tau  still above its 1.0 floor     :', bool(jnp.all(net.exc.tau.value() > 1.0)))
print('V_th still above its 0.1 floor     :', bool(jnp.all(net.exc.V_th.value() > 0.1)))
print('synaptic weights still non-negative:', bool(jnp.all(net.syn_ee.w.value() >= 0.0)))

loss        : 0.036572 -> 0.002921

firing rate : 0.2000 -> 0.1000   (target 0.1)
exc.tau mean: 20.0000 -> 21.5448
tau  still above its 1.0 floor     : True
V_th still above its 0.1 floor     : True
synaptic weights still non-negative: True

The network reached the target rate, and exc.tau moved to get there — so the constraints were under real pressure rather than holding by default. None of them were enforced by clipping, projection, or a post-update repair step: the domain is guaranteed by the transform itself, so the optimizer never has to be corrected.

7. Summary#

Param carries two orthogonal facilities on top of ParamState: a transform fixing where a parameter may live, and a regularization expressing which values are preferred. value() returns the constrained value, val is the array the optimizer updates, and Const marks a value as never trained.

Once a model has more than a handful of parameters, four Module methods do the tree-level work:

Method

Answers

named_param_modules()

What parameters does this model have, and what are they called? Yields (dotted path, param); the path doubles as a checkpoint key

param_modules()

The same walk without names, for bulk operations

reg_loss()

What is the total penalty? Sums every regularizer in the tree in one call

param_precompute()

How do I avoid recomputing transforms repeatedly? A scoped cache: warmed on entry, cleared on exit, exception-safe

All four accept only allowed_hierarchy=(min_depth, max_depth) — a closed interval, with root parameters at depth 1. They take no filter function; filter with a list comprehension.

Two behaviours are worth remembering. Traversal includes Const, so check .fit when trainability is what you mean. And value() never populates the cache on its own — only param_precompute() or Param.cache() does — which is what keeps a value traced inside jit from leaking outside that trace.