Kai9987kai/supermix-v56-latent-state-reasoner

Supermix v56 — Latent State Reasoner

An 808,626-parameter model that learns a 4-step composed modular-arithmetic algorithm, reaching 0.9740 held-out accuracy where the previous architecture on the identical task reached 0.1710 against a 0.1430 majority-class floor.

Scope, before any number. This model solves one synthetic task. It takes a fixed 128-dimensional vector encoding a start digit and four (operation, operand) pairs, and returns a digit 0–9. It is not a language model — it has no tokenizer and cannot process text. It is a research artifact about inductive bias in composed-function learning, and nothing here transfers to any other task.

The task, and why it is hard

x encodes a start digit 0–9 and four operations (add / multiply / subtract, operands 1–9); the label is the composed result mod 10. There are 5,314,410 distinct inputs, so a 12,000-example training set covers 0.23% of them and only 4 of the 1,000 canonical test items appear in training. It is a generalisation test.

The task is close to non-decomposable. Bayes-optimal accuracy measured over 4,000,000 samples, by which part of the input a predictor may see:

information given	Bayes accuracy	Bayes cross-entropy
nothing (majority class)	0.1415	2.2843
the last operation only	0.1720	2.1240
the last two operations	0.2013	1.9810
the last three operations	0.2389	1.8406
all four operations, no start digit	0.4105	1.3162
start + first three operations	0.2521	1.9742
start + all four operations	1.0000	0.0000

Every proper subset is worth almost nothing and the complete input is worth everything, so gradient descent gets no partial credit and has no path from a partial solution to the exact one. A flat MLP, a slot-tokenised transformer, a GRU recurrence and a latent state machine all stall between 0.24 and 0.27 — and raising the training set from 12,000 to 96,000 examples does not move it.

Read any accuracy against this table. 0.4105 is the ceiling for any model that ignores the start digit, so clearing it is the first evidence a model composes the whole chain rather than reading its tail.

Architecture

Reuses the v53 MiMoMix components — hybrid sliding-window/global attention with learnable sinks, auxiliary-loss-free sparse MoE, and a recursive thinking core with ACT halting and a calibrated verifier — and adds an explicit latent state machine between the trunk and the answer:

• Slot tokenisation. The same flat vector is read as a short sequence. No input dimension is discarded; the unused tail is projected into a context token.
• Position equivariance. All operator slots share one encoder, turning n_blocks × n_operations maps to learn into n_operations.
• Row-stochastic transitions. The state is a distribution over 10 latent states; each slot emits a 10×10 row-stochastic matrix, composed in log space.
• Identity initialisation. A product of near-uniform stochastic matrices mixes to uniform and starves the gradient reaching the first operator, so operator logits are biased toward the identity at init.
• A crispness prior. The maps being composed are deterministic functions, so an entropy penalty pushes transition rows toward one-hot.

Multi-token prediction and speculative decoding are deliberately absent: the model emits one answer per input, so there is no next token to draft.

Measured results

Evaluated on the untouched held-out set, make_chained_task(1000, seed=52):

model	params	training data	held-out accuracy
majority-class constant	—	—	0.1430
previous architecture (v51)	2,245,715	12,000 × 4 epochs	0.1710
v56, matched protocol	808,626	same 12,000 × 4 epochs	0.2410
v56, curriculum (this checkpoint)	808,626	160,000	0.9740

The matched row changes only the model — identical examples, seed and epoch budget, with 2.8× fewer parameters. The curriculum row also changes the training recipe and is reported separately for that reason. Calibration on this checkpoint: NLL 0.0680, ECE 0.0194.

Paired promotion gate

40,000 paired samples over 20 fresh cohorts, no seed reused from the training split or any earlier evaluation (v56_promotion_gate.json):

arm	accuracy	95% Wilson CI
v56 (this checkpoint)	0.9762	[0.9747, 0.9776]
previous architecture	0.1718	[0.1681, 0.1755]
majority-class floor	0.1393	—

20 seed wins to 0 (sign test p = 1.9e-6). McNemar on discordant pairs: 32,288 candidate-only versus 112 baseline-only, p below floating-point resolution.

What moved the number

Ablation at an identical budget:

change	accuracy
first curriculum (identity ops pinned to a trailing prefix)	0.9220
+ active operations spread over random slots	0.9480
+ operator-entropy prior (this checkpoint)	0.9740
+ no positional embeddings (full equivariance)	0.8210
+ strictly slot-local operator	0.6920

The prefix curriculum starved the last slot — 83% of that model's errors first diverged at exactly the final step. Note the last two rows: removing positional embeddings for full equivariance hurt, the opposite of what the design argument predicted. The measurement is recorded over the argument.

Usage

The architecture is custom, not a transformers model, so the modules ship with the weights:

pip install torch huggingface_hub
python example.py

import torch, reasoner_curriculum as rc
from mimomix_reasoner import load_reasoner

model, payload = load_reasoner("v56b_randslots_entropy.pt")

# ((((7*3)+8)-5)*6) mod 10 — ops: 0=add, 1=multiply, 2=subtract
x, truth = rc.encode_chain(
    torch.tensor([7]), torch.tensor([[1, 0, 2, 1]]), torch.tensor([[3, 8, 5, 6]]),
    generator=torch.Generator().manual_seed(0),
)
with torch.no_grad():
    out = model(x)
print(int(out.logits.argmax(-1)), "expected", int(truth[0]))

Reading the reasoning trace

out.state_trace holds the latent state after every operation, and out.operator_log_probs the learned transition matrix per slot. example.py prints what the model would answer at each step:

((((7 * 3) + 8) - 5) * 6) mod 10
  prediction 4   true 4   correct   confidence 1.000
    after multiply 3  would answer 1   true 1   ok
    after add      8  would answer 9   true 9   ok
    after subtract 5  would answer 4   true 4   ok
    after multiply 6  would answer 4   true 4   ok

Two cautions on reading it. Latent state indices are not answers — the class head is a linear map over the whole distribution, not a per-state lookup, so no "state s means answer a" table would be truthful; the only faithful decoding is running the real head on the real state, which is what the example does. And the start row is not graded: the head is trained on final states, so reading it on the initial state is out of distribution.

Limitations

• One synthetic task. Not language, not general arithmetic, not transferable.
• Fixed 128-dimensional input with a fixed 4-operation structure. Operands outside 1–9 or a start outside 0–9 are unrepresentable.
• Residual errors concentrate at the final step. Of the errors that remain, 93.2% first diverge at the last operation. Mean confidence separates correct from wrong at 0.9734 versus 0.6504, so what is left is mostly uncertain rather than confidently wrong.
• The curriculum is a recipe difference. Only the matched protocol (0.2410 vs 0.1710) isolates the architecture.

What this does not prove

• That the architecture is better in general. The latent state machine has an inductive bias matched to composed maps over a small state set. That is the right bias for this task family and says nothing about any other.
• That the verifier is semantic. p(correct) is supervised against the model's own correctness on this task. It ranks this task's answers and is not a general self-knowledge signal.
• That the state trace is an explanation. It is the model's arithmetic, faithfully reported — a test pins that the displayed trace composes to the displayed answer. Faithful is not the same as interpretable.
• That reported latency is a product latency. CPU-only, batched, one machine.

Citation and source

Part of the Supermix project, built on the v53 MiMoMix line. Design notes: docs/V56_LATENT_STATE_REASONER.md.

Companion model: supermix-v57-mimomix-talk, the same MiMoMix stack trained as a language model.

MIT licensed.