Reparametrise a weight as direction times magnitude (Salimans & Kingma 2016).
Replaces module.<name> with the derived tensor
where g is a learnable scale (one entry per slice along dim)
and v is a learnable direction sharing the original weight's
shape. Decoupling magnitude from direction often accelerates
convergence and improves conditioning, especially in deep
convolutional and recurrent networks.
Parameters
moduleModuleThe module whose parameter is to be reparametrised. Mutated in
place — the original leaf parameter is removed and replaced with
a derived tensor recomputed before every forward call.
namestr= 'weight'Attribute name of the parameter to normalise. Default
"weight".dimint= 0Axis along which the scale
g has its own entry; norms are
reduced over every other axis. 0 (the default) is correct
for Linear and Conv* (one scale per output channel).
Use -1 for output-row Linear layouts.Returns
ModuleThe same module, now carrying <name>_g and <name>_v as
leaf parameters in place of the original <name>.
Notes
Two new parameters are registered:
<name>_g— broadcast-shaped tensor (1along every axis exceptdim) holding the per-slice magnitudes.<name>_v— same shape as the original weight, holding the unnormalised direction.
A forward pre-hook recomputes
so subsequent autograd / inference sees the same shape as before.
Invert with remove_weight_norm.
Examples
>>> import lucid.nn as nn
>>> from lucid.nn.utils import weight_norm
>>> layer = weight_norm(nn.Linear(128, 64))
>>> layer.weight_g.shape, layer.weight_v.shape
((64, 1), (64, 128))