Investigate MaxSim acceleration strategies for PostgreSQL

Problem

The populate_binary_maxsim_cache() and populate_approach2_maxsim_cache() functions have O(N²) complexity where N = number of pages in a corpus. For each page, we compute MaxSim against every other page to find the nearest neighbor.

MaxSim computes: for each query token, find the max similarity to any target token, then average.

MaxSim(Q, T) = (1/|Q|) × Σ max(cos_sim(q, t)) for all q in Q, t in T

This requires |Q| × |T| similarity computations per page pair. With 100 tokens per page and 173 pages:

• Per comparison: 100 × 100 = 10,000 operations
• Full cache build: 173 × 172 × 10,000 = ~297 million operations

Current Performance

| Corpus Size | Approximate Time |
|-------------|------------------|
| 50 pages | ~5-30 seconds |
| 100 pages | ~30 seconds - 2 minutes |
| 500 pages | ~10-30 minutes |
| 1000 pages | ~1-2 hours |
| 5000 pages | Many hours |

No standard index (HNSW, IVFFlat) can accelerate MaxSim because it's not a simple nearest-neighbor search—it requires computing token-level interactions.

---

Potential Acceleration Strategies

1. SIMD-Optimized C Extension

PostgreSQL extensions can use SIMD (AVX2/AVX-512) for parallel vector operations.

// Pseudocode for SIMD MaxSim
float maxsim_simd(float* query_tokens, int q_len, 
                  float* target_tokens, int t_len, int dim) {
    float sum = 0;
    for (int q = 0; q < q_len; q++) {
        float max_sim = -1;
        for (int t = 0; t < t_len; t++) {
            // AVX-512 can process 16 floats at once
            float sim = dot_product_avx512(&query_tokens[q*dim], 
                                           &target_tokens[t*dim], dim);
            max_sim = fmax(max_sim, sim);
        }
        sum += max_sim;
    }
    return sum / q_len;
}

Expected speedup: 8-16x over scalar code

2. GPU Acceleration (CUDA/OpenCL)

MaxSim is embarrassingly parallel—perfect for GPUs.

• Load all page embeddings into GPU memory
• For each query page, compute similarity matrix Q×T in parallel
• Reduce to find max per row, then average

Expected speedup: 100-1000x for large corpora

Implementation options:

• pgvector is exploring GPU support
• Custom extension using CUDA + PostgreSQL foreign data wrapper
• External service (compute MaxSim outside Postgres, store results)

3. Approximate MaxSim with Token Clustering

Instead of comparing all tokens, cluster tokens and compare cluster representatives:

Preprocessing (once per page):
1. Cluster tokens into K groups (K << |T|)
2. Store cluster centroids + token-to-cluster mapping

Query time:
1. For each query token, find nearest cluster centroid
2. Only compare against tokens in that cluster + neighboring clusters

Expected speedup: ~10-50x (trades some accuracy for speed)

4. Inverted Index for Tokens (ColBERT v2 / PLAID style)

This is how ColBERT v2 and PLAID achieve fast retrieval at scale:

Build inverted index:
- Quantize token embeddings to centroids (e.g., 64K centroids)
- Store: centroid_id -> list of (page_id, token_position, residual)

Query:
1. For each query token, find top-k nearest centroids
2. Only score pages that have tokens in those centroids
3. Use residual vectors for precise scoring

Expected speedup: 100-1000x for large corpora (sublinear scaling!)

5. Optimized Binary MaxSim

We already have binary quantized vectors (bq_vectors). The binary MaxSim uses Hamming distance which is very fast (single CPU cycle for POPCNT).

Potential improvement: Ensure the max_sim() function for bit arrays is fully optimized with SIMD.

6. Matryoshka/Truncated Embeddings

Use shorter embeddings for initial filtering:

1. Store embeddings at multiple dimensions: 128, 64, 32, 16
2. First pass: MaxSim with 16-dim (very fast)
3. Re-rank top candidates with full 128-dim

Expected speedup: 4-8x for initial filtering

---

Recommended Approach

A custom PostgreSQL extension combining several techniques:

// pg_maxsim extension

// 1. Optimized storage: pack token embeddings contiguously
CREATE TYPE token_matrix AS (
    num_tokens int2,
    embeddings float4[]  -- flat array, row-major
);

// 2. SIMD-optimized MaxSim function
CREATE FUNCTION maxsim_fast(
    query token_matrix,
    target token_matrix
) RETURNS float4
AS 'pg_maxsim', 'maxsim_fast'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

// 3. Batch version for multiple targets
CREATE FUNCTION maxsim_batch(
    query token_matrix,
    targets token_matrix[]
) RETURNS float4[]
AS 'pg_maxsim', 'maxsim_batch'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

Quick Wins (Before Custom Extension)

1. Store embeddings as flat arrays instead of separate rows per token
2. Optimize unnest usage in current PL/pgSQL implementation
3. Ensure max_parallel_workers_per_gather is configured for parallel query execution

---

References

• ColBERT v2 Paper - Efficient passage retrieval with late interaction
• PLAID - Fast ColBERT retrieval with deferred interaction
• pgvector - Vector similarity for PostgreSQL