Note

This tutorial can be run as an IPython notebook.

Visualize Token Probabilities and Model Confidence in LLM Predictions

In this tutorial we show how eli5.explain_prediction() can be used to visualize LLM predictions, highlighting tokens proportionally to the log probability. In many cases, this can help to see where LLM is uncertain about its predictions:

LLM token probabilities visualized with eli5.explain_prediction

LLM token probabilities visualized with eli5.explain_prediction

1. OpenAI models

To follow this tutorial you need the openai library installed and working.

import eli5
import openai

client = openai.OpenAI()

First let’s define our task: we’ll be extracting specific product properties from a free-form product description:

product_description = """\
Stay is designed by Danish Maria Berntsen and is functional and beautiful lighting
in one and the same design. Stay has several nice features, where the lamp's flexible
swivel mechanism makes it possible to direct the light and create optimal lighting
conditions. Furthermore, the switch is discreetly placed behind the lamp head,
making it easy and convenient to turn the light on and off. Thus, Stay combines great
flexibility and elegant appearance in a highly convincing way.

Stay table lamp is highly functional, as the arm and head of the lamp can be adjusted
according to wishes and needs, which make it ideal as task lighting in the office.
Furthermore, the silky matte grey metal support the modern and minimalistic expression.\
"""
prompt = f"""\
Product description:
{product_description.strip()}

Extract the following properties from product description:
'materials' (list of strings),
'type' (string, e.g. 'LED'),
'color' (string),
'price' (non-zero float with 2 decimal places),
'summary' (string, a very short summary).
Respond with JSON dict with above keys. Always make sure to return the price or it's estimate.
"""
completion = client.chat.completions.create(
    messages=[{"role": "user", "content": prompt}],
    model="gpt-4o",
    logprobs=True,
    temperature=0,
)
print(completion.choices[0].message.content.strip('`'))

json
{
    "materials": ["metal"],
    "type": "task lighting",
    "color": "silky matte grey",
    "price": 150.00,
    "summary": "Stay table lamp with adjustable arm and head for optimal task lighting."
}

If you examine the prompt and description above, you can see that not all attributes are equally clear: - “material” is quite clear and explicitly mentioned - “color” is also clear, but its mention is more ambiguous - “type” is underspecified, we don’t extractly specify what we want, - “summary” is quite clear, but there is no single correct summary, - “price” is not present at all, but we ask a model to make a guess.

You may have noticed that above we included logprobs=True, in a call to the model, this allows us to get a log-probability for each token, and then we can visualize them with eli5.explain_prediction():

eli5.explain_prediction(completion)

```json { "materials": ["metal"], "type": "table lamp", "color": "silky matte grey", "price": 150.00, "summary": "Stay is a flexible and elegant table lamp designed by Maria Berntsen." } ```

We can clearly see that the model is very confident in the material – if you hover over the prediction, you can see a probability for each token – and less confident about type and color. The confidence in price is a lot lower, while summary, being a longer piece of text, is harder to interpret – we can see that some words follow more obviously.

We can also obtain the same result by passing client and prompt to eli5.explain_prediction(), in this case it would call the client, and we can pass extra keyword arguments – here we’ll pass n=2 to obtain two different predictions, and would leave temperature at default.

explanation = eli5.explain_prediction(client, prompt, model='gpt-4o', n=2)
explanation

```json { "materials": ["metal"], "type": "", "color": "silky matte grey", "price": 199.99, "summary": "Stay table lamp with flexible swivel mechanism for optimal task lighting." } ```

```json { "materials": ["metal"], "type": "table lamp", "color": "silky matte grey", "price": 165.00, "summary": "A functional and elegant adjustable table lamp designed by Maria Berntsen." } ```

We can obtain the original prediction from the explanation object via explanation.targets[0].target, e.g. use explanation.targets[0].target.message.content to get the prediction text.

2. Limitations

Even though above the model confidence matched our expectations, it’s not always the case. For example, if we use “Chain of Thought” (https://arxiv.org/abs/2201.11903) reasoning, asking the model first to think about the price estimate, it would be much more confident in the price in its final output, but that does not reflect the real confidence of the model, as it’s smeared over CoT:

prompt_cot = prompt + """
Before outputting the JSON with extracted results, provide analysis of the most likely price.
"""
eli5.explain_prediction(client, prompt_cot, model='gpt-4o', temperature=0)

To determine the most likely price for the Stay table lamp, we need to consider several factors: 1. **Design and Brand**: The lamp is designed by Danish designer Maria Berntsen, which suggests a focus on high-quality design and potentially a higher price point due to the designer's reputation. 2. **Functionality and Features**: The lamp has several advanced features, such as a flexible swivel mechanism and adjustable arm and head, which add to its functionality and likely increase its cost. 3. **Materials and Finish**: The description mentions a "silky matte grey metal," indicating the use of quality materials that contribute to a modern and minimalistic aesthetic. This choice of materials can also influence the price. 4. **Market Positioning**: Given the emphasis on both functionality and elegant appearance, the lamp is likely positioned in the mid to high-end market segment. Based on these factors, a reasonable estimate for the price of the Stay table lamp would be in the range of $150.00 to $300.00. For the purpose of this exercise, I will choose a mid-point estimate of $225.00. Now, I will provide the JSON with the extracted properties: ```json { "materials": ["metal"], "type": "table lamp", "color": "silky matte grey", "price": 225.00, "summary": "A functional and elegant table lamp with adjustable features." } ```

We can see that the model has already committed to a specific price point as part of it’s analysis, and it’s condfidence is very high in a particular prediction, but this is not indicative of the true confidence.

Finally, an interesting point is that if we leave the temperature at its default value of 1, the analysis would show up a lot of less condfient predictions, which is expected given the sampling performed at non-zero temperatures:

eli5.explain_prediction(client, prompt_cot, model='gpt-4o')

Analysis of Price: The product description does not provide an explicit price for the Stay table lamp. To estimate the price, it's important to consider factors such as the design origin (Danish design by Maria Berntsen), functionality (flexible swivel mechanism and adjustable arm and head), material quality (silky matte grey metal), and intended use (task lighting in modern, minimalistic settings). Based on these characteristics, the Stay table lamp is likely positioned as a mid-to-high-end product in the market. Danish-designed lighting products known for combining aesthetics with functionality often range from approximately $150 to $400. Given these aspects, a reasonable estimate for the price could be around $250.00. Now, I will provide the JSON dict with all extracted information: ```json { "materials": ["metal"], "type": "table lamp", "color": "silky matte grey", "price": 250.00, "summary": "A flexible and elegant table lamp designed by Maria Berntsen." } ```

3. Open Source and other models

If an API endpoint can provide logprobs in the right format, then it should work. However few APIs or libraries do provide it, even for open source models. One library which is know to work is mlx_lm (Mac OS only), e.g. if you start the server like this:

mlx_lm.server --model mlx-community/Mistral-7B-Instruct-v0.3-4bit

Then you can explain predictions with a custom client:

client_custom = openai.OpenAI(base_url="http://localhost:8080/v1", api_key="dummy")
eli5.explain_prediction(
    client_custom,
    prompt + ' Price should never be zero.',
    model="mlx-community/Mistral-7B-Instruct-v0.3-4bit",
)

{ "materials": ["silky matte grey metal"], "type": "Not specified in the description", "color": "Not specified in the description", "price": 99.99, "summary": "Stay is a flexible and beautiful Danish-designed table lamp with a discreet switch and adjustable arm and head, ideal for office task lighting." }

Consider checking other libraries which support explaining predictions of open source LLMs: