pyjuice.parameterizations.BlockScale
- class pyjuice.parameterizations.BlockScale(block_size: int | None = None, ch_block_size: int | None = None, apply_z_correction: bool = True, tie_external: bool = False)
A per-sample multiplicative gate on each parameter block, supplied externally at call time.
Every gate owns one scalar per sample. Writing g(n,c) for the gate covering the edge from node n to child c, the effective parameters for batch element b are
theta_b[n, c] = phi[b, g(n,c)] * theta_shared[n, c] / Z_b[n] Z_b[n] = sum_c phi[b, g(n,c)] * theta_shared[n, c]
normalized over the children of n. Factoring out each gate’s share of the node’s mass, sigma[g,n] = sum_{c in g} theta_shared[n,c], shows what the gate actually does:
theta_b[n,c] = [ phi[b,g] sigma[g,n] / sum_{g'} phi[b,g'] sigma[g',n] ] * ( theta_shared[n,c] / sigma[g,n] )– it reweights the gate-level mixture weights per sample and leaves each gate’s internal conditional untouched. Two consequences follow directly:
it is an exact no-op when a node has only one gate, since the reweighting then cancels against the normalizer. Whether a node has several gates is entirely the caller’s choice of gate sizes; a dense transition with one edge block and the default gate size does nothing at all.
a gate that does not depend on the parent node is equivalent to multiplying the child’s value, i.e. to soft evidence on the children, which PyJuice already provides. The capability that is new here is a gate that depends on the parent group and the child group.
Why it is cheap. Because phi is constant within a gate it pulls out of the block matmul, so in log space it folds into the child values as an elementwise add – the dense per-sample parameter matrix is never materialized, and the tensor is [B, num_gates] rather than anything shaped like the parameters. For a 1024-state layer at block_size = 32 that is 1 MB at batch 256.
Tensor layout. One float32 tensor of LOG gates:
phi: [B, Nk, Ck], a dense grid with Nk = ns.num_nodes / block_size gates along the node axis and Ck = ns.num_ch_nodes / ch_block_size along the child axis. phi[b, i, j] scales the parameters from child gate-block j into node gate-block i, for sample b.
The grid is indexed by the GATE’s own block sizes and the layer’s node counts – it does not depend on ns.block_size, so asking for gates of a given size gives a grid of exactly that many whatever blocking pyjuice happens to use inside. Entries for (node block, child block) pairs the layer does not connect are ignored, and their gradients come back zero.
Log space keeps the gate positive for free and makes the fold an add rather than a multiply.
- Note:
the gate tile is deliberately NOT forced to equal the parameter block. Refining along the CHILD axis is free (the fold is elementwise in the child index), while refining along the NODE axis costs one reduction per subgroup, because a single staged tile can only serve nodes that share a gate. Shrinking ns.block_size to obtain a finer gate would instead tax the whole layer: the node tile of the standard kernels collapses with it, below the 16 rows an efficient MMA wants, and the compiled index tensors grow as K^2 / block_size.
- Parameters:
block_size (Optional[int]) – the GATE’s block size – how many nodes one gate spans, NOT the node’s own block_size. None means the node’s, i.e. one gate per parameter block, and it must divide the node’s. Only the full block size is supported so far; splitting the node axis needs a separate kernel.
ch_block_size (Optional[int]) – the GATE’s child block size – how many children one gate spans. None means the node’s own. Any divisor is allowed and costs nothing, which is why this is the axis to refine.
apply_z_correction (bool) – make EM exact under a live gate. The gate makes the normalizer Z = sum_c phi * theta depend on theta, so pyjuice’s plain M-step theta <- normalize(F+) solves a stationarity condition that is missing the term sum_b f_b * theta_b. With this set, the backward also accumulates that term as a DENOMINATOR flow F- and the M-step becomes the conditional theta <- normalize(theta * F+ / F-), whose fixed point is the true one (and whose F+ - F- is d LL / d log theta). Costs a second flow accumulator – kept in gate space, ch_block_size times smaller than param_flows – and ~1% of an EM step. On by default; with False the layer allocates nothing and takes the plain M-step. Not implemented for a gate finer than ns.block_size along the node axis (the backward raises): pass False there. Supported by FullBatchEM / MiniBatchEM / Anemone; the gradient optimizers warn, since their descent direction is uncorrected.
tie_external (bool) – share one gate tensor across every copy of a tied node, instead of one per copy.