RAG and Batch Inference: Reducing LLM Costs in Production

Adaptive Retrieval: A New Approach for LLMs

In the field of artificial intelligence, retrieval-augmented generation, or RAG, is emerging as a key architecture for production applications. Unlike an approach that relies solely on a model's training data, RAG integrates relevant information from external knowledge bases to enrich the generated responses. The goal is to provide the model with pertinent context, thereby increasing the accuracy of the answers.

However, a major challenge arises: determining the exact amount of context to retrieve for each query.

The Hidden Cost of Fixed Retrieval

Consider the example of an internal chatbot designed for a company. When a user asks a simple question like "What are your opening hours?", the system retrieves 10 documents from a vector database, as the pipeline is configured this way:

• documents = vectorstore.similarity_search(query, k=10)

These documents may include various information such as the employee manual, HR policies, or safety guidelines. Yet, only one sentence is needed to answer the question. Imagine this process repeating for 50,000 queries per day. Most of the retrieved tokens are unnecessary for the final answer, but they are still billed.

Tailoring Retrieval Size to Each Query

Not all questions require the same amount of context. Let's consider two different queries:

• Query 1: "What are your opening hours?" A limited number of relevant documents is sufficient.
• Query 2: "Compare our healthcare reimbursement policy with last year's financial guidelines." This request requires multiple documents from various sources.

Although both queries are important, they should not result in the same volume of information being retrieved.

Towards Adaptive Retrieval

Rather than treating each query uniformly, production systems first assess the complexity of the question. Simple questions retrieve limited context, while more complex questions require more data. Here’s a simplified implementation:

if query_type == "simple":
    top_k = 2
elif query_type == "medium":
    top_k = 5
else:
    top_k = 10
documents = vectorstore.similarity_search(query, k=top_k)

While this logic is straightforward, it can significantly reduce the number of unnecessary tokens processed across millions of queries.

Reranking: Prioritizing Quality

Retrieving more documents does not guarantee a better quality of response. Production systems often incorporate a reranking step, which filters the retrieved documents to pass only the most relevant ones to the model:

• Retriever → Top 10 Documents → Reranker → Top 3 Relevant Documents → LLM

The reranker evaluates the relevance of each document and only transmits the best matches. This presents two major advantages:

• The model processes a reduced number of tokens.
• It receives higher-quality context.

In many cases, a smaller number of documents improves the quality of responses, as the model is less distracted by irrelevant information.

Production Insight

Many teams invest time experimenting with more efficient embedding models. However, a significant improvement can result from a simpler approach: avoiding sending superfluous documents to the model. A more restricted and clearer context can often enhance accuracy and reduce costs.

Batch Inference: Optimizing Query Processing

So far, optimization has focused on what is sent to the language model. However, the efficiency of query processing is equally crucial. This is especially true when your AI application must handle a high volume of documents, emails, product descriptions, or customer reviews every hour. At this scale, processing requests one by one can become costly.

The Hidden Cost of Sequential Processing

Imagine a document search system. Before users can perform searches, each document must be converted into an embedding and stored in a vector database. Suppose you have 1,000 PDF documents to index. A simple approach might be:

for document in documents:
    embedding = embedding_model.embed(document)
    vector_db.insert(embedding)

While this works, it results in 1,000 distinct API calls. Each request involves:

• Network latency
• Authentication overhead
• Request initialization
• Response processing

The model spends almost as much time managing requests as it does generating embeddings.

A More Efficient Approach

Rather than sending documents individually, production systems process them in batches:

batch_size = 100
for i in range(0, len(documents), batch_size):
    batch = documents[i:i + batch_size]
    embeddings = embedding_model.embed(batch)
    vector_db.insert(embeddings)

Thus, instead of making 1,000 API calls, you only make 10. The number of tokens remains nearly the same, but the infrastructure becomes much more efficient.

Why Batching Improves Performance

Imagine ordering coffee for your team. Would you prefer to make twenty trips to order coffee each time, or gather all the orders and make a single visit? Both approaches yield the same result, but one wastes much less time. Batch inference works similarly. By processing multiple inputs together, the system reduces overhead and improves throughput.

Where Batch Inference is Most Effective

Batching is particularly effective for tasks that do not require an immediate response. A few examples include:

• Generating embeddings for large collections of documents
• Indexing knowledge bases
• Classifying customer feedback
• Offline summarization
• Processing support tickets
• Content moderation

These tasks occur in the background, where processing speed is more important than immediate user interaction. For real-time chatbots, however, batching is often less suitable, as users expect immediate responses.