Active Learning for RAG: Improving Retrieval Over Time

Active Learning for RAG: Improving Retrieval Over Time

Retrieval-Augmented Generation (RAG) systems have become the backbone of production LLM applications, but most deployments suffer from a silent problem: their retrieval quality stagnates. The embeddings, chunking strategy, and vector index are set up once, and then the system runs unchanged — even as real users expose weaknesses that could be fixed. Active Learning for RAG addresses this by turning user interactions into a feedback loop that continuously improves retrieval performance over time.

What Is Active Learning for RAG?

Active Learning is a machine learning paradigm where a model identifies which unlabeled examples would be most valuable to label, so a human (or another model) can annotate only those high-value samples. Applied to RAG, it means selectively collecting feedback on the queries and retrieved passages where the system is most uncertain or most likely to fail, then using that feedback to refine the retriever, the embeddings, or the document corpus.

Instead of manually auditing thousands of queries, you let the system surface the few dozen that actually matter. Those become training signal for fine-tuning embeddings, adjusting chunk boundaries, adding missing documents, or rewriting queries.

Why It Matters

• Retrieval is the bottleneck. Most RAG failures are retrieval failures, not generation failures. If the right chunk isn't in the top-k context, no amount of prompt engineering will save the answer.
• User traffic is free signal. Every query, click, follow-up, and reformulation tells you something about what your retriever is missing.
• Annotation budget is finite. Labeling every query is impractical. Active learning focuses human effort where it has the highest marginal value.
• Distribution shift is real. What users ask in month six is different from month one. Static indexes decay; active learning keeps the system aligned with actual usage.

Core Components of an Active Learning Loop for RAG

A practical active learning pipeline for RAG has four stages: uncertainty estimation, sample selection, annotation, and model update. Let's walk through each with code.

1. Logging Retrieval Metadata

Before you can select informative samples, you need to record what happened during retrieval. At minimum, log the query, the top-k retrieved chunks with their scores, and any downstream signal (user feedback, answer rating, follow-up queries).

import sqlite3
import json
from datetime import datetime

class RetrievalLogger:
    def __init__(self, db_path="rag_logs.db"):
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS retrieval_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                query TEXT,
                retrieved_chunks TEXT,
                scores TEXT,
                timestamp TEXT,
                user_feedback TEXT,
                answer TEXT
            )
        """)
        self.conn.commit()

    def log(self, query, chunks, scores, answer="", feedback=None):
        self.conn.execute(
            "INSERT INTO retrieval_events (query, retrieved_chunks, scores, timestamp, user_feedback, answer) VALUES (?, ?, ?, ?, ?, ?)",
            (
                query,
                json.dumps([c["id"] for c in chunks]),
                json.dumps(scores),
                datetime.utcnow().isoformat(),
                feedback,
                answer,
            ),
        )
        self.conn.commit()

    def fetch_all(self):
        cur = self.conn.execute("SELECT * FROM retrieval_events")
        return cur.fetchall()

2. Estimating Uncertainty

The key question in active learning is: which queries should we review? Several heuristics work well for RAG systems.

Score margin: The difference between the top-1 and top-2 retrieval scores. A small margin means the retriever is torn between candidates — a good sign the ranking may be wrong.

Low top-1 confidence: If even the best chunk has a low similarity score, the query may be out of distribution or the corpus may be missing relevant content.

Embedding distance to known queries: Queries whose embeddings are far from anything in your training set are exploratory and worth reviewing.

import numpy as np

def score_margin(scores):
    """Small margin between top-1 and top-2 indicates uncertainty."""
    if len(scores) < 2:
        return 0.0
    return abs(scores[0] - scores[1])

def top1_confidence(scores):
    """How confident is the best retrieval?"""
    return scores[0] if scores else 0.0

def uncertainty_score(scores, alpha=0.6, beta=0.4):
    """
    Combine margin and confidence into a single uncertainty score.
    Higher = more uncertain = more worth labeling.
    """
    margin = score_margin(scores)
    confidence = top1_confidence(scores)
    # Normalize margin to [0,1] range (assume cosine similarity)
    normalized_margin = 1.0 - min(margin, 1.0)
    normalized_confidence = 1.0 - max(min(confidence, 1.0), 0.0)
    return alpha * normalized_margin + beta * normalized_confidence

# Example usage on a batch of logged events
def rank_events_by_uncertainty(events):
    ranked = []
    for event in events:
        _, query, chunk_ids, scores_json, ts, feedback, answer = event
        scores = json.loads(scores_json)
        if feedback is not None:
            continue  # skip already-labeled
        u = uncertainty_score(scores)
        ranked.append((u, event))
    ranked.sort(key=lambda x: x[0], reverse=True)
    return ranked

3. Selecting and Annotating Samples

Once you've ranked events by uncertainty, select the top N for annotation. Annotation can be done by domain experts, by an LLM-as-judge pipeline, or by a hybrid where an LLM proposes labels and a human reviews them.

def select_samples_for_annotation(events, batch_size=50):
    ranked = rank_events_by_uncertainty(events)
    return [event for _, event in ranked[:batch_size]]

def llm_judge_annotation(query, retrieved_chunks, answer, judge_llm):
    """
    Use a strong LLM to judge whether the retrieved context
    was sufficient to answer the query.
    """
    prompt = f"""You are evaluating a RAG retrieval step.

Query: {query}

Retrieved context:
{chr(10).join([c['text'] for c in retrieved_chunks])}

Generated answer: {answer}

Respond with JSON:
{{
  "retrieval_sufficient": true/false,
  "missing_information": "description of what was missing, if anything",
  "better_query_rewrite": "an improved version of the query, if helpful"
}}
"""
    response = judge_llm.complete(prompt)
    return json.loads(response)

The judge produces structured labels you can act on: whether retrieval was sufficient, what was missing, and a rewritten query that might retrieve better.

4. Updating the Retriever

Annotation is only useful if it feeds back into the system. There are several update strategies, each addressing a different failure mode.

Strategy A: Fine-Tune Embeddings

If you have query-document pairs labeled as relevant or irrelevant, you can fine-tune your embedding model with contrastive loss to push relevant pairs closer together and irrelevant pairs apart.

import torch
import torch.nn.functional as F

def contrastive_loss(anchor, positive, negative, margin=0.5):
    """
    anchor: embedding of the query
    positive: embedding of a relevant chunk
    negative: embedding of an irrelevant chunk
    """
    distance_pos = 1.0 - F.cosine_similarity(anchor, positive)
    distance_neg = 1.0 - F.cosine_similarity(anchor, negative)
    losses = F.relu(distance_pos - distance_neg + margin)
    return losses.mean()

def build_training_pairs(annotated_events, embed_model, chunk_store):
    """
    Convert annotated events into (query, positive_chunk, negative_chunk) triples.
    """
    pairs = []
    for event in annotated_events:
        query, chunk_ids, scores_json, ts, feedback, answer = event[1:]
        label = json.loads(feedback)
        if not label.get("retrieval_sufficient"):
            # Use the judge's rewrite as an augmented positive query
            query = label.get("better_query_rewrite", query)
        query_emb = embed_model.encode(query)
        # positive: highest-scoring chunk the judge deemed relevant
        # negative: a chunk the judge flagged as irrelevant
        pos_emb = embed_model.encode(chunk_store.get(chunk_ids[0]))
        neg_emb = embed_model.encode(chunk_store.get(label.get("irrelevant_chunk_id")))
        pairs.append((query_emb, pos_emb, neg_emb))
    return pairs

Strategy B: Query Rewriting and Expansion

Not every problem requires retraining. Sometimes the fix is a better query transformation layer. Collect the judge's rewritten queries and look for patterns — maybe users consistently use abbreviations your retriever doesn't expand, or ask in a style that doesn't match document prose.

class QueryRewriteBank:
    def __init__(self):
        self.rewrites = []  # (original, rewrite) pairs

    def add(self, original, rewrite):
        self.rewrites.append((original, rewrite))

    def train_classifier(self):
        """
        Train a lightweight seq2seq or rule-based rewriter
        from collected (original -> rewrite) pairs.
        """
        # In practice, fine-tune a small T5/BART model here
        # or build a lookup table for common reformulations.
        pass

    def rewrite(self, query):
        # Apply learned rewrites at inference time
        return self._apply_rules(query)

Strategy C: Corpus Gaps

If the judge consistently reports missing information that isn't in your corpus at all, that's a signal to add new documents. Track these gaps systematically:

def detect_corpus_gaps(annotated_events, threshold=3):
    """
    If the same missing topic appears across multiple queries,
    flag it as a corpus gap to fill.
    """
    gap_counter = {}
    for event in annotated_events:
        feedback = json.loads(event[5]) if event[5] else {}
        missing = feedback.get("missing_information", "")
        if missing:
            gap_counter[missing] = gap_counter.get(missing, 0) + 1

    return [gap for gap, count in gap_counter.items() if count >= threshold]

Putting the Loop Together

The full active learning cycle runs as a scheduled job: pull recent logs, rank by uncertainty, send a batch for annotation, apply updates, and redeploy. Here's a simplified orchestrator:

def active_learning_cycle(logger, judge_llm, embed_model, chunk_store, batch_size=50):
    # 1. Fetch unlabeled events
    events = logger.fetch_all()
    unlabeled = [e for e in events if e[5] is None]
    if not unlabeled:
        print("No new events to process.")
        return

    # 2. Select most uncertain samples
    samples = select_samples_for_annotation(unlabeled, batch_size)

    # 3. Annotate with LLM judge
    annotated = []
    for event in samples:
        _, query, chunk_ids, scores_json, ts, _, answer = event
        chunks = [chunk_store.get(cid) for cid in json.loads(chunk_ids)]
        label = llm_judge_annotation(query, chunks, answer, judge_llm)
        annotated.append((event, label))
        # Persist the label
        logger.conn.execute(
            "UPDATE retrieval_events SET user_feedback=? WHERE id=?",
            (json.dumps(label), event[0]),
        )
    logger.conn.commit()

    # 4. Detect corpus gaps
    gaps = detect_corpus_gaps(
        [(e[0], e[1], e[2], e[3], e[4], json.dumps(label), e[6]) for e, label in annotated]
    )
    if gaps:
        print(f"Corpus gaps detected: {gaps}")

    # 5. Build training pairs and fine-tune embeddings
    pairs = build_training_pairs(
        [(e[0], e[1], e[2], e[3], e[4], json.dumps(label), e[6]) for e, label in annotated],
        embed_model,
        chunk_store,
    )
    if pairs:
        fine_tune_embeddings(embed_model, pairs, epochs=3)

    print(f"Cycle complete. Processed {len(annotated)} samples.")

Best Practices

• Start with heuristics before ML. Score-margin and top-1 confidence are cheap and surprisingly effective. Don't build a learned acquisition function until you've exhausted simple heuristics.
• Diversify your batch. Don't select 50 nearly identical uncertain queries. Use clustering on query embeddings and sample across clusters to maximize coverage.
• Combine explicit and implicit feedback. Thumbs-up/down is gold, but so are implicit signals like query reformulation (a user rephrasing their question usually means the first retrieval failed).
• Version your retriever. Every time you fine-tune embeddings or update the index, snapshot it. Run a held-out evaluation set before deploying so you can detect regressions.
• Watch for selection bias. Active learning optimizes for the queries you've seen. Keep a random sampling stream alongside the uncertainty-based stream to catch blind spots.
• Set a labeling budget. Active learning is about efficiency. Decide upfront how many samples per week you can afford to judge, and let the acquisition function fill that budget.
• Don't forget the generator. Some failures look like retrieval problems but are actually prompt or generation issues. Have the judge distinguish between "context was insufficient" and "context was fine but answer was wrong."

Measuring Improvement

To know your active learning loop is working, maintain a fixed evaluation set of queries with known-good answers and track retrieval metrics over time: Recall@k, MRR, and nDCG. Also track business metrics like answer acceptance rate and follow-up query frequency. A healthy active learning pipeline should show steady improvement in retrieval recall on the eval set and a declining rate of user reformulations in production.

def evaluate_retriever(eval_set, retriever, k=5):
    """
    eval_set: list of {query, relevant_chunk_ids}
    """
    recalls = []
    rr_scores = []
    for item in eval_set:
        results = retriever.search(item["query"], top_k=k)
        retrieved_ids = [r["id"] for r in results]
        relevant = set(item["relevant_chunk_ids"])

        hit = any(rid in relevant for rid in retrieved_ids)
        recalls.append(1.0 if hit else 0.0)

        # MRR
        rr = 0.0
        for i, rid in enumerate(retrieved_ids):
            if rid in relevant:
                rr = 1.0 / (i + 1)
                break
        rr_scores.append(rr)

    return {
        "recall_at_k": sum(recalls) / len(recalls),
        "mrr": sum(rr_scores) / len(rr_scores),
    }

Conclusion

Active Learning for RAG transforms a static retrieval system into a self-improving one. By logging retrieval metadata, estimating uncertainty, selectively annotating high-value queries, and feeding those labels back into embeddings, query rewriting, and corpus management, you create a compounding return on every user interaction. The upfront investment is modest — a logger, an acquisition function, and a judge — but the payoff is a RAG system that gets sharper with every query it serves, staying aligned with real user needs long after the initial deployment.