Decily-MLX-4bit

Decily — calibrated candidate-scorer decision model

Decily scores a runtime-provided set of candidates against a state and a question, and returns a calibrated probability distribution over those candidates. It does not generate text. Decily is built for decisions where you want a probability you can threshold, not a sentence you have to parse.

Decily.decide(
    state="Our app crashes on startup after the latest update.",
    question="What is the customer's intent?",
    options=["technical", "billing", "shipping", "returns"],
)  # -> {"technical": 0.9999, "billing": 0.0001, "returns": 0.0, "shipping": 0.0}
# measured with the released weights (bf16, T=0.45)

This repository

Apple MLX 4-bit build — the on-device variant. The tower is quantized (group-64 affine, 968 MB); the scoring head ships separately in decision_head.safetensors (kept in fp32 for calibration).

import mlx.core as mx, mlx_lm
from transformers import AutoTokenizer

model, _ = mlx_lm.load(".")                       # this repository
tok = AutoTokenizer.from_pretrained(".")
head = dict(mx.load("decision_head.safetensors"))

def _gelu(x):
    return 0.5 * x * (1.0 + mx.erf(x / mx.sqrt(mx.array(2.0))))

def decide(state, question, options, temperature=0.45):
    prefix = (tok(state, truncation=True, max_length=256)["input_ids"]
              + tok(question, truncation=True, max_length=96)["input_ids"])
    seqs = [prefix + tok(o, truncation=True, max_length=64)["input_ids"] for o in options]
    width = max(len(s) for s in seqs)
    ids, mask = mx.zeros((len(seqs), width), dtype=mx.int32), mx.zeros((len(seqs), width))
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = mx.array(s, dtype=mx.int32)
        mask[i, :len(s)] = 1.0
    h = model.model(ids)                          # [B, L, d]
    scores = (h @ head["pool.proj.weight"].T).squeeze(-1) + head["pool.proj.bias"]
    scores = mx.where(mask > 0, scores, mx.array(-1e9, dtype=scores.dtype))
    w = mx.softmax(scores, axis=-1)
    pooled = (h * w[..., None]).sum(axis=1).astype(mx.float32)
    ln_w, ln_b = head["head.0.weight"], head["head.0.bias"]
    mu = pooled.mean(axis=-1, keepdims=True)
    var = ((pooled - mu) ** 2).mean(axis=-1, keepdims=True)
    x = (pooled - mu) / mx.sqrt(var + 1e-5) * ln_w + ln_b
    x = _gelu(x @ head["head.1.weight"].T + head["head.1.bias"])
    logits = (x @ head["head.4.weight"].T + head["head.4.bias"]).squeeze(-1) / temperature
    probs = mx.softmax(logits, axis=-1)
    mx.eval(probs)
    return dict(zip(options, probs.tolist()))

print(decide("Our app crashes on startup after the latest update.",
             "What is the customer's intent?",
             ["technical", "billing", "shipping", "returns"]))

Model family

This card covers all four released artifacts; they are the same model in different containers.

Model details

• Model type: cross-encoder candidate scorer. The backbone encodes state + question + candidate jointly; an attention-pooling layer over all tokens feeds a LayerNorm/GELU MLP that outputs one logit per candidate. Probabilities come from the softmax over the candidates you pass in.
• Backbone: Qwen/Qwen3-1.7B-Base (1.72B parameters, Apache-2.0), with an attention-pooling scoring head (decision_head recorded in config.json; pool/head weights live in the checkpoint).
• Input limits (training configuration): state ≤256 tokens, question ≤96 tokens, each candidate ≤64 tokens, 2–16 candidates per call. Longer inputs are truncated — for long documents, chunk the state and score chunks separately (the cross-encoder logits are candidate-set independent, so chunked shortlisting then re-ranking is lossless; see the repository's two-stage inference utility).
• Precision: bf16 (PyTorch/MLX), 4-bit affine group-64 (on-device), int8 (ONNX).
• Temperature: raw logits are over-confident out of distribution. Fitted temperatures observed: ≈1.1 (in-domain), ≈0.45 on unseen label sets, ≈1.35 on the zero-shot fair suite. Always report/calibrate the temperature for your data (see Evaluation).

Training summary

Full, reproducible recipe: TRAINING.md in the training repository arczhi/decily (design rationale in DESIGN.md, all experiments and ablations in TRAIN-REPORT.md).

Key recipe properties:

• Route B (explicit candidate scorer) beats the LM-head-letter-logits route (Route A) by +10.6 pt held-out accuracy in a controlled same-backbone test.
• Ensemble distillation: the 4-model ensemble reaches NLL 1.455; the distilled single model reaches 1.451 at 1× inference cost.
• Consumer hardware: the whole pipeline trains on a single RTX 5090 32 GB in well under a day.

Evaluation

Protocol: accuracy / NLL / ECE after temperature fitting (raw T=1 numbers are over-confident). "In-task" = the 24 trained task families; "held-out" = unseen label sets; "fair suite" = 8 tasks unseen by both this model and the external baseline.

Selective prediction (unseen label sets): taking only the most-confident 5% of predictions gives 93.3% accuracy (an SFT baseline with the same protocol gives 79%); abstention thresholds can be calibrated to a target error rate (empirically 4.7% achieved at a 5% target).

Honest reading of the numbers:

• Decily leads the larger, 95-task decider-2b by +5.2 pt on tasks both were trained on, at 1.7B vs 2B parameters and ~1/4 of the task coverage.
• On tasks neither model has seen, decider-2b leads by 4.6 pt; the gap is concentrated in knowledge-heavy tasks (sciq / pubmedqa / quality).
• The zero-shot gap is an efficiency result, not a ceiling: with this recipe, task coverage — not parameter count — is the documented main lever. Scaling the 24-task mixture toward 60–95 tasks with the same pipeline is expected to close and exceed the baseline (see TRAINING.md).

Uses

Direct use

• Classification / ranking / selection when the label set is known at runtime and may change per call (intents, topics, sentiment, NLI-style relations, multiple-choice answers, routing, moderation labels, tool selection).
• Calibration-sensitive automation: threshold the probability, auto-handle the confident head and escalate the rest to a human or a larger model.
• On-device / private inference: 968 MB 4-bit build, no GPU required.
• Synthetic decision data generation and reward/verifier scoring for LLM pipelines (it returns probabilities, not text).

Out of scope

• Text generation, chat, summarization.
• Open-ended label spaces (the candidate set must be provided at call time).
• High-stakes decisions without a calibrated threshold and human review.
• Very long inputs beyond the 256-token state window without chunking.

Quick start

Apple Silicon (MLX) — the standalone reference implementation ships with the training repository (mlx/decision_mlx.py, pure mlx.core, no torch):

python mlx/decision_mlx.py --model-dir <Decily-MLX dir> \
  --state "The invoice was paid twice..." \
  --question "What is the customer's intent?" \
  --options "billing,technical,shipping,returns" \
  --temperature 0.45

4-bit on-device: load the quantized tower with mlx_lm.load and apply the attention-pooling head from decision_head.safetensors (both files are in the repository; config.json records the head spec).

PyTorch: load model.safetensors + config.json with the training repository's model class; the checkpoint keeps backbone and head under their training key names (model.*, pool.*, head.*).

Limitations

• Task coverage: trained on 24 task families; zero-shot behavior on unseen domains trails a 95-task baseline (see Evaluation). Validate on your domain before trusting the probabilities.
• Calibration drift: the model needs temperature scaling; the fitted value depends on the data distribution (≈0.45–1.35 in our measurements).
• Input truncation: state is truncated at 256 tokens; long documents must be chunked (importance within a chunk, then combine).
• Language: training data is predominantly English (some multilingual intent data); other languages are untested.
• No safety layer: outputs are probabilities over the candidates you provide; the model does not refuse or filter inputs. Do not expose it as an autonomous decision-maker in safety-critical settings.
• Inherited biases: as a fine-tune of Qwen3-1.7B-Base on public datasets, it can reproduce biases present in those datasets (toxicity, sentiment and bias-classification tasks were part of the mixture, which reduces but does not eliminate this).

Environmental impact

Training used a single RTX 5090 32 GB: teacher ≈4 h, SFT ≈1 h, members ≈1 h, distillation ≈1.5 h (plus data conversion). Total well under 10 GPU-hours; no model was trained more than once per stage.

Citation and acknowledgements

• Backbone: Qwen3-1.7B-Base (Apache-2.0).
• The Route B formulation and the calibration/abstention pipeline build on the Decily line of work; the external baseline in the tables is Mapika/decider.
• If you use Decily, please cite this model card and link the training repository (TRAINING.md).