Laya-TypedDecisions-MXFP8

Laya-Typed-Decisions-MXFP8

This is the Typed-Decisions encoder backbone of convaiinnovations/laya (ModernBERT-large, 421M params, 1024 hidden size, 28 layers, 1024 token context), quantized to MXFP8 for MLX on Apple Silicon.

What's included -- and what's not

Laya is not a plain encoder: it is ModernBERT-large (this repo) plus a small from-scratch decision head (a 2-layer torch.nn.TransformerEncoder, an option-marker scorer, and an act/escalate head) defined in Laya's own rl_common.py, trained separately with RL (RLCD). Only the shared bidirectional encoder backbone is quantized and published here -- the decision head is a tiny (~15M param) uncompiled torch module that gains nothing from MLX quantization, and porting it to MLX would not make Laya's laya.load(...) / RLAgent API usable anyway (that API expects a torch model directory as-is).

To use Laya's actual decision-making API, keep using the upstream convaiinnovations/laya repo. This repo is for people who want the ModernBERT/mmBERT encoder only (e.g. as a quantized general-purpose feature extractor, or to build a custom MLX head on top) at a fraction of the memory footprint.

• Size on disk: ~407 MB
• Architecture: ModernBertModel (via mlx-embeddings)
• Quantized with: mlx-embeddings (nn.quantize(..., mode="mxfp8"), group size 32)

Other quantizations

• English MXFP4
• English MXFP8
• Multilingual MXFP4
• Multilingual MXFP8
• Typed-Decisions MXFP4

Use with mlx-embeddings

pip install mlx-embeddings

import mlx.core as mx
import mlx.nn as nn
from safetensors import safe_open
from transformers import AutoTokenizer
from mlx_embeddings.models.modernbert import ModelArgs, ModernBertModel
import json

mlx_dir = "sahilchachra/Laya-Typed-Decisions-MXFP8"  # or a local snapshot_download() path
with open(f"{mlx_dir}/config.json") as f:
    cfg = json.load(f)
args = ModelArgs(model_type=cfg["model_type"], vocab_size=cfg["vocab_size"], hidden_size=cfg["hidden_size"],
    num_hidden_layers=cfg["num_hidden_layers"], intermediate_size=cfg["intermediate_size"],
    num_attention_heads=cfg["num_attention_heads"], max_position_embeddings=cfg.get("max_position_embeddings"),
    norm_eps=cfg.get("layer_norm_eps", 1e-5), attention_bias=cfg.get("attention_bias", False),
    global_attn_every_n_layers=cfg.get("global_attn_every_n_layers", 3), local_attention=cfg.get("local_attention", 128))
rp = cfg.get("rope_parameters")
if rp:
    args.global_rope_theta = rp["full_attention"]["rope_theta"]
    args.local_rope_theta = rp["sliding_attention"]["rope_theta"]

model = ModernBertModel(args)
weights = {}
with safe_open(f"{mlx_dir}/model.safetensors", framework="numpy") as f:
    for k in f.keys():
        weights[k] = mx.array(f.get_tensor(k))
qcfg = cfg["quantization"]
nn.quantize(model, group_size=qcfg["group_size"], bits=qcfg["bits"], mode=qcfg["mode"],
            class_predicate=lambda p, m: hasattr(m, "to_quantized") and f"{p}.scales" in weights)
model.load_weights(list(weights.items()), strict=True)
mx.eval(model.parameters())

tok = AutoTokenizer.from_pretrained(mlx_dir)
enc = tok("The quick brown fox jumps over the lazy dog.", return_tensors="np")
out = model(mx.array(enc["input_ids"]), attention_mask=mx.array(enc["attention_mask"]))
print(out["last_hidden_state"].shape)

Verification

Verified against the bf16/fp32 torch reference (answerdotai/ModernBERT-large / mmBERT-base weights loaded via the original rl_common.build_model) in two stages:

1. Port correctness (unquantized). The unquantized MLX encoder's last_hidden_state was compared directly to the torch reference on held-out text: max abs diff 0.045, mean abs diff 0.0025 (hidden states have mean abs magnitude ~0.63) -- confirms the mlx-embeddings ModernBERT port itself is numerically correct, independent of quantization.

2. End-to-end decision quality (quantized). The MLX-quantized encoder's last_hidden_state was fed into Laya's original, unmodified torch decision head (same weights, loaded from model.safetensors) and compared against the full torch reference pipeline, on real choice / score / noul-type questions (moderation and sentiment-style prompts):

Honest caveat: this is a small (~421M), LayerNorm-heavy encoder, not an LLM -- it is noticeably more sensitive to block quantization than modern decoder-only models. Hidden-state drift from MXFP4/MXFP8 is meaningfully larger in relative terms (~35-45% mean abs perturbation) than typical MXFP4 LLM quantization (~1-5%), and one of three test questions (a typed-decisions MXFP8 score question with two closely-scored options) flipped its top-1 answer. Top-1 answers matched in 5/6 checkpoint x mode combinations, but if you need calibrated, tie-breaking-sensitive decisions, prefer the original bf16/fp32 encoder or treat quantized confidence scores as directional only, not exact.

LM Studio

Not applicable. Laya is a non-autoregressive encoder + custom scoring head, not a chat/completion model -- there is no generate() path or chat template, so LM Studio (which only runs chat/completion-style models) cannot load or serve this repo.