I've seen this in a UI. They went a step further and you could select a word (well token but anyway) and "regenerate" from that point by selecting another word from the token distribution. Pretty neat. Had the heatmaps that you mentioned, based on probabilities returned by the LLM.

This should also be pretty cheap (just one pass through the LLM).