MIPS -> L2 Reduction

Background

Maximum Inner-Product Search (MIPS) isn't directly compatible with L2-nearest-neighbour structures because MIPS isn't a metric — there's no triangle inequality on dot products. MIPS-to-L2 reduction (Bachrach et al. 2014) augments each vector with one extra dimension chosen so that on the unit sphere MIPS-argmax equals L2-argmin. After the augmentation, any off-the-shelf L2 ANN structure (HNSW, IVF, KD-tree) handles MIPS correctly.

Problem statement

Implement two functions:

• mips_augment(db_vecs) — return (augmented_db, M) where $M = \max_i \|x_i\|$ and each augmented row is $\bigl(\tfrac{x_i}{M},\ \sqrt{1 - \|x_i/M\|^2}\bigr)$ (unit norm).
• mips_augment_query(q) — return $\bigl(\tfrac{q}{\|q\|},\ 0\bigr)$ (unit norm too).

In the augmented space, $\arg\max_i \langle q, x_i\rangle$ (MIPS) equals $\arg\min_i \|q_{\text{aug}} - x_i^{\text{aug}}\|_2$ (L2-NN).

Input

• db_vecs — np.ndarray shape $(n, d)$.
• q — 1-D np.ndarray of length $d$.

Output

• mips_augment returns (augmented_db: (n, d+1), M: float).
• mips_augment_query returns np.ndarray of length $d + 1$.

Examples

Example 1 — augmented shape

Input:  db_vecs shape=(5, 4)
Output: augmented shape=(5, 5)

Example 2 — augmented db rows have unit norm

Input:  any non-zero db
Output: every row's L2 norm == 1.0 ± 1e-6

Example 3 — MIPS = L2 in augmented space

Input:  random db, query
Output: argmax(db @ q) == argmin(||aug_db - aug_q||^2)

Constraints

• Handle M = 0 (all-zero db) by setting M = 1.
• Use np.maximum(0, ...) before the square root to avoid floating-point issues that produce tiny negatives.
• Augmented vectors have exactly one extra dimension.

Notes

• Why query has zero in the extra slot. The reduction relies on $\langle q_{\text{aug}}, x_{\text{aug}}\rangle = \langle q, x\rangle / M$; the extra dimension contributes nothing because $q_{\text{aug, last}} = 0$.
• Asymmetric variant. Other reductions augment query and db differently (asymmetric LSH) — same effect, slightly different geometry.