UEmbed-2B-MLX-4bit

4-bit MLX quantization of the dense text embedding tower of Alibaba-NLP/UEmbed-2B (Qwen3.5 text backbone), packaged to run natively on Apple Silicon with mlx-lm. It produces L2-normalized dense text embeddings via the upstream chat-template contract and was gated against the MLX bf16 reference before release.

At a glance

Source	Alibaba-NLP/UEmbed-2B @ e7501a4d1be34ac4c7f8d1565cbeaa5b3f5b41b3
Format	MLX (safetensors)
Quantization	4-bit, group size 32 (affine)
Disk size	1.3 GB

Variants

Variant	Size	Status
UEmbed-2B-MLX-8bit	2.0 GB	public
UEmbed-2B-MLX-4bit (this repo)	1.3 GB	public

Verification (2026-08-11, M-series Mac)

• Quantization gate (this pack): worst cosine-vs-bf16 0.9369; retrieval ordering PASS (matched 0.78/0.76 vs distractors ~0.0).
• Reference parity (bf16 baseline vs upstream): MLX bf16 vs upstream torch (transformers 5.4.0, shipped Qwen35Embedder, fp32): cosine >= 0.99998 on all 6 fixed parity inputs; retrieval sims match (matched pairs 0.77/0.74, distractors ~0).

Fidelity note: at 4-bit the embeddings drift measurably from the bf16 reference (worst-case cosine 0.9369 on the gate inputs). Retrieval ordering is preserved with wide margins, but if you need maximum embedding fidelity, prefer the 8-bit pack above.

Usage

Requires: mlx-lm with qwen3_5 support (git main as of 2026-08-11; newer than the 0.31.3 release).

from mlx_lm import load
import mlx.core as mx
model, tok = load("majentik/UEmbed-2B-MLX-4bit")
def embed(text, instruction="Represent the user's input."):
    ids = tok.apply_chat_template(
        [{"role": "system", "content": instruction}, {"role": "user", "content": text}],
        add_generation_prompt=True, enable_thinking=False)
    h = model.model(mx.array([ids]))[0, -1, :].astype(mx.float32)
    return h / mx.sqrt((h * h).sum())

Embedding contract: messages=[system(instruction), user(text)]; tokenizer.apply_chat_template(add_generation_prompt=True, enable_thinking=False); dense embedding = final-norm hidden state at the LAST prompt token, L2-normalized. (Upstream appends 16 EOS and pools at last-16 — identical under causal attention.) Default instruction: "Represent the user's input."

DENSE TEXT embeddings v1: the pack contains the Qwen3.5 text tower only (the upstream repo's visual tower and SPLADE sparse heads are NOT included; use upstream for image/video/sparse). Checkpoint keys normalized from the upstream tripled prefix (model.language_model.language_model.language_model.*) to the standard text layout.

Reproduce

Prerequisite: the upstream checkpoint stores the text tower under a tripled key prefix (model.language_model.language_model.language_model.*); normalize the keys to the standard text layout (dropping the visual tower and SPLADE sparse heads) before conversion. Then:

mlx_lm.convert --hf-path <normalized-uembed-2b-text-dir> --mlx-path UEmbed-2B-MLX-4bit -q --q-bits 4 --q-group-size 32

License

CC-BY-4.0, inherited from upstream Alibaba-NLP/UEmbed-2B. Attribution: this repository redistributes a quantized derivative of UEmbed-2B by Alibaba-NLP, used under CC BY 4.0, with modifications (text-tower extraction, key normalization, 4-bit group-size-32 MLX quantization).