techtheist/laya-onnx

Laya — ONNX exports with a dynamic sequence length (int4 / int8)

Unofficial ONNX exports of Laya by Convai Innovations (Apache-2.0), made for Engram Alpha, which uses them as an optional contradiction judge. The weights are the upstream checkpoints, unchanged apart from quantization. Any credit for the model belongs upstream.

path	source checkpoint	quantization	size
en/model_int4.onnx	convaiinnovations/laya (ModernBERT-large, 421M)	int4 MatMulNBits (block 32, symmetric) on every MatMul and the embedding Gather	262 MB
en/model_int8.onnx	convaiinnovations/laya	int8 dynamic (MatMulInteger, per-channel); embeddings stay fp32	554 MB
multilingual/model_int8.onnx	convaiinnovations/laya-multilingual (mmBERT-base, 322M)	int8 dynamic; the 256k-vocab embedding table stays fp32	873 MB

Each folder also holds the files the runtime needs: tokenizer.json, tokenizer_config.json (which names the special tokens: [CLS]/[SEP]/[MASK]/[PAD] for English, <bos>/<eos>/<mask>/<pad> for multilingual) and rl_agent_config.json (max_len, head_max_len, fitted temperatures).

Why another export

The sequence length here is truly dynamic. The exports were made with torch.onnx.export(dynamo=True, dynamic_shapes=…): batch up to 64, sequence 8–8192, markers up to 255. A TorchScript-traced export can record dynamic axes and still hard-code the length it was traced at inside the decision head's attention reshape. Short inputs then fail, or have to be padded to full length on every call. Before quantization, the fp32 export matches PyTorch to ≤ 2e-5 on the logits at sequence lengths 80, 222 and 429. scripts/ holds the export and quantization scripts.

Contract

Inputs: input_ids [batch, seq] int64, attention_mask [batch, seq] int64, marker_pos [batch, k] int64, marker_mask [batch, k] bool, qtype [batch] int64 (choice 0, score 1, noul 2). Outputs: logits [batch, k] (one logit per option marker) and act_logits. Build the sequence exactly as laya.common.build_sequence does: [CLS] <type> question: <instructions> [SEP] [MASK] opt₀ [MASK] opt₁ … [SEP] <state> [SEP]. Then divide the logits by the temperature for the (type, option-count) bucket in rl_agent_config.json, clamped to [0.5, 5.0], and softmax.

Measured (zero-shot contradiction detection, Engram's eval)

Engram's engram-eval --shapes bench: 300 generated note pairs across 15 contradiction and trap shapes, judged with Laya's own XNLI question. AUROC is how well contradictions separate from traps; context means title + first body sentence + up to 7 related notes. For comparison, the deberta-v3-small-tasksource-nli baseline scores 0.68 on titles alone.

model	AUROC, titles only	AUROC, with context	CPU per call (Apple M1 Pro, titles only)
English fp32 (not published)	0.71	0.79	231 ms
en/model_int4.onnx	0.72	0.77	122 ms
en/model_int8.onnx	0.63	0.68	97 ms
multilingual/model_int8.onnx	0.63	0.58	40 ms

int4 tracks fp32 closely; dynamic int8 loses more. The multilingual checkpoint ships without fitted temperatures, so treat its probabilities as uncalibrated. Receipts and method: engram eval/.

License

Apache-2.0, as upstream.