Writing
jevdecision modeltrainingLLMs

jevlite - Experience report in training the Decision Model

In the previous post we talked about how jevlite makes decisions using the underlying adapter. In this blog post, I want to talk about the dataset selection, synthetic dataset generation and how I almost ruined it and the actual training run. Having trained smaller LLMs before, this model was unusually pleasant to train and I full expect a whole bunch of the decision model clones to pop up.

Code | Model | Dataset

Data used for training

The goal of this decision model is to take decisions given a state of the world. The state and the questions are both arbitrary. So the training data has to teach it a proper skill and not just teach about categories. This is what makes it different from just a classification model. The base LLMs already have this generality so we just need to point it in the right direction.

So we need diversity in terms of states and we need to split by tasks for the train test split and not just by the rows. This will help us test if the model generalizes across domains. The data set generation code takes care of this. The task level data is held back for evals. See the dataset link below to for the synth subset. But this was done for all data.

The dataset is derived from SNI Super-NaturalInstructions from AI2 (Allen Institute for AI), MNLI, ANLI, BoolQ, RACE, Yelp and SST2. The data generation script was written to remove non-classification problems and downsampled tasks where a large portion of the answers were skewed to one choice. The most important part was synthesizing new data for the model to behave better.

Data Synthesis for business domains

The above data set did not have representative data for support tickets, SIEM alerts and so on. So I had to generate those. I used the Qwen-3.6-35B-A3B via llama.cpp to generate synthetic data. The teacher.py script in the repo generates the training data across domains like customer support, SIEM, agent transcripts etc. Here is a sample of one of the rows.

{
  "type": "choice",
  "state": "Invoice #INV-2023-8842 from Vendor 'Acme Supplies' dated 2023-10-27 for $500.00 lacks a Purchase Order number. The system only has internal ticket ID #TKT-991 for this engagement, and no PO was ever issued. Payment is frozen due to compliance policy requiring a valid PO.",
  "question": "What is the status of the associated internal request?",
  "options": ["Approved", "Pending", "Closed"],
  "criteria": {
    "Approved": "An internal ticket exists and has been marked as approved for purchase.",
    "Pending": "An internal ticket exists but is still waiting for manager review.",
    "Closed": "The internal ticket was rejected or cancelled without a PO being generated."
  },
  "label": [0.19158, 0.70574, 0.10268],
  "teacher_entropy": 0.7248,
  "task": "synth:a vendor invoice with line items",
  "source": "synth",
  "_key": "synth-2-1"
}

I used different tones like “terse and technical”, “rambling and emotional”, “frustrated” etc along with different lengths like “3-4 sentences”, “a long detailed account with timestamps and ids” etc. This generated the variety of data that the model needed.

I also added the probability distribution generated by the teacher. This is a 35B model used to drive a 8B model. I also had to turn off reasoning. I actually let the first version run overnight and the probability scores were all wrong. Everything snapped to 1.0. Turns out reasoning for this is a wrong idea. So I turned off reasoning and took the log probs of what the teacher thought. This cost me 13 hours of overnight data run. I also added a bit of shuffling during training. I also enriched the labels on the data from the other training set like SNI since it only had labels. I gave it descriptions for the set so that it behaves better. It turned out some of the results of SNI were a direct contradiction and plain Qwen 3.5 35B A3B was a better judge. The SNI has a bunch of mislabeled data that Qwen correctly identified and fixed. As a follow up, I should check this data against a panel of judges just to ensure that jevlite doesn’t just mimic Qwen.

Scoring and Confidence

I first tried normalized entropy as the scoring method as we are doing anomaly detection. Shannon Entropy has a nice property that concentrated entropy gives it confidence but spread out distribution gives uncertainity. But normalized entropy was the wrong thing to score against. The probabilities for choices [0.08, 0.85 and 0.07] ended up with [0.202, 0.138, 0.186] for entropy contributions per choice. The sum of two of the smaller choices is higher than the sure choice. It was a ln(1/p_i) (surprisal) to me. The simplest answer was the right one. I just used top probability. I was anyway going to use brier and ECE that need these probabilities. I did a softmax on the logits to get the probabilities of the choices and derived the confidence from there.

def confidence(kind, probs, score=None):
    if kind == "score":
        return sum(p for i, p in enumerate(probs) if abs(i - score) <= 0.5)
    return max(probs)

What the training actually optimizes

The training is a PEFT technique using qlora. I used gemma4-e4b-it which is a 8B parameter with 4.5B active params. I used bitsandbytes to train against nf4 quant and the adapter is bf16.

The loss is soft-label cross-entropy over the option letters. This is different from how loss is trained for a single LLM. You just do a cross entropy loss only on the token and not on the confidence scores across the options. This penalizes the model for being too certain as well as being wrong. Options are shuffled for each epoch so a model cannot learn that “A is usually right”.

# log-probs over this row's letters
lp = option_logprobs(model, enc, row, options)
# target is a distribution, not an index
loss = -(target * lp).sum()          

After fine-tuning, 94.9% of the model’s probability mass at the answer position is for option letters now. But without the adapter, the base model for the same prompt puts only 2.5% there.

Beyond accuracy

Looking at acc or accuracy data alone for this model is not enough. We also want to look at nll (Negative Log Likelihood), brier (Squared Error on probabilities) and ece (Expected Calibration Errors). Gemma 4 is a strong model and actually started at a good place. The whole point of this training is to observe these numbers reducing. This asserts that the model is building a better internal representation is more confident about its choices.

step acc nll brier ece
250 0.7993 0.5316 0.2684 0.0855
500 0.7993 0.4608 0.2425 0.0311
750 0.8008 0.4505 0.2347 0.0411
1000 0.8056 0.4529 0.2354 0.0525
1250 0.8140 0.4359 0.2258 0.0291
final 0.8156 0.4297 0.2222 0.0185

The above is the eval progression against tasks held back during training. NLL down 19%, Brier down 17%, accuracy up 1.6 points, ECE down substantially but noisy across checkpoints.

Why this problem is unusually well suited to fine-tuning

This problem is like catnip for ML practitioners. I am sure there are a ton of others like me who will be nerdsniped into building something like this. I was sitting on the bylines without access to enough compute to run let alone train models that are in the zeitgeist.

We are literally reading off the logits and manipulating them. This is what the state of the art was in 2019 and 2020 well before ChatGPT. But we did not have these awesome models. BERT was very weak. RoBERTa made it a bit better. Currently jevlite is text only but there is no reason for it. In fact one of the reasons why I chose gemma-4-E4B as the backbone is that it is multi modal.

This is also a very unique ML problem. There is no divergence from training and inference. This is almost close to the idea. With normal LLMs you have things like inference params like top-p, top-k, repetition penalties, stopping criteria etc. But we are literally reading logits. So this is a perfect model usecase.

What I’d do differently

Check the teacher before the overnight run. The reasoning-entropy problem and the truncation bug would both have been caught by labelling twenty rows and reading them. The labeling took 13 hours that I had to throw away.

Question the base model choice earlier. I built it on gemma-4-E4B-it, an instruction-tuned model, then fed it raw text that ignores its chat template. It was an oversight but it still worked out okay. One thing to verify is if the base (non-instruct) model might be a better starting point.

Get more score data. I could use yet stronger models to synthesize more data and evaluate against a panel of judges.


Code: github.com/vagmi/jevlite Huggingface Adapter: vagmi/jev-lite Dataset: vagmi/jevlite_dataset