pyjuice.parameterizations.LowRank
- class pyjuice.parameterizations.LowRank(rank: int, form: str = 'exp_sum', apply_z_correction: bool = True, tie_external: bool = False, flow_precision: str = 'bf16')
A per-sample multiplicative low-rank scale on the sum node’s parameters, supplied externally at call time.
A node vector of M nodes over N children takes two tensors per sample b, A_b of [M, k] and B_b of [k, N]; they make a scale g_b[i, j] of every (node, child) pair, and the effective parameters are
theta_b[i, j] = theta_shared[i, j] * g_b[i, j] / Z_b[i] Z_b[i] = sum_j theta_shared[i, j] * g_b[i, j]
normalized over the children of node i, as
BlockScale’ are. form says how the factors make the scale:"exp_sum": g_b[i, j] = exp( sum_r A_b[i, r] * B_b[r, j] ) # the log-scale is the rank-k product "sum_exp": g_b[i, j] = sum_r exp( A_b[i, r] ) * exp( B_b[r, j] ) # a mixture of k rank-one scales
“exp_sum” gives every pair an exponent of its own, so it costs one exp per (sample, node, child) and the samples share no matmul; “sum_exp” is k matmuls with the shared parameters. Both are defined on the pairs the node vector’s edge_ids connect; a pair without parameters has none here either.
Tensor layout. Two float32 tensors per node vector, in this order:
A: [B, ns.num_nodes, k]
B: [B, k, ns.num_ch_nodes]
whatever blocks the node vector is stored in. A registered group (TensorCircuit.register_external_params_group) concatenates both along axis 1: A over the members’ nodes, B over their ranks.
Training. The backward writes the element and parameter flows itself, and d LL / d A, d LL / d B into the gradient buffers. With apply_z_correction it also accumulates the denominator flow’s statistic, and the M-step is the conditional theta <- normalize(theta * F+ / F-) (
compute_em_correction()): Z_b depends on theta, so normalize(F+) alone is not EM here, as for BlockScale. What is kept isW[i, j] = sum_b f_b[i] * g_b[i, j] / Z_b[i] F-[i, j] = theta[i, j] * W[i, j]
in the parameter flows’ layout (it has no smaller form here), FREE OF THE PARAMETERS as BlockScale’s gate-space statistic is, and multiplied by them when F- is read. That is what an optimizer averaging the flows over steps needs (Anemone’s momentum): the update is then normalize(avg F+ / avg W), the M-step of the averaged statistics. Kept as F- itself, the average would carry the parameters of past steps against the current ones in theta * F+ / F- – a parameter that grew is pushed further – and the training is another: MEASURED, a low-rank scale built to equal a BlockScale gate, Anemone with momentum 0.9, parameters against BlockScale’s after 16 steps 3e-2 apart that way, and the same to rounding this way and without momentum either way.
- Note:
the passes: “exp_sum” through its fused kernels (lowrank/kernels/c/lowrank_exp_sum.cu: the scale made tile by tile on tensor cores and used where it is made, never stored) where a node vector’s parameters lie block-major in one run and the rank is at most 64 (lowrank.forward.exp_sum_fused) – under the default precision the factors’ products in bf16 (A B from bf16 hi + lo, ~2^-16 of each product; the gradients’ products one bf16 term each), under precision = “fp32” in three-term TF32 (~fp32) –, “sum_exp” through the plain sum-layer kernels on a virtual batch (lowrank.virtual) where it serves; elsewhere plain torch – dense per node vector, exp_sum over [B, M, N] in chunks of samples, sum_exp by float64 matmuls (a sample’s child values spread over hundreds of nats, more than float32 holds beside one shift a sample). Correct on every structure.
- Parameters:
rank (int) – the rank k
form (str) – “exp_sum” or “sum_exp” (above)
apply_z_correction (bool) – accumulate the denominator flow and take the conditional M-step (default); False leaves the plain normalize(F+) and allocates nothing
tie_external (bool) – share one pair of tensors across every copy of a tied node, instead of one per copy
flow_precision (str) – “bf16” (default) or “fp8”, an OPT-IN for “sum_exp”: the pass’s end’s two products – the parameter flows’ theta * (Ud Vd^T) and the denominators’ Un Vt^T, over every (copy, sample, rank) of the pass – from e4m3 operands (a power-of-2 scale a row, fp32 accumulation) where they are bf16 by default (precision “auto”, the products large enough: lowrank.virtual); about twice the bf16 products’ speed, the statistics’ entries about 16x less accurate (MEASURED against float64 on an HMM of 1024 states, rank 16, batch 64: the M-step’s F+ / F- a median 3e-4 .. 1.4e-2 and at most 2e-3 .. 0.18 off as the factors’ spread goes from 0.3 to 4 nats, against bf16’s 2e-5 .. 8e-4 and 1e-4 .. 1e-2; no entry made 0). Under precision “fp32” the products stay fp32. The gradients are not affected.