Convex Optimization for Retrieval-Augmented Generation Pipeline Tuning

A study of convex and non-convex optimization structure in RAG retrieval signal weighting. We formulate the problem of learning retrieval weights as a constrained Quadratic Program (QP), characterize where convexity holds and where it breaks down, and empirically measure the relaxation gap on TriviaQA.

---

Overview

Retrieval-Augmented Generation (RAG) pipelines combine multiple retrieval signals — dense cosine similarity, sparse BM25 lexical matching, and MMR diversity — whose weights are typically set by hand. This project asks: what is the true optimization structure of this weighting problem, and how much does a convex formulation lose compared to a non-convex one?

We show that:

• The full problem (joint selection + weight learning) is non-convex and NP-hard
• The weight learning subproblem (with fixed features) is a strongly convex QP with a unique global minimum
• The convex relaxation gap is 2–3% on Recall@k — small but measurable
• Both optimizers drive the MMR diversity weight to zero on single-answer factual QA, which is a genuine empirical finding about when diversity signals add value

---

Results Summary

System	Recall@1	Recall@5	MRR
Convex QP	0.635	0.892	0.933
Non-Convex	0.655	0.922	0.955
Equal Weights	0.500	0.685	0.776
Cosine Only	0.705	0.929	0.990
BM25 Only	0.580	0.779	0.850

End-to-end generation (Convex QP + Groq LLaMA-3 8B): Exact Match = 0.667 | F1 = 0.715

Learned weights:

Convex QP  : cosine=0.412  bm25=0.588  diversity=0.000
Non-Convex : cosine=0.547  bm25=0.452  diversity=0.001

All result figures are in results/figures/.

Lambda sensitivity (fig8): main.py automatically sweeps λ ∈ {10⁻⁴ … 10¹} during Stage 4b, saves results/lambda_results.json, and generate_plots.py produces fig8_lambda_sensitivity.png.

---

Project Structure

RAG-Convex-Optimization/
├── main.py                    # Pipeline orchestrator — run this
├── generate_plots.py          # Generates all result figures from results/results.json
├── lambda_sensitivity.py      # Standalone lambda sweep script (alternative to main.py Stage 4b)
├── config.yaml                # All parameters — nothing is hardcoded
├── requirements.txt
├── .env.example               # Copy to .env and add your API key(s)
├── src/
│   ├── data_loader.py         # TriviaQA loading and corpus construction
│   ├── retriever.py           # BM25 (sparse) + FAISS dense retrieval
│   ├── features.py            # Feature extraction: cosine, BM25, MMR diversity
│   ├── optimizer.py           # Convex QP (CVXPY) + Non-Convex (scipy) + baselines
│   ├── evaluator.py           # Recall@k, Precision@k, MRR, Exact Match, F1
│   └── llm_generator.py       # Groq / Gemini API integration
└── results/
    ├── results.json           # Full experiment output
    ├── lambda_results.json    # Lambda sensitivity sweep output (auto-generated by main.py)
    ├── X_train.npz            # Saved feature matrices (for lambda_sensitivity.py reruns)
    └── figures/               # Publication-quality plots (300 DPI)
        ├── fig1_recall_at_k.png
        ├── fig2_precision_at_k.png
        ├── fig3_mrr_bar.png
        ├── fig4_weights.png
        ├── fig5_relaxation_gap.png
        ├── fig6_generation.png
        ├── fig7_summary_heatmap.png
        └── fig8_lambda_sensitivity.png

---

Setup

1. Clone and install dependencies

git clone https://github.com/Lord2709/RAG-Convex-Optimization.git
cd RAG-Convex-Optimization
pip install -r requirements.txt

2. Configure API keys

cp .env.example .env
# Open .env and add your key(s):
#   GROQ_API_KEY   — get from https://console.groq.com/keys
#   GEMINI_API_KEY — get from https://aistudio.google.com/app/apikey

3. Choose your LLM provider

In config.yaml, set:

llm:
  provider: "groq"     # or "gemini"

---

Running the Pipeline

Full run (retrieval + LLM generation evaluation):

python main.py

Retrieval only (no API key needed — good for testing):

python main.py --skip-llm

Custom config:

python main.py --config config.yaml

Generate plots (after running main.py):

python generate_plots.py

Lambda sensitivity only (if main.py has already run once, uses saved X_train.npz):

python lambda_sensitivity.py

Lambda sensitivity with custom grid:

python lambda_sensitivity.py --lambdas 0.0001 0.001 0.01 0.1 1.0 10.0

Lambda sensitivity from scratch (no prior main.py run needed — recomputes features):

python lambda_sensitivity.py --recompute

For a fast smoke test, set num_queries: 20 in config.yaml — completes in under 2 minutes.

---

Configuration

All parameters are controlled via config.yaml. Key settings:

Parameter	Default	Description
dataset.num_queries	200	Total queries loaded from TriviaQA
dataset.num_train_queries	150	Queries used for weight learning
dataset.num_test_queries	50	Held-out queries for evaluation
retrieval.top_k	10	Documents retrieved per query
retrieval.embedding_model	all-MiniLM-L6-v2	Sentence-transformers model
optimization.lambda_reg	0.01	L2 regularization for convex QP
optimization.solver	OSQP	CVXPY solver
llm.provider	groq	LLM backend: groq or gemini
llm.num_eval_queries	30	Queries used for generation eval

---

Lambda (Regularization) Sensitivity Analysis

A key contribution of this project is a systematic study of how the regularization parameter λ affects the QP solution. We sweep λ ∈ {10⁻⁴, 3×10⁻⁴, 10⁻³, 3×10⁻³, 10⁻², 3×10⁻², 10⁻¹, 3×10⁻¹, 1, 3, 10} and record:

• Regularization path: how each learned weight [w_cosine, w_bm25, w_diversity] changes with λ
• Objective value: cost ‖Xw−y‖² + λ‖w‖² at the optimum as a function of λ
• Recall@5: training retrieval quality at each λ (peak identifies the best regularization strength)
• Weight composition: stacked-area view of how λ blends the three retrieval signals

This is saved as fig8_lambda_sensitivity.png and results/lambda_results.json. The analysis is automatically run in Stage 4b of main.py.

Key finding: The MMR diversity weight remains near zero across all λ values on TriviaQA single-answer QA, confirming that diversity signals add no value for this task regardless of regularization strength. The cosine/BM25 split evolves with λ: at extreme λ, regularization forces weights toward [1/3, 1/3, 1/3]; at small λ, the less-regularized solution allocates weight more aggressively based on empirical signal strength.

---

Optimization Formulation

Convex subproblem (weight learning):

Given pre-computed feature matrix X ∈ ℝⁿˣ³ with columns [cosine_similarity, BM25, MMR_diversity] and binary relevance labels y, we solve:

minimize   ‖Xw − y‖² + λ‖w‖²
subject to w ≥ 0,  Σwᵢ = 1

This is a strongly convex QP (λ > 0) with constraints on the probability simplex. Solved via CVXPY + OSQP with a guaranteed unique global optimum.

Why the full problem is non-convex:

1. Selecting the optimal subset of k documents from n candidates is combinatorial (NP-hard)
2. The downstream generation loss is non-differentiable with respect to w when selection is a function of w

Non-convex comparison: We also optimize w directly on Recall@k using Nelder-Mead (gradient-free), which maximizes a non-differentiable, non-convex objective. The gap between the two approaches is measured as the relaxation gap.

---

Systems Compared

System	Description
Convex QP	Weights learned via constrained QP (our main contribution)
Non-Convex	Weights learned by directly maximizing Recall@k (Nelder-Mead)
Equal Weights	Uniform w = [1/3, 1/3, 1/3] — naive baseline
Cosine Only	w = [1, 0, 0] — dense retrieval only
BM25 Only	w = [0, 1, 0] — sparse retrieval only

---

Dataset

• TriviaQA (rc.wikipedia split) via HuggingFace Datasets
• Wikipedia evidence pages used as the retrieval corpus
• Documents from all queries are pooled together, making retrieval non-trivial (each query must find its evidence among other queries' documents acting as distractors)
• Relevance labels derived from exact-match / token-F1 between document content and gold answers

---

Dependencies

• cvxpy — convex QP solver
• sentence-transformers — dense document embeddings
• faiss-cpu — approximate nearest neighbour search
• rank-bm25 — BM25 sparse retrieval
• datasets — TriviaQA loading
• groq / google-genai — LLM generation
• scipy — non-convex optimization (Nelder-Mead)
• matplotlib, seaborn — result visualizations

---

License

MIT License — see LICENSE.