Add probability distribution to choices

[andy-zhou]

Presentation of the new feature

It would be very helpful to include the probability distribution of the different options (both log probabilities and real probabilities) present in outlines.generate.choice(). This is useful for evaluating the certainty of the model for any given classification.

We use it as a pre-filter step for deciding if we should generate more expensive reasoning (for example, CoT) to arrive at a more certain classification.

Two areas of complexity that I'm aware of:

1. If you ever implement token-healing, the probabilities need to include the healed tokens
2. OpenAI doesn't include logits in their non-deprecated models

Are you willing to open a PR?

Yes, though I would need pointers on where to start.

[wdhitchc]

Greate Idea. Here are some quick thoughts on how we might be able to implement this although I'm not completely sure if this would work..

On in line 75 of serve.py we have:

   text_outputs = [prompt + output.text for output in request_output.outputs]

https://github.com/outlines-dev/outlines/blob/main/outlines/serve/serve.py

Request_output.output should be a VLLM CompletionOutput?, which has log probs as an argument. If that was the case you could just add an optionl to return that as well.

https://github.com/vllm-project/vllm/blob/main/vllm/outputs.py

[rlouf]

There's a small subtlety here. There may be several combinations of tokens that lead to either of the choices. In this case do we need to return the logprob that corresponds to all possible paths or only the path that was sampled?

If we go with "all possible paths", the basic idea is to find all paths in the FSM that lead to either choice and pass the corresponding token ids + prompt through the model to get the corresponding logprobs.

[HerrIvan]

Maybe off-topic: What about doings this for the openai part. For that we should give (optional) access to the logprobs values (https://cookbook.openai.com/examples/using_logprobs).

Of course, there is the subtlety that a generation may be the output of "n" api-calls, so there may be some decisions to be made on how to returns the aggregates.

A pre-condition for this may be using the pickleserializer for the persistence, so that we can save the whole api response. But that change would affect also the "transformers" part of the module.

(Should this go to a separate issue?)

[rlouf]

I am still not sure what the API would look like, especially since we still want outlines.generate.choices to return one of the choices.

[dnhkng]

Can we have a new method, "probabilities"?

Also, can you point out where in the codebase the actual decision on which class to select is made?

[lapp0]

We might consider returning a GenerationResult object. e.g.

>>> result = outlines.generate.choice(model, ["Positive", "Negative"], sampler=BeamSampler(2))`
>>> result.text
"Positive"
>>> result.relative_probs  # probability relative to actual generations
{"Positive": 0.6, "Negative": 0.4}
>>> result.absolute_probs  # un-normalized probabilities, valuable other `outlines.generate` functions
{"Positive": 0.3, "Negative": 0.2}

Note that if choices have multiple tokens, we aren't guaranteed we know the probabilities. However, with a beam search sampler we can guarantee we know the relative_probs of N samples. For generate.choice, N would have to equal the number of choices, since we want P("Choice"∣choices) as opposed to P("Choice"∣entire set of possible generations).

@dnhkng you might run into issues implementing this before beam search is available. The actual decision on which class to select is determined by the language model, not based on post-processing. https://github.com/outlines-dev/outlines/blob/main/outlines/generate/choice.py

[dnhkng]

Ahh, ok. I thought the selection was done by post-processing the probabilities. Otherwise, you might select categories with high initial token probability, but with a beam search you would find the overall most likely category.

I have an interesting use case that would require the probabilities.

[lapp0]

With beam search, you are guaranteed to explore all legal paths given that the number of legal paths is equal to the number of beams. This is why I suggest beam search.

Although, reconsidering - there may be multiple legal paths for each choices, e.g. ["Pos", "itive"] and ["Posit", "ive"]. We would need to guarantee all choices are generated through other means.

I agree that this is a valuable and interesting use case. Here are a few steps that would need to be done to accomplish this:

• 1. Ensure SequenceGenerator can return logits
• 2. Create outlines.generate.logits which returns the result and logits of a prompt and uses Greedy sampler by default.
• 3. Create outlines.generate.probabilities which calls outlines.generate.logits with each choice.

[rlouf]

With beam search, you are guaranteed to explore all legal paths given that the number of legal paths is equal to the number of beams. This is why I suggest beam search.

This is overkill

Although, reconsidering - there may be multiple legal paths for each choices, e.g. ["Pos", "itive"] and ["Posit", "ive"]. We would need to guarantee all choices are generated through other means.

You can walk the FSM created when calling RegexFSM, get all the possible token combinations, run them in one batch through the model and sum the path probabilities.

[dnhkng]

Sum of the average probability per token of each combination?

Some care needs to be taken with the target categories. Imagine a character level LLM, and we want the probabilities of 'yes' or 'no' for some prompt question. Not only are there more letters in 'yes', but there are also many more words that start with 'no', biasing the selection.

In this case, although we want just 'yes' or 'no' we should use something like 'yes.' or 'yes ', as the probability on the ' ' or '.' will compensate the letters when we average over all characters.

[lapp0]

You can walk the FSM created when calling RegexFSM, get all the possible token combinations, run them in one batch through the model and sum the path probabilities.

I'm concerned about the number of combinations of tokens, it would have exploding growth. Is there something I'm missing here?

>>> generate_substring_combinations("foo")
[['f', 'o', 'o'], ['f', 'oo'], ['fo', 'o'], ['foo']]
>>> generate_substring_combinations("foo1")
[['f', 'o', 'o', '1'], ['f', 'o', 'o1'], ['f', 'oo', '1'], ['f', 'oo1'], ['fo', 'o', '1'], ['fo', 'o1'], ['foo', '1'], ['foo1']]
>>> generate_substring_combinations("foobar")
[['f', 'o', 'o', 'b', 'a', 'r'], ['f', 'o', 'o', 'b', 'ar'], ['f', 'o', 'o', 'ba', 'r'], ['f', 'o', 'o', 'bar'], ['f', 'o', 'ob', 'a', 'r'], ['f', 'o', 'ob', 'ar'], ['f', 'o', 'oba', 'r'], ['f', 'o', 'obar'], ['f', 'oo', 'b', 'a', 'r'], ['f', 'oo', 'b', 'ar'], ['f', 'oo', 'ba', 'r'], ['f', 'oo', 'bar'], ['f', 'oob', 'a', 'r'], ['f', 'oob', 'ar'], ['f', 'ooba', 'r'], ['f', 'oobar'], ['fo', 'o', 'b', 'a', 'r'], ['fo', 'o', 'b', 'ar'], ['fo', 'o', 'ba', 'r'], ['fo', 'o', 'bar'], ['fo', 'ob', 'a', 'r'], ['fo', 'ob', 'ar'], ['fo', 'oba', 'r'], ['fo', 'obar'], ['foo', 'b', 'a', 'r'], ['foo', 'b', 'ar'], ['foo', 'ba', 'r'], ['foo', 'bar'], ['foob', 'a', 'r'], ['foob', 'ar'], ['fooba', 'r']]

>>> choices = ("She is at home", "She is at the store")
>>> len(generate_substring_combinations(choices[0]))
6930
>>> len(generate_substring_combinations(choices[1]))
203513

I don't think we can explore all tokenization paths for a given choice. It seems the best we can do is calculate the probability the best path for each choice (via greedy for now, beam later) and compare, OR strictly limit the size of probabilistic choices.

[dnhkng]

Although the number of combinations feels n^2, I think the paths overlap, and it resolves to n. Feels like a dynamic programming coding interview question 😅

Break down the input string into subchunks recursively, and then do a batch on an LLM to get the logits, and fill in the graph. Finally, calculate all the paths based on the probabilities, calculate the average probability per token per path, and sum them?

[rlouf]

The simplest here would still be approximate by taking multiple samples once #533 is merged. SMC on the roadmap should give better results.

[dnhkng]

Yes, monte carlo might be fine ;)

BTW, can someone tell me what FSM stands for? Finite state machine maybe?

[rlouf]

Finite State Machine indeed.

[LouisHernandez17]

Is anyone still actively working on this ? @dnhkng ? If not, I can give it a try myself, I also need it.

[dnhkng]

No, not working on this feature.

[aaronsnoswell]

+1 for this feature - this would be very useful!

[LouisHernandez17]

For BeamSearch Sampler, do you think it would be a satisfying approximation to consider the weights returned by the sampler as the log probabilities?

By default, BeamSearch with choice already returns one prediction per beam, ordered by beam weight. We can then easily get the probability by applying an exp, and, finally, group the beam prediction by final output, and sum the probabilities.

Pros:

• Almost no extra computation, as we use the weights already computed by beam-search
• Actually shows what's going on when sampling, since the probabilities are the same as the ones used to filter and sort the topk.

Cons:

• Only works for BeamSearch sampler
• The probabilities don't always sum to one, in particular when many allowed path are thrown away by topk.

I implemented this in a PR I just submitted (#895)

[cpfiffer]

This is part of a wider set of requests for being able to view the logprobs of each token (#614).

I think the API here should probably come in two parts. The simplest API is just to view the next-token distribution, rather than the combination of all tokens that map to a given choice, particularly since obtaining log probabilities of entire sequences is quite difficult (as noted).

One API I could imagine using without too much fuss is a SequenceGeneratorAdapter extension.

generator = outlines.generate.choice(model, ["a", "b"])
result = generator.logprobs("Pick a or b.")

which would return an object (SampleResult or something) with

• The sampled choice (a or b)
• A mapping from token to probabilities

For an arbitrarily sampled sequence, we'd have a list of SampleResult that anyone using Outlines can hack on. This would be awesome for researchers working on structured text, misc probability problems, etc.

This is pretty easy to work with if you use generate_substring_combinations("foo") and beam search. Multiple beams would provide a ragged array or something.

This way you'd get a rough approximation of sequence probabilities you can work with, and run all the math yourself.

[sigjhl]

I also think this would be a good feature.

There is another repo that provides this functionality, but I'm not proficient enough to tell if this can be implemented in outlines as well. Maybe someone can take a look?

https://github.com/kddubey/cappr

[lapp0]

I'm not proficient enough to tell if this can be implemented in outlines as well.

No, it currently isn't.

cappr's implementation makes sense. It doesn't generate, it simply calculates the logprobs of each choice.

Happy to provide guidance to anyone seeking to tackle this issue.

The easiest approach is likely beam search since Outlines has interfaces for generation, but not for forward.

The cleanest approach is to implement forward for all models, which returns the logits of all tokens passed.