I have always been bothered with the multinomial layer that picks the next token. It seems like you should be able to stick a classifier in there to detect if the probabilities align with a satisfactory answer and whether the model is unsure of the response. If that is the case, then it should branch out to a bigger model which would compute that single token that it is unsure about, and from that the rest of the response is obvious for the smaller model.

This way you could combine e.g. a 7B parameter model with a 70B parameter model, and get the quality of the larger model while most of the time you are only running the small model.

Edit: You could also store the full probabilities for each token, and then the classifier could detect if it had gone down a bad path, and then unwind the tokens and pick a different path.