Build a Jev-Style Decision Model from an LLM
Part 1 ended with a limit. Logit bias and guided choice both get a decision out of an ordinary LLM in one pass. Both read their scores from the vocabulary. That is where label mass comes from. Label mass is the share of the model’s probability that lands on your labels. The model may never have wanted your labels and the label-only softmax hides this. The last recommendation in Part 1 was a native scoring endpoint. It scores options directly instead of borrowing the vocabulary head.
This series builds the core of that endpoint from scratch and logs what went wrong along the way.
The design does not read the vocabulary at all. We write the question and the options into one token sequence in a fixed layout. We run one forward pass. This is a single trip through the model without generating text. Then we read the model’s internal vector at the token where each option ends. A small trained head scores each option against a final “decide” vector. A softmax then turns the scores into probabilities. Three things from Part 1 stop applying. Labels do not have to be single tokens. There is no top-N window to fall out of. Label mass does not exist, because no vocabulary softmax is involved. One thing replaces them: the head has to be trained and training can go wrong in new ways. That is the subject of Parts 3 and 4.
The idea is not ours. It follows a published decision-model design (called Jev) and an open-source reimplementation. What we did was rebuild every piece on a small open model. The goal was to learn what is inside such a system, where it breaks and what each design decision costs.
The setting is deliberately modest: one laptop, one 8 GB GPU and a model with 0.85 billion parameters that also understands images. Images matter later. The goal is a decision model that can look at a screenshot and answer typed questions about it.
The software that calls the model asks three kinds of question. Pick one from a list (choice). Answer yes or no (noul). Or place a level on an ordered scale such as calm, frustrated, furious (score). All three are the same primitive: a list of options and one probability for each. A yes/no is a choice with two options. A score is a choice whose options have an order.
The series in three parts, each readable on its own:
The whole system in one picture. This article covers the first three steps.
Setup used throughout the series:
A vision-language model sounds like one thing. Loaded into Python, it is a tree of modules. The first job was to walk that tree and read it. One term will come up often. A hidden vector is the list of 1,024 numbers the model holds for one token at one stage of processing.
The first surprise was what was missing. lm_head did not show up in the count at all. It is a real linear layer with its own name in the module tree. The explanation is weight tying. The input uses a matrix to turn a token id into a vector. The output uses a matrix to turn a hidden vector into vocabulary scores. These two matrices are the same tensor (one array of numbers stored once). PyTorch's parameter listing skips duplicates, so the shared tensor was counted once, under whichever module appeared first. A one-line identity check confirmed it and came back True.
Nearly a third of the model is one lookup table that serves as both the front door and the back door. This table is called the embedding matrix: one row of 1,024 numbers for each token in the vocabulary. Keep that in mind for the section on what to remove.
A normal transformer repeats one block 24 times. This model repeats a pattern of four:
layers 0, 1, 2 : Gated DeltaNet (a "linear attention" layer)layer 3 : Gated Attention (classic softmax attention)layers 4, 5, 6 : Gated DeltaNetlayer 7 : Gated Attention... (six repetitions)
We checked this three ways, because documentation is easy to misread. The model card said so. The printed list of layer classes showed it. The configuration listed layer_types explicitly. Every layer, of either kind, is followed by the same feed-forward block (a small network applied to each token on its own). Each layer is also wrapped in RMSNorm layers, which rescale the numbers to a steady size.
Three layers of one kind, then one of the other. The rest of this article is about what that ratio forbids.
If your deep learning background is convolutional networks, here is the shortest path to attention.
A convolution applies the same fixed filter at every token. What a token sees is decided by the architecture. It is a small neighbourhood with weights that are learned once and then frozen.
Attention throws away the fixed neighbourhood. Every token produces three vectors: a query (what am I looking for), a key (what do I contain) and a value (what do I hand over if you pick me). A token scores how well each key matches its query. A softmax turns those scores into shares. The token then takes an average of the values, weighted by those shares. In symbols:
Q, K and V hold the queries, keys and values of all tokens. d_k is the size of a key vector. M is the mask described below.
The attention shares are computed fresh for every input, from the content. The learned weights of the network are different because they stay fixed after training. That is the conceptual leap: a filter that depends on the data.
Here is the same idea with plain numbers. Imagine a support ticket where a customer says they were billed twice and wants their money back. It contains the phrase “charged twice, refund”. The model is working on the word “refund”. For each word it can look at, it asks how useful that word is for understanding “refund” here. The answer is one match score per word. Every word also hands over a note, which in this toy is one number: how urgent that word sounds, from 0 to 10.
Softmax turns the scores into shares that add up to 1. Each score is exponentiated, giving e² = 7.39, e¹ = 2.72 and e⁰ = 1.00. Each result is then divided by their total (7.39 + 2.72 + 1.00 = 11.11) resulting in share as charged=0.67, twice=0.24 and refund=0.09. The result for “refund” is a blend of the notes. Each note counts in proportion to its share: 0.67 x 9 + 0.24 x 5 + 0.09 x 1, which is about 7.3. Put a different sentence in and the scores change, so the shares change. That is the filter that depends on the data.
During generation a token must not look at the future and the mask enforces that. It holds minus infinity for every forbidden pair. Since the exponential of minus infinity is zero, those pairs get a share of exactly 0. If a word “soon” followed “refund”, its score at this token would be minus infinity. Its share would be 0 and it could not change the 7.3. This is the same idea as the causal convolutions used for audio, applied to attention shares. An attention mask is nothing more than an include/exclude list for pairs of tokens.
At generation time the model emits one token at a time. Past keys and values never change, so they are stored in a KV cache instead of recomputed. Think of that cache as a notebook with one page per past token. It grows with the sequence and each page can be individually included or excluded by a mask. Remember that property. It decides a lot later.
The Gated DeltaNet layer solves a different problem. Attention’s cost grows with sequence length, because every token looks at every earlier one. A DeltaNet layer instead carries a fixed-size memory through the sequence and updates it token by token. The closest older idea is the hidden state of an LSTM or GRU, two recurrent networks that were common before transformers. The layer runs several heads in parallel. For each head the memory is a small matrix, updated by:
S is the memory, a small matrix. k, v and q are the key, value and query of the current token. beta is the write strength and alpha is the forget gate.
In words: the term in brackets erases what the memory currently stores for key k. This is the “delta rule” part. It corrects the memory by the prediction error, the way a perceptron update corrects a classifier by its mistake. The second term writes the new value v under that key. The forget gate is the “Gated” part and works like an LSTM gate.
The smallest possible memory makes this concrete. It is one slot holding one number, the customer’s mood on a scale from 0 (calm) to 10 (furious). With a single slot the key is just 1 and the formula shrinks to:
The slot currently holds 2. A new message suggests a mood of 8. Leave the forget gate open (alpha = 1) and vary the write strength:
The rule behind the table is: new = old + strength x (new value minus old). For a strength of 0.5 that is 2 + 0.5 x (8–2) = 5. The write strength works like a learning rate for the memory. The forget gate is a separate dial that fades the old content before anything is written. With alpha = 0.5 and a write strength of 0, the slot becomes 0.5 x 2 = 1.
The real memory is a 128 by 128 grid of such numbers per head. Its keys and values are vectors instead of single numbers. The idea is unchanged and so is the size: 16,384 numbers (128 x 128) per head after two tokens or after two million. A short causal convolution sits in front of the recurrence. It mixes the last few tokens locally before the memory is touched.
The practical difference between the two layer types is this. A KV cache is the notebook from earlier, so it grows with the sequence and a mask can cross out any page. A DeltaNet memory is a whiteboard that gets rewritten as the text goes by, so it has a constant size and has already blended everything it has seen.
Diagrams are easy to believe. So we ran a 28-token prompt through the model and inspected the cache that comes back. It is a list of 24 entries. Their class names split along the same 3-and-1 line: LinearAttentionLayer for the memory layers and DynamicLayer for the attention layers. Here is what came back:
The number 28 appears in the attention shapes and nowhere in the DeltaNet shapes. Here is how the odd-looking numbers decode, all confirmed against the configuration:
Now count what each layer type has to store as the input grows. An attention layer stores 2 x 2 x 256 = 1,024 numbers per token: keys and values for 2 heads of size 256. A DeltaNet layer stores 16 x 128 x 128 + 6144 x 4 = 262,144 + 24,576 = 286,720 numbers at any length.
In notebook terms, the attention notebook gains a page per token. The DeltaNet whiteboard never changes size. At 28 tokens the “constant” whiteboard is ten times bigger than the notebook (286,720 numbers against 28,672). The two cross at 280 tokens, where both hold 286,720. Across the whole model at maximum length, the 6 attention layers would hold 1.61 billion numbers (6 x 268,435,456). That is 3.2 GB in bf16. All 18 DeltaNet layers together hold only 5.2 million numbers (18 x 286,720). That is about 10 MB. That is the reason the design exists: long inputs stay affordable. On short inputs, like a support ticket, there is no benefit to see.
An honest list, because these are the moments where understanding was slow:
A support ticket usually carries several questions: which team, is it urgent, how upset is the customer. An efficient design would write the ticket once, then add the three questions after it in one sequence. A mask would let each question see the ticket and itself, but not the other questions. That is a block-causal mask and it computes all three answers in one pass. In a pure-attention model it works, because the mask is just an include/exclude list.
For this model it cannot work and we proved it from the source instead of trusting a claim. The DeltaNet layer’s code does accept an attention_mask, which looked promising. Reading it, the mask goes to a helper that zeroes the vectors at padding tokens before the recurrence. It can say "this token does not exist". It cannot say "this token exists, but the third question may not read it". By the time the third question is processed, the memory has already blended in everything before it. The attention layers, by contrast, feed the mask straight into the score computation.
In the notebook and whiteboard picture from earlier, you can cross out any page of the notebook before reading. You can ask the whiteboard to ignore blank space. You cannot ask it to forget question 2 and keep the rest, because question 2 is already mixed into what is written.
Same parameter name, completely different job. The consequence: questions about the same ticket must be computed in separate rows, each as “ticket plus one question”. We checked that this rows mode matches the reference computation exactly (torch.equal returns True). So it is the baseline everything else is compared against.
Yes, in principle. Run the ticket once, save the cache, copy it per question and continue each copy with that question’s tokens. Nothing has to be masked because the copies never see each other. Compare the work with round numbers: a ticket of 100 tokens and three questions of 10 tokens each.
The longer the ticket compared with the questions, the more the repeat costs. Two details matter. Use a real copy of the cache. A lightweight view of the same memory would let one question overwrite another’s state. And check that tokenizing the ticket alone gives the same tokens as the first part of the tokenized ticket plus question. Text can tokenize differently at the join. That is one reason we later used dedicated marker tokens.
We ran the experiment and compared the cached version against the plain full pass. The final-layer vector differed by up to 0.375. That felt too large to be rounding noise. So we went hunting. We checked position ids (the index given to each token) and rotary encoding (how the model uses those indexes). We also checked cache flags, copy semantics and the convolution code. Each came back clean. Passing explicit position ids changed nothing, because the model derives them from the cache length anyway. Feeding tokens one at a time changed nothing either.
A separate investigation ran the identical code in 32-bit floats. The difference collapsed. There was no bug.
Part 1 met this effect from the other side, when a bias of 100 was added to scores. The rule is the same: bf16 has 8 bits of precision, so its rounding step depends on the size of the number. Our final-layer vector has entries up to about 30 and this is what the spacing looks like at that scale:
A difference of 0.25 to 0.375 is therefore one to three rounding steps. Two mathematically equivalent computations round differently along the way. One runs everything as one chunk. The other continues from a saved state.
Think of a kitchen scale that, for anything above 32 kg, can only show multiples of 0.25 kg. You start with 33 kg and add two scoops of 0.1 kg. Add them one at a time and the scale rounds each time. The first scoop makes 33.1, which shows as 33.0. The second scoop makes 33.1 again, which shows as 33.0 again. So the scale still says 33.0. Add the two scoops together first (0.2 kg) and then put them on: 33.2 shows as 33.25.
The true total is 33.2. Same additions, different order, readings one full step (0.25) apart. Our two paths differ in order in exactly this way, thousands of times per layer.
A repeat of the identical call gave exactly 0.0 difference, so the model is deterministic. The gap was systematic rounding, not randomness. Against an fp32 ground truth, the cached path was no further from the truth than the plain path. The bf16 baseline sat 0.242 and 0.203 away. The cache-reuse path sat 0.231 and 0.163 away. The cosine similarity between the two bf16 paths was 0.9999 (relative error 1.4%).
Our own reasoning was the problem. We judged “0.375 is big” without asking “big compared to what, in this number format, at this magnitude?”. Our rule now is to compare in fp32 before concluding anything is broken or to use scale-aware measures. Cache reuse itself is correct. We kept separate rows as the simple reference and the reuse path as a verified optimization.
We never generate text, so the vocabulary head is dead weight. But the tying matters: its weights are the embedding matrix, still needed at the input. So removing the head saves no parameters. It saves activations. Running the head over every token creates a tensor with one score per vocabulary entry per token.
All of that is allocated before the loss even starts. We call the inner model directly and never build it.
The distinction is worth remembering. Parameters are the stored weights. Activations are the temporary tensors a forward pass creates. Removing a layer can change both or only one of them. We also confirmed that the vector we read is the final one. The last entry of the returned hidden states equals the model’s final output. So the final normalization is already applied.
The model has to produce a vector we can score. We design the input so that specific tokens play specific roles, using five special marker tokens. Here is the layout of a department question:
<STATE> ...ticket text... <QUESTION> Which team should handle this?<OPT> billing <END> <OPT> technical <END> <OPT> sales <END> <DECIDE>
The model produces one 1,024-number vector for every token. We keep four of them and throw the rest away:
Scores will come from comparing vectors 1 to 3 with vector 4. The markers must each be exactly one token. If a marker splits into several ordinary tokens, there is no single token to read.
This is also why an option can be a whole phrase. Part 1 needed single-token labels because the score was a token’s logit, the raw score the model gives that token. Here the score comes from the vector at the end marker, so “billing: charges, invoices and refunds” works as well as “billing”.
We needed five tokens that mean nothing to the model. The options were to add five new tokens, which start with random embeddings or to reuse rarely used existing special tokens. This tokenizer has 33 special tokens and none of them is a blank spare. Several belong to features unused here: fill-in-the-middle markers and bounding-box markers. We reused those, which avoids growing the embedding matrix. In practice the end of an option is the bounding-box end token. The decide token is a fill-in-the-middle suffix token. The two strings are <|box_end|> and <|fim_suffix|>. Tokens that mark reasoning or tool calls would have been a poor choice, since the model has strong habits around them.
Our first attempt got this wrong in an instructive way. We built a lookup from marker names to token strings, then accidentally tokenized the names instead of the strings. Names such as <state> and <q> are ordinary text and split into three tokens each. The check meant to catch it was verifying the wrong side of the dictionary. A marker map has to go from concept to the exact registered token string. You tokenize the string and you assert that one id comes out.
Text from users goes into the same sequence as our markers. If a customer writes the literal marker string inside their message, can they forge a decision boundary? We expected the tokenizer’s add_special_tokens=False setting to prevent it. It does not. Fast tokenizers match registered special-token strings in raw text regardless of that flag. A test confirmed that user text containing a marker string produced the marker's id.
The fix has two layers. Rewrite any <|...|> pattern in user text to a look-alike with different bar characters before tokenizing. Then drop any marker ids that still appear. After that, the same test shows no leak. It is the same class of problem as SQL injection, in a new setting.
A record can now be written into one sequence with exactly one-token markers at known places and user text cannot forge them. One forward pass gives a 1,024-number vector for every token. We need four of them: one per option and one for the decision.
What we have not done is look at what those vectors contain or turn them into scores. That is the next article and it starts with a surprise. Take a real ticket where a customer was charged twice, with three options. The vectors are 110 to 125 units long. The option vectors point almost the same way. Their cosines, a measure where 1 means exactly the same direction, are between 0.82 and 0.98. That decides the shape of the scoring head.
Check out Part 3, where the model surgery begins: we attach a small pointer head to the vectors we read at the markers, train it on support tickets and test whether its probabilities can be trusted. The LLM gets converted into a Jev.
Build a Jev-Style Decision Model from an LLM, Part 2: Opening the Model was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.