Implementation of logit threshold sampler and confidence breaker #657
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Edit 23/10: Simplified the implementation of logit threshold sampler, removing the previous logit_min_threshold parameter and replacing it with a new temperature parameter applied only to logits above logit_temp_threshold. This should make the sampler easier to tune across different models, with conservative settings (logit_high_temp=3.0; logit_temp_threshold=16.0) improving text diversity while maintaining coherence across a number of tested models (Gemma 9b, Llama 3.1 8B, Llama 3.2 3B, Mistral Nemo 12B). The below explanations and example outputs have been updated to reflect these changes.
Implementation of two new samplers enabling coherent diverse text generation.
The key innovation is to use absolute logit values instead of softmax probabilities to retain the model’s raw confidence in each token. By using the raw logits to make decisions, we preserve the model's absolute confidence in each token, allowing for more nuanced sampling based on how confident the model is about each token.
Explanation: Logit threshold sampler
In language model text generation, adjusting the temperature parameter affects the randomness of the output:
Leading approaches to addressing these issues, such as min-p, rely upon the manipulation of token probabilities after softmax, however this can only go so far, as the softmax operation discards absolute logit values, so that we can no longer distinguish between model uncertainty in cases where probabilities are flat due to many good choices vs no good choices. This despite the fact that ideally we want to broaden our options in the first case and narrow them in the second.
On the other hand, Logit threshold temperature allows us to avoid the diversity-coherence trade-off to a far greater extent through selectively applying high temperature only to high-confidence tokens. This enhances creativity among the tokens the model is most confident about, without introducing incoherence. Those remaining tokens below this threshold are not eliminated, but neither do they increase their existing probabilities.
Primary Applications
Parameters
Suggested initial conservative starting values: logit_high_temp=3.0; logit_temp_threshold=16.0 - if this is a little too wild then bump up the threshold in increments of 1-2 and observe outputs for coherence.
As with all samplers, you may need to adjust these settings based on intended use case and model.
Explanation: Confidence breaker sampler
Current generation language models are well known for producing certain cliched phrases, which would not necessarily be problematic in a single instance, but are known to be produced repeatedly in response to varied prompts. This is the so-called ai-slop problem. In 'deterministic' use cases this is usually not an issue, as we are simply looking for the one correct answer. In scenarios where engaging, diverse language choices are valued, ai-slop represents a significant limitation.
Current approaches to resolving this issue involve user defined lists of banned strings. This can be reasonably effective, however fail to address the deeper issue, which is the tendency of language models to funnel their responses into unintentionally learned 'tram-track' token sequences, where token choices are strongly conditioned by their immediate predecessors in a manner which is not logically or grammatically implied by the prompt or by those preceding choices.
This is a far deeper and more pervasive problem than the well known handful of phrases commonly associated with the idea of ai-slop. These tram-tracks exist within trained models because the tokens within them are good predictors for text completion. Nevertheless a user passing either the same or similar prompts repeatedly, looking for diverse outputs, will quickly observe that model responses which seemed initially impressive are in fact tram-track patterns, and the apparent diversity of outputs is an illusion.
Using the confidence breaker to jump tracks
The confidence breaker sampler addresses the issue by looking for logit patterns that signal we have entered a tram-track, and extending exllamav2's banned string functionality to roll-back to the token directly before we entered the tram-tracks. The tram-tracked tokens are then discarded and replaced by a novel generation, which will take us down a different, less traveled path.
Based on empirical observation of logits and the conditions for the appearance of tram-tracks, the pattern that the confidence breaker looks for to identify these tram-tracks is a sequence of mid-high valued logits, logits which have been nudged higher than a good score by over-training, but which are not yet so guaranteed as logical or grammatical necessity.
Empirical observation also validates the decision to return back and alter the token before the tram-track, rather than the first tram-track. This is a token which had a reasonable range of viable alternatives, but once the model settled on the decision it made, the tram-track became nearly inevitable. So, this is the error we have to correct.
Parameters
Note on use
Output comparison
Prompt: "Write at least 500 words, beginning with and in the same style as the following:\nOnce upon a time,"
Model: gemma-2-9b-it-exl2-6bpw
Note that the trials below are cut-off after the first paragraph block.
Baseline output for comparison
Logit threshold sampler
Confidence breaker
Logit threshold sampler + Confidence breaker