Principled and Scalable Diversity-Aware Retrieval via Cardinality-Constrained Binary Quadratic Programming

1. Problem Statement

The paper addresses a critical bottleneck in Retrieval-Augmented Generation (RAG): the trade-off between relevance and semantic diversity in retrieved passages.

• The Issue: Standard top-$k$ retrieval maximizes relevance but often returns semantically redundant passages (e.g., near-duplicate news articles), wasting the limited context window of Large Language Models (LLMs).
• Limitations of Existing Methods:
  • Maximal Marginal Relevance (MMR): A greedy algorithm that lacks theoretical approximation guarantees for non-monotone submodular functions. Its runtime scales linearly with the number of retrieved passages ($k$), becoming a bottleneck for large $k$.
  • Determinantal Point Processes (DPP): A probabilistic framework where finding the Maximum A Posteriori (MAP) subset is NP-hard. Existing approximations have super-linear time complexity in $k$ and suffer from difficult-to-interpret trade-off parameters.
• Goal: Develop a method that explicitly balances relevance and diversity with theoretical guarantees, interpretable parameters, and sub-linear scalability in $k$.

2. Methodology

A. Problem Formulation (CCBQP)

The authors formulate diversity-aware retrieval as a Cardinality-Constrained Binary Quadratic Programming (CCBQP) problem.

• Objective: Select a subset of $k$ passages (represented by binary vector $x$) to maximize:
  $$\max_{x \in \{0,1\}^n} \theta \cdot (k-1) \cdot c^T x + (1-\theta) x^T (I - EE^T) x$$
  • $c$: Relevance scores (cosine similarity to query).
  • $E$: Embedding matrix of candidate passages.
  • $x^T (I - EE^T) x$: Diversity term (penalizes pairwise cosine similarity).
  • $\theta \in [0,1]$: Interpretable trade-off parameter.
  • The coefficient $(k-1)$ aligns the scales of the relevance and diversity terms.

B. Theoretical Analysis & Relaxation

Since CCBQP is NP-hard, the authors propose a non-convex tight continuous relaxation:

1. Relaxation: Relax the binary constraint $x \in \{0,1\}^n$ to the convex hull $x \in [0,1]^n$ while maintaining the cardinality constraint $\mathbf{1}^T x = k$.
2. Tightness Guarantee: They prove that if a penalty parameter $\lambda \ge 2$ is added to the objective, the global maximum of the relaxed problem coincides with the binary solution (Theorem 2.1).
3. Integral Solutions: Under the realistic assumption that passage embeddings are not perfectly anti-aligned (Assumption 2.2), they prove that all global maximizers of the relaxed problem are integral (Corollary 2.3). This eliminates the need for post-hoc rounding.
4. Optimization Landscape: They establish a "benign" landscape where all non-optimal stationary points are strict saddles with positive curvature (Theorem 2.4). This ensures gradient-based methods can escape local minima/saddles easily.

C. Optimization Algorithm (Frank-Wolfe)

The authors design a specialized Frank-Wolfe algorithm to solve the relaxed problem efficiently:

• Exact Line Search: Instead of using Lipschitz-based adaptive step sizes (which are slow due to the large eigenvalues of dense embedding matrices), they derive a closed-form solution for the optimal step size $\gamma$ along the update direction.
• Sparsity Acceleration: To avoid $O(nd)$ matrix-vector multiplications in every iteration, they exploit the sparsity of the search direction. The update $E^T d^{(t)}$ is computed as $E^T s^{(t)} - E^T x^{(t)}$, where $s^{(t)}$ is a $k$-sparse vector. This reduces the complexity of this step to $O(kd)$.
• Complexity: The total per-iteration complexity is dominated by $O(nd)$, but the wall-clock time scales sub-linearly with $k$ because the heavy lifting depends on $k$ only via the sparse term.

3. Key Contributions

1. Principled Formulation: Introduced the CCBQP framework for RAG, offering an interpretable trade-off parameter $\theta$ that balances relevance and diversity.
2. Theoretical Guarantees: Proved the tightness of a non-convex continuous relaxation and established that the relaxed problem's global optima are always integral (no rounding needed). Provided convergence guarantees for the proposed algorithm.
3. Scalable Algorithm: Developed a Frank-Wolfe variant with exact line search and sparsity-aware updates, achieving sub-linear scaling in $k$, a significant improvement over the linear (MMR) and super-linear (DPP) scaling of baselines.
4. Empirical Dominance: Demonstrated superior performance on the relevance-diversity Pareto frontier across multiple datasets and retrieval sizes.

4. Experimental Results

The method was evaluated on ASQA and QAMPARI datasets with candidate pools $n > 2 \times 10^6$ and retrieval sizes $k \in \{25, 50, 100\}$.

• Retrieval Quality (Pareto Frontier): The proposed method consistently dominates MMR and DPP, achieving higher Recall for a given Intra-List Average Distance (ILAD) diversity metric. It globally optimizes the trade-off rather than making greedy local decisions.
• Efficiency (Speedup):
  • For $k=100$, MMR and DPP require $>850$ ms.
  • The proposed method reduces latency to 40–180 ms.
  • Speedup: Achieves 2.4× to 22.9× speedup over MMR in the practical regime ($\theta \ge 0.5$).
• Generation Quality: End-to-end experiments using Llama-3.1-8B showed that diversity-aware retrieval improves generation quality (F1 score) compared to standard top-$k$.
  • Note: Excessive diversity ($\theta$ too low) can hurt performance by removing grounding facts, but moderate diversity ($\theta \approx 0.8-0.9$) yields consistent improvements.
  • Increasing $k$ from 25 to 50 improved recall but dropped precision due to information overload; the method adapts well to these constraints.

5. Significance

This work provides a theoretically grounded and computationally efficient solution to a major limitation in modern RAG systems.

• Scalability: As LLM context windows grow, the ability to retrieve large numbers of diverse passages ($k$) without linear or super-linear time penalties is crucial. This method makes large-scale diversity retrieval feasible in real-time.
• Robustness: By proving that the continuous relaxation yields integral solutions, the method avoids the heuristic rounding steps that often degrade performance in other optimization-based retrieval approaches.
• Practicality: The interpretable parameter $\theta$ allows practitioners to easily tune the system based on specific application needs (e.g., prioritizing factual coverage vs. broad topic coverage).

In summary, the paper transforms diversity-aware retrieval from a heuristic, slow process into a principled optimization problem with strong theoretical guarantees and state-of-the-art efficiency.