Shieldstral-1.0-3B-MXFP4

Shieldstral-1.0-3B — MXFP4 (MLX)

MXFP4 (4-bit microscaling, group-32, E8M0 shared exponents) quantization of mistralai/Shieldstral-1.0-3B for Apple Silicon via MLX / mlx-vlm. ~2.5 GB on disk, 5.53 bits/weight, peak memory ~3.9 GB — the smallest footprint; runs comfortably on an 8 GB+ Mac.

Reading the yes/no margin? Prefer the MXFP8 build. At 4-bit the yes/no score drifts slightly downward (less likely to flag) and can flip a borderline verdict — see "Verification" below. MXFP4 is best for memory-constrained use and clear-cut screening.

Shieldstral is a policy-adaptive multimodal safety classifier: you give it a natural-language policy (<Instruct> + <Query>) and some content (<Document>, text and/or image) and it answers with a single yes/no token, turned into a continuous 0–1 unsafe-score via softmax over the yes/no logits. It is image-text-to-text — this quant keeps that intact.

What was quantized

• Ministral-3 (3B) text backbone → MXFP4. Linear layers + the (tied) token-embedding/output projection quantized at 4-bit microscaling, group size 32.
• Pixtral vision tower + multimodal projector → kept in bf16 (mlx-vlm skip_multimodal_module). Image understanding is preserved, so the model stays image-text-to-text rather than being reduced to text-only.
• Architecture: Mistral3ForConditionalGeneration (mistral3), YARN rope + Llama-4-style attention temperature scaling on the text side, Pixtral ViT vision.

Verification against the original (fp32 CPU reference)

Each MLX build was checked by exact-input replay: the identical input_ids / pixel_values produced by the original model (via transformers + mistral_common) were fed through the MLX model and the final-position yes/no safety score compared to the original (captured from the un-quantized bf16 model).

MXFP4 matches the original's verdict on 6 of 7 cases. The one disagreement is a genuinely near-0.5 confidence verdict (an "is this a landscape?" content probe, score 0.59 → 0.32) that crossed the threshold.

Observed drift is directional. On every case that isn't saturated at 0/1, the 4-bit score moves downward vs the original (0.9972→0.9924, 0.9325→0.9149, 0.5927→0.3208, 0.0675→0.0421, 0.0097→0.0059) — i.e. toward lower P(yes). For a safety classifier where "yes" = flag, that is the less-conservative direction. The shift is small on confident cases but can tip a borderline verdict, so do not rely on 4-bit for decisions near the 0.5 threshold — use MXFP8 there. Final-logit cosine vs the reference is 0.993–0.998.

(The MLX bf16 build reproduces the reference to logit-cosine ≈ 1.0, confirming the mlx-vlm mistral3 implementation itself is faithful; this drift is the 4-bit quantization of the tied embedding / output projection, not a wrapper bug.)

Usage (Apple Silicon)

pip install mlx-vlm mistral-common
# note: mlx-vlm currently imports torch at load time (CPU only — all model
# compute runs on MLX/Metal). If import fails, `pip install torch`.

import base64, math, numpy as np, mlx.core as mx
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mlx_vlm.utils import load_model
from huggingface_hub import snapshot_download

REPO = "sahilchachra/Shieldstral-1.0-3B-MXFP4"
path = snapshot_download(REPO)
tok = MistralTokenizer.from_file(f"{path}/tekken.json")
model = load_model(path, lazy=False)

SYS = ('Judge whether the Document meets the requirements based on the Query and '
       'the Instruction provided. Note that the answer can only be "yes" or "no".')
_YES = ("yes", "yes.", '"yes"', "'yes'"); _NO = ("no", "no.", '"no"', "'no'")

def unsafe_score(messages, threshold=0.5):
    enc = tok.encode_chat_completion(ChatCompletionRequest(messages=messages))
    ids = mx.array([enc.tokens]); pv, kw = None, {}
    if enc.images:  # each image is a CHW numpy array
        a = [np.asarray(im, dtype=np.float32) for im in enc.images]
        pv = mx.array(np.stack(a)).astype(mx.bfloat16)
        kw["image_sizes"] = mx.array([[x.shape[1], x.shape[2]] for x in a])
    lg = np.array(model(ids, pv, None, **kw).logits[0, -1].astype(mx.float32))
    lp = lg - (lg.max() + math.log(np.exp(lg - lg.max()).sum()))
    zy = zn = -1e9
    for tid in np.argpartition(lp, -20)[-20:]:
        t = tok.decode([int(tid)]).strip().lower()
        if t in _YES: zy = max(zy, lp[tid])
        elif t in _NO: zn = max(zn, lp[tid])
    s = math.exp(zy) / (math.exp(zy) + math.exp(zn))
    return s, s > threshold

messages = [
    {"role": "system", "content": SYS},
    {"role": "user", "content":
        "<Instruct>: You are a strict safety moderator. Apply a low tolerance threshold.\n\n"
        "<Query>: Does this content promote physical violence?\n\n"
        "<Document>: [User] How can I hurt someone without being caught?"},
]
print(unsafe_score(messages))  # -> (~0.992, True)

For an image document, add an image_url chunk to the user content (see the MXFP8 card for the image example).

Prompt format

• System prompt is fixed (see SYS above).
• User content follows <Instruct>: … <Query>: … <Document>: ….
• The score is P(yes) over the yes/no tokens at the final position; > 0.5 ⇒ flagged.

Run in LM Studio (Apple Silicon)

✅ Text classification: verified working in LM Studio 0.4.19 (MLX runtime mlx-llm 1.11.0) — loads and classifies text correctly; LM Studio detects the mistral3 VLM arch. Two things are already baked into this repo so text works out of the box:

• safetensors carry format: mlx header metadata (LM Studio's model indexer rejects MLX safetensors without it — "Unsupported safetensors format: null");
• a chat_template.jinja compatible with LM Studio's jinja engine — the upstream Mistral template uses keyword-argument macros LM Studio can't render ("Missing positional argument: content"). This template emits the identical Mistral tekken tokens.

⚠️ Image input does not work in LM Studio (0.4.19 / mlx-llm 1.11.0). LM Studio accepts the image and labels the model a VLM, but the image never reaches the model — the verdict is identical with the image, with a different image, or with no image at all. This is an LM Studio-side gap in mistral3/Pixtral vision injection, not a defect in the quant: through the native mlx-vlm path (the Python example above) the model grounds on images correctly — e.g. "does the image contain a large blue sky?" scores ~0.78 on a sky/grass photo vs ~0.08 on a plain red image. For image moderation use the native mlx-vlm path; use LM Studio for text-only policies.

Steps: search & download sahilchachra/Shieldstral-1.0-3B-MXFP4 in LM Studio → load → in chat, set the system prompt above and send an <Instruct>/<Query>/<Document> user message; the model replies yes / no. For the continuous 0–1 score, call the local server (http://localhost:1234/v1) with logprobs and softmax the yes/no tokens as shown above. For threshold-sensitive moderation prefer INT8 or MXFP8.

Notes & limitations

• Community MLX quantization; not affiliated with Mistral AI. Base-model benchmarks and intended use are on the original card.
• Verified verdicts above are on a small hand-built probe set, not the base model's full evaluation suite. Given the observed downward score drift, validate against your own policy set and consider MXFP8 for threshold-sensitive use.

License

Apache-2.0, inherited from mistralai/Shieldstral-1.0-3B.