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

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.