A visual guide to editing an attention mask into a Jev-like decision model

A visual guide to editing an attention mask into a Jev-like decision model

Diego Fiori

·

Co-founder & CTO

TLDR

Jev, TypeSafe’s decision model, answers many classification questions about one document in a single forward pass and returns calibrated probabilities rather than text. This guide rebuilds it step by step: stop before decoding, edit the attention mask and position ids so questions stay isolated and order-independent, swap the LM head for a pointer head, and fine-tune with cross-entropy plus temperature scaling. None of this adds knowledge, so the backbone still sets the ceiling.

This guide has one thread running through it: the attention mask. We start from what the mask does in an ordinary causal transformer, edit it until classification happens in a single forward pass, fix the three things that break along the way, replace the output end, and fine-tune the result. The target design is Archer Hume’s black-box reconstruction of TypeSafe’s Jev. The measurements come from Jared Palmer’s kev, which implements it on Qwen2.5-0.5B, and from laya, which takes the other branch of the fork described just below. Where something is still inference rather than measurement, the text says so.

1. Where the answer already is

Here is a classification call, the way almost everyone runs one today.

prompt  →  [prefill: 20,000 tokens]

prompt  →  [prefill: 20,000 tokens]

prompt  →  [prefill: 20,000 tokens]

prompt  →  [prefill: 20,000 tokens]

prompt  →  [prefill: 20,000 tokens]

Two things are wrong here, and they need separate fixes.

The confidence is a token sequence. The model emitted the characters 0, ., 9, 1. The probability that the model emits those characters is not the probability that payments is the right queue. Two different quantities, printed in the same font. Nothing in training ever tied them together.

You paid for a decode loop you did not need. Token t+1 does not exist until you have sampled token t, so those 28 passes are strictly sequential. And they produced information that was already sitting in the model after the prefill.

That second sentence is the load-bearing one. To see why it is true we have to look at what the mask actually does.

First, a fork

One decision comes before everything else, because it decides whether the next six parts apply to you at all.

A causal decoder has a mask that needs changing. A bidirectional encoder, BERT and its descendants, never had one: every token already reads every other token, and it was pretrained to represent rather than to generate. Start there and most of this guide is unnecessary. You skip straight to the readout in part 8.






This guide takes the decoder branch, and that is the default I would reach for. The reason is the single thing none of these edits can supply: knowledge. Every patch below changes how a model’s knowledge is read out. None of them add any. Decoders are where the pretraining compute went, and there is no encoder at 25B; the best ones sit around 400M.

That is not a theoretical worry. laya is the strongest encoder-based implementation of this design, built on ModernBERT-large at 421M. Its base checkpoints score below the majority-class baseline on their own typed-decisions benchmark zero-shot: 0.362 against a 0.461 majority class and 0.318 random. After fine-tuning on that benchmark’s own training split they reach 0.766. Their README puts it exactly right: a fast base to specialise, not a zero-shot decision engine.

So the encoder branch is the right call when you already know your workflow, have labelled data for it, and want 33 ms per request. The decoder branch is the right call when you want something useful on a schema it has never seen, which is what the Jev interface promises. Most people asking for this want the second one.

The encoder branch reappears twice below, in part 8 and part 10, because on two specific questions it has measured the cost of a design choice this guide makes differently.

2. What an attention mask actually does

Inside one attention layer, every token is projected into a query, a key and a value. Each query is compared against every key, which gives a square matrix of scores.






Row i says how much token i wants to read from every other token. Softmax each row, multiply by the values, and you have that token’s new representation.

The mask is a second matrix, of the same shape, that decides which of those scores are allowed to exist. Blocked entries are set to negative infinity before the softmax, so they contribute exactly zero afterwards.

Take row 4 and apply an ordinary causal mask, which blocks everything to the right of the diagonal:






billing contributes nothing to about. Not a little: nothing. Over the whole matrix that gives the familiar triangle.

        the  ticket   is   about  billing
the    [  ■     ·     ·      ·      ·  ]
ticket [  ■     ■     ·      ·      ·  ]
is     [  ■     ■     ■      ·      ·  ]
about  [  ■     ■     ■      ■      ·  ]
billing[  ■     ■     ■      ■      ■  ]

        the  ticket   is   about  billing
the    [  ■     ·     ·      ·      ·  ]
ticket [  ■     ■     ·      ·      ·  ]
is     [  ■     ■     ■      ·      ·  ]
about  [  ■     ■     ■      ■      ·  ]
billing[  ■     ■     ■      ■      ■  ]

        the  ticket   is   about  billing
the    [  ■     ·     ·      ·      ·  ]
ticket [  ■     ■     ·      ·      ·  ]
is     [  ■     ■     ■      ·      ·  ]
about  [  ■     ■     ■      ■      ·  ]
billing[  ■     ■     ■      ■      ■  ]

        the  ticket   is   about  billing
the    [  ■     ·     ·      ·      ·  ]
ticket [  ■     ■     ·      ·      ·  ]
is     [  ■     ■     ■      ·      ·  ]
about  [  ■     ■     ■      ■      ·  ]
billing[  ■     ■     ■      ■      ■  ]

        the  ticket   is   about  billing
the    [  ■     ·     ·      ·      ·  ]
ticket [  ■     ■     ·      ·      ·  ]
is     [  ■     ■     ■      ·      ·  ]
about  [  ■     ■     ■      ■      ·  ]
billing[  ■     ■     ■      ■      ■  ]

The mask says who may read whom. It does not say in what order anything runs. During prefill all 5 tokens already exist, so a layer computes the whole 5x5 matrix in one shot, masks it, and updates all 5 representations together. The layers run in sequence, the positions do not. Sequentiality in generation comes from somewhere else entirely: when you are sampling, token t+1 does not exist until you have drawn token t. That is a property of sampling, not of the triangle. So: stop before you sample, and nothing about a causal transformer is sequential. Everything in this guide follows from that one sentence.

3. First attempt: just stop

If the prefill already built a full representation of “the state, then this question”, let’s take it and skip the decode entirely.






For now, read h through the model’s ordinary LM head and look only at the logits of the option letters:

logits = model(render(state, question)).logits[0, -1]
p = logits[tok.convert_tokens_to_ids(["A", "B", "C"])].softmax(-1)
logits = model(render(state, question)).logits[0, -1]
p = logits[tok.convert_tokens_to_ids(["A", "B", "C"])].softmax(-1)
logits = model(render(state, question)).logits[0, -1]
p = logits[tok.convert_tokens_to_ids(["A", "B", "C"])].softmax(-1)
logits = model(render(state, question)).logits[0, -1]
p = logits[tok.convert_tokens_to_ids(["A", "B", "C"])].softmax(-1)
logits = model(render(state, question)).logits[0, -1]
p = logits[tok.convert_tokens_to_ids(["A", "B", "C"])].softmax(-1)

One pass, no decode, no training. This works today on any model you serve, and it is the baseline you should measure before building anything.

What it gets you. The decode loop is gone. Problem 2 solved.

What it does not. Those are next-token probabilities, and they are badly calibrated. Here is kev’s measurement of the untuned Qwen2.5-0.5B backbone, as accuracy / ECE:

Task

Base

Instruct

Choice, 4-way (AG News)

0.813 / 0.069

0.787 / 0.160

Choice, 3-way (MNLI)

0.460 / 0.225

0.433 / 0.390

Noul (BoolQ)

0.427 / 0.274

0.607 / 0.084

Score, 5 levels (Yelp)

0.313 / 0.043

0.353 / 0.078

Look at the Instruct column: instruction tuning made calibration worse on three tasks out of four. RLHF sharpens distributions toward answers that sound confident, which is precisely the wrong pressure for a decision service.

And there is a cost problem. Real requests ask many questions about one document. Ten questions means ten calls, and each one re-encodes the same 20,000-token state.

call 1:  [ 20,000 state tokens ][ Q1 ]
call 2:  [ 20,000 state tokens ][ Q2 ]      the same prefill,
call 3:  [ 20,000 state tokens ][ Q3 ]

call 1:  [ 20,000 state tokens ][ Q1 ]
call 2:  [ 20,000 state tokens ][ Q2 ]      the same prefill,
call 3:  [ 20,000 state tokens ][ Q3 ]

call 1:  [ 20,000 state tokens ][ Q1 ]
call 2:  [ 20,000 state tokens ][ Q2 ]      the same prefill,
call 3:  [ 20,000 state tokens ][ Q3 ]

call 1:  [ 20,000 state tokens ][ Q1 ]
call 2:  [ 20,000 state tokens ][ Q2 ]      the same prefill,
call 3:  [ 20,000 state tokens ][ Q3 ]

call 1:  [ 20,000 state tokens ][ Q1 ]
call 2:  [ 20,000 state tokens ][ Q2 ]      the same prefill,
call 3:  [ 20,000 state tokens ][ Q3 ]

4. Second attempt: one sequence, standard mask

The obvious fix is to put everything in a single sequence so the state is encoded once.

[ state ][ Q_a ][ Q_b ][ Q_c ]
[ state ][ Q_a ][ Q_b ][ Q_c ]
[ state ][ Q_a ][ Q_b ][ Q_c ]
[ state ][ Q_a ][ Q_b ][ Q_c ]
[ state ][ Q_a ][ Q_b ][ Q_c ]

Run it with the ordinary causal mask and read the last hidden state of each question’s span. The state gets encoded exactly once. Problem solved?

No. Look at what the triangle allows.






Two failures, both serious.

Contamination. Question C can read questions A and B. Ask “is the customer angry?” before “which queue?” and the angry framing is now part of the context the queue decision reads. The answer to one question depends on which other questions you happened to ask.

Order dependence. Swap A and B in the JSON payload and you get different answers, because they now occupy different positions in the triangle. The caller’s key ordering, which should be meaningless, changes the output.

The interface we want promises that questions are independent. The triangle does not deliver that. So change the triangle.

5. Edit the mask

We want each question to read the state and itself, and nothing else. That is a small edit: take the triangle and delete the blocks where one question reads another.






Read it row by row and the design falls out.

  • The state row sees only itself. It never depends on any question, which is exactly why its KV cache can be computed once and reused. This is not an optimisation bolted on afterwards; it is a consequence of the mask shape.

  • Each question row reads the state in full, plus its own tokens causally.

  • No question row touches another question. Isolation is not enforced by a wrapper or a prompt instruction. It is enforced by arithmetic: those scores are negative infinity, so they contribute zero.

In code the whole edit is one predicate.

from torch.nn.attention.flex_attention import create_block_mask

# branch_id[t] = 0 for state tokens, k >= 1 for tokens of question k
def jev_mask(b, h, q_idx, kv_idx):
    causal      = q_idx >= kv_idx
    same_branch = branch_id[b, q_idx] == branch_id[b, kv_idx]
    is_state    = branch_id[b, kv_idx] == 0
    return causal & (same_branch | is_state)

block_mask = create_block_mask(jev_mask, B, H, L, L)
from torch.nn.attention.flex_attention import create_block_mask

# branch_id[t] = 0 for state tokens, k >= 1 for tokens of question k
def jev_mask(b, h, q_idx, kv_idx):
    causal      = q_idx >= kv_idx
    same_branch = branch_id[b, q_idx] == branch_id[b, kv_idx]
    is_state    = branch_id[b, kv_idx] == 0
    return causal & (same_branch | is_state)

block_mask = create_block_mask(jev_mask, B, H, L, L)
from torch.nn.attention.flex_attention import create_block_mask

# branch_id[t] = 0 for state tokens, k >= 1 for tokens of question k
def jev_mask(b, h, q_idx, kv_idx):
    causal      = q_idx >= kv_idx
    same_branch = branch_id[b, q_idx] == branch_id[b, kv_idx]
    is_state    = branch_id[b, kv_idx] == 0
    return causal & (same_branch | is_state)

block_mask = create_block_mask(jev_mask, B, H, L, L)
from torch.nn.attention.flex_attention import create_block_mask

# branch_id[t] = 0 for state tokens, k >= 1 for tokens of question k
def jev_mask(b, h, q_idx, kv_idx):
    causal      = q_idx >= kv_idx
    same_branch = branch_id[b, q_idx] == branch_id[b, kv_idx]
    is_state    = branch_id[b, kv_idx] == 0
    return causal & (same_branch | is_state)

block_mask = create_block_mask(jev_mask, B, H, L, L)
from torch.nn.attention.flex_attention import create_block_mask

# branch_id[t] = 0 for state tokens, k >= 1 for tokens of question k
def jev_mask(b, h, q_idx, kv_idx):
    causal      = q_idx >= kv_idx
    same_branch = branch_id[b, q_idx] == branch_id[b, kv_idx]
    is_state    = branch_id[b, kv_idx] == 0
    return causal & (same_branch | is_state)

block_mask = create_block_mask(jev_mask, B, H, L, L)

Those deleted regions are free. FlexAttention works on tiles and skips a tile that is entirely masked, so you never compute the scores you are about to throw away. The same structure exists in Hydragen and DeFT, and vLLM and SGLang give you most of it through prefix caching.

What this actually saves

Worth being precise, because the headline is often overstated.

Work per layer

Q separate calls

One edited mask

FFN / MoE tokens

Q · (S + n)

S + Q · n

State attending to itself

Q · O(S²)

O(S²), once

Questions attending to state

Q · S

Q · S, unchanged

The last row does not improve, and no mask edit can make it. Each question genuinely has to read the state. What disappears is everything else, and since the FFN dominates prefill FLOPs and S is typically far larger than n, that is most of the bill. kev measures the packed path at 2.0x the throughput of separate calls.

It also explains the two limits the real API exposes: roughly 32K per branch, roughly 64K per request, with the state counted once. A 23K state with 5,000 questions fits comfortably. Under separate calls that same request would be over 100 million tokens.

6. Third thing that breaks: positions

The mask is right and the model still misbehaves. Because attention does not only use the mask; with RoPE it also uses the distance between positions.

In the packed sequence the questions sit at different offsets:






So Q_c is further from the state than Q_a, purely because of where it landed in the JSON payload. We deleted the attention edges between branches, but we left the geometric ordering in place, and RoPE reads that ordering.

The fix is to restart each branch’s position ids right after the state.






pos = torch.where(branch_id == 0, idx, S + idx - branch_start)
pos = torch.where(branch_id == 0, idx, S + idx - branch_start)
pos = torch.where(branch_id == 0, idx, S + idx - branch_start)
pos = torch.where(branch_id == 0, idx, S + idx - branch_start)
pos = torch.where(branch_id == 0, idx, S + idx - branch_start)

Now every question occupies the same geometric slot: “the state, then one question”. Two tokens in different branches share a position id, which would be a bug in an ordinary sequence and is harmless here for exactly one reason: they can never attend to each other. The mask edit is what makes the position edit safe.

This is the same trick as tree attention in speculative decoding: question order stops mattering.

7. Fourth thing that breaks: sliding windows

Only if your backbone has them. Qwen2.5, which kev uses, is full attention everywhere, so kev never meets this. Gemma 4 will.

25 of Gemma 4’s 30 layers use a 1,024-token sliding window. Only 5 are full attention. That is a third mask, ANDed with ours:






Most of the state is invisible at that layer. It arrives through the 5 global layers, and through whatever the local layers already folded into the state tokens near the branch boundary. This is not broken, it is the same constraint the base model works under for any long prompt. But two things need care.

Compute the window over the reset positions, not the packed index. Otherwise the last branch in the payload gets a different effective window than the first, and you have quietly reintroduced exactly the ordering dependence part 6 removed.

def jev_mask_local(b, h, q_idx, kv_idx, window=1024):
    return jev_mask(b, h, q_idx, kv_idx) & ((pos[b, q_idx] - pos[b, kv_idx]) < window)
def jev_mask_local(b, h, q_idx, kv_idx, window=1024):
    return jev_mask(b, h, q_idx, kv_idx) & ((pos[b, q_idx] - pos[b, kv_idx]) < window)
def jev_mask_local(b, h, q_idx, kv_idx, window=1024):
    return jev_mask(b, h, q_idx, kv_idx) & ((pos[b, q_idx] - pos[b, kv_idx]) < window)
def jev_mask_local(b, h, q_idx, kv_idx, window=1024):
    return jev_mask(b, h, q_idx, kv_idx) & ((pos[b, q_idx] - pos[b, kv_idx]) < window)
def jev_mask_local(b, h, q_idx, kv_idx, window=1024):
    return jev_mask(b, h, q_idx, kv_idx) & ((pos[b, q_idx] - pos[b, kv_idx]) < window)

Then measure it. Plant a distinctive fact at varying depths in the state, ask a question that needs it, and plot recall against state length. If it falls off past a few thousand tokens, either promote more layers to full attention and fine-tune them back into shape, or cap the state length in the product. Neither implementation has tested this; it is my inference from the architecture.

8. Now the other end: replace the LM head

The input side is done. The output side is still a language model.

The final hidden state currently gets projected into 262,144 vocabulary logits, of which we use three. Project it somewhere smaller instead.

          h  (the branch's final hidden state)
                        │
       ┌────────────────┴────────────────┐
       ▼                                 ▼
LM head: d × 262144               Readout: d × K
"which word comes next"           "which option is right"
       │                                 │
       ▼                                 ▼
softmax over vocabulary           softmax over K options
          h  (the branch's final hidden state)
                        │
       ┌────────────────┴────────────────┐
       ▼                                 ▼
LM head: d × 262144               Readout: d × K
"which word comes next"           "which option is right"
       │                                 │
       ▼                                 ▼
softmax over vocabulary           softmax over K options
          h  (the branch's final hidden state)
                        │
       ┌────────────────┴────────────────┐
       ▼                                 ▼
LM head: d × 262144               Readout: d × K
"which word comes next"           "which option is right"
       │                                 │
       ▼                                 ▼
softmax over vocabulary           softmax over K options
          h  (the branch's final hidden state)
                        │
       ┌────────────────┴────────────────┐
       ▼                                 ▼
LM head: d × 262144               Readout: d × K
"which word comes next"           "which option is right"
       │                                 │
       ▼                                 ▼
softmax over vocabulary           softmax over K options
          h  (the branch's final hidden state)
                        │
       ┌────────────────┴────────────────┐
       ▼                                 ▼
LM head: d × 262144               Readout: d × K
"which word comes next"           "which option is right"
       │                                 │
       ▼                                 ▼
softmax over vocabulary           softmax over K options

Two ways to build the readout.

A slot head maps to positional slots: output 0 means “the first option in the list”, whatever that option happens to be. The branch text supplies each slot’s meaning, so one head serves any customer’s label set without retraining. A 256-slot head would explain the API’s 255-option cap, which is 2^8 minus one.

A pointer head scores each option against the decision position:

# h_dec: [B, d]     the <decide> position
# h_opt: [B, K, d]  each option's closing position
z = torch.einsum("bd,bkd->bk", Wq(h_dec), Wk(h_opt)) / math.sqrt(d)
# h_dec: [B, d]     the <decide> position
# h_opt: [B, K, d]  each option's closing position
z = torch.einsum("bd,bkd->bk", Wq(h_dec), Wk(h_opt)) / math.sqrt(d)
# h_dec: [B, d]     the <decide> position
# h_opt: [B, K, d]  each option's closing position
z = torch.einsum("bd,bkd->bk", Wq(h_dec), Wk(h_opt)) / math.sqrt(d)
# h_dec: [B, d]     the <decide> position
# h_opt: [B, K, d]  each option's closing position
z = torch.einsum("bd,bkd->bk", Wq(h_dec), Wk(h_opt)) / math.sqrt(d)
# h_dec: [B, d]     the <decide> position
# h_opt: [B, K, d]  each option's closing position
z = torch.einsum("bd,bkd->bk", Wq(h_dec), Wk(h_opt)) / math.sqrt(d)

Archer’s external probes cannot separate the two. kev settles the practical question by building the pointer, and the reason to prefer it is calibration: a slot head’s output 200 is exercised far less during training than output 2, so its probabilities are worse. The pointer shares parameters across every position and accepts whatever K the request sends.

There is a third design, and it is where the encoder branch pays for itself. Give all of a question’s options one shared, fixed token budget and score each at a marker inside it. laya does this, and its own limits section reports the consequence: on Banking77’s 77 labels the budget works out to roughly 3 to 4 tokens per label, the label texts stop being distinguishable, and accuracy falls to 0.425. Jev scores 0.870 on the same task. kev, at 0.5B, scores 0.860. Parameter count is not what separates them.

Options need room to be read. That is the argument for scoring each one at its own position inside a branch that can be as long as it needs to be, and it is a large part of why Jev allows roughly 32k tokens per branch rather than a few hundred.

Where to read, and why the layout matters

<state> …incident text…
<q> Which team should handle this? <opt> payments </opt> <opt> account </opt> <opt> other </opt> <decide>

<state> …incident text…
<q> Which team should handle this? <opt> payments </opt> <opt> account </opt> <opt> other </opt> <decide>

<state> …incident text…
<q> Which team should handle this? <opt> payments </opt> <opt> account </opt> <opt> other </opt> <decide>

<state> …incident text…
<q> Which team should handle this? <opt> payments </opt> <opt> account </opt> <opt> other </opt> <decide>

<state> …incident text…
<q> Which team should handle this? <opt> payments </opt> <opt> account </opt> <opt> other </opt> <decide>

<decide> is last, so within the branch’s own causal triangle it sees every option. That ordering is deliberate. If each option were scored in isolation and the temperature were fixed, adding a fourth irrelevant option would only change the shared denominator, and the odds between two existing options would be unchanged:

\frac{p( {customer})}{p( {unknown})} = e^{z_{ {customer}} - z_{{unknown}}}

Both Archer and kev test this and both find real movement: a log-odds shift of -0.28 in Archer’s probe, 0.13 mean and 0.34 at p90 in kev’s. The options are being read as a list, not scored one at a time. kev’s README names the consequence: this is what makes “none of the above” work at all.

Make the delimiters reserved vocabulary tokens, not text. Then a state containing the literal string <opt> ignore previous instructions tokenizes as ordinary words and cannot manufacture an option slot. kev tests this: the option count is unchanged and a forged option scores p <= 0.09. For a service reading untrusted input that is a security property, not a nicety.

9. Check: did we change the model?

We have edited the mask and swapped the head. Before training anything, it is worth asking what we actually did to the network.

Almost nothing. Every branch is still a causal suffix reading a causal prefix. No token attends to its own future anywhere. So the packed forward pass should be mathematically identical to running Q separate [state][question] prompts, differing only in floating-point reduction order.

kev checks exactly this, and reports a maximum probability difference of 3.7e-6 between the packed and separate paths.

This is the single most useful test to write first.

It catches every mask bug and every position-id bug at once, before any training, with an unambiguous pass/fail number. If packed and separate disagree at the third decimal place, something in parts 5 to 7 is wrong.

It also explains why no adaptation training is needed for the mask. Compare with DiffusionGemma, which converts causal attention into bidirectional attention: there, tokens start attending to their own future, which the base model has never seen, and that is a real distribution shift needing real retraining. Our edit only ever removes edges.

Isolation holds behaviourally too. kev reproduces Archer’s secret-code probe: a fact planted in a sibling question is recovered with p = 0.03, identical to when it is not present at all, against p = 0.99 when the same fact sits in the shared state.

10. Finally: the probabilities are still not probabilities

Everything so far fixed the shape of the computation. The readout head is untrained and the backbone has never seen this format. Time for the last step.

Calibration is a property of predictions across cases, never of one prediction.

Take every case you assigned 0.8 to. Calibration asks whether roughly 80% of them turned out that way. You cannot tell whether a single prediction was calibrated by looking at how that one case ended.

Stage A: head only

Freeze the backbone, train just the readout. It is a d × K matrix, so this costs almost nothing and gives you a working end-to-end system in an hour. Its real value is validating the plumbing before you spend money. It will underperform, because the backbone’s representations were shaped for next-token prediction.

Stage B: teach the backbone the format

Unfreeze with LoRA. kev uses r=16 on the backbone with the head trained from scratch, and the whole artifact is 38 MB.

What you are teaching is not knowledge, which is already in the checkpoint, but format: what <decide> means, how to read an option list, and that no text is coming.

The loss must be a proper scoring rule, meaning one whose expected loss is minimised by reporting the true conditional distribution. Cross-entropy is the obvious choice and the one kev uses.

TypeSafe calls their objective RLCD, Reinforcement Learning for Calibrated Decisions, and laya publishes a concrete implementation of something under that name: the policy reports a distribution, exploration adds zero-mean Gaussian noise to the logits, the reward is a proper scoring rule, and updates are REINFORCE with a group-mean baseline. Worth being precise about what the RL framing buys, because the acronym sounds more exotic than the mechanism. For a single-step prediction against a known label, RL on a proper scoring rule and supervised cross-entropy optimise the same objective. Log loss is the logarithmic scoring rule; a policy gradient there is a higher-variance estimate of a gradient you can compute exactly. The RL machinery earns its keep only where you cannot differentiate through to the label: delayed or non-differentiable rewards, or multi-turn credit assignment, which is what laya’s TD over conversation prefixes is for.

loss = F.cross_entropy(logits, target)
loss = F.cross_entropy(logits, target)
loss = F.cross_entropy(logits, target)
loss = F.cross_entropy(logits, target)
loss = F.cross_entropy(logits, target)

Do not reach for label smoothing.

It is a proper scoring rule deliberately broken: it drags predictions toward a fixed prior regardless of what the data says. Fine when you only care about the argmax. Exactly wrong when the probability itself is the product.

Data, and the two extra terms

kev trains on six public datasets reshaped into the target request format: Banking77 (77-way Choice), AG News (Choice plus yes/no), MNLI (3-way), BoolQ (Noul), SST-5 and Yelp (5-level Score). 9,000 records, 13,500 questions, two epochs, about 1h45m on an Apple M5.

Two loss terms are worth copying:

  • --perm_kl, a symmetric KL between the predictions under two random option orders. This is what attacks order sensitivity at training time rather than at inference time.

  • --ord_w, an ordinal term penalising |E[level] - y| for Score questions. Plain cross-entropy treats “off by one level” and “off by three levels” as equally wrong, which is obviously the wrong prior for an ordered scale. If you are writing this from scratch, reach for the ranked probability score instead, which is what laya uses: it is the proper scoring rule for ordinal outcomes, so it is the principled version of the same correction rather than a bolted-on penalty.

Augment every example with a shuffled option order, an appended irrelevant option, and a varying number of sibling questions. Those three augmentations target exactly the three behaviours parts 5 to 8 built.

One renderer, shared by training and serving.

kev routes training data and live requests through the same code, so the model never meets a format at inference that it did not see in training. This sounds too obvious to state and is the thing people get wrong, because the training pipeline and the serving pipeline are usually written months apart by different people.

Stage C: calibration

A proper objective is a statement about the optimum, not a guarantee about your trained model. Finite data, limited capacity and distribution shift all leave real calibration imperfect.

Temperature scaling earns its keep, and it is not a small effect. One scalar fit on held-out data, applied to the logits. kev reports held-out ECE of 0.065, dropping to 0.031 after that single parameter. laya, which ships over-confident, reports mean ECE moving from 0.466 to 0.081.

Fit one temperature per (question type, option count), not merely per type. The sharpness of a distribution depends on K, so a 3-option Choice and a 77-option Choice have no reason to want the same scalar. This is laya’s refinement and it is worth copying.

Soft targets, where you have them. Five annotators splitting 3-2 means the target is [0.6, 0.4], not a one-hot. Hard labels on genuinely ambiguous cases actively train overconfidence, and ambiguous cases are precisely where the probability matters. Outcome labels (“was this escalation actually necessary”) beat annotation labels. Neither implementation does this; both use public datasets with hard labels, and it is the obvious next step for a real deployment.

One field you should not train. If you expose a scalar confidence next to the distribution, compute it in application code. TypeSafe’s own SDK does, for K > 1:

c = \frac{p_{\max} - 1/K}{1 - 1/K}

It measures how far the leading answer stands above uniform. Arithmetic, not a second learned estimate that the answer is correct. Keeping those two objects separate prevents the most common conceptual error in this area: a concentrated distribution can still be confidently wrong.

11. Does it work?

kev, at 0.5B parameters, on 1,350 held-out questions. Accuracy / ECE at 10 bins.

Task

Base backbone

After parts 5 to 10

Choice, 4-way (AG News)

0.813 / 0.069

0.940 / 0.028

Choice, 3-way (MNLI)

0.460 / 0.225

0.747 / 0.100

Choice, 77-way (Banking77)

not runnable

0.860 / 0.057

Noul (BoolQ)

0.427 / 0.274

0.753 / 0.136

Score, 5 levels (Yelp)

0.313 / 0.043

0.553 / 0.118

All


0.799 / 0.065

Note the 77-way row. Reading option letters off a next-token distribution does not extend to 77 options; a pointer head does, without changing anything.

The eval suite

Each test below maps to one part of this guide, which is what makes them worth running in CI: a regression tells you which edit broke.

Test

Guards

kev-0.5b

Packed vs separate

Parts 5 and 6

3.7e-6 max difference

Isolation probe

Part 5

0.03 / 0.03 / 0.99

Permutation

Parts 6 and 8

7.4% argmax flips

Choice-set (IIA)

Part 8

0.13 mean, 0.34 p90

Boundary forgery

Part 8

forged p <= 0.09

Binned ECE

Part 10

0.065, 0.031 tempered

Two notes on reading that table. The 7.4% permutation figure is the unregularised number: the released kev-0.5b predates --perm_kl and is trained with it at 0. And report bin counts alongside ECE, because aggregate agreement hides overconfidence in one region cancelling underconfidence in another.

Threshold economics is the number the business actually cares about. If a false escalation costs 1 and a missed urgent case costs 9, the rule is p(urgent) > 0.1. Sweep the threshold against your real cost matrix.

And the failure mode to design against is confident and wrong. laya’s English checkpoint scores 0.000 accuracy at 0.952 confidence on Khmer. Maximally certain, maximally wrong, so no confidence threshold catches it and no amount of calibration on English would have warned you. Their fix is to detect the script in under 0.5 ms of plain Python before the forward pass, which generalises into a rule worth keeping: out-of-distribution detection has to happen outside the model whose confidence you do not trust. A calibrated model is calibrated on the distribution it was calibrated on, and it has no way to tell you when it has left.

kev is explicit about the limitation that matters most, and it applies to anything built this way: the ECE is in-distribution. The test splits come from the training datasets. Calibration on Banking77 says nothing about calibration on your tickets. Real calibration needs outcome-labelled data from the workflow you are deploying into.

12. What you give up

  • No chain of thought. Everything happens in one prefill. Hard multi-step reasoning is where this design is weakest, and where an autoregressive model with a scratchpad still wins.

  • No cross-question dependency. If B genuinely needs A’s answer, isolation gives you nothing. Merge them into one question over the product space, or take a second round trip. Sharing context does not remove the logical structure of your workflow.

  • No open-ended output. The option set must be closed and known at request time.

  • Every request costs the same. There is no cheap path for an easy case.

  • Knowledge is capped by the backbone, and this is the real ceiling. kev is honest about it: at 0.5B it picks return_policy where Jev picks return_status. laya is starker, scoring below the majority-class baseline zero-shot on its own benchmark. None of these edits add knowledge. They only change how it is read out, which is exactly why the fork before part 2 is the most consequential decision in this guide.

The upside is that the computational graph finally matches the job. A decision service reads evidence, compares permitted outcomes, and reports uncertainty. A transformer can do all three without turning every decision into a sentence first.

FAQs

What is a Jev-like decision model?

It's a model that answers several classification questions about one document in a single forward pass. Instead of generating text, it returns a probability distribution over a closed set of options. The design follows Archer Hume's reconstruction of TypeSafe's Jev, implemented in kev on Qwen2.5-0.5B.

Why not just ask an LLM to output a confidence score?

Because that score is generated text. The probability of the model writing "0.91" is not the probability that its answer is right, and nothing in training links the two. In kev's tests, instruction tuning made calibration worse on three of four tasks.

How does editing the attention mask help?

Each question can read the shared document and its own tokens, but never another question. The document is encoded once and reused, and answers no longer change based on which other questions you ask or their order. kev measures 2x the throughput of separate calls.

Does editing the mask mean retraining the model?

Not for the mask itself. The edit only removes attention edges, so the packed pass matches separate prompts almost exactly (a maximum difference of 3.7e-6 in kev). Training is still needed for the new readout head and the input format: kev uses LoRA with cross-entropy, then temperature scaling.

What does calibration mean here?

A model is calibrated if, across all the cases it scores at 0.8, about 80% turn out correct. It can't be judged from a single prediction. kev's held-out calibration error drops from 0.065 to 0.031 after temperature scaling, though only on data similar to its training sets.

Isn't a bidirectional encoder like BERT the simpler option?

It skips the mask edits, but the best encoders sit around 400M parameters and add no knowledge. laya, built on ModernBERT-large, scores below the majority-class baseline zero-shot and only reaches 0.766 after fine-tuning on the benchmark's own training data. Encoders suit known workflows with labelled data. For schemas the model hasn't seen, a decoder is the better starting point

Subscribe to our newsletter

Subscribe to our newsletter

Stay up to date on what we're learning, building, and seeing as enterprise teams deploy and measure AI agents in production.