If you want to try a _legit_ Jev implementation that matches (at least in my evals), the vLLM patch to turn DiffusionGemma into Jev is available.
On my DGX Spark I get very similar latency numbers, and it matches my evals + or - a few points on each test (DG wins some, Jev wins some, both show low confidence when wrong).
I ran the same evals against a Qwen36 and it clearly lost to both of them, so you are leaving both knowledge and instinctual reasoning on the table with any smaller models, FWIW.
DebertaV3's architecture and noising should be even better as a basis because it had a couple inductive biases (cross encoder, disentangled attention and RTD corruptions) that enabled it to have unmatched weight performance ratio on such tasks.
My gut tells me that a better approach to a calibrated 0-shot classifier than shoehorning DiffusionGemma would be starting from another gemma, T5GemmaV2. Take its encoder and do continual training on an RTD objective and a large relational synthetic data mix. Then finetuning (multi-annotator data will help calibration) on as many proper NLI datasets as possible. That still lacks the DebertaV3 disentangled attention's inductive bias, however.
Jev also has its calibrated predictions component which is important. Temperature scaling is probably the easiest first pass. But there's lots of sensible options to improve on that.
ModernBERT might be the easier, more stable starting point than T5Gemma though.
I suspect you could get interesting results, but DiffusionGemma has a lot of knowledge that may be challenging to train into the smaller models. The advantage of pulling a fully-trained diffusion model off the shelf is that it already knows all of this, has been trained as a MoE, etc.
What I think these models actually need is a structured decision thinking mode. As it stands now, the only way to think about the answers with DiffusionGemma is to diffuse a thinking block, but giving the model an auto-regressive thinking space to reason, even just lightly, could drastically improve performance.
This is very interesting! Seems like a promising direction.
I wonder though if it supports the same claims as Jev: answers are not impacted by other answers to the same questions, nor the existance of other questions? It seems by sharing KV cache all questions will be visible. And I think the diffusion causes the answers to attend to eachother?
Also I think without fine-tuning we can not say that the probabilities it output are actually probabilities. Maybe fine tuning using Brier scoring would do the trick?
Maybe some kind of hierarchical structure of the KV cache can make questions independent, and with smaller diffusion canvas’ generated in batch can be a way to make answers generate independently?
I believe you get better answers by diffusing each question together, but the PR's server gives you finer control over that. If each question is independent, you can get better parallelism.
> It seems by sharing KV cache all questions will be visible
Yeah, this is partially why in my approach I've done this instead, and not used diffusion model:
1. Convert the state into one shared prompt.
2. Run that shared prompt through the model once.
3. Fork the model’s internal state once per question.
4. Add a different question to each fork.
5. Ask each fork for its next-token scores.
6. Calculate only 64 possible label scores—not the whole vocabulary.
Basically ... skip decode.
Won't be as fast as doing diffusion model parallel across a pile of questions at once, but:
a) let's you use pretty much any existing text model (with some modifications). I've got qwen3.6 moe and qwen3.8 flash next running, am getting gemma4 working now
b) the problem you identified
It's possible I'm getting high on my own supply and misunderstand entirely the whole thing, but it seems to work?
https://github.com/rdaum/eider/commit/9c2d5c049068c33da2a48b...
I don't have the chutzpah to go creating PRs for vLLM to do the same.
Yes that seems sensible for isolating questions/answers.
> 6. Calculate only 64 possible label scores—not the whole vocabulary.
This I don’t understand though, could you expand this please?
So the model normally (like during a normal decode) takes its final hidden vector and multiplies it by the entire vocabulary head -- so like about 250K rows -- to produce one logit per possible next token (and then so on and so on...)
Instead I use a fixed set of up to only 64 single-token labels. At model load, I gather only those 64 rows from the vocabulary head into a small matrix. Each question maps its permitted answers onto some of those labels.
It is a probability distribution conditional on the allowed labels. Calibration is a separate problem that I have to solve still and will be model specific :-) But I do seem to get reasonable answers right now.
So "64" is just in the end the endpoint’s maximum answer-label set. Most questions use only two or three of those rows. And, yeah, some calibration required. WIP on that
why not mask attention and do it all in one forward pass ? tokens belonging to a question can just see that question and the main prompt
Ohh this is really cool.
So one could pack all of common state, every question, every answer in the same prefill, using attention mask to only let them attend to their logical parent.
Then additionally do position encoding for token based on their logical position rather than physical.
Then the diffusion step also applies an attention mask to prevent bidirectional attention between answers.
I think you're right. Better. I will have to think through if it would be faster or slower. If understand what you're getting at with this.. broken analogy...
What I described was -- we have a bunch of orders to the kitchen, all of which have the same "base" meal but different topics.
> Cook the "base" meal once, divide servings onto multiple plates, then add different toppings to each plate.
vs what you suggest:
> Put the base meal and every topping through the kitchen together, but use some kind of dividers so the toppings never mix up together.
Except.. ok, that analogy is confusing lol.
fwiw, w/ gemma4 -- non-diffusion -- I get about 170ms for a single question -> answer and then an additional ~33ms on adding more. While I see people reporting 300ms for this vLLM PR on same hardware (Spark.)
So I don't see the advantage to their approach until you're up beyond 6 or 7 questions?
Latest commits added gemma4 and instructions. I'll work on making a version of all of this that is standalone and not specific to DGX Spark.
This PR is interesting but it's making the assumption that what Jev has done is based on a diffusion model or that a diffusion model is superior for this work. Which may or may not be the case.
If I understand it though it does mean you can evaluate a bunch of questions simultaneously, which is an advantage.
Also: While I think it's expected/normal to see LLM-generated programs... there's a lot of LLM written comments in that PR, which is sad to see. Auto-human.