LLM/trail: Topic Clustering for Chat Interfaces

Introduction

LLM chat interfaces make it easy to start with one query and quickly branch into related (and unrelated!) topics. Unfortunately, a long conversation is hard to navigate when all you can do is scroll. I built LLM/trail to cluster your chat history by topic, letting you jump directly to messages via the sidebar. Try it at llmtrail.com, or explore a live demo with sample data at llmtrail.com/demo.

TLDR: How It Works

To organize your chat history, LLM/trail uses an Anchored Recency Window algorithm. When you send a query, the app uses Gemini to generate a summary of the query and response, which is converted into a vector embedding. The app calculates the cosine similarity of the new turn's embedding and the average embedding of the last three messages in each group (the recency window), for historical context. The turn is added to the best scoring group if it exceeds a threshold and is still similar enough to the very first message in the group (the anchor), to prevent long-term topic drift.

For all the details, keep reading!

Design Constraints

LLM/trail's main affordance is an outline of your chat history, organized by topic. To build this view, I couldn't use unstable clustering methods like k-means because each new message might cause noticeable shifts in the groupings, which would be confusing for a user. To get a stable ordering, I chose an incremental approach where each message was assigned to a group once.

Preprocessing: Clustering Summaries vs. Raw Text

A note on terminology: "Query" refers to what the user types into LLM/trail for their chats. "Prompt" refers to the instructions provided to the LLM on the backend for summarization and categorization. "Turn" refers to a single user query and response pair.

I didn't categorize the raw text of the query and the response. Instead, I generated summaries of the query and the response and ran the clustering algorithms on the summaries. I found that the queries and responses often contained conversational filler and transitional phrases favored by LLMs. Using the summaries focused the text, reduced context pollution, and improved clustering quality.

The querySummary prompt was: "The core topic of the user's query. Ensure any implicit contexts, entities, or references are explicitly resolved. Must be written as a descriptive noun phrase."

• "Ensure any implicit contexts, entities, or references are explicitly resolved" is there to handle the case when the query is ambiguous, such as "Yes", "No,", "Tell me about the other one"
  
• "Descriptive noun phrase" is there to prevent the model from generating a summary like "The user requested more information about ocean exploration". When all the query summaries had the prefix "The user requested…," the models seemed to cluster turns together more often because it thought all the queries were requests.

As a result, a query like "Tell Me More" is summarized as "Deep Sea Exploration History" as opposed to "The user requested more information about deep sea exploration history."

The responseSummary prompt was: "A concise summary of the core topic and main points of the assistant response, no more than 3 sentences long. Ensure any implicit contexts, entities, or references are explicitly resolved."

Approach 1: Text Generation Models

Being a project about LLMs, it seemed natural to use the LLM to do the categorization. I tried two different approaches.

LLM as Similarity Scorer

My first approach was to use the LLM to generate similarity scores.

1. On a new turn, the LLM generates a 3-5 word summary that represents both the query and the response.
   
2. If this is the first turn, that turn's category is the 3-5 word summary. We are done.
   
3. If it is not the first turn, provide the LLM with a list of existing categories, e.g. [category1, category2, ..., categoryN]
   
4. Prompt the LLM to score the new turn against the list of existing categories and have it return a similarity score from 0 to 1. The prompt looked something like: For each category in the "List of existing categories", assign a semantic similarity score from 0.0 to 1.0, representing how closely this Q&A turn relates to that category. Use these scoring guidelines. 0.9 to 1.0: Direct, immediate continuation of a thread (e.g., clarifying a term just used, expanding a specific bullet point from the previous answer). 0.7 to 0.8: Very close follow-up within the same topic. 0.4 to 0.6: Shift to a new sub-topic, task. 0.0 to 0.3: Completely unrelated domains or subjects.
   
5. If the new turn's score is above a threshold, assign it to that category. Otherwise the turn is assigned to a new category.

The results seemed promising, but it was hard to iterate on because it was challenging to write a prompt that improved the scoring rubric. It wasn't clear to me how to describe the difference between a "direct continuation" vs. a "very close followup." It was hard to provide concrete examples in a way that would generalize across multiple types of conversations and still have a concise prompt.

I realized that this approach was naive. If I wanted similarity scores, it would be better to use a vector embedding model, so I decided to try a different approach.

LLM as Router

As an alternative, I used the LLM to decide whether a new turn was a member of an existing category or should be a new category. I thought this might be better than the previous approach because it didn't have to produce scores, it just had to pick a category.

1. On a new turn, the LLM generates a query summary and a response summary.
   
2. If it is the first turn in the session, assign that turn a group label, such as G1. We are done.
   
3. If it is not the first turn, provide the LLM with a list of the previous turns and their group labels. For example, in a chat with 3 turns and 2 groups, it would look like: [(querySummary1, responseSummary1, G1), (querySummary2, responseSummary2, G2), (querySummary3, responseSummary3, G1)]
   
4. Provide the LLM with the new turn's summaries: (newQuerySummary, newResponseSummary)
   
5. Prompt the LLM with: "You are classifying which topic group a new conversation turn belongs to. Decide whether this new turn continues the conversation topic of one of the existing groups, or starts a new topic. If it continues an existing topic, return that group's label. If it starts a new topic, return \NEW_GROUP`."`

This approach seemed to produce better results than LLM as Similarity Scorer. I could look at the output after the fact and rationalize why the LLM produced a certain set of groups. However, sometimes I thought the LLM would create too many groups and sometimes too few. I found it hard to concisely describe in a prompt how I wanted the clustering to be different. When I provided specific examples, it seemed to overfit.

Approach 2: Text Embedding Models

I switched to using embeddings because it seemed easier to work with vectors and adjust numerical thresholds instead of editing a prompt. I computed cosine similarities between a new turnEmbedding and the existing groupEmbedding. If the group with the highest similarity score exceeded a threshold, then the turn would be assigned to that group. If no group met the threshold, a new group was created.

Local Embeddings vs API-based

I started with local embeddings for cost reasons. I thought I would need to generate a lot of embeddings, so I wanted something that was not API based. In addition to generating embeddings for turns, I needed to generate embeddings for highlighted text. In LLM/trail, you can highlight text and save it to the sidebar for later. Once in the sidebar, you can tap on the highlighted text to return to the message it appears in.

In the early prototypes, message history and saved highlights were distinct sections in the sidebar. I thought users might come across similar concepts in different messages, so it could make sense to cluster those by topic. Since highlighting could happen quite often, I thought it would be too slow if we had to make a network call every time a user highlighted some text.

I started with an English-only model, Xenova/all-MiniLM-L6-v2, which is a 23MB download. The embeddings worked well, and the output of the groupings were fine. But to get multiple language support, I had to switch to a different model such as multilingual-e5-small, which was a 118MB download. This seemed like it would be too much for users on slower devices or networks.

To minimize download size, I switched to the Gemini embedding API which handles multiple languages. However, by switching away from a local model, I had to reduce the number of embedding calls. As a result, I decided to change the product so highlights no longer needed to be clustered. Instead of them being in a separate section, I showed the highlights below their source message. After some use, I realized that this was a better design because the message provides context for the highlighted text.

Group Embeddings

I explored a few different ways to represent a group's embedding vector:

• Centroid, which averages all the turn embeddings in a group. As a group got bigger, its embedding vector got more diluted, becoming more like the "average vector." It made the algorithm under-split and resulted in too few groups.
  
• Anchor-First, which uses the first turn's embedding. If the new message is too far away from the first topic, it creates a new group. This seemed to be too sensitive to the gradual evolution that happens in a conversation. It made the algorithm over-split, creating too many groups.
  
• Recency Window, which averages the most recent K turns. When K is small, like K=3, it doesn't get too diluted like Centroid. This approach performed better than Centroid and Anchor-First, but can still allow some drift from the original topic if the K is big enough.

Anchored Recency Window

In production, we use an algorithm that combines the Recency Window with Anchor-First.

1. On a new turn, the LLM generates a query summary and a response summary.
   
2. Let newTurnEmbedding = average(querySummaryEmbedding, responseSummaryEmbedding)
   
3. If it is the first turn, assign the turn to a new group. We are done.
   
4. If it is not the first turn, compare newTurnEmbedding against each candidate group. For each group, let groupEmbedding be the average embedding of the most recent K=3 turns, and compute groupScore = cosineSim(groupEmbedding, newTurnEmbedding)
   
5. Find the candidate group G which maximizes groupScore
   
6. Let firstNodeEmbedding be the embedding of the first node in group G. Calculate firstNodeScore = cosineSim(firstNodeEmbedding, newTurnEmbedding)
   
7. If groupScore >= SIMILARITY_THRESHOLD and firstNodeScore >= ANCHOR_THRESHOLD, then assign the new turn to group G. Otherwise, assign the new turn to a new group.

With gemini-embedding-2 at 768 dimensions, we are currently using a SIMILARITY_THRESHOLD of 0.85 and an ANCHOR_THRESHOLD of 0.78.

This approach seemed to work well because it allowed for some topic drift that feels natural. When K=1, the window feels a bit small; if your message doesn't match against the previous message, you create a new group. With K=3, you allow matches against a larger set of turns, which extends the algorithm's referential lookback. Adding the anchor allows us to split the group if the newest message gets too far away from the first one.

Evaluation

To evaluate these clustering methods, I had a small dataset consisting of my own exported chats, which ranged in length from about 10 to 40 turns, on a variety of different topics. As I tried different algorithms and variants, I would run them against these conversations to compare the groupings that were generated. Once I had a working prototype of LLM/trail, I would use it to get a qualitative sense of how the clustering behaved. When I found conversations that broke the algorithm, I would export those chats and add them to the dataset.

Once I had settled on Anchored Recency Window, I built some tools that allowed me to tweak the clustering thresholds directly in the prototype, which helped me to get a better sense of how the similarity and anchor thresholds interacted.

With n=1, I can't claim a generalizable solution, but it seems to work well for me and how I tend to interact with LLMs. I'm curious to see how well it works for others and would love to get your feedback.

Conclusion

Thank you for reading to the end! Please let me know if you have any comments or questions by sending me an email.