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