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 when sigmoid ≥ 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. See temperature_calibration in 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.

Downloads last month

-

Downloads are not tracked for this model. How to track
Safetensors
Model size
0.4B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for attentionai/laya-extended

Finetuned
(80)
this model

Space using attentionai/laya-extended 1