UEmbed-9B-MLX-4bit

UEmbed-9B-MLX-4bit

4-bit (group-size 32) MLX quantization of the dense text-embedding tower of Alibaba-NLP/UEmbed-9B, packaged for Apple silicon. UEmbed-9B is built on the Qwen3.5 text architecture (hybrid linear attention); this pack runs it natively with mlx-lm and produces L2-normalized dense embeddings.

At a glance

Variants

<!-- sibling-table -->

Smaller family members: UEmbed-4B-MLX, UEmbed-2B-MLX (note: 2B uses a different chat-template tail — enable_thinking=False).

Verification (2026-08-11, M-series Mac, stock mlx-lm git main)

• Quantization gate (this pack): worst cosine-vs-bf16 0.9451, worst cosine-vs-torch 0.9458; retrieval ordering PASS (margins >0.7). Typical 4-bit drift for this size class — prefer the 8-bit pack (0.9997) when embedding-space interchangeability matters.
• Port parity (shared basis): MLX vs upstream torch (transformers 5.4.0, shipped Qwen35Embedder, fp32): cosine >= 0.999976 on all 6 fixed parity inputs.

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-9B-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)
    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) — do NOT pass enable_thinking=False: UEmbed-9B's upstream embedder uses the template-default open <think> tail (this differs from UEmbed-2B); 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."

Scope — 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 prefix (model.language_model.language_model.language_model.*) across sharded safetensors; normalize the keys to the standard Qwen3.5 text layout first, then convert the normalized checkpoint:

mlx_lm.convert --hf-path <normalized-UEmbed-9B> \
  --mlx-path UEmbed-9B-MLX-4bit \
  -q --q-bits 4 --q-group-size 32

License

CC-BY-4.0, inherited from the upstream Alibaba-NLP/UEmbed-9B. Attribution: this pack is a derivative of UEmbed-9B by Alibaba-NLP, redistributed with modifications (text-tower extraction, key normalization, 4-bit MLX quantization) under the terms of the CC-BY-4.0 license.