Self-RAG and Adaptive Retrieval

Standard RAG always retrieves, even when the model already knows the answer. This wastes compute and can introduce irrelevant information. Self-RAG (Asai et al., 2023) trains the model to make retrieval decisions dynamically.

Reflection Tokens

Token Types

Implementation

class SelfRAG(nn.Module):
    def __init__(self, base_model, vocab_size):
        super().__init__()
        self.base_model = base_model
        hidden_size = base_model.config.hidden_size
        
        # Reflection token embeddings
        self.retrieve_head = nn.Linear(hidden_size, 2)  # Yes/No
        self.isrel_head = nn.Linear(hidden_size, 2)     # Relevant/Irrelevant
        self.issup_head = nn.Linear(hidden_size, 3)     # Fully/Partial/No
        self.isuse_head = nn.Linear(hidden_size, 2)     # Useful/Not Useful
    
    def forward(self, input_ids, attention_mask=None):
        outputs = self.base_model(input_ids, attention_mask=attention_mask)
        hidden_states = outputs.last_hidden_state
        
        # Generate reflection logits
        retrieve_logits = self.retrieve_head(hidden_states[:, -1, :])
        isrel_logits = self.isrel_head(hidden_states[:, -1, :])
        issup_logits = self.issup_head(hidden_states[:, -1, :])
        isuse_logits = self.isuse_head(hidden_states[:, -1, :])
        
        return {
            "retrieve": retrieve_logits,
            "isrel": isrel_logits,
            "issup": issup_logits,
            "isuse": isuse_logits,
            "hidden": hidden_states
        }

Training Self-RAG

Retrieval Decision Training

def train_retrieval_decision(model, data, optimizer):
    """Train the retrieval decision head."""
    total_loss = 0
    
    for batch in data:
        # Forward pass
        outputs = model(batch["input_ids"])
        
        # Retrieval decision loss
        retrieve_loss = F.cross_entropy(
            outputs["retrieve"],
            batch["retrieve_labels"]  # 0=No, 1=Yes
        )
        
        # Relevance prediction loss
        isrel_loss = F.cross_entropy(
            outputs["isrel"],
            batch["isrel_labels"]  # 0=Irrelevant, 1=Relevant
        )
        
        # Support prediction loss
        issup_loss = F.cross_entropy(
            outputs["issup"],
            batch["issup_labels"]  # 0=No, 1=Partial, 2=Full
        )
        
        loss = retrieve_loss + isrel_loss + issup_loss
        total_loss += loss
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    return total_loss / len(data)

Adaptive Retrieval Strategy

Dynamic Retrieval Threshold

def adaptive_retrieve(query, model, retriever, confidence_threshold=0.7):
    """Retrieve only when model confidence is low."""
    # Get model's initial prediction
    with torch.no_grad():
        outputs = model(query)
        retrieve_prob = F.softmax(outputs["retrieve"], dim=-1)[0, 1].item()
    
    # Decide whether to retrieve
    if retrieve_prob > confidence_threshold:
        # Model is uncertain, retrieve
        documents = retriever.retrieve(query, top_k=5)
        
        # Filter by relevance
        relevant_docs = []
        for doc in documents:
            with torch.no_grad():
                rel_outputs = model(query, doc)
                isrel_prob = F.softmax(rel_outputs["isrel"], dim=-1)[0, 1].item()
                if isrel_prob > 0.5:
                    relevant_docs.append(doc)
        
        return relevant_docs
    else:
        # Model is confident, no retrieval needed
        return []

Iterative Retrieval

def iterative_self_rag(query, model, retriever, max_iterations=3):
    """Perform iterative retrieval and generation."""
    context = ""
    generated = ""
    
    for i in range(max_iterations):
        # Generate with current context
        full_prompt = f"Query: {query}\nContext: {context}\nAnswer:"
        outputs = model(full_prompt)
        
        # Check if more retrieval is needed
        retrieve_prob = F.softmax(outputs["retrieve"], dim=-1)[0, 1].item()
        
        if retrieve_prob < 0.3 or i == max_iterations - 1:
            # No more retrieval needed
            generated = decode(outputs["hidden"])
            break
        
        # Retrieve more information
        new_docs = retriever.retrieve(query + " " + generated, top_k=3)
        
        # Filter by relevance
        for doc in new_docs:
            rel_outputs = model(query, doc)
            if F.softmax(rel_outputs["isrel"], dim=-1)[0, 1].item() > 0.5:
                context += f"\n{doc}"
    
    return generated

Self-RAG vs Standard RAG

CRAG: Corrective RAG

def crag_retrieve(query, retriever, evaluator, web_searcher):
    """Corrective RAG with quality-aware retrieval."""
    # Retrieve initial documents
    documents = retriever.retrieve(query, top_k=5)
    
    # Evaluate retrieval quality
    quality_scores = evaluator.score_documents(query, documents)
    avg_quality = sum(quality_scores) / len(quality_scores)
    
    if avg_quality > 0.8:
        # High quality: use documents directly
        return documents, "direct"
    elif avg_quality > 0.4:
        # Uncertain: combine with web search
        web_docs = web_searcher.search(query, top_k=3)
        return documents + web_docs, "combined"
    else:
        # Low quality: web search only
        web_docs = web_searcher.search(query, top_k=5)
        return web_docs, "web_only"

Practice Exercises

1. Self-RAG Training: Train a Self-RAG model on a dataset with retrieval labels. How does the retrieval decision accuracy compare to a fixed threshold?

2. Adaptive Threshold: Implement a dynamic retrieval threshold that adjusts based on query complexity. How does this affect retrieval frequency and answer quality?

3. Iterative Retrieval: Compare single-round vs iterative Self-RAG on questions requiring multiple pieces of evidence. How many iterations are typically needed?

4. CRAG Implementation: Implement CRAG with a quality evaluator. How does the quality threshold affect the balance between direct use and web search?

Key Takeaways

What to Learn Next