Laya-Extended — Recurrent Latent Reasoning
A drop-in extension of Laya that adds a weight-tied recurrent reasoning core with adaptive halting on top of a (mostly frozen) ModernBERT-large backbone. Laya is a fast, non-autoregressive decision engine: given a state (any text, ticket, email or JSON) and typed questions, it returns calibrated answers in a single parallel forward pass. This checkpoint keeps that speed but lets the model iterate on hard problems — thinking in latent space for up to 3 steps and stopping early when confident.
Honest framing. This is not a leaderboard SOTA result. On the repo's own benchmark it is at accuracy parity with base Laya (0.81 vs 0.82) with better calibration (ECE 0.078 vs 0.179), plus a demonstrably working reasoning loop. The contribution is reasoning capability added at no accuracy cost and improved calibration, verified by ablation — not a claim of beating the field.
Metrics
Published benchmark — mteb/amazon_massive_intent, English, 20-option choice, adaptive halt
(same harness/protocol as the base-model baseline: research/eval/laya_eval.py, per-lang 100 rows):
| model | accuracy | ECE ↓ | macro-F1 |
|---|---|---|---|
base convaiinnovations/laya |
0.82 | 0.179 | 0.788 |
| laya-extended (this) | 0.81 | 0.078 | 0.762 |
Accuracy is at parity (inside the ±0.02 regression tolerance); calibration is halved. Accuracy is temperature-invariant — recalibration only moved ECE, not accuracy.
Reasoning gains — in-domain held-out decision set (agentic1200, synthetic, n=180). Forcing exactly K core
steps shows the loop measurably corrects hard cases, and the learned halt policy captures it:
| slice | K=1 | K=2 | K=3 | as-halted |
|---|---|---|---|---|
| overall (n=180) | 0.511 | 0.556 | 0.572 | 0.567 |
| "flip" subset — step-0 answer wrong (n=154) | 0.461 | 0.513 | 0.532 | 0.526 |
multi-option choice (n=72) |
0.417 | 0.569 | 0.583 | — |
| non-flip (already right at step 0) | 0.808 | 0.808 | 0.808 | 0.808 |
+7.1 pts on hard flip items and +16.6 pts on multi-option choice from a single→triple reasoning step,
with no regression on easy items. Mean steps actually taken at inference: 2.87. The ordinal
score question type does not benefit from depth.
Architecture
- Encoder:
answerdotai/ModernBERT-large(28 layers, d=1024). Lower 21 layers frozen; top 7 fine-tuned so reasoning has capacity to adapt features. - Readout head: 2-layer Transformer encoder (runs once) → initial reasoning state
s₀. - Recurrent core: 1 weight-tied Transformer layer, looped up to
max_steps=3(min_steps=1). Each step re-injects the encoded problem:s ← block(norm(s) + inject(e)). - Heads:
scorer(→ per-option logits → softmax),halt_head(→ P(stop); halts whensigmoid ≥ 0.5),act_head. Confidence is read from the final halted distribution. - Calibration: per-
(qtype, option-count)temperature scaling (NLL-min, bounded [0.5, 5]), refit on this checkpoint's own logits. Seetemperature_calibrationin the config.
Question types: choice (pick one), score (0..N ordinal), noul (true/false).
⚠️ Requires the reasoning-enabled laya code
The recurrent core is newer than the released laya package (0.3.20) — a stock pip install laya
has only the single-pass forward and cannot load or run this checkpoint. Use the reasoning-enabled
source, which is vendored for reference in the companion demo Space. Loading with matching code:
import laya
agent = laya.load("attentionai/laya-extended") # add token=... while private
out = agent.system_one(
"p99 latency on checkout jumped to 4.2s; Redis cache hit-rate fell 94%→11% one minute earlier.",
{"cause": {"type": "choice",
"instructions": "Most likely root cause?",
"criteria": {"deploy": "the recent deploy",
"database": "the database is overloaded",
"cache": "the cache eviction caused a cold-cache stampede",
"load_balancer": "load balancer misrouting"}}},
)
ans = out["answers"]["cause"]
print(ans["choice"], ans["answer_confidence"])
Training
Process supervision with deep supervision over the K-step loop: a per-step answer loss (a strictly
proper scoring rule: log + spherical + RPS), a cosine distillation loss pulling each step's latent
toward a teacher "thought" embedding, and a small BCE halting loss. See the research/ docs in the
Laya repository for the full method.
Intended use & limitations
Fast, calibrated typed decisions (triage, routing, classification, guardrail-style true/false).
Tested in English only here; the reasoning gain is strongest on multi-option choice and weakest
(none) on ordinal score. Parity — not SOTA — against the base model.
Base model: convaiinnovations/laya · encoder: answerdotai/ModernBERT-large.
Model tree for attentionai/laya-extended
Base model
convaiinnovations/laya