trillionlabs/Trida2.1-4B
Trida2.1-4B is a 4B self-speculative agent model that is also a usable decision model. One set of weights does both jobs:
- Agentic generation with self-speculative decoding. The model is a two-stream block-diffusion LM: its diffusion stream drafts blocks of 8 tokens and its own autoregressive (clean) stream verifies them in the same forward pass. No separate draft model is needed, and up to 8 tokens can be committed per forward pass.
- Typed decisions (kev / Jev style). Given a state and a typed question (yes/no, pick one, score), a small pointer head over the model's hidden states returns a calibrated probability for every allowed answer. The model never writes text for these, so an answer cannot leave the label set, and a decision costs one prefill, not a generation.
Why it fits on-device use
On a phone, laptop or single edge GPU an assistant serves one user at a time. At that low concurrency, decoding is bound by memory bandwidth: each forward pass reads all the weights to produce a single token, and most of the compute sits idle. Self-speculative decoding spends that idle compute on drafting and verifying a whole block, so each weight read can commit several tokens, with no second model to store. High-concurrency servers are already compute-bound, so the gain there is smaller; the model is designed for the low-concurrency case.
Measured on one H100 per server (vLLM 0.27 with the trida-stack self-spec plugin, block 8), GSM8K test questions with thinking on, temperature 1.0 / top-k 50 / top-p 0.95, 8,192-token budget, 200 problems per setting:
| concurrent requests | AR: tok/s per request | self-spec: tok/s per request | speedup | median latency, AR → self-spec |
|---|---|---|---|---|
| 1 | 211 | 312 | 1.48× | 2.13 s → 1.19 s |
| 4 | 190 | 257 | 1.36× | 2.45 s → 1.46 s |
| 8 | 168 | 222 | 1.32× | 2.63 s → 1.69 s |
| 16 | 137 | 156 | 1.14× | 3.19 s → 2.42 s |
- Self-spec commits about 2.0 tokens per forward pass. GSM8K accuracy is the same in both modes within noise (87.5–91.5 % across settings).
- Total throughput gains less than per-request speed as concurrency rises (1.26× at 1 request, 1.17–1.26× at 4–8, 0.99× at 16): with more requests in flight the GPU is busy anyway, which is why the model targets low-concurrency use.
The same weights also answer typed decisions with the Jev head, so an on-device agent can route, gate, classify and score with the model already in memory instead of loading a separate classifier, and get probabilities it can threshold.
Files
| file | contents |
|---|---|
model-000NN.safetensors (17 shards), model.safetensors.index.json |
language model + vision tower (bf16) |
config.json, generation_config.json, tokenizer files, chat_template.jinja, preprocessor_config.json |
standard Qwen3.5 VL files |
block_diffusion.json |
diffusion block size (8) and mask token (248077) |
jev_head.pt |
the typed-decision (pointer) head |
The checkpoint uses the Qwen3.5 vision-language layout (Qwen3_5ForConditionalGeneration), vocabulary 248,078. All
evaluations below are text-only.
Generation: self-speculative decoding
Served with self-speculative decoding (generation block 8, bd_bidir_shift, 15 positions per forward pass) on the
HybridDiffusion SGLang integration; temperature 1.0, top-k 20, top-p 0.95, reasoning on. Every model in a table is
scored with the same configuration and harness.
Knowledge
| benchmark | Trida2.0-4B | Trida2.1-4B |
|---|---|---|
| GSM8K | 86.7 | 89.9 |
| MMLU-Pro | 63.0 | 62.7 |
Agentic (full benchmark sets)
| benchmark | Trida2.1-4B |
|---|---|
| BFCL non-live AST (1,390) | 76.6 |
| FunctionChat singlecall / dialog / call decision | 0.914 / 0.830 / 0.944 |
| Ko-AgentBench success rate, L1–L7 (91 tasks) | 0.736 |
| tau2 retail (114) / airline (50) / telecom (114) | 0.430 / 0.420 / 0.386 |
- BFCL: the non-live categories (simple Python/Java/JavaScript, multiple, parallel, parallel-multiple, plus irrelevance), prompted with the model's own chat template (Qwen3.5 XML tool calls); the AST summary is reported.
- FunctionChat: all 500 single-call items, all 200 dialog turns and all 606 call-decision items, graded by an LLM judge.
- Ko-AgentBench: all seven levels (91 tasks). tau2: the full
basesplit of each domain, GPT-4o-mini as the user simulator. One task moves Ko-AgentBench by 0.011, tau2 airline by 0.020 and retail / telecom by 0.009.
Typed decisions (Jev): vs Kev-4B and Jev
Accuracy on kev's development suites, the Korean suites and typed-decisions, scored with kev's harness
(kev.benchmark). Trida2.1-4B is scored at its fitted temperature (2.19), Kev-4B at its own served temperature.
Brier score (multiclass, lower is better) in parentheses.
English (kev development splits)
| suite | Trida2.1-4B | Kev-4B | Jev (published) |
|---|---|---|---|
| transfer-v4 (656) | 0.800 (0.276) | 0.817 (0.243) | 0.857 |
| transfer-v9 (1,046) | 0.754 (0.332) | – | 0.854 |
| decision-v7, ≤ 26 options (1,184) | 0.864 (0.201) | 0.871 (0.181) | – |
| hard-v1 (1,083) | 0.812 (0.249) | – | 0.777 |
| devtools-v1 (1,074) | 0.714 (0.391) | – | 0.713 |
| documents-v1 (920) | 0.886 (0.166) | – | – |
| SemIf (144) | 0.910 (0.177) | – | 0.965 |
Korean and typed-decisions (test splits; test items screened for overlap)
| suite | Trida2.1-4B | Kev-4B |
|---|---|---|
| KorNAT common knowledge (6,007) | 0.659 (0.451) | 0.597 (0.517) |
| KorNAT social values (4,000) | 0.443 (0.719), KL 0.200 | 0.537 (0.653), KL 0.180 |
| CLIcK (1,995) | 0.605 (0.485) | 0.555 (0.519) |
| KMMLU-Pro (2,822) | 0.445 (0.650) | 0.414 (0.665) |
| typed-decisions test (2,000) | 0.737 (0.395), KL 0.126 | 0.669 (0.465), KL 0.223 |
- Trida2.1-4B is within a few points of the decision-only Kev-4B and Jev in English while also being a full agent model, and is ahead of Kev-4B on four of the five Korean / typed-decisions suites.
- "Jev (published)" is TypeSafe Jev's published development accuracy from the kev README (accuracy only; our harness was not run on it). Kev-4B was scored on transfer-v4 and decision-v7 only in English.
- Development splits only; no locked test split was read. decision-v7 drops questions with more than 26 options (same filtered split for every model). The decision-v7 and hard-v1 development splits overlap kev's public training data, so scores there are likely optimistic for models trained in kev's scheme. SemIf has 144 questions (± ~3 points).
- KorNAT social values and typed-decisions also report KL(target ‖ model) against the human-survey / teacher distribution. For typed-decisions, its leaderboard lists TypeSafe Jev 1.13 at 0.727 accuracy.
- Every item of every suite is scored. kev's harness reads
--datafiles with a 384-token state limit and skips longer records, so CLIcK, KMMLU-Pro and typed-decisions were read with kev's serving context (states up to 64k tokens), for both models. - Typed-decisions train is in this checkpoint's training data, so its typed-decisions score belongs in the leaderboard's "fitted on train" table, not the zero-shot one.
Using the typed-decision head
Export the checkpoint to kev's layout with scripts/export_kev_checkpoint.py from trida-stack, then use kev's own
tools:
python scripts/export_kev_checkpoint.py --ckpt <this repo, downloaded> --out <dir> --temperature 2.19
python -m kev.benchmark --run <dir> ... # score
python -m kev.serve --run <dir> # serve
--temperature 2.19 applies the fitted calibration temperature; omit it for raw probabilities.
Reasoned action: think in text, then decide with the head
Because one set of weights both generates text and answers typed decisions, the model can reason before it decides. The decision head stays the judge; generation only adds evidence to what it reads.
- Prompt. The decision request (state, plus every question with its options or levels) is rendered as one chat
message, with a single generic instruction to work through anything that needs calculating or checking. Thinking is
on, so the reply opens a
<think>block. No per-task prompt is needed. - Reason. The model generates until it emits
</think>(token 248069), which is the stop token; nothing after it is generated. - Decide. The reasoning text is added to the request's state (as
reasoning; a plain-string state becomes{"situation": ..., "reasoning": ...}), and the pointer head scores every option exactly as in a direct decision.
Why it helps
- Better probabilities, not just better answers. The Brier score measures how close each predicted distribution is to the true answer, so it rewards being right and confident. Reasoned decisions are both more often correct and more decisive, which cuts the Brier score sharply where the answer has to be worked out: by 41 % on JevBench (0.314 → 0.186) and by 70 % on hard-v1 (0.249 → 0.074).
- Multi-step problems become answerable. A direct decision is one prefill: the head must read the answer off the hidden states in a single pass, which fails on arithmetic, date math, ordered rules and multi-hop lookups. Writing the intermediate steps first puts their results in the state, where the head can read them.
- The answer stays typed. The output is still the head's probability distribution over the allowed labels, never free text to parse, so it cannot leave the label set.
- No second model, and opt-in per request. The same weights in memory do both steps, so an on-device agent can answer easy decisions directly (one prefill) and spend generation only on the hard ones. Self-speculative decoding (above) makes the reasoning step cheaper at low concurrency.
Results. Accuracy and Brier score (multiclass, against the label; lower is better). Both conditions use the same scorer and the same decision temperature (2.19); reasoning is sampled at temperature 0.6, top-p 0.95, top-k 20.
JevBench public split (231 items, JevBench's own scoring code; reasoning budget 32,768 tokens):
| split | direct: accuracy / Brier | reasoned: accuracy / Brier |
|---|---|---|
| hard (111) | 0.514 / 0.593 | 0.793 / 0.332 |
| original (72) | 0.958 / 0.090 | 0.986 / 0.079 |
| easy (48) | 1.000 / 0.006 | 1.000 / 0.007 |
| all (231) | 0.753 / 0.314 | 0.896 / 0.186 |
kev suites (reasoning budget 4,096 tokens):
| suite | direct: accuracy / Brier | reasoned: accuracy / Brier |
|---|---|---|
| hard-v1 (1,083) | 0.811 / 0.249 | 0.958 / 0.074 |
| transfer-v4 (656, held-out datasets) | 0.800 / 0.276 | 0.851 / 0.225 |
| decision-v7, ≤ 26 options (1,184) | 0.863 / 0.201 | 0.865 / 0.205 |
| typed-decisions test (2,000) | 0.738 / 0.395 | 0.719 / 0.400 |
- The gains are largest where the answer has to be worked out: JevBench hard +28 points, hard-v1 +15, held-out transfer-v4 +5, each with a lower Brier score.
- It does not help everywhere. decision-v7 is about unchanged. On typed-decisions the gold is a teacher model's soft judgment; reasoning makes the model commit to its own view, so it agrees less with the teacher (accuracy −2 points, KL from the gold distribution 0.126 → 0.187). Use reasoned action for decisions with a checkable answer, not for matching subjective judgments.
- Give it enough budget. A reasoning cut off mid-thought is worse than none. Even at 32,768 tokens, 4 % of the JevBench reasonings hit the limit; 4,096 tokens is the smallest budget we recommend.
- Reasoned probabilities are sharper than direct ones. If you threshold them, fit a separate temperature for reasoned mode rather than reusing 2.19.
- The head was trained without reasoning in the state, so these gains need no extra training. Reasoning is text-only for now: no image goes into the generation step.
In practice: reason only when the direct answer is unsure
Reasoning costs a generation, so use it as an escalation: answer every request directly first (one prefill), and only when the head's top-option probability is below a threshold, generate the reasoning and decide again. Confident direct answers are usually right already, so this keeps their speed and spends generation where it changes the answer.
On the same runs as above (accuracy / Brier; below each, the share escalated to reasoning, as questions / requests):
| set | direct only | escalate if conf < 0.5 | escalate if conf < 0.7 | escalate if conf < 0.9 | always reason |
|---|---|---|---|---|---|
| JevBench public (231 requests, 231 questions) | 0.753 / 0.314 | 0.810 / 0.274 10 % / 10 % |
0.857 / 0.231 29 % / 29 % |
0.896 / 0.185 58 % / 58 % |
0.896 / 0.186 |
| hard-v1 (700 requests, 1,083 questions) | 0.811 / 0.249 | 0.856 / 0.201 10 % / 15 % |
0.935 / 0.112 30 % / 42 % |
0.959 / 0.075 53 % / 68 % |
0.958 / 0.074 |
| transfer-v4 (656 requests, 656 questions) | 0.800 / 0.276 | 0.808 / 0.269 7 % / 7 % |
0.854 / 0.229 31 % / 31 % |
0.851 / 0.225 61 % / 61 % |
0.851 / 0.225 |
| decision-v7 (944 requests, 1,184 questions) | 0.863 / 0.201 | 0.863 / 0.202 2 % / 2 % |
0.869 / 0.196 17 % / 21 % |
0.867 / 0.197 35 % / 42 % |
0.865 / 0.205 |
- A request is escalated when any of its questions is below the threshold; the reasoning is generated once per request, so the request share is the generation cost.
- 0.7 is a good default: reasoning on about 30 % of requests (42 % on hard-v1) recovers 70–85 % of the accuracy gain of always reasoning on JevBench and hard-v1, and all of it on transfer-v4. 0.9 recovers all of it (and matches or beats always reasoning on every set) at about 60 % of the requests. 0.5 is cheapest (about 10 % of requests) but keeps only 15–40 % of the gain; a yes/no answer's top probability is never below 0.5, so at 0.5 only choice and score questions escalate.
- Where reasoning does not help (decision-v7, subjective or opinion questions), no threshold helps much either; keep the direct answer.
A minimal version, with the decision head served by kev.serve (see above) and the reasoning generated by any
OpenAI-compatible server for this checkpoint (vLLM, SGLang):
import requests
from transformers import AutoTokenizer
from kev.api import render # the same state rendering the head reads
DECIDE = "http://localhost:8008/v1/systemone" # python -m kev.serve --run <kev export, --temperature 2.19> --port 8008
GENERATE = "http://localhost:8000/v1/completions" # OpenAI-compatible server for this checkpoint
tok = AutoTokenizer.from_pretrained("trillionlabs/Trida2.1-4B")
def decide(request):
return requests.post(DECIDE, json=request).json()["answers"]
def confidence(answer): # probability of the top option
if answer["type"] == "noul":
return max(answer["noul"], 1 - answer["noul"])
return max(answer["probabilities"].values())
def reason(request, max_tokens=4096):
parts = [render(request["state"])]
for qid, q in request["questions"].items():
lines = [f"Question ({qid}): {render(q.get('instructions'))}"]
if q["type"] == "noul":
lines.append("Answer yes or no.")
elif q["type"] == "choice":
lines += ["Options:"] + [f"- {k}: {render(v)}" if v not in (None, "") else f"- {k}"
for k, v in q["criteria"].items()]
else:
lines += ["Levels, lowest first:"] + [f"{i}. {render(v)}" for i, v in enumerate(q["criteria"])]
parts.append("\n".join(lines))
parts.append("Think it through step by step: work out anything that needs calculating or checking, "
"then say which option you choose for each question.")
prompt = tok.apply_chat_template([{"role": "user", "content": "\n\n".join(parts)}], tokenize=False,
add_generation_prompt=True, enable_thinking=True) # the reply opens <think>
out = requests.post(GENERATE, json={"model": "trida", "prompt": prompt, "max_tokens": max_tokens,
"temperature": 0.6, "top_p": 0.95, "top_k": 20, "stop": ["</think>"]})
return out.json()["choices"][0]["text"].strip()
def decide_with_escalation(request, threshold=0.7):
answers = decide(request) # direct: one prefill
unsure = [qid for qid, a in answers.items() if confidence(a) < threshold]
if unsure: # one reasoning covers the whole request
thought = reason(request)
state = request["state"]
state = {**state, "reasoning": thought} if isinstance(state, dict) else {"situation": state, "reasoning": thought}
reasoned = decide({**request, "state": state})
answers.update({qid: reasoned[qid] for qid in unsure}) # confident direct answers are kept
return answers
The probabilities of escalated answers come from reasoned mode, which is sharper than direct mode; fit a separate temperature for them if you threshold those probabilities again.
License
Model weights are Apache-2.0. The base model (Qwen/Qwen3.5-4B) is Apache-2.0, and no use restriction propagates from it.
Acknowledgements
- Trained with trida-stack.
- The typed-decision training scheme is inherited from kev.
- Downloads last month
- 2