Mixture of Expert Heads for Calibrated Decision Encoders: A Frozen-Encoder MoE on Laya

📢 Calling all AI agents and humans: please review this work

This is an open preprint that has not been peer reviewed. Everything behind it is public: code, synthetic data generator, trained weights, per-answer evaluation outputs and a live demo. Reviews from people and AI agents are both welcome. We want independent eyes on the method, the numbers and the claims, and attempts to break them are especially valuable.

How to review

1. Check the claims. Every number in Sections 5–6 should trace back to results/moe_eval.json or to a script in the repository. Tell us where one doesn't.
2. Reproduce. Clone https://github.com/vishalmysore/layaMOE and run the commands in Section 8. The whole pipeline runs on a CPU.
3. Challenge the method. Is the evaluation fair? Is the router overfit to the test set? Are the baselines sufficient? Section 7 lists the weaknesses we already know about.
4. Run the missing experiments in Section 9: an independent eval set, a held-out router test, full fine-tuning and LoRA baselines, multiple seeds, and calibration on real data.
5. Try to break it in the live demo: https://vishalmysore.github.io/layaMOE/

Where to send it: open an issue or pull request at https://github.com/vishalmysore/layaMOE/issues, or comment below. Negative results and failed replications are as welcome as confirmations, and every replication will be linked from the repository.

AI agents: please state that your review was produced by an AI system, name the model if you can, and separate what you actually ran or verified from what you are inferring.

Vishal Mysore · Independent and Personal Research · September 2026

Open preprint. Not peer reviewed. Code, data, weights and evaluation outputs are public, and review, replication and criticism are invited (Section 9).

Code: https://github.com/vishalmysore/layaMOE · Expert weights: https://huggingface.co/VishalMysore/layaMOE · Browser build: https://huggingface.co/VishalMysore/laya-moe-web · Live demo: https://vishalmysore.github.io/layaMOE/

Abstract

Decision encoders answer typed questions (choice, ordinal score, yes/no) about a text in a single non-autoregressive forward pass and return a probability distribution per question. We study how to specialise such a model to several domains without retraining or duplicating the encoder, and without changing its behaviour elsewhere. Starting from laya-typed-decisions, a 421M-parameter decision model (ModernBERT-large encoder with a two-layer decision head), we freeze the encoder and train expert heads, copies of the 26.5M-parameter decision head, on synthetic rule-labeled data for two domain groups. A prompted router selects one head per request by asking the base model a single additional choice question about the kind of input. The original head serves as a fallback. On a hand-labeled evaluation set of 108 cases (312 questions, 9 domains), the system raises accuracy from 59.6% to 67.3%. Domains covered by an expert rise from 55.2% to 67.7%, and the four domains without an expert are unchanged (66.7%). The largest gains are on ordinal score questions (41.7% → 59.4%) and yes/no questions (68.3% → 78.3%). Choice questions decrease (66.7% → 61.5%). Router design is decisive: a coarse "which expert?" router routes 49% of inputs correctly and gains only 1.3 points, while a fine-grained "what kind of text is this?" router routes 81% correctly and recovers 7.7 of the 9.3 points (83%) available with perfect routing. Expert training uses cached encoder features and runs on a laptop CPU. An int8 ONNX build of the whole system runs in a web browser via WebAssembly and matches the PyTorch accuracy (67.9%). We report negative and cautionary findings, including high-confidence errors inherited from the synthetic labeling rules, and we list the threats to validity that we consider most serious, including a small evaluation set and a router revised after inspecting that set.

---

1. Introduction

Generative language models are commonly used to make operational decisions such as routing a ticket, flagging a risky agent action or rating urgency. The decision is extracted from generated text, which requires parsing and validation, and does not naturally yield a calibrated probability. Decision encoders take a different approach. The input text (the state) and a set of typed questions are encoded together, and the model returns one distribution per question over options defined at request time. Laya [1] is an open example, trained with strictly proper scoring rules so that honest probabilities maximise the reward.

A general decision model is uneven across domains. On our evaluation set, laya-typed-decisions ranges from 83.3% (patient-message routing) to 39.6% (judging the risk of AI-agent actions). The usual remedy, fine-tuning, has two costs in deployment settings where Laya is attractive (self-hosted, edge, in-browser):

1. Size. One fine-tuned 421M-parameter model per domain multiplies memory and download size.
2. Interference. A single model fine-tuned on several new domains can change its behaviour on the domains it already handled well.

Mixture-of-experts (MoE) models [2, 3, 4] address a related problem in large language models by routing tokens to specialised feed-forward experts. That design is poorly matched to decision encoders. Per-token experts specialise in whatever reduces the loss, and nothing makes them align with human-meaningful domains. The parameter count also grows with the number of experts in every layer, and common browser runtimes provide no sparse MoE kernel.

We instead study a sequence-level mixture of expert heads. The encoder is shared and frozen, each expert is a copy of the small decision head, and a router picks one head per request. Our contributions are:

• C1. Architecture. A frozen-encoder mixture of decision heads with a prompted router and a fallback to the original head. By construction, domains routed to the fallback keep the base model's outputs exactly.
• C2. Training recipe. Cached encoder features plus head-only training on rule-generated synthetic data, feasible on a CPU. This includes a learning-rate finding that is easy to get wrong.
• C3. Router study. An empirical comparison showing that routing by fine-grained input kind is far more accurate than routing by expert description when the router is the base decision model itself.
• C4. Evaluation and deployment. Evaluation against perfect-routing ("oracle") and no-routing baselines on hand-labeled data, and an int8 browser deployment whose accuracy matches PyTorch.
• C5. Open artifacts and an explicit call for review, with the specific experiments we believe are needed to confirm or refute these results.

2. Background and related work

Laya. Laya [1] combines a ModernBERT-large encoder 5 with a decision head: a learned type embedding per question type, two pre-norm transformer layers, and a scorer MLP applied at one [MASK] marker per option, followed by a softmax over that question's options. The input is [CLS] <type> question [SEP] [MASK] opt1 [MASK] opt2 … [SEP] state [SEP], so the options are part of the input and new schemas need no retraining. Training uses reinforcement learning against strictly proper scoring rules 6. Probabilities are post-hoc calibrated with temperature scaling [8] per (question type, option count) bucket. laya-typed-decisions is fine-tuned from the base English checkpoint on four synthetic workflows and reports 0.766 accuracy on that benchmark, against 0.362 zero-shot for the base checkpoint [1].

Mixture of experts. The idea of gating between specialised experts dates to Jacobs et al. [2]. Sparsely-gated MoE layers [3] and their simplified top-1 variant [4] scale language models by routing tokens to experts within each layer, with auxiliary load-balancing losses and later router z-losses [9] for stability. Sparse upcycling [10] initialises such experts from a dense checkpoint. Our design differs in three ways. Routing is per sequence. Experts are whole decision heads, not feed-forward blocks. And the router is not a learned gate but the base model answering a question.

Parameter-efficient specialisation. Adapters [11] and LoRA [12] add small trainable modules to a frozen backbone. Task-specific heads on a shared encoder are a standard multi-task pattern. Expert heads can be viewed as full-rank "adapters" placed only after the encoder, which permits caching encoder outputs during training and sharing one encoder pass across heads at inference.

3. Method

3.1 Architecture

Let E be the frozen encoder, and H₀ the original decision head (type embedding, two transformer layers, scorer; 26,512,131 parameters). For a state s and a question q with options o₁…o_k, Laya computes hidden states h = E(x(s, q)) and logits z = H(h) at the k marker positions. An expert Hᵢ is initialised as a copy of H₀ and fine-tuned. At inference:

1. Build the input sequences for the router question q_r and all user questions q₁…q_n, and run E once on the batch.
2. Apply H₀ to the router row to get a distribution over input kinds. Map the most probable kind to an expert via a fixed table. If no expert is mapped, or the probability is below a threshold τ, use H₀.
3. Apply the chosen head to the user-question rows, then that head's own temperatures and a softmax.

Equivalence check. Before training, we verified that E followed by a copied H₀ reproduces the original model's logits exactly (maximum absolute difference 0.0), so the decomposition is exact. A consequence is that inputs routed to H₀ receive exactly the base model's answers.

3.2 Router

The router is a single choice question answered by the base model. We evaluated two designs:

• R1 (expert-level): "What kind of input is this?" with three options: an agent action or a comment needing moderation (→ safety), a support ticket, delivery problem or work email (→ customer_ops), and anything else (→ general). τ = 0.6.
• R2 (kind-level): "What kind of text is this?" with ten concrete kinds (agent action, public comment, support ticket, delivery issue, work email, IT alert, patient message, product review, sales inquiry, other) and a fixed kind → expert table. Unmapped kinds go to general. τ = 0.3.

R2 was designed after R1 failed on the evaluation set (Section 7).

3.3 Expert training

Feature caching. Because E is frozen, E(x(s, q)) is computed once per training item and stored in fp16. The largest deviation from fp32 logits was 3.3 × 10⁻⁴. Training then touches only the head.

Objective. Cross-entropy over the options, plus the ranked probability score on ordinal (score) questions, with weight 1:

L = CE(z, y) + 1[score] · RPS(softmax(z), y), where RPS = Σ_j (F_j − 1[j ≥ y])² / (k − 1)

and F is the predicted cumulative distribution. The RPS term penalises probability mass far from the correct level.

Optimisation. AdamW (weight decay 0.01), batch size 32, linear warm-up over 6% of steps followed by linear decay, gradient clipping at 1.0. Length-bucketed batches. The epoch with the best validation accuracy is kept.

Calibration. After training, one temperature per (type, option-count) bucket with at least 20 validation items, and one per type as a fallback, is fitted by grid search (0.30–3.00) to minimise validation negative log-likelihood.

3.4 Synthetic training data

No public labeled corpus exists for most of our target decisions, such as whether an agent's planned action needs human approval. We therefore generate cases compositionally. Each case is assembled from slots, and its labels are computed from the slots by explicit rules. For agent guardrails the slots are the operation class (read, low-impact write, internal or external communication, delete or overwrite, security change, payment), the environment (production, staging, development, sandbox), the safety net (verified backup, versioning, --dry-run, none) and the scope. For example, a delete in production with no safety net is labeled destructive, needing a human, and critical risk. The same command with --dry-run is non-destructive, safe, and low risk. Analogous generators cover content moderation, support tickets, delivery exceptions and email triage.

Question wordings are paraphrased (three instructions per question and two or three option or level texts), and choice options are shuffled, so a head cannot memorise a single prompt or option position. We generated 500 distinct states per domain, which gives 3,500 question items for the safety expert (agent guardrails and content moderation) and 4,500 for customer_ops (support tickets, delivery exceptions, email triage). About 10% of states, split by state, are held out for validation. Any generated state sharing a word 5-gram with an evaluation case is discarded (101 guardrail and 57 support-ticket states were dropped this way).

4. Experimental setup

Evaluation data. 108 hand-written cases, 12 per domain across 9 domains, with 312 labeled questions (96 choice, 96 score, 120 yes/no). The cases were written and labeled by the author before this work, for an earlier browser port of Laya, and were never used for training. Some labels are marked "hard" (judgement calls). We report all labels.

Metrics. Choice: top option equals the label. Score: most probable level equals the label (exact). Yes/no: P(yes) ≥ 0.5 matches the label.

Systems.

• general: laya-typed-decisions unchanged (equivalently, H₀ for every input).
• oracle: the expert that owns the case's domain. This is the upper bound for perfect routing and equals general for domains without an expert.
• moe: the full system with a router.

Hardware. Intel Core i9-13900H laptop, CPU only, PyTorch 2.14 (fp32). Browser: ONNX Runtime Web 1.30 (WASM, 4 threads) in Microsoft Edge and Chromium on the same machine.

5. Results

5.1 Main result

Table 1. Accuracy on the 312 hand-labeled questions (PyTorch fp32, router R2).

Domain	Expert	general	oracle	moe	routed (R2)
agent-guardrails	safety	39.6	66.7	66.7	safety 12/12
content-moderation	safety	75.0	83.3	75.0	safety 3, general 8, customer_ops 1
delivery-exceptions	customer_ops	47.2	63.9	63.9	customer_ops 11, general 1
email-triage	customer_ops	52.8	58.3	50.0	customer_ops 5, general 5, safety 2
support-tickets	customer_ops	66.7	80.6	83.3	customer_ops 11, general 1
it-incidents	none	61.1	61.1	61.1	general 12/12
patient-messages	none	83.3	83.3	83.3	general 12/12
product-reviews	none	77.8	77.8	77.8	general 9, customer_ops 3
sales-leads	none	41.7	41.7	41.7	general 12/12
covered domains (192)		55.2	70.3	67.7
uncovered domains (120)		66.7	66.7	66.7
All (312)		59.6	68.9	67.3

Table 2. Accuracy by question type.

Type	n	general	oracle	moe
choice	96	66.7	66.7	61.5
score (exact)	96	41.7	59.4	59.4
yes/no	120	68.3	78.3	78.3

5.2 Router comparison

Table 3. Router R1 (expert-level) vs R2 (kind-level). Routing accuracy counts a case as correct when it reaches its domain's expert, or the general head for uncovered domains.

	R1	R2
Routing accuracy, all 108 cases	49.1%	80.6%
covered-domain cases reaching their expert	15/60	42/60
uncovered-domain cases kept on general	38/48	45/48
moe accuracy, all	60.9%	67.3%
moe accuracy, uncovered domains	63.3%	66.7%
share of the oracle gain recovered (9.3 points)	14%	83%

With R1, the base model assigned the safety option only 17% average probability on moderation comments, and routed 6 of 12 patient messages to customer_ops, which reduced patient-message accuracy from 83.3% to 70.8%. R2 kept all patient messages and IT alerts on the general head.

5.3 Expert training

Table 4. Validation accuracy on held-out synthetic data (not the evaluation set).

Expert	train / val items	base head	expert	lr, epochs	CPU time
safety	3,120 / 380	58.4	73.4	6e-4, 6	≈ 45 min
customer_ops	4,011 / 489	57.3	75.9	6e-4, 6	≈ 100 min

Per question type on validation, the largest changes were guardrail yes/no (55.4 → 76.8), guardrail risk score (21.4 → 35.7), email-triage score (48.6 → 91.9), and delivery and email yes/no (77.6 → 100 and 62.2 → 100). Delivery action (choice) barely moved (41.8 → 43.3), which suggests the head cannot separate some of the generator's action classes from frozen features.

Learning rate. Head-only training is sensitive to the learning rate. For the safety expert:

lr	epochs	best val. accuracy
3e-5	4	59.5 (base 58.4)
2e-4	3	65.3
6e-4	6	73.4 (still rising at epoch 6)

At 3e-5 the loss barely changed, which initially looked like a bug. A 64-example overfitting test separated the two explanations: at 3e-4 the head reached 95% training accuracy within 40 steps, at 1e-4 it reached 89%, and at 3e-5 it reached 61%. The head is trainable but needs step sizes an order of magnitude larger than typical full fine-tuning.

5.4 Browser deployment

The encoder and each head are exported as separate ONNX graphs with weight-only int8 quantization: a block-wise 8-bit MatMulNBits operator for weight matrices and per-row int8 for the embedding table. That gives an encoder of about 390 MiB and heads of about 32 MiB each, split into 24 MiB parts with SHA-256 hashes.

Table 5. int8 ONNX vs PyTorch on the full evaluation set.

	PyTorch fp32	ONNX int8
general head	59.6	58.7
moe (R2)	67.3	67.9
routing decisions identical		312/312

The largest per-probability deviation on a four-question spot check was 0.014 for the general head, 0.095 for customer_ops and 0.134 for safety. Fine-tuned heads were more sensitive to quantization than the original head.

In the browser (WASM, 4 threads), a request with three questions and 325–433 input tokens took 3.8–7.0 s for the encoder pass and 0.3–0.8 s for the router plus the chosen head and the general head together. The encoder dominates. Running an additional head costs a small fraction of a full pass.

6. Analysis

Where the experts help. The gains are concentrated in ordinal and binary questions. The base model tends to place its score predictions in the middle of the scale (it rated all twelve agent actions "Medium" risk). The RPS term together with domain data moves mass toward the extremes when the text warrants it.

Why choice accuracy fell. In covered domains, the lost choice answers come from two sources: cases routed to the wrong head (email to safety, moderation to customer_ops) and cases where the synthetic labeling rule disagrees with the human label. An example of the latter is the email generator's rule for when to "delegate" versus "reply".

High-confidence errors. Fine-tuning on rule-generated labels transfers the rules' blind spots, and often with higher confidence. On an input where a parcel had not moved for 12 days and the customer had written three times, the customer_ops expert answered "customer is upset" = no with probability 0.989, while a similar evaluation case is labeled "upset". The generator marked customers as upset only when explicit angry language was present, and the expert learned exactly that rule. Temperature fitting on synthetic validation data cannot detect this, because the validation data shares the same rule.

A safety-relevant miss. For an agent plan to run DROP TABLE on a production database with no backup, the safety expert gives P(destructive) = 0.65–0.68 but P(needs human approval) = 0.23–0.32 (PyTorch and int8 respectively). The confidence gate (0.90) sends this case to a person, but the underlying probability is wrong in a way that matters.

Interference. Because uncovered inputs are routed to the unchanged head, interference is limited to routing errors. With R2, three product reviews were routed to customer_ops and still scored the same in aggregate (77.8%). Measuring interference with a single jointly fine-tuned model is left open (Section 9).

7. Limitations and threats to validity

1. Evaluation size. There are 12 cases per domain. One case is three or four answers, so per-domain differences below about 10 points are within noise. The 312 answers are also not independent: questions about the same case are correlated.
2. Router designed after seeing the test set. R2's kind list was written after R1's failure on the evaluation set. The descriptions are generic and nothing was tuned per case, but this remains test-set adaptation and probably inflates R2's routing accuracy.
3. Single annotator. The evaluation labels were written by the author and were not independently reviewed. Some are explicitly marked as judgement calls.
4. Synthetic training data. The experts learn the author's labeling rules, including their blind spots (Section 6). Validation accuracy on synthetic data (Table 4) says little about real-world accuracy.
5. Single seeds. Each expert was trained once (seed 0). No variance across seeds or data samples is reported.
6. Calibration not evaluated on real data. The temperatures were fitted on synthetic validation data. We report no ECE on the evaluation set.
7. One base model and one hardware setup. All results use laya-typed-decisions on one laptop and one browser configuration.

8. Reproducibility

Everything needed to reproduce the results is public. It runs on a CPU and requires no GPU.

git clone https://github.com/vishalmysore/layaMOE && cd layaMOE
python -m venv .venv && .venv/Scripts/activate      # or source .venv/bin/activate
pip install -r requirements.txt
python scripts/eval_baseline.py                      # general head on data/eval
python scripts/gen_train_data.py                     # synthetic data (seed 1234)
python scripts/cache_features.py --expert safety
python scripts/train_expert.py  --expert safety      # lr 6e-4, 6 epochs (defaults)
python scripts/cache_features.py --expert customer_ops
python scripts/train_expert.py  --expert customer_ops
python scripts/eval_moe.py --out results/moe_eval.json
python scripts/export_web.py && python scripts/eval_onnx.py

Per-answer outputs for all three systems, including router probabilities, are in results/moe_eval.json. Trained heads and the browser build are on Hugging Face (links at the top).

9. Call for peer review

This is an open preprint by an independent author, and it has not been peer reviewed. We would value review of the method, the evaluation and the claims, and especially attempts to break them. The experiments we consider most important:

1. A larger, independently labeled evaluation set, ideally several hundred cases per domain from real traffic, labeled by more than one annotator, with inter-annotator agreement reported. This is the single most important check on every number in this paper.
2. Held-out router evaluation. Freeze R2 and evaluate it on inputs its author never saw. Compare it with a trained router (a small head on E's [CLS] features, trained on labeled input kinds).
3. Baselines we did not run:
   • a single head fine-tuned on all five domains (does one head suffice?);
   • a full fine-tune of the whole model on the same data (the interference question);
   • LoRA [12] on the upper encoder layers, per expert;
   • the base laya checkpoint as the starting point instead of laya-typed-decisions.
4. Calibration on real data. ECE and reliability diagrams of expert heads on held-out human-labeled data, before and after refitting temperatures on that data.
5. Seeds and data variance. Several training seeds and several generator seeds, with confidence intervals on Table 1.
6. Synthetic-data artifacts. A systematic search for high-confidence errors caused by generator rules, like the "upset" case in Section 6, and methods to reduce them (label noise, paraphrased or LLM-rewritten states, mixing in real labeled data).
7. Scaling the number of experts. Router accuracy and uncovered-domain stability as experts are added. Kind-level routing may degrade as kinds multiply.

How to contribute. Please open an issue or pull request at https://github.com/vishalmysore/layaMOE with your results, including negative ones. If you run a replication, we will link to it from the repository, whether or not it agrees with ours.

10. Conclusion

A frozen-encoder mixture of decision heads is a cheap way to specialise a calibrated decision encoder. Experts cost 26.5M parameters each and train on a laptop CPU, and domains without an expert keep the base model's exact behaviour. On a small hand-labeled evaluation, two experts and a prompted router raise accuracy from 59.6% to 67.3%, with most of the gain on ordinal and binary questions, and the system runs unchanged in a browser at int8. The router is the component that decides whether the approach works. Routing by concrete input kind recovered most of the available gain, and routing by expert description did not. The evaluation is small, the router was revised after seeing it, and the experts inherit the blind spots of their synthetic labels. We present the results as a starting point for independent verification.

---

Disclaimer. All figures come from controlled experimental runs on a laptop CPU and a single browser configuration. Accuracy is measured on a small hand-labeled evaluation set (108 cases, 312 questions, 9 domains, labeled by the author and not independently reviewed), and the expert heads were trained on synthetic, rule-generated data. Figures quoted for Laya itself come from its model cards. These are experimental results, not production guarantees.

Acknowledgements. Laya and laya-typed-decisions are by Convai Innovations (Apache-2.0). ModernBERT is by Answer.AI and LightOn (Apache-2.0). This work is unofficial and not affiliated with either. Code and experts are released under Apache-2.0.

References

[1] Convai Innovations. Laya: calibrated, non-autoregressive System 1 decision models (model cards and code). https://huggingface.co/convaiinnovations/laya · https://github.com/NandhaKishorM/laya, 2026.

[2] R. A. Jacobs, M. I. Jordan, S. J. Nowlan, G. E. Hinton. Adaptive Mixtures of Local Experts. Neural Computation, 3(1), 1991.

[3] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, J. Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR, 2017.

[4] W. Fedus, B. Zoph, N. Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR, 2022.

[5] B. Warner et al. Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (ModernBERT). arXiv:2412.13663, 2024.

[6] T. Gneiting, A. E. Raftery. Strictly Proper Scoring Rules, Prediction, and Estimation. Journal of the American Statistical Association, 102(477), 2007.

[7] E. S. Epstein. A Scoring System for Probability Forecasts of Ranked Categories. Journal of Applied Meteorology, 8(6), 1969.

[8] C. Guo, G. Pleiss, Y. Sun, K. Q. Weinberger. On Calibration of Modern Neural Networks. ICML, 2017.

[9] B. Zoph et al. ST-MoE: Designing Stable and Transferable Sparse Expert Models. arXiv:2202.08906, 2022.

[10] A. Komatsuzaki et al. Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints. ICLR, 2023.

[11] N. Houlsby et al. Parameter-Efficient Transfer Learning for NLP. ICML, 2019.

[12] E. J. Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. ICLR, 2022.

---

Appendix A. Demo screenshots

The live demo shows the routing decision and the expert's answers next to the general head's.

Figure A1. DROP TABLE on production: the router reads "agent_action" (74%) and the safety expert answers. All confidences are below 0.90, so the case is escalated. P(needs human) is wrong (Section 6).

Figure A2. A support ticket: the customer_ops expert is confident and correct, while the general head put 55% on "angry" for a calm message.

Figure A3. A high-confidence error inherited from the synthetic labeling rules: P(upset) = 0.011 for a customer who has written three times about a lost parcel.

Figure A4. A sales lead: no expert is mapped, so the unchanged general head answers.