Choose a parameter transform#
This how-to is the transform catalog: which transform maps onto which constrained domain, how to compose them, and how to choose. Four related documents cover different ground:
Document |
What it gives you |
|---|---|
The next step — combining a transform with a regularization penalty |
|
The guided tour, with three worked models |
|
The why behind the design |
|
Exact signatures |
import jax
import jax.numpy as jnp
import brainunit as u
import brainstate
import brainstate.nn as nn
brainstate.random.seed(0)
brainstate.__version__
'0.5.2'
1. What a transform is, and why not clipping#
A transform is a bijection between an unconstrained space — all of \(\mathbb{R}\), where gradient descent is well behaved — and a constrained space such as the positives, an interval, or the probability simplex. The optimizer moves the unconstrained value; the model reads the constrained one.
The tempting alternative is to clip after each update. It breaks learning at exactly the place you care about: below the bound the clipped output is constant, so its gradient is zero and the optimizer gets no signal to climb back.
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 = {float(theta)}')
print(f' clip -> value {float(jnp.clip(theta, 0.0, None)):.6f}, gradient {float(clip_grad):.6f}')
print(f' SoftplusT-> value {float(soft.forward(theta)):.6f}, gradient {float(soft_grad):.6f}')
at theta = -3.0
clip -> value 0.000000, gradient 0.000000
SoftplusT-> value 0.048587, gradient 0.047426
The clipped parameter is stuck: a zero gradient means the optimizer cannot move it, and any momentum it had is wasted. The transformed parameter has a small but non-zero gradient everywhere, so it can always recover.
2. The catalog, by constrained domain#
Constrained domain |
Transforms |
|---|---|
\((\text{lower}, \infty)\) — positive quantities |
|
\((-\infty, \text{upper})\) — negative quantities |
|
\((\text{lower}, \text{upper})\) — a bounded interval |
|
Non-negative, summing to one — probabilities |
|
Unit \(L_2\) norm |
|
Monotonically increasing entries |
|
Reparameterizations |
|
Composition and masking |
|
IdentityT() is the default: no constraint at all.
Note that TanhT and SoftsignT take explicit lower and upper bounds — they are bounded
interval maps, not fixed \((-1, 1)\) squashers.
2.1 Every transform round-trips#
inverse() undoes forward(). That is what lets Param store an unconstrained value while
you read and write constrained ones. The check below runs the whole catalog through a
round-trip.
CASES = [
('IdentityT()', nn.IdentityT(), jnp.array([-1.0, 0.0, 2.0])),
('SoftplusT(0.0)', nn.SoftplusT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),
('ExpT(0.0)', nn.ExpT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),
('PositiveT()', nn.PositiveT(), jnp.array([0.5, 2.0, 8.0])),
('ReluT()', nn.ReluT(), jnp.array([0.5, 2.0, 8.0])),
('LogT(0.0)', nn.LogT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),
('NegSoftplusT(0.0)', nn.NegSoftplusT(upper=0.0), jnp.array([-0.5, -2.0, -8.0])),
('NegativeT()', nn.NegativeT(), jnp.array([-0.5, -2.0, -8.0])),
('SigmoidT(0, 1)', nn.SigmoidT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),
('ScaledSigmoidT(0, 1)', nn.ScaledSigmoidT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),
('ClipT(0, 1)', nn.ClipT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),
('TanhT(-1, 1)', nn.TanhT(lower=-1.0, upper=1.0), jnp.array([-0.5, 0.0, 0.5])),
('SoftsignT(-1, 1)', nn.SoftsignT(lower=-1.0, upper=1.0), jnp.array([-0.5, 0.0, 0.5])),
('SimplexT()', nn.SimplexT(), jnp.array([0.2, 0.3, 0.5])),
('UnitVectorT()', nn.UnitVectorT(), jnp.array([0.6, 0.8])),
('OrderedT()', nn.OrderedT(), jnp.array([-1.0, 0.5, 2.0])),
('AffineT(2, 1)', nn.AffineT(scale=2.0, shift=1.0), jnp.array([-1.0, 0.0, 3.0])),
('PowerT()', nn.PowerT(), jnp.array([0.5, 2.0, 8.0])),
]
print(f'{"transform":<22} {"round-trip":<12} max abs error')
for name, t, x in CASES:
back = u.get_magnitude(t.forward(t.inverse(x)))
err = float(jnp.max(jnp.abs(back - x)))
print(f'{name:<22} {str(bool(err < 1e-4)):<12} {err:.1e}')
transform round-trip max abs error
IdentityT() True 0.0e+00
SoftplusT(0.0) True 3.0e-08
ExpT(0.0) True 0.0e+00
PositiveT() True 0.0e+00
ReluT() True 0.0e+00
LogT(0.0) True 0.0e+00
NegSoftplusT(0.0) True 3.0e-08
NegativeT() True 3.0e-08
SigmoidT(0, 1) True 0.0e+00
ScaledSigmoidT(0, 1) True 0.0e+00
ClipT(0, 1) True 0.0e+00
TanhT(-1, 1) True 0.0e+00
SoftsignT(-1, 1) True 0.0e+00
SimplexT() True 3.0e-08
UnitVectorT() True 0.0e+00
OrderedT() True 0.0e+00
AffineT(2, 1) True 0.0e+00
PowerT() True 0.0e+00
Every transform recovers the original constrained value to within floating-point tolerance.
Note the u.get_magnitude(...) call. Most transforms return a plain JAX array, but a few —
ExpT and LogT among them — return a dimensionless brainunit.Quantity. get_magnitude
normalizes both to a plain array. See the pitfalls in section 6.
3. Composing transforms#
ChainT applies its transforms in order, so you can reach a domain no single transform covers.
MaskedT applies a transform to only the entries its mask selects, leaving the rest
unconstrained.
chained = nn.ChainT(nn.AffineT(scale=2.0, shift=1.0), nn.SoftplusT(lower=0.0))
x = jnp.array([0.5, 2.0, 8.0])
print('ChainT round-trip max error:',
float(jnp.max(jnp.abs(u.get_magnitude(chained.forward(chained.inverse(x))) - x))))
# Constrain entries 0 and 2 to be positive; leave entry 1 free.
masked = nn.MaskedT(jnp.array([1.0, 0.0, 1.0]), nn.SoftplusT(lower=0.0))
y = jnp.array([0.5, -2.0, 8.0])
print('MaskedT keeps the unmasked entry negative:',
u.get_magnitude(masked.forward(masked.inverse(y))))
ChainT round-trip max error: 2.9802322387695312e-08
MaskedT keeps the unmasked entry negative: [ 0.49999997 -2. 8. ]
4. Choosing one#
If you need |
Use |
|---|---|
A rate, conductance, or time constant that must stay positive |
|
The same, but spanning orders of magnitude |
|
A mixing weight or probability in \([0, 1]\) |
|
A categorical distribution over \(k\) outcomes |
|
A direction, with magnitude handled separately |
|
Sorted thresholds or bin edges |
|
To constrain only some entries of an array |
|
A domain no single transform covers |
|
SoftplusT is the default choice for positivity: it is gentler near the bound than ExpT,
which grows exponentially and can turn a moderate unconstrained value into a very large
constrained one.
5. Attaching a transform to a parameter#
Pass it as t=. From then on value() returns the constrained value, and set_value()
accepts a constrained value and applies the inverse for you.
tau = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0))
print('value() :', float(tau.value()))
print('stored in val:', float(tau.val.value))
tau.val.value = jnp.array(-50.0) # an aggressive optimizer step
print('after a large negative update:', float(tau.value()), '(still above the 1.0 floor)')
value() : 5.0
stored in val: 3.9815146923065186
after a large negative update: 1.0 (still above the 1.0 floor)
Note what the printout above says about the constructor: Param(5.0, t=SoftplusT(lower=1.0))
gives value() == 5.0. The value you pass in is interpreted in constrained space, and the
inverse transform is applied once to derive the unconstrained number stored in val — here
3.98. You never have to work out the unconstrained value yourself.
Param also accepts precompute=, a callable applied to the constrained value after the
transform. Use it when the model needs a derived quantity rather than the constrained value
itself — for instance a decay factor computed from a time constant.
plain = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0))
decay = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0),
precompute=lambda tau: jnp.exp(-1.0 / tau))
print('without precompute, value() is the time constant :', float(plain.value()))
print('with precompute, value() is the derived factor :', float(decay.value()))
print('check exp(-1 / tau) :',
float(jnp.exp(-1.0 / plain.value())))
without precompute, value() is the time constant : 5.0
with precompute, value() is the derived factor : 0.8187307715415955
check exp(-1 / tau) : 0.8187307715415955
6. Pitfalls#
Saturation. SigmoidT and TanhT flatten far from the origin. If the unconstrained value
drifts to \(\pm 100\), the gradient is effectively zero and the parameter stops moving. Keep
initial values near the middle of the range.
Initializing exactly on the bound. SoftplusT(lower=0.0) cannot represent 0.0 — the
inverse diverges. Initialize strictly inside the domain.
Inverting outside the domain. inverse() is defined only on the constrained domain. Calling
SimplexT().inverse() on a vector that does not sum to one is undefined behaviour, not an error.
Quantity vs. plain array. forward() does not return the same type for every transform:
ExpT and LogT return a dimensionless brainunit.Quantity, while SoftplusT returns a plain
JAX array. Passing a Quantity to a raw jnp function raises a TypeError. Wrap the result in
u.get_magnitude(...) when you need a plain array.
print('SoftplusT forward type:', type(nn.SoftplusT(lower=0.0).forward(jnp.array(1.0))).__name__)
print('ExpT forward type:', type(nn.ExpT(lower=0.0).forward(jnp.array(1.0))).__name__)
try:
jnp.abs(nn.ExpT(lower=0.0).forward(jnp.array(1.0)))
except TypeError as e:
print('raw jnp on a Quantity ->', type(e).__name__)
print('with get_magnitude ->',
float(jnp.abs(u.get_magnitude(nn.ExpT(lower=0.0).forward(jnp.array(1.0))))))
SoftplusT forward type: ArrayImpl
ExpT forward type: Quantity
raw jnp on a Quantity -> TypeError
with get_magnitude -> 2.7182817459106445
7. Summary#
A transform is a bijection between unconstrained \(\mathbb{R}\) and a constrained domain. The optimizer works unconstrained;
value()returns the constrained value.Clipping zeroes the gradient at the bound; a transform never does.
Pick by constrained domain using the table in section 4.
SoftplusTis the default for positivity.Compose with
ChainT; restrict to part of an array withMaskedT.Watch for saturation, initialization on a bound, out-of-domain inverses, and the
Quantity-vs-array difference between transforms.