mixle.engines.lns_nn module¶
Neural-network operations in a logarithmic number system.
Built on mixle.engines.lns.LogNumberSystem. Two families of op benefit from LNS, both because
they fight transcendentals (exp/log), not BLAS:
softmax / cross-entropy / log-softmax: the normalizer is a
logsumexpover logits (the LM head over the vocabulary, a classifier over classes, an attention/MoE-router softmax). The integerlogsumexpreplaces theexp``+``logwith integermax+ a LUT (~2x measured). The model’s logits are quantized by the same log step; only softmax-back-to-linear still needs anexp.sum-product circuits / probabilistic circuits: the whole forward pass is sums and products of probabilities, so in LNS every product node is an integer ADD and every sum node an integer
logadd. The entire network runs in integer log-space, not only the normalizer.
These are inference and scoring operations. The gradient path stays in floating autograd, while LNS provides a compact integer representation for log-space math.
- log_softmax(logits, lns, axis=-1)[source]
Compute log-softmax through the integer log-partition.
- softmax(logits, lns, axis=-1)[source]
Softmax with the normalizer computed in LNS; back-to-linear needs one
exp(for attention.V etc.).Renormalized (MXR-080-0141): exponentiating the LNS-approximate log-softmax does not, by itself, sum to exactly 1 – each class’s rounding error is independent of the others, so they do not cancel (100 equal logits summed to
1.00518pre-fix), which breaks the basic probability-simplex contract a function namedsoftmaxhas to satisfy. Dividing by the actual computed sum costs one reduction and trades a sliver of the raw approximation’s “purity” for an exactly-normalized guarantee.
- cross_entropy(logits, targets, lns, axis=-1)[source]
Mean negative log-likelihood
mean(logsumexp(logits) - logit[target])via the integer normalizer.The LM / classifier loss: the log-partition over the vocab/classes is an integer
logsumexp; the target logit is gathered from the same quantized logits, so the loss is integer until the final scale.Targets must be exact, in-range integer class indices (MXR-080-0141: no negative-index wraparound, no silent fractional truncation) – validated identically whether or not the compiled fused path is taken (MXR-080-0140: the compiled path has boundscheck disabled, so an unvalidated target is a genuine out-of-bounds read, not merely a wrong answer). A target whose own logit quantized to the reserved “log of an exact zero” sentinel has a mathematically infinite loss – and so does the batch mean – so that is returned directly rather than ever subtracting the sentinel as a plain code (MXR-080-0157: that subtraction is exactly the kind of extreme-code-range operation that overflows int64, in both this path and the compiled one).
- class SumProductCircuit(nodes)[source]
Bases:
objectA probabilistic circuit evaluated entirely in integer log-space (product=add, sum=logadd).
nodesis a topologically ordered list (children before parents), each a tuple:("leaf", leaf_id)– an input whose log-value is supplied at evaluation,("product", [child indices])– log-output = sum of children (integer ADD),("sum", [child indices], [log_weights])– log-output = logsumexp of weighted children.
The root is the last node.
evaluate_lnsruns the whole forward pass on integers;evaluate_floatis the float64 reference. Leaf values may be scalars or arrays (the forward broadcasts).- evaluate_lns(lns, leaf_values)[source]
Evaluate the circuit with log-number-system arithmetic.