mlboydaisuke / bekko-embedding-v1-a25m-LiteRT

bekko-embedding-v1-a25m — LiteRT

hotchpotch/bekko-embedding-v1-a25m converted to LiteRT ( .tflite ) for on-device inference. An ultra-compact multilingual embedding model (100+ languages, strong Japanese) for retrieval and RAG: 123M parameters of which only 25M are active per token — the rest is the embedding table — producing 384-dimensional L2-normalized vectors, fully offline, on CPU.

Mean pooling and L2 normalization are inside the graph : one call in, one finished embedding out. No prefixes — bekko is trained without query: / passage: instructions; encode raw text on both sides.

Both are task-lossless against the PyTorch reference on every gate below; embt8 is the exact analog of the upstream repo's own default ONNX/OpenVINO artifacts (embedding table int8, transformer float). Peak RAM measured on an M4 Max Mac, 8 threads, all four signatures loaded and invoked.

Signatures

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

Pad the token ids into the smallest signature that fits and set the mask accordingly. Padding is fully masked inside the graph, so the result is independent of which signature you route through : the same text through embed_64 / 128 / 256 / 512 returns bitwise identical vectors, and pad-region token ids cannot influence the output at all.

Embeddings are L2-normalized, so cosine similarity is a dot product. Matryoshka truncation works as documented upstream: slice the first 256 / 128 / 64 dimensions and re-normalize (verified on JSTS for every variant, table below).

Usage (Python)

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

PAD_ID = 0
tok = AutoTokenizer.from_pretrained("hotchpotch/bekko-embedding-v1-a25m")
it = Interpreter(model_path="bekko-embedding-v1-a25m_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"embed_{s}") for s in LENS}

def embed(text):
    ids = tok(text)["input_ids"][:LENS[-1]]
    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
    return list(runners[S](input_ids=x, attention_mask=m).values())[0][0]

q = embed("日本で一番高い山は？")
d = embed("Mount Fuji is the highest mountain in Japan, at 3,776 meters.")
print("cosine:", float(q @ d))

Texts longer than 512 tokens must be chunked (the upstream model accepts 8192, but a static on-device graph at that length is not practical; chunk-and-average or chunk-and-max is the usual approach).

Quality

Four independent checks, each run on every variant against the PyTorch fp32 reference.

1. The base card's own quickstart examples. Every variant ranks the right document first in both documented examples (6/6 argmax), and reproduces the documented cross-lingual similarities to within 0.003 (e.g. int8: 0.454 vs 0.457, 0.564 vs 0.563).

2. JSTS (Japanese semantic similarity, JGLUE v1.3 validation, 300 pairs, Spearman), including Matryoshka truncation:

3. JSQuAD retrieval (Japanese question → Wikipedia paragraph, 800-paragraph corpus, 150 questions): PyTorch nDCG@10 0.9179 / hit@1 0.840; wi8fc 0.9310 / 0.873; embt8 0.9179 / 0.840. The corpus is subsampled, so absolute numbers are not comparable to published benchmarks; at 150 queries the gate cannot resolve differences below a couple of percent — the supported claim is "not worse than the reference".

4. Cross-variant retrieval — the RAG deployment shape. Documents encoded with the PyTorch model, queries with the int8 artifact (i.e. index built on a server, queried on device): nDCG@10 0.9186 vs the 0.9179 all-PyTorch control. The quantized embedding space is compatible with an upstream-built index.

Multilingual spot-check (STS17 Spearman, 100 pairs each): en-en 0.886 / ko-ko 0.885 / es-en 0.767 / en-ar 0.733 for wi8fc, all within 0.0025 of the PyTorch reference.

Speed

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

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

Memory

The interpreter allocates and XNNPACK-packs every signature subgraph at creation time , whether or not you call it — so peak RAM scales with the signatures present in the file, not the ones you use. The wi8fc build holds all four signatures in 225 MiB. An fp16 variant was built and measured but is not shipped: XNNPACK expands fp16 weights to fp32 per signature, so it peaked at 2.4 GiB for identical task scores — strictly dominated by embt8 .

Conversion

Encoder lane — a direct multi-signature litert_torch trace of the HF model (not an LLM export), with the ModernBERT attention masks built by hand inside the traced wrapper. Two things worth knowing if you reproduce it:

• ModernBERT alternates full attention (every 3rd layer) with ±64-token sliding-window attention. Under right padding, a pad query whose whole window is pads produces an all-masked attention row → NaN → 0 × NaN poisons the mean pool. The exported graph always allows the diagonal, which is a no-op for real tokens and keeps pad rows finite.
• ModernBertModel.forward accepts a prebuilt {layer_type: bias} mask dict, which bypasses the mask-construction path that would otherwise specialize away the attention mask under torch.export when traced with an unpadded sample. Trace with padded samples and gate on pad-content invariance either way.

Script and full notes: hf-to-litertlm .

License

MIT, inherited from the base model by Yuichi Tateno (hotchpotch).

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