Layers

Multiplicative nn.Module building blocks: the dense SigmaPiLinear, the convolutional ConvSigmaPi2d block, and the low-rank bilinear PolyLinear.

polyweave.layers

nn.Module building blocks that wrap the low-level ops.

PolyLinear

PolyLinear(in_features: int, out_features: int | None = None, *, rank: int = 8, bias: bool = True, symmetric: bool = False, quad_scale_init: float = QUAD_SCALE_INIT)

Bases: Module

Linear + low-rank quadratic (degree-2 factorization-machine) layer.

Parameters:

Name	Type	Description	Default
in_features	int	size of each input sample's last dimension.	required
out_features	int | None	size of each output's last dimension. Defaults to in_features (channels-preserving).	None
rank	int	number of rank-1 bilinear factors in the quadratic branch. 0 disables the quadratic branch entirely (pure linear). Higher rank = more multiplicative capacity; controls the parameter budget.	8
bias	bool	whether the linear branch carries a bias term (default True).	True
symmetric	bool	if True the two factor matrices are tied (V = U), giving a true symmetric quadratic form sum_r (u_r · x)^2; if False (default) U and V are independent, giving a general (asymmetric) bilinear form.	False
quad_scale_init	float	initial per-output gate (default -2.0).	QUAD_SCALE_INIT

Parameter count: out*in + (bias?out:0) + 2*rank*in + out*rank + out (with symmetric=True the 2*rank*in term becomes rank*in).

Source code in polyweave/layers/poly_linear.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    *,
    rank: int = 8,
    bias: bool = True,
    symmetric: bool = False,
    quad_scale_init: float = QUAD_SCALE_INIT,
) -> None:
    super().__init__()
    out_features = in_features if out_features is None else out_features
    if rank < 0:
        raise ValueError(f"rank must be >= 0, got {rank}")
    self.in_features = in_features
    self.out_features = out_features
    self.rank = rank
    self.symmetric = symmetric

    self.linear = nn.Linear(in_features, out_features, bias=bias)
    if rank > 0:
        self.U = nn.Parameter(torch.empty(rank, in_features))
        self.V = self.U if symmetric else nn.Parameter(torch.empty(rank, in_features))
        self.mix = nn.Parameter(torch.empty(out_features, rank))
        self.quad_scale = nn.Parameter(torch.full((out_features,), float(quad_scale_init)))
        self.reset_quadratic_parameters()
    else:
        self.register_parameter("U", None)
        self.register_parameter("mix", None)
        self.register_parameter("quad_scale", None)

quad_scale_mean

quad_scale_mean() -> float

Polynomial-recruitment diagnostic exp(quad_scale).mean().

Returns 0.0 when the quadratic branch is disabled (rank == 0), matching "no multiplicative branch".

Source code in polyweave/layers/poly_linear.py

@torch.no_grad()
def quad_scale_mean(self) -> float:
    """Polynomial-recruitment diagnostic ``exp(quad_scale).mean()``.

    Returns ``0.0`` when the quadratic branch is disabled (``rank == 0``),
    matching "no multiplicative branch".
    """
    if self.rank == 0:
        return 0.0
    return self.quad_scale.exp().mean().item()

ConvSigmaPi2d

ConvSigmaPi2d(channels: int, kernel_size: int = 3, padding: int = 1, pi_scale_init: float = PI_SCALE_INIT, max_exponent: float = DEFAULT_MAX_EXPONENT, max_log: float = DEFAULT_MAX_LOG, signed_products: bool = False, center_product: bool = False, eps: float = DEFAULT_EPS)

Bases: Module

Sigma-Pi convolutional block with a genuine geometric-product pi branch.

Parameters:

Name	Type	Description	Default
channels	int	number of input == output channels.	required
kernel_size	int	conv kernel size for both branches (default 3).	3
padding	int	conv padding (default 1, i.e. "same" for kernel 3).	1
pi_scale_init	float	initial value of the per-channel pi_scale gate (default -2.0 so exp(pi_scale) ~= 0.135).	PI_SCALE_INIT
max_exponent	float	bound on each log-space exponent; the pi conv weight is max_exponent * tanh(raw) (default 0.5).	DEFAULT_MAX_EXPONENT
max_log	float	the log-space accumulation is clamped to [-max_log, max_log] before exponentiation (default 6.0).	DEFAULT_MAX_LOG
signed_products	bool	if True, multiply the magnitude product by a sign-vote surrogate (a flagged ablation; default False is magnitude-only).	False
center_product	bool	if True the pi branch uses expm1(u) (= product - 1) instead of exp(u), so it starts at the multiplicative identity and the pi_scale gate recovers a clean "volume knob" meaning (default False). Caveat (conv): this silent-init property only holds when u ~= 0, which needs every exponent ~0. With a wide receptive field (channels * k * k terms) the accumulated u is far from zero even at init, so expm1 starts large anyway and this flag is a near-no-op here — it matters mainly for narrow-fan-in dense layers. (Verified null on the Q/K diagnostic, 2026-06-07.)	False
eps	float	stabiliser inside log(|x| + eps) (default 1e-8).	DEFAULT_EPS

Source code in polyweave/layers/sigmapi_conv.py

def __init__(
    self,
    channels: int,
    kernel_size: int = 3,
    padding: int = 1,
    pi_scale_init: float = PI_SCALE_INIT,
    max_exponent: float = DEFAULT_MAX_EXPONENT,
    max_log: float = DEFAULT_MAX_LOG,
    signed_products: bool = False,
    center_product: bool = False,
    eps: float = DEFAULT_EPS,
) -> None:
    super().__init__()
    self.channels = channels
    self.padding = padding
    self.max_exponent = float(max_exponent)
    self.max_log = float(max_log)
    self.signed_products = bool(signed_products)
    self.center_product = bool(center_product)
    self.eps = eps

    # Sigma: signed additive branch (carries sign information).
    self.sigma_conv = nn.Conv2d(channels, channels, kernel_size, padding=padding)
    # Pi: raw log-space exponent kernel, squashed to (-max_exponent, max_exponent).
    # No bias (a log-space offset is a multiplicative constant -> the gate's job).
    self.pi_weight_raw = nn.Parameter(
        torch.zeros(channels, channels, kernel_size, kernel_size)
    )
    nn.init.normal_(self.pi_weight_raw, std=0.1)
    # Per-channel learnable amplitude gate / recruitment diagnostic.
    self.pi_scale = nn.Parameter(torch.full((channels, 1, 1), float(pi_scale_init)))
    self.bn = nn.BatchNorm2d(channels)
    # Runtime ablation switch: when False the pi branch is zeroed in ``forward``
    # (the block degrades to ``relu(bn(sigma))``). Lets us measure the pi branch's
    # *functional* contribution by toggling it off on an already-trained block.
    self.pi_enabled: bool = True

pi_weight

pi_weight() -> torch.Tensor

Bounded signed log-space exponent kernel max_exponent * tanh(raw).

Source code in polyweave/layers/sigmapi_conv.py

def pi_weight(self) -> torch.Tensor:
    """Bounded signed log-space exponent kernel ``max_exponent * tanh(raw)``."""
    return self.max_exponent * torch.tanh(self.pi_weight_raw)

exponent_abs_mean

exponent_abs_mean() -> float

Recruitment metric A — mean(|exponent|) over the pi weights.

The geometric product is prod |x| ** w; w = 0 means that factor is |x| ** 0 = 1 (a no-op). So the mean absolute exponent measures how far the learned product departs from doing nothing — it starts ~0 and grows as the block recruits multiplicative structure. Unlike pi_scale this is meaningful even when the product starts at the identity (exp(u) form), because it reads the product's SHAPE, not its amplitude/volume.

Source code in polyweave/layers/sigmapi_conv.py

@torch.no_grad()
def exponent_abs_mean(self) -> float:
    """**Recruitment metric A** — ``mean(|exponent|)`` over the pi weights.

    The geometric product is ``prod |x| ** w``; ``w = 0`` means that factor is
    ``|x| ** 0 = 1`` (a no-op). So the mean absolute exponent measures *how far the
    learned product departs from doing nothing* — it starts ~0 and grows as the
    block recruits multiplicative structure. Unlike ``pi_scale`` this is meaningful
    even when the product starts at the identity (``exp(u)`` form), because it reads
    the product's SHAPE, not its amplitude/volume.
    """
    return self.pi_weight().abs().mean().item()

branch_energy

branch_energy(x: Tensor) -> dict

Recruitment metric B — how much each branch moves the output.

Returns {"sigma_rms", "pi_rms", "pi_share", "pi_effect_postbn"}.

• pi_share = pi_rms / (sigma_rms + pi_rms) is the pre-BN scale share of the pi branch. This OVERSTATES the branch's functional role: the geometric product exp(u) is heavy-tailed, so its RMS is inflated by outliers that BatchNorm then renormalises away. Read with caution.
• pi_effect_postbn is the honest, post-normalisation measure: the relative L2 change in the block's actual output (relu(bn(.)), in eval mode so running stats are used and untouched) when the pi branch is removed, ||relu(bn(sigma+pi)) - relu(bn(sigma))|| / ||relu(bn(sigma+pi))||. Because BatchNorm IS a z-score, this reflects how much the pi branch moves the activations the rest of the network sees — ~0 = pi functionally idle.

Source code in polyweave/layers/sigmapi_conv.py

@torch.no_grad()
def branch_energy(self, x: torch.Tensor) -> dict:
    """**Recruitment metric B** — how much each branch moves the output.

    Returns ``{"sigma_rms", "pi_rms", "pi_share", "pi_effect_postbn"}``.

    * ``pi_share = pi_rms / (sigma_rms + pi_rms)`` is the *pre-BN* scale share of
      the pi branch. This OVERSTATES the branch's functional role: the geometric
      product ``exp(u)`` is heavy-tailed, so its RMS is inflated by outliers that
      BatchNorm then renormalises away. Read with caution.
    * ``pi_effect_postbn`` is the honest, *post-normalisation* measure: the
      relative L2 change in the block's actual output (``relu(bn(.))``, in eval
      mode so running stats are used and untouched) when the pi branch is removed,
      ``||relu(bn(sigma+pi)) - relu(bn(sigma))|| / ||relu(bn(sigma+pi))||``.
      Because BatchNorm IS a z-score, this reflects how much the pi branch moves
      the activations the rest of the network sees — ~0 = pi functionally idle.
    """
    sigma, pi = self._branches(x)
    sigma_rms = sigma.float().pow(2).mean().sqrt().item()
    pi_rms = pi.float().pow(2).mean().sqrt().item()
    denom = sigma_rms + pi_rms

    # Post-BN ablation effect: use eval-mode BN so the SAME running stats apply to
    # both the with-pi and without-pi outputs (and no stats are updated here).
    was_training = self.bn.training
    self.bn.eval()
    y_on = F.relu(self.bn(sigma + pi)).float()
    y_off = F.relu(self.bn(sigma)).float()
    if was_training:
        self.bn.train()
    on_rms = y_on.pow(2).mean().sqrt()
    effect = ((y_on - y_off).pow(2).mean().sqrt() / on_rms).item() if on_rms > 0 else 0.0

    return {
        "sigma_rms": sigma_rms,
        "pi_rms": pi_rms,
        "pi_share": (pi_rms / denom) if denom > 0 else 0.0,
        "pi_effect_postbn": effect,
    }

pi_scale_mean

pi_scale_mean() -> float

Current value of the pi-branch diagnostic exp(pi_scale).mean().

Source code in polyweave/layers/sigmapi_conv.py

@torch.no_grad()
def pi_scale_mean(self) -> float:
    """Current value of the pi-branch diagnostic ``exp(pi_scale).mean()``."""
    return self.pi_scale.exp().mean().item()

SigmaPiLinear

SigmaPiLinear(in_features: int, out_features: int | None = None, *, bias: bool = True, pi_scale_init: float = PI_SCALE_INIT, max_exponent: float = DEFAULT_MAX_EXPONENT, max_log: float = DEFAULT_MAX_LOG, signed_products: bool = False, center_product: bool = False, eps: float = DEFAULT_EPS)

Bases: Module

Sigma-Pi fully-connected layer with a genuine geometric-product pi branch.

Parameters:

Name	Type	Description	Default
in_features	int	size of each input sample's last dimension.	required
out_features	int | None	size of each output sample's last dimension. Defaults to in_features (channels-preserving, like the conv block).	None
bias	bool	whether the sigma (additive) branch carries a bias term (default True). The pi branch has no bias — an additive offset in log space would be a multiplicative constant, which the pi_scale gate already supplies.	True
pi_scale_init	float	initial value of the per-output-feature pi_scale gate (default -2.0 so exp(pi_scale) ≈ 0.135).	PI_SCALE_INIT
max_exponent	float	bound on the magnitude of each log-space exponent; weights are max_exponent * tanh(raw) (default 0.5).	DEFAULT_MAX_EXPONENT
max_log	float	the log-space accumulation is clamped to [-max_log, max_log] before exponentiation (default 6.0).	DEFAULT_MAX_LOG
signed_products	bool	if True, multiply the magnitude product by prod_i sign(x_i) (a true signed product) — a flagged ablation. The default False is magnitude-only, with sign carried by sigma.	False
center_product	bool	if True the pi branch uses expm1(u) (= product - 1) instead of exp(u), so it starts at the multiplicative identity (silent) and the pi_scale gate recovers a clean "volume knob" meaning (default False). The silent-init property holds only when the accumulated u is near zero at init, i.e. for narrow fan-in; a wide in_features sums many small exponents into a non-negligible u and weakens the effect.	False
eps	float	stabiliser inside log(|x| + eps) (default 1e-8).	DEFAULT_EPS

Source code in polyweave/layers/sigmapi_linear.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    *,
    bias: bool = True,
    pi_scale_init: float = PI_SCALE_INIT,
    max_exponent: float = DEFAULT_MAX_EXPONENT,
    max_log: float = DEFAULT_MAX_LOG,
    signed_products: bool = False,
    center_product: bool = False,
    eps: float = DEFAULT_EPS,
) -> None:
    super().__init__()
    out_features = in_features if out_features is None else out_features
    self.in_features = in_features
    self.out_features = out_features
    self.max_exponent = float(max_exponent)
    self.max_log = float(max_log)
    self.signed_products = bool(signed_products)
    self.center_product = bool(center_product)
    self.eps = eps

    # Sigma: signed additive branch (also carries sign information).
    self.sigma = nn.Linear(in_features, out_features, bias=bias)
    # Pi: raw log-space exponent weights, squashed to (-max_exponent, max_exponent).
    # Init near zero so exponents start ~0 (product ~ 1) and the branch is gentle.
    self.pi_weight_raw = nn.Parameter(torch.zeros(out_features, in_features))
    nn.init.normal_(self.pi_weight_raw, std=0.1)
    # Per-output learnable amplitude gate / recruitment diagnostic.
    self.pi_scale = nn.Parameter(torch.full((out_features,), float(pi_scale_init)))

pi_weight

pi_weight() -> torch.Tensor

Bounded signed log-space exponents max_exponent * tanh(raw).

Source code in polyweave/layers/sigmapi_linear.py

def pi_weight(self) -> torch.Tensor:
    """Bounded signed log-space exponents ``max_exponent * tanh(raw)``."""
    return self.max_exponent * torch.tanh(self.pi_weight_raw)

exponent_abs_mean

exponent_abs_mean() -> float

Recruitment metric A — mean(|exponent|) over the pi weights.

The geometric product is prod_i |x_i| ** w_i; w_i = 0 means that factor is |x_i| ** 0 = 1 (a no-op). So the mean absolute exponent measures how far the learned product departs from doing nothing — it starts ~0 and grows as the layer recruits multiplicative structure. Unlike pi_scale, this is meaningful even when the product starts at the identity (exp(u) form), because it reads the product's SHAPE, not its amplitude/volume.

Source code in polyweave/layers/sigmapi_linear.py

@torch.no_grad()
def exponent_abs_mean(self) -> float:
    """**Recruitment metric A** — ``mean(|exponent|)`` over the pi weights.

    The geometric product is ``prod_i |x_i| ** w_i``; ``w_i = 0`` means that
    factor is ``|x_i| ** 0 = 1`` (a no-op). So the mean absolute exponent measures
    *how far the learned product departs from doing nothing* — it starts ~0 and
    grows as the layer recruits multiplicative structure. Unlike ``pi_scale``,
    this is meaningful even when the product starts at the identity (``exp(u)``
    form), because it reads the product's SHAPE, not its amplitude/volume.
    """
    return self.pi_weight().abs().mean().item()

branch_energy

branch_energy(x: Tensor) -> dict

Recruitment metric B — how much each branch moves the output.

Returns {"sigma_rms", "pi_rms", "pi_share"} where pi_share = pi_rms / (sigma_rms + pi_rms) on the given batch: the fraction of the output's scale carried by the multiplicative branch. ~0 = pi branch idle, toward 1 = output dominated by the product. A direct, if noisier, complement to metric A (which reads the weights; this reads the actual activations).

Source code in polyweave/layers/sigmapi_linear.py

@torch.no_grad()
def branch_energy(self, x: torch.Tensor) -> dict:
    """**Recruitment metric B** — how much each branch moves the output.

    Returns ``{"sigma_rms", "pi_rms", "pi_share"}`` where ``pi_share =
    pi_rms / (sigma_rms + pi_rms)`` on the given batch: the fraction of the
    output's scale carried by the multiplicative branch. ~0 = pi branch idle,
    toward 1 = output dominated by the product. A direct, if noisier, complement
    to metric A (which reads the weights; this reads the actual activations).
    """
    sigma, pi = self._branches(x)
    sigma_rms = sigma.float().pow(2).mean().sqrt().item()
    pi_rms = pi.float().pow(2).mean().sqrt().item()
    denom = sigma_rms + pi_rms
    return {
        "sigma_rms": sigma_rms,
        "pi_rms": pi_rms,
        "pi_share": (pi_rms / denom) if denom > 0 else 0.0,
    }

pi_scale_mean

pi_scale_mean() -> float

Current value of the pi-branch diagnostic exp(pi_scale).mean().

Source code in polyweave/layers/sigmapi_linear.py

@torch.no_grad()
def pi_scale_mean(self) -> float:
    """Current value of the pi-branch diagnostic ``exp(pi_scale).mean()``."""
    return self.pi_scale.exp().mean().item()