mLateOn

mLateOn — LiteRT

lightonai/mLateOn converted to LiteRT (.tflite) for on-device inference. LightOn's multilingual late-interaction (ColBERT-style) retriever: one 128-dimensional unit vector per token, scored with MaxSim — fully offline, on CPU.

The projection head and per-token L2 normalization are inside the graph; MaxSim scoring is a few lines host-side (below).

Both are task-lossless against the PyTorch reference on every gate below; the converted reference reproduces the base card's own published MaxSim scores to all four printed decimals.

Why late interaction on device

The base model's headline is long-document retrieval (MLDR 87.69 NDCG@10, ~9 points ahead of the next model) and generalization to languages it was never retrieval-trained on — it is trained on nine European languages plus Arabic, yet scores strongly on Japanese, Cyrillic, and other unseen scripts. Both survive conversion; the unseen-language claim is measured below on Japanese.

Host contract (read this before using the output)

Reimplemented from the PyLate reference and verified against it:

• Pad with token id 4 (<mask>) — that is what PyLate pads with. Note config.json's pad_token_id: 0 is not what the reference stack uses; the vendor's own onnx_config.json says 4. (With padding masked out, the graph output on real tokens is provably independent of the pad id — but use 4 anyway.)
• Queries: tokenize normally (<bos> … <eos>), then insert [Q] (id 256000) at position 1 and a matching 1 in the attention mask. No query expansion, no skiplist — every valid position is scored.
• Documents: same with [D] (id 256001). Keep the vectors at valid (mask=1) positions.
• Score = MaxSim: for each query vector, take the max dot product over the document's vectors, then sum over query vectors. Vectors are unit length, so dot product = cosine.

Signatures

Batch-1, right-padded static shapes: input_ids int32 [1, S], attention_mask int32 [1, S] (1 = real token, then 0s), for S in 32 / 128 / 256 / 512.

Route each text to the smallest signature that fits and keep only the valid positions of the output. Padding is masked inside the graph, so the same text returns bitwise identical vectors through every signature, and pad-region token ids cannot influence the valid positions at all.

Usage (Python)

import numpy as np
from ai_edge_litert.interpreter import Interpreter
from transformers import AutoTokenizer

PAD_ID, Q_ID, D_ID = 4, 256000, 256001
tok = AutoTokenizer.from_pretrained("lightonai/mLateOn")
it = Interpreter(model_path="mLateOn_wi8fc.tflite", num_threads=8)

LENS = sorted(int(n.split("_")[1]) for n in it.get_signature_list())
runners = {s: it.get_signature_runner(f"encode_{s}") for s in LENS}

def encode(text, prefix_id, max_len=512):
    e = tok(text, truncation=True, max_length=max_len - 1)["input_ids"]
    ids = [e[0], prefix_id] + e[1:]          # insert [Q]/[D] after <bos>
    S = next(s for s in LENS if len(ids) <= s)
    x = np.full((1, S), PAD_ID, np.int32)
    m = np.zeros((1, S), np.int32)
    x[0, :len(ids)] = ids
    m[0, :len(ids)] = 1
    v = list(runners[S](input_ids=x, attention_mask=m).values())[0][0]
    return v[:len(ids)]                       # keep valid positions only

def maxsim(q, d):
    return float((q @ d.T).max(axis=1).sum())

q = encode("Which planet is the Red Planet?", Q_ID)
docs = ["Mars, known for its reddish appearance, is often referred to as the Red Planet.",
        "Venus is often called Earth's twin because of its similar size and proximity."]
print([maxsim(q, encode(d, D_ID)) for d in docs])

Documents longer than 512 tokens must be chunked (the upstream model accepts 8192, but a static on-device graph at that length is not practical).

Quality

Five independent checks against the PyTorch fp32 reference.

1. The base card's own usage example (Red Planet query × Mars-en/fr/de + Venus): the converted fp32 reference reproduces the card's published MaxSim scores [9.6029, 9.5838, 9.5877, 9.4578] to all four printed decimals; every variant ranks all three Mars documents (including French and German — the cross-lingual property) above Venus.

2. NanoSciFact retrieval (English, 600-doc corpus, 50 claims, MaxSim): PyTorch nDCG@10 0.8882 / hit@1 0.820; int8 0.8951 / 0.840; fp16 identical to PyTorch.

3. JSQuAD retrieval — an unseen language (Japanese question → Wikipedia paragraph, 500-paragraph corpus, 150 questions). Japanese is not among the nine training languages; the base card's generalization claim, measured on this artifact: PyTorch nDCG@10 0.9710 / hit@1 0.9467; int8 0.9719 / 0.9467; fp16 identical to PyTorch.

4. Cross-variant retrieval — the deployment shape. Document banks encoded with the PyTorch model, queries with the int8 artifact (index built on a server, queried on device): JSQuAD identical to the all-PyTorch control in every metric (nDCG@10 0.9710); NanoSciFact within 0.0015.

5. Mechanics: cross-signature outputs bitwise identical; pad-content invariance exactly 0; all positions finite in fp32 at the longest signature (the sliding-window edge case is guarded inside the graph); per-token norms within 6e-8 of 1.

Corpora are subsampled, so absolute numbers are not comparable to published benchmarks.

Speed

CPU/XNNPACK, median of 10 runs, 75%-full signatures:

That is ~4,600 tokens/s of document encoding at encode_512 (int8) — late interaction pays its multi-vector cost at scoring time, not at encoding time.

A static signature computes all S positions regardless of how many are real, so route each text to the smallest signature that fits.

Memory

Measured at load+invoke on an M4 Max Mac (8 threads), all four signatures loaded: the int8 file peaks at ~680 MiB. The fp16 file peaks at ~6 GiB (XNNPACK expands fp16 weights to fp32 per signature subgraph) — it is a desktop artifact; use int8 on device.

Conversion

Encoder lane — a direct multi-signature litert_torch trace of the HF ModernBERT backbone (not an LLM export), with the full-attention and sliding-window (±64) masks built by hand inside the traced wrapper. The three PyLate Dense modules (768→1536→768→128, all-linear with residual branches) fold exactly into a single 128×768 projection, verified against the module stack. Gated on: bitwise agreement with the vendor's own mask path, pad-content invariance, all-position finiteness (sliding-window edge rows), cross-signature bitwise agreement, and the base card's published scores.

Script and full notes: hf-to-litertlm.

License

Apache 2.0, inherited from the base model by LightOn.

Modification notice: these files are converted, not original. The weights were exported to LiteRT and quantized (int8 dynamic-range / fp16); the projection head and per-token L2 normalization were folded into the graph. No fine-tuning or weight modification beyond quantization was performed.