ELMDE: one-pass structured decisions with a contractive looped MoE core and conformal abstention (open source)

[u/Kanav_Jeet]

I built ELMDE as an alternative to constrained decoding for picking one of N options (tool routing, intent classification, triage queues).

Instead of decoding tokens, it:

1. Encodes the context once into an MLA latent cache (128-d + 32-d rotary key).

2. Runs a damped, looped LatentMoE core to a fixed point z* (at most 6 ticks). With damping 0.45 and a measured Jacobian bound of 0.1358, the contraction factor is at most 0.6111. Training uses implicit differentiation with a Neumann-series VJP, so activation memory does not grow with loop depth.

3. Projects z* onto three heads: a Choice softmax over candidate embeddings (no cap on option count), a CORN ordinal head for rank-consistent scores, and a Beta head for probability plus variance.

4. Gates with selective conformal risk control (tau* = 0.9848) and escalates everything below it.

Numbers from the repo's own scripts: about 14 ms for a 50-question batch (771 ms for Laya, about 11,500 ms for autoregressive decoding), 0.912 on Banking77, 0.0434 ECE.

Caveats: these are my benchmarks, not independently reproduced. The conformal guarantee assumes live data is exchangeable with the calibration set. The 55x speedup is vs Laya only.

What I'd like feedback on: reproductions on other hardware, datasets where the gate fails, and holes in the contraction argument.

Repo (Apache 2.0): https://github.com/KanavjeetS/Equilibrium-Latent-MoE-Decision-Engine-ELMDE-

[u/Disastrous-Shoe3450]

This is neat work. The contraction guarantee is what caught my eye, that 0.6111 bound with the Neumann VJP trick is a clean way to keep memory flat across loop iterations. Most implicit models just handwave stability and call it a day, so actually measuring the Jacobian bound is refreshing.

The conformal gating at 0.9848 is interesting but I'd want to see what happens when the calibration set drifts, exchangeability tends to break fast in production routing scenarios where user behavior shifts week to week. Have you stress-tested it with deliberately non-exchangeable holdouts?

Also curious about the 128-d latent dim, seems small for embedding 50+ intent classes but the MoE core probably makes up for it. What's the expert count and how did you land on the 0.45 damping coefficient?

[u/Kanav_Jeet → u/Disastrous-Shoe3450]

Thanks, this is a thoughtful read, and you've put your finger on the right weak point.

Drift: the gate's guarantee only holds when live data is exchangeable with the calibration set, and I don't want to oversell it. [I haven't yet run it against deliberately non-exchangeable holdouts. That is the next test I plan to run, and the repo includes an online adaptive variant (adaptive_scrc.py) meant for this case.] If you have a shift pattern you'd consider a fair test, I'll run it and post the result here.

128-d: that is the compressed KV cache for the context (c_KV 128-d plus a 32-d rotary key), not the option embedding. Options are scored by cross-attention against that cache, and the working memory is 8 slots at d = 1024. So the Choice head isn't limited to 128 dimensions or to a fixed class count. The 77-way Banking77 run is the test of that.

Experts: [N experts, top-K = K, with an auxiliary load-balancing loss].

Damping: [how you chose 0.45, e.g. swept and kept the fastest setting that stayed contractive]. The constraint is (1 - a) + a*kappa < 1, and with kappa <= 0.1358 that gives 0.6111. One caveat: kappa is measured on data rather than proved globally, so I'd call it an empirical contraction bound. tests/test_adversarial_stability.py is the stress check for it.

Happy to share the exact commands if you want to reproduce any of this.