Increasing the Accuracy of Embedding-Based Retrieval

We've run into some interesting challenges with utterance matching in certain use cases of Parlant's heavier enterprise users.

For a bit of background, utterance matching is a unique Parlant feature designed for regulated enterprise customers who deploy customer-facing AI agents at an extremely large scale (hundreds of thousands to millions of monthly conversations). In those cases, even a hallucination rate of 0.01% is unacceptable (since, with hallucinations, it's not just a matter of their frequency—the bigger problem is their magnitude, which is often completely unbounded).

Utterance matching works by allowing the LLM to generate a fluid draft message, and then matching it to the best candidate out of a set of pre-defined utterance templates. This pre-defined set is created offline, normally by a conversation designer or business expert, and overrides the LLM's draft message to ensure that the agent's responses are safe and consistent with the company's policies and tone of voice.

The bigger the set of templates, the more likely it is that the LLM will generate a draft message that matches one of them, and the more likely it is that the agent's response will be appropriate for the conversation.

As with most things in software (and perhaps life), you start out thinking the problem is simple, but as you dig deeper, you realize that a robust solution is more complex than you initially thought.

In our case, the problem manifests once you have a large number of utterances, and you need to ensure that your ideal utterance is returned within the K most relevant results to consider for a given draft message. We do need to limit it to some small constant K, because otherwise the LLM will have too many options to choose from, and it will be less likely to pick the right one.

Now, it's important to note that, when looking at your utterance store, you usually know which one is the best fit for a given draft message—either because it's obvious, or because a conversation designer or business expert has told you that directly!

The naive approach is to use a vector database to index all of the utterance templates according to their content, and then find the closest template to the query. But it turns out this doesn't work well, and, just as importantly, it's unpredictable and inconsistent. More concretely, there are two big problems we've found with this approach:

1. Many templates include fields in them. For example, I'm sorry to hear about {{generative.customer_issue}}. Please consider these solutions: {% for sm in suggested_mitigations %}\n- {{sm}}{% endfor %}. This field syntax causes cosine similarity to be smaller and often unexpected compared to the fluid (and fully rendered) draft message.
2. Sometimes, the desired template is not the closest one to the draft message. It's particularly relevant when you just don't agree with how the LLM chose to phrase the draft message, and you want to reply with something more appropriate. For example, the draft message is, Certainly, I can help you book a flight to Paris. Please provide me with the dates and times you prefer. and the desired template is the simpler Sure, when are you going?, which may easily be further away in the vector space than other templates.

This means that the most fitting templates might appear outside of the top K results returned by the vector database. And the larger your set of templates is, the more likely this is to happen.

So, what can we do about it?

The first thing we had to understand was that matching the right utterance depends on a human label. In other words, having a person say, "This template is the best for for this draft message."

But how do you get an arbitrarily-chosen best-fit template to return in the top K results using an off-the-shelf embedding model that you don't want to fine-tune yourself?

This is where we came up with a new concept in Parlant, which we call queries. Queries are strings that are attached to retrievable elements in Parlant—such as utterance templates. With queries, each utterance template has N>=1 embedding vectors associated with it: one for the content of the template itself, and one for each query that is attached to it.

When we retrieve utterance templates, we retrieve the top K templates based on the cosine similarity between the draft message's embedding and the embeddings of all of its queries, as well as the embedding of the template itself.

Get in touch with us

How this works

1. A conversation designer notes, "This template is the best fit for this draft message."
2. They then create a query that contains the draft message, and attach it to the utterance template. This is essentially the label that says "This is a great fit for this draft message."
3. When the LLM generates a draft message, we retrieve the top K embeddings, containing both template as well as query embeddings.
4. Since the added query is specifically designed to align with the draft message, it greatly increases the chances of retrieving it within the top K results.

But here's the tricky part: we don't simply want to retrieve the top K closest embeddings to a draft message; rather, we want to retrieve the top N templates.

For example, say that each template has a total of 2 embeddings associated with it: one for the template itself, and another one for an associated query. If we then want to retrieve the top 2 templates (note: not vectors, but the templates themselves), we can't just retrieve the top 2 embeddings, because we could easily end up with two embeddings that belong to the same template, which is not what we want.

And it's an even bigger problem if we have more than 2 embeddings per template (on average). In practice, we would have hundreds if not thousands of templates, and potentially dozens of queries per template.

The naive solution

An easy approach would be to calculate the cosine similarity between our draft message's embedding and all of our stored embeddings, sort them by similarity, and then extract the top K templates from them.

However, this is not efficient, and doesn't scale well. If we have a large number of templates, this approach will take a long time to compute, and it will be very resource-intensive.

Toward an efficient solution

This is where we came up with a new and simple approach that is both efficient and effective.

The idea is to utilize the vector database's KNN search capabilities to retrieve the top K embeddings, which is much more efficient than retrieving all of the embeddings, computing their similarity, and sorting them manually.

The problem with that approach is that, with vector databases, you need to tell them, "Get the top K closest vectors to this input vector." So what do we choose for K when we actually want the N templates (where N<=K)?

Concrete example

Suppose we want to retrieve the top 2 templates, and that each template has between 1 to 10 embedding vectors associated with it.

As I said before, if we want to retrieve the top 2 templates (N=2), setting K=2 will not work, as we may get just one template whose entire 10 vectors are at the top of the similarity ranking. This won't work, so we need to set K to a larger value than N.

But how much larger should K be? That's the real question we have to answer. If we have that number, and that number is relatively small (and ideally upper-bounded by a de facto constant), then our problem is solved, and our retrieval will be efficient and correct.

The algorithm

This is the solution we came up with:

1. Pre-compute the number of embeddings per template.
2. Pre-compute a max-heap where each node in the heap represents a template, and the value of the node is how many embeddings are associated with that template.
3. Find the N largest nodes in the heap, representing the N templates with the most embeddings.
4. Set K to the sum of the values of those N nodes—that's it.

Understanding the algorithm

In the case above, if N=2 (number of templates to retrieve), setting K=13 (number of vectors to retrieve) would work.

Why 13? Because it guarantees that we retrieve enough embeddings to cover the templates with the most embeddings, while still being efficient.

Here's the analysis.

Suppose that the top 2 templates are indeed the ones with the most embeddings, and that they have 7 and 6 embeddings, respectively. In this case, we will retrieve all of their embeddings, which is 13 in total. So this case works.

But if the top 2 templates are not the two with the most embeddings, we will still retrieve enough embeddings to cover the top 2 templates.

For example, if the top 2 templates have 5 and 4 embeddings, respectively, we will retrieve 12>=9 embeddings, which is still more than enough to cover the top 2 templates.

As long as we retrieve enough embeddings to cover the worst case, which is when the top N templates have the most embeddings, we are guaranteed to retrieve the top N templates.

And this is how we achieve efficiency. We use KNN (literally, to retrieve the K nearest neighbors) to retrieve the top K embeddings, and, using our max-heap, we ensure that K is large enough to cover the top N templates with the most embeddings.

We then extract the top N templates from the fairly small set of K retrieved embeddings, which is efficient and correct.

Finally, here's the relevant implementation of finding K in Python:

import heapq

K = sum(heapq.nlargest(N, [count_vectors_per_utterance(u) for u in stored_utterances]))

Get in touch with us