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.).
- 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.
- 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.