Training?

[xdevfaheem]

@noamgat is it possible to train a model with constrained output tokens?

[noamgat]

While I have not done this personally, it should be possible. You could devise a "legal token discipline loss" in which you look at the output that the LLM generates, and penalize it according to the weights it assigns to the illegal tokens.

So you would do something like:

• Give the LLM a prompt that specifies what schema the output should be in
• Let it generate output logits
• Run the LMFE on the tokens it generates
• Create a loss between the original distribution and the filtered distribution (if the LLM originally gave -inf scores to illegal tokens, the loss will be 0)
• Sample a legal token from the filtered distribution (This is akin to teacher forcing) and continue generation.

The -inf kind of scares me, so you might want to replace the -inf with a negative number (like -100 which is the maximum in OpenAI API logits bias). This is all theoretical though, I haven't tried it in practice.

[xdevfaheem]

Cool, this makes sense. Is there any scripts available somewhere as of you your knowledge to get s quick overview of the implementation?

[xdevfaheem]

I tried, training the model with preferred json output for the input without any format enforcers. it performs good. Eventhough, i wants to use LMFE to increase it's reliability. Am i wrong?

[noamgat]

There are no scripts that try to do this as far as I'm aware, this is brainstorming... Indeed, normal SFT training can work as well, and then you don't need LMFE in the training loop