Upload model checkpoints and code

README.md

@@ -1,3 +1,69 @@
----
-license: apache-2.0
----
+# Multi-Manifold Retrieval: Proof of Concept
+
+A proof-of-concept implementation of the Multi-Manifold Retrieval defense against spectral poisoning attacks (GeoPoison-RAG) on Retrieval-Augmented Generation systems.
+
+## Core Idea
+
+Standard RAG systems use a single shared embedding space for queries and documents, making the **document geometry identical to the retrieval geometry**. GeoPoison-RAG exploits this by computing the spectral structure (Fiedler vector) of the document graph Laplacian to find optimal adversarial placement.
+
+Multi-Manifold Retrieval **decouples** these geometries by using:
+- Separate query and document manifolds (M_Q and M_D)
+- A non-decomposable cross-manifold relevance operator R(q, d)
+
+This breaks the attack because the Laplacian the attacker computes (document space) no longer predicts the Laplacian governing retrieval (cross-manifold).
+
+## Project Structure
+
+```
+multi_manifold_retrieval/
+├── models/
+│   ├── cross_manifold_operator.py   # Construction C: Attention-Geometric Hybrid
+│   ├── encoders.py                  # Sentence-transformer wrapper
+│   └── baseline.py                  # Standard cosine similarity baseline
+├── training/
+│   ├── train.py                     # Training loop
+│   ├── data.py                      # MS MARCO data loading
+│   └── losses.py                    # Contrastive loss
+├── evaluation/
+│   ├── spectral_analysis.py         # L_D, L_R, spectral discrepancy, Fiedler alignment
+│   ├── retrieval_metrics.py         # MRR@10, Recall@100
+│   └── attack_simulation.py         # GeoPoison-RAG simulation
+proofs/
+├── proof_theorem_4_3.tex            # Spectral Decoupling theorem
+└── proof_theorem_6_1.tex            # Query Complexity Lower Bound theorem
+configs/
+└── default.yaml                     # Hyperparameters
+run_experiment.py                    # End-to-end pipeline
+```
+
+## Setup
+
+```bash
+pip install -r requirements.txt
+```
+
+## Running
+
+Full experiment (train + evaluate + spectral analysis + attack):
+```bash
+python run_experiment.py --config configs/default.yaml
+```
+
+Skip training and load from checkpoint:
+```bash
+python run_experiment.py --skip-train --checkpoint checkpoints/best_operator.pt
+```
+
+## Key Metrics
+
+| Metric | Baseline (expected) | Multi-Manifold (expected) |
+|--------|-------------------|--------------------------|
+| Spectral discrepancy δ | ≈ 0 | > 0 (significant) |
+| Fiedler alignment cos(θ) | ≈ 1 | < 0.5 |
+| ASR@10 | > 0.8 | Significantly lower |
+| MRR@10 | Reference | ≥ 80% of baseline |
+
+## Formal Proofs
+
+- `proofs/proof_theorem_4_3.tex`: Proves that non-decomposable R with positive cross-manifold curvature guarantees spectral decoupling δ ≥ Ω(κ_R · λ_2(L_D)).
+- `proofs/proof_theorem_6_1.tex`: Proves that an adaptive adversary needs Ω(Vol(M_Q) / V_{d_Q}(ε/κ_R)) oracle queries to reconstruct R.

checkpoints/best_operator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c63df76c2941f40d9f840057dcdc54d506b74cba5dc71d329117c965fdef783
+size 6969493

checkpoints/final_operator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa9be37ddd7e178a56f93eeb0a3259b8409399ee8be89118d84b53b8533318b1
+size 6969525

configs/default.yaml

@@ -0,0 +1,63 @@
+# Multi-Manifold Retrieval - Default Configuration
+
+seed: 42
+
+# Encoder settings
+encoder:
+  model_name: "sentence-transformers/all-MiniLM-L6-v2"
+  embedding_dim: 384
+  freeze: true  # Freeze pretrained encoders
+
+# Cross-manifold operator (Construction C)
+cross_manifold:
+  num_heads: 4
+  head_dim: 96  # embedding_dim / num_heads
+  value_mlp_hidden: 256
+  value_mlp_layers: 2
+  dropout: 0.1
+
+# Training
+training:
+  batch_size: 64
+  learning_rate: 2.0e-4
+  weight_decay: 1.0e-2
+  epochs: 5
+  warmup_steps: 500
+  max_train_samples: 100000
+  num_negatives: 7
+  max_seq_length: 128
+  fp16: true
+  gradient_accumulation_steps: 1
+  log_every: 100
+  eval_every: 2000
+  save_dir: "checkpoints"
+
+# Evaluation
+evaluation:
+  max_eval_queries: 5000
+  metrics:
+    - mrr@10
+    - recall@100
+
+# Spectral analysis
+spectral:
+  num_documents: 1000
+  num_queries: 500
+  k_neighbors: 20  # For sparse Laplacian (optional)
+
+# Attack simulation
+attack:
+  target_domain: "medical"
+  num_target_queries: 100
+  top_k: 10
+  medical_keywords:
+    - "health"
+    - "medical"
+    - "doctor"
+    - "patient"
+    - "treatment"
+    - "disease"
+    - "symptom"
+    - "diagnosis"
+    - "medicine"
+    - "clinical"

multi_manifold_retrieval/evaluation/attack_simulation.py

@@ -0,0 +1,332 @@
+"""Simplified GeoPoison-RAG attack simulation.
+
+Realistic threat model (matching GeoPoison-RAG Phase 1):
+- Attacker has shadow queries approximating target query distribution.
+- Attacker has access to document embeddings.
+- Attacker builds bipartite query-document graph using COSINE SIMILARITY
+  (their model of how retrieval works).
+- Attacker computes Fiedler vector and places adversarial doc at the
+  spectral-optimal position in document space.
+
+Defense argument:
+- Baseline (cosine sim): attacker's model is correct → high ASR.
+- Multi-manifold (R(q,d)): attacker's model is wrong because R ≠ cosine → lower ASR.
+"""
+
+import numpy as np
+import torch
+from scipy.sparse.linalg import eigsh
+from sklearn.metrics.pairwise import cosine_similarity
+
+from multi_manifold_retrieval.evaluation.spectral_analysis import compute_document_laplacian
+
+
+def select_domain_documents(
+    passages: list[str],
+    keywords: list[str],
+    max_docs: int = 200,
+) -> tuple[list[int], list[str]]:
+    """Select documents belonging to a target domain by keyword matching."""
+    indices = []
+    texts = []
+    for i, text in enumerate(passages):
+        text_lower = text.lower()
+        if any(kw in text_lower for kw in keywords):
+            indices.append(i)
+            texts.append(text)
+            if len(indices) >= max_docs:
+                break
+    return indices, texts
+
+
+def build_bipartite_fiedler_placement(
+    query_embs: np.ndarray,
+    doc_embs: np.ndarray,
+    t_nn: int = 20,
+) -> tuple[np.ndarray, dict]:
+    """GeoPoison-RAG Phase 1: bipartite spectral placement (cosine-based).
+
+    The attacker:
+    1. Builds bipartite query-document graph using cosine similarity.
+    2. Computes Fiedler vector of the normalized Laplacian.
+    3. Extracts document component of Fiedler vector.
+    4. Places adversarial doc at Fiedler-weighted centroid of documents.
+
+    The placement is in DOCUMENT SPACE — the attacker optimizes where to
+    place a document, guided by the query-document spectral structure.
+    But the attacker assumes retrieval = cosine similarity.
+    """
+    nq = query_embs.shape[0]
+    nd = doc_embs.shape[0]
+
+    # Cosine similarity between queries and documents (attacker's model)
+    S = cosine_similarity(query_embs, doc_embs)  # (nq, nd)
+
+    # Sparsify: keep top-t per query
+    t = min(t_nn, nd - 1)
+    S_sparse = np.zeros_like(S)
+    for i in range(nq):
+        top_idx = np.argpartition(S[i], -t)[-t:]
+        S_sparse[i, top_idx] = S[i, top_idx]
+
+    # Build bipartite adjacency: A = [[0, S], [S^T, 0]]
+    n = nq + nd
+    A = np.zeros((n, n))
+    A[:nq, nq:] = S_sparse
+    A[nq:, :nq] = S_sparse.T
+
+    # Normalized Laplacian: L = I - D^{-1/2} A D^{-1/2}
+    degrees = A.sum(axis=1)
+    degrees[degrees == 0] = 1.0
+    D_inv_sqrt = np.diag(1.0 / np.sqrt(degrees))
+    L = np.eye(n) - D_inv_sqrt @ A @ D_inv_sqrt
+
+    # Fiedler vector (2nd smallest eigenvector)
+    k = min(3, n - 1)
+    eigenvalues, eigenvectors = eigsh(L, k=k, which="SM")
+    sorted_idx = np.argsort(eigenvalues)
+    fiedler_vec = eigenvectors[:, sorted_idx[1]]
+    fiedler_val = eigenvalues[sorted_idx[1]]
+
+    # Extract document component and use as weights
+    doc_component = fiedler_vec[nq:]
+    weights = np.abs(doc_component)
+    weights = weights / (weights.sum() + 1e-12)
+
+    # Fiedler-weighted centroid of documents
+    adv_embedding = (weights[:, None] * doc_embs).sum(axis=0)
+
+    # L2-normalize
+    norm = np.linalg.norm(adv_embedding)
+    if norm > 0:
+        adv_embedding = adv_embedding / norm
+
+    info = {
+        "method": "bipartite_fiedler",
+        "fiedler_eigenvalue": float(fiedler_val),
+        "weight_entropy": float(-np.sum(weights * np.log(weights + 1e-12))),
+        "max_weight": float(weights.max()),
+        "adv_mean_cos_to_queries": float(
+            cosine_similarity(adv_embedding.reshape(1, -1), query_embs).mean()
+        ),
+        "adv_mean_cos_to_docs": float(
+            cosine_similarity(adv_embedding.reshape(1, -1), doc_embs).mean()
+        ),
+    }
+
+    return adv_embedding, info
+
+
+def compute_doconly_fiedler_placement(doc_embs: np.ndarray) -> tuple[np.ndarray, dict]:
+    """Document-only Fiedler placement (no query access).
+
+    Weaker attacker that only has document embeddings.
+    Uses document-space Laplacian L_D directly.
+    """
+    n = doc_embs.shape[0]
+    if n < 3:
+        centroid = doc_embs.mean(axis=0)
+        return centroid / np.linalg.norm(centroid), {"method": "centroid_fallback"}
+
+    L_D, _ = compute_document_laplacian(doc_embs)
+
+    k = min(3, n - 1)
+    eigenvalues, eigenvectors = eigsh(L_D, k=k, which="SM")
+    sorted_idx = np.argsort(eigenvalues)
+    fiedler_vec = eigenvectors[:, sorted_idx[1]]
+    fiedler_val = eigenvalues[sorted_idx[1]]
+
+    weights = np.abs(fiedler_vec)
+    weights = weights / (weights.sum() + 1e-12)
+
+    adv_embedding = (weights[:, None] * doc_embs).sum(axis=0)
+    norm = np.linalg.norm(adv_embedding)
+    if norm > 0:
+        adv_embedding = adv_embedding / norm
+
+    return adv_embedding, {
+        "method": "doconly_fiedler",
+        "fiedler_eigenvalue": float(fiedler_val),
+    }
+
+
+def compute_asr_threshold(
+    query_embeddings: torch.Tensor,
+    corpus_embeddings: torch.Tensor,
+    adv_embedding: torch.Tensor,
+    operator,
+    top_k: int = 10,
+    device: str = "cpu",
+    batch_size: int = 50,
+) -> tuple[float, dict]:
+    """Compute ASR@k using per-query threshold (oracle-style).
+
+    For each query, the k-th highest corpus score is the threshold.
+    Attack succeeds if the adversarial doc's score >= threshold.
+    Mirrors gp_rag/plan_single.py oracle check.
+    """
+    num_queries = query_embeddings.shape[0]
+    corpus_emb = corpus_embeddings.to(device)
+    adv_emb = adv_embedding.to(device)
+
+    operator.eval()
+    successes = 0
+    margins = []
+
+    with torch.no_grad():
+        for start in range(0, num_queries, batch_size):
+            end = min(start + batch_size, num_queries)
+            q_batch = query_embeddings[start:end].to(device)
+            bs = q_batch.shape[0]
+
+            # Score adversarial document
+            adv_expanded = adv_emb.unsqueeze(0).expand(bs, -1)
+            adv_scores = operator(q_batch, adv_expanded)
+
+            # Score corpus documents
+            corpus_scores = operator.compute_pairwise(q_batch, corpus_emb)
+
+            # k-th highest corpus score = threshold
+            topk_vals, _ = torch.topk(corpus_scores, top_k, dim=1)
+            thresholds = topk_vals[:, -1]
+
+            for j in range(bs):
+                margin = float(adv_scores[j].item() - thresholds[j].item())
+                margins.append(margin)
+                if adv_scores[j] >= thresholds[j]:
+                    successes += 1
+
+    asr = successes / num_queries
+    margins_arr = np.array(margins)
+    info = {
+        "mean_margin": float(margins_arr.mean()),
+        "median_margin": float(np.median(margins_arr)),
+        "p25_margin": float(np.percentile(margins_arr, 25)),
+        "fraction_positive_margin": float((margins_arr >= 0).mean()),
+    }
+
+    return asr, info
+
+
+def run_attack_simulation(
+    encoder,
+    operator,
+    baseline_operator,
+    passages: list[str],
+    passage_embeddings_torch: torch.Tensor,
+    target_query_texts: list[str],
+    medical_keywords: list[str],
+    top_k: int = 10,
+    max_domain_docs: int = 200,
+    device: str = "cpu",
+) -> dict:
+    """Run GeoPoison-RAG attack simulation.
+
+    Tests two attacker models:
+    1. Bipartite Fiedler (realistic): attacker has shadow queries + docs,
+       builds cosine-based bipartite graph, optimizes in document space.
+    2. Doc-only Fiedler (weaker): attacker has only document embeddings.
+
+    Both assume cosine similarity governs retrieval. The defense breaks
+    this assumption via the cross-manifold operator R.
+    """
+    print("\n=== Attack Simulation ===", flush=True)
+
+    # Step 1: Select target domain documents
+    domain_indices, domain_texts = select_domain_documents(
+        passages, medical_keywords, max_domain_docs
+    )
+    print(f"Selected {len(domain_indices)} domain documents.", flush=True)
+
+    if len(domain_indices) < 5:
+        print("Warning: Too few domain documents found.")
+        return {"error": "insufficient domain documents"}
+
+    domain_embs_np = passage_embeddings_torch[domain_indices].cpu().numpy()
+    domain_corpus = passage_embeddings_torch[domain_indices]
+
+    # Step 2: Encode target queries (attacker's shadow queries)
+    print(f"Encoding {len(target_query_texts)} target queries...", flush=True)
+    query_embeddings = encoder.encode_queries(target_query_texts, show_progress=False)
+    q_np = query_embeddings.cpu().numpy()
+
+    # Step 3a: Bipartite Fiedler placement (realistic attacker)
+    print("\nComputing bipartite Fiedler placement (attacker has shadow queries)...", flush=True)
+    adv_bipartite_np, bp_info = build_bipartite_fiedler_placement(
+        q_np, domain_embs_np, t_nn=min(20, len(domain_indices) - 1)
+    )
+    adv_bipartite = torch.tensor(adv_bipartite_np, dtype=torch.float32)
+    print(f"  Fiedler eigenvalue: {bp_info['fiedler_eigenvalue']:.6f}", flush=True)
+    print(f"  Adv mean cos to queries: {bp_info['adv_mean_cos_to_queries']:.4f}", flush=True)
+    print(f"  Adv mean cos to docs:    {bp_info['adv_mean_cos_to_docs']:.4f}", flush=True)
+
+    # Step 3b: Doc-only Fiedler placement (weaker attacker)
+    print("\nComputing doc-only Fiedler placement (no query access)...", flush=True)
+    adv_doconly_np, do_info = compute_doconly_fiedler_placement(domain_embs_np)
+    adv_doconly = torch.tensor(adv_doconly_np, dtype=torch.float32)
+
+    # Step 4: Measure ASR for bipartite attack
+    print(f"\n--- Bipartite Fiedler Attack (realistic GeoPoison-RAG) ---", flush=True)
+
+    asr_bp_base, bp_base_info = compute_asr_threshold(
+        query_embeddings, domain_corpus, adv_bipartite,
+        baseline_operator, top_k, device
+    )
+    print(f"  Baseline ASR@{top_k}:        {asr_bp_base:.4f} (mean margin: {bp_base_info['mean_margin']:.4f})", flush=True)
+
+    asr_bp_mm, bp_mm_info = compute_asr_threshold(
+        query_embeddings, domain_corpus, adv_bipartite,
+        operator, top_k, device
+    )
+    print(f"  Multi-manifold ASR@{top_k}:  {asr_bp_mm:.4f} (mean margin: {bp_mm_info['mean_margin']:.4f})", flush=True)
+
+    # Step 5: Measure ASR for doc-only attack
+    print(f"\n--- Doc-only Fiedler Attack (weaker attacker) ---", flush=True)
+
+    asr_do_base, do_base_info = compute_asr_threshold(
+        query_embeddings, domain_corpus, adv_doconly,
+        baseline_operator, top_k, device
+    )
+    print(f"  Baseline ASR@{top_k}:        {asr_do_base:.4f} (mean margin: {do_base_info['mean_margin']:.4f})", flush=True)
+
+    asr_do_mm, do_mm_info = compute_asr_threshold(
+        query_embeddings, domain_corpus, adv_doconly,
+        operator, top_k, device
+    )
+    print(f"  Multi-manifold ASR@{top_k}:  {asr_do_mm:.4f} (mean margin: {do_mm_info['mean_margin']:.4f})", flush=True)
+
+    # Summary
+    results = {
+        "bipartite_attack": {
+            "baseline_asr": asr_bp_base,
+            "multi_manifold_asr": asr_bp_mm,
+            "baseline_margins": bp_base_info,
+            "multi_manifold_margins": bp_mm_info,
+            "placement_info": bp_info,
+        },
+        "doconly_attack": {
+            "baseline_asr": asr_do_base,
+            "multi_manifold_asr": asr_do_mm,
+            "baseline_margins": do_base_info,
+            "multi_manifold_margins": do_mm_info,
+            "placement_info": do_info,
+        },
+        "num_domain_docs": len(domain_indices),
+        "num_target_queries": len(target_query_texts),
+        "top_k": top_k,
+        # For backward compat with summary printing
+        "baseline_asr": asr_bp_base,
+        "multi_manifold_asr": asr_bp_mm,
+    }
+
+    def _reduction(base, mm):
+        return (1 - mm / max(base, 1e-9)) * 100
+
+    print(f"\n=== Attack Results Summary ===", flush=True)
+    print(f"                         Baseline    Multi-Manifold    Reduction", flush=True)
+    print(f"  Bipartite (realistic): {asr_bp_base:.4f}      {asr_bp_mm:.4f}"
+          f"            {_reduction(asr_bp_base, asr_bp_mm):.1f}%", flush=True)
+    print(f"  Doc-only (weaker):     {asr_do_base:.4f}      {asr_do_mm:.4f}"
+          f"            {_reduction(asr_do_base, asr_do_mm):.1f}%", flush=True)
+
+    return results

multi_manifold_retrieval/evaluation/retrieval_metrics.py

@@ -0,0 +1,99 @@
+"""Retrieval evaluation metrics: MRR@10 and Recall@100."""
+
+import numpy as np
+import torch
+
+
+def compute_retrieval_metrics(
+    query_embeddings: torch.Tensor,
+    doc_embeddings: torch.Tensor,
+    operator,
+    query_texts: list[str],
+    positive_passages: list[list[str]],
+    all_passages: list[str],
+    passage_embeddings: torch.Tensor,
+    device: str = "cpu",
+    batch_size: int = 32,
+) -> dict:
+    """Compute MRR@10 and Recall@100.
+
+    Args:
+        query_embeddings: (num_queries, d) query embeddings.
+        doc_embeddings: Not used directly (passage_embeddings used instead).
+        operator: Relevance operator (cross-manifold or baseline).
+        query_texts: List of query strings.
+        positive_passages: List of lists of positive passage texts per query.
+        all_passages: Flat list of all candidate passages.
+        passage_embeddings: (num_passages, d) embeddings for all_passages.
+        device: Computation device.
+        batch_size: Batch size for scoring.
+
+    Returns:
+        Dict with mrr@10 and recall@100.
+    """
+    num_queries = query_embeddings.shape[0]
+    num_passages = passage_embeddings.shape[0]
+
+    # Build positive passage index: for each query, which passage indices are relevant
+    passage_to_idx = {text: idx for idx, text in enumerate(all_passages)}
+    positive_indices = []
+    for pos_list in positive_passages:
+        indices = set()
+        for text in pos_list:
+            if text in passage_to_idx:
+                indices.add(passage_to_idx[text])
+        positive_indices.append(indices)
+
+    passage_embeddings = passage_embeddings.to(device)
+    operator.eval()
+
+    mrr_sum = 0.0
+    recall_100_sum = 0.0
+    valid_queries = 0
+
+    with torch.no_grad():
+        for i in range(0, num_queries, batch_size):
+            end = min(i + batch_size, num_queries)
+            q_batch = query_embeddings[i:end].to(device)  # (bs, d)
+
+            # Score all passages: (bs, num_passages)
+            scores = operator.compute_pairwise(q_batch, passage_embeddings)
+            scores_np = scores.cpu().numpy()
+
+            for j in range(scores_np.shape[0]):
+                query_idx = i + j
+                pos_set = positive_indices[query_idx]
+                if not pos_set:
+                    continue
+
+                # Rank by score (descending)
+                ranked = np.argsort(-scores_np[j])
+
+                # MRR@10
+                rr = 0.0
+                for rank, doc_idx in enumerate(ranked[:10]):
+                    if doc_idx in pos_set:
+                        rr = 1.0 / (rank + 1)
+                        break
+                mrr_sum += rr
+
+                # Recall@100
+                top_100 = set(ranked[:100].tolist())
+                recall = len(pos_set & top_100) / len(pos_set)
+                recall_100_sum += recall
+
+                valid_queries += 1
+
+    mrr_at_10 = mrr_sum / valid_queries if valid_queries > 0 else 0.0
+    recall_at_100 = recall_100_sum / valid_queries if valid_queries > 0 else 0.0
+
+    results = {
+        "mrr@10": mrr_at_10,
+        "recall@100": recall_at_100,
+        "num_queries": valid_queries,
+    }
+
+    print(f"MRR@10: {mrr_at_10:.4f} | Recall@100: {recall_at_100:.4f} "
+          f"({valid_queries} queries)")
+
+    return results

multi_manifold_retrieval/evaluation/spectral_analysis.py

@@ -0,0 +1,205 @@
+"""Spectral analysis: compute L_D, L_R, spectral discrepancy δ, and Fiedler alignment cos(θ)."""
+
+import numpy as np
+import torch
+from scipy import sparse
+from scipy.sparse.linalg import eigsh
+
+
+def compute_document_laplacian(doc_embeddings: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the document-space graph Laplacian L_D.
+
+    Args:
+        doc_embeddings: (n, d) array of L2-normalized document embeddings.
+
+    Returns:
+        L_D: (n, n) graph Laplacian.
+        W_D: (n, n) cosine similarity matrix.
+    """
+    # Cosine similarity (embeddings are already L2-normalized)
+    W_D = doc_embeddings @ doc_embeddings.T
+
+    # Clip to [0, 1] to ensure non-negative weights
+    W_D = np.clip(W_D, 0, None)
+    np.fill_diagonal(W_D, 0)  # No self-loops
+
+    # Degree matrix and Laplacian
+    degrees = W_D.sum(axis=1)
+    D_D = np.diag(degrees)
+    L_D = D_D - W_D
+
+    return L_D, W_D
+
+
+def compute_retrieval_laplacian(
+    doc_embeddings: torch.Tensor,
+    query_embeddings: torch.Tensor,
+    operator,
+    device: str = "cpu",
+    batch_size: int = 50,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the retrieval Laplacian L_R.
+
+    (W_R)_{ij} = (1/|Q|) * sum_q R(q, d_i) * R(q, d_j)
+
+    Args:
+        doc_embeddings: (n, d) document embeddings tensor.
+        query_embeddings: (m, d) query embeddings tensor.
+        operator: Cross-manifold operator or baseline operator.
+        device: Computation device.
+        batch_size: Number of queries to process at once.
+
+    Returns:
+        L_R: (n, n) retrieval Laplacian.
+        W_R: (n, n) retrieval similarity matrix.
+    """
+    n = doc_embeddings.shape[0]
+    m = query_embeddings.shape[0]
+
+    doc_embeddings = doc_embeddings.to(device)
+    query_embeddings = query_embeddings.to(device)
+
+    # Accumulate W_R = (1/m) * R^T R where R_{ki} = R(q_k, d_i)
+    W_R = np.zeros((n, n), dtype=np.float64)
+
+    operator.eval()
+    with torch.no_grad():
+        for start in range(0, m, batch_size):
+            end = min(start + batch_size, m)
+            q_batch = query_embeddings[start:end]  # (bs, d)
+
+            # Compute R(q, d) for all docs: (bs, n)
+            scores = operator.compute_pairwise(q_batch, doc_embeddings)
+            scores_np = scores.cpu().numpy().astype(np.float64)  # (bs, n)
+
+            # Outer product accumulation: W_R += scores^T @ scores
+            W_R += scores_np.T @ scores_np
+
+    W_R /= m
+
+    # Ensure non-negative and zero diagonal
+    W_R = np.clip(W_R, 0, None)
+    np.fill_diagonal(W_R, 0)
+
+    # Laplacian
+    degrees = W_R.sum(axis=1)
+    D_R = np.diag(degrees)
+    L_R = D_R - W_R
+
+    return L_R, W_R
+
+
+def compute_spectral_discrepancy(L_D: np.ndarray, L_R: np.ndarray,
+                                  num_eigenvalues: int = 50) -> float:
+    """Compute spectral discrepancy δ = ||σ(L_D) - σ(L_R)||_2.
+
+    Uses the smallest num_eigenvalues eigenvalues (normalized).
+
+    Args:
+        L_D: Document-space Laplacian.
+        L_R: Retrieval Laplacian.
+        num_eigenvalues: Number of eigenvalues to compare.
+
+    Returns:
+        δ: Spectral discrepancy.
+    """
+    n = L_D.shape[0]
+    k = min(num_eigenvalues, n - 2)
+
+    # Compute smallest eigenvalues (Laplacians have smallest eigenvalue = 0)
+    eigs_D = eigsh(L_D, k=k, which="SM", return_eigenvectors=False)
+    eigs_R = eigsh(L_R, k=k, which="SM", return_eigenvectors=False)
+
+    # Sort
+    eigs_D = np.sort(eigs_D)
+    eigs_R = np.sort(eigs_R)
+
+    # Normalize so max eigenvalue = 1
+    max_D = eigs_D[-1] if eigs_D[-1] > 0 else 1.0
+    max_R = eigs_R[-1] if eigs_R[-1] > 0 else 1.0
+    eigs_D_norm = eigs_D / max_D
+    eigs_R_norm = eigs_R / max_R
+
+    delta = np.linalg.norm(eigs_D_norm - eigs_R_norm)
+    return delta
+
+
+def compute_fiedler_alignment(L_D: np.ndarray, L_R: np.ndarray) -> float:
+    """Compute Fiedler vector alignment cos(θ) = |v_2(L_D)^T v_2(L_R)| / (||v_2(L_D)|| * ||v_2(L_R)||).
+
+    Args:
+        L_D: Document-space Laplacian.
+        L_R: Retrieval Laplacian.
+
+    Returns:
+        cos(θ): Absolute cosine of angle between Fiedler vectors (1 = aligned, 0 = orthogonal).
+    """
+    # Compute the two smallest eigenvalues/vectors
+    _, vecs_D = eigsh(L_D, k=2, which="SM")
+    _, vecs_R = eigsh(L_R, k=2, which="SM")
+
+    # Fiedler vector = eigenvector for 2nd smallest eigenvalue (index 1 after sorting)
+    v2_D = vecs_D[:, 1]
+    v2_R = vecs_R[:, 1]
+
+    # Normalize
+    v2_D = v2_D / np.linalg.norm(v2_D)
+    v2_R = v2_R / np.linalg.norm(v2_R)
+
+    # Absolute cosine similarity
+    cos_theta = np.abs(np.dot(v2_D, v2_R))
+    return cos_theta
+
+
+def run_spectral_analysis(
+    doc_embeddings_np: np.ndarray,
+    doc_embeddings_torch: torch.Tensor,
+    query_embeddings_torch: torch.Tensor,
+    operator,
+    baseline_operator,
+    device: str = "cpu",
+) -> dict:
+    """Run full spectral analysis for both multi-manifold and baseline.
+
+    Returns dict with all metrics.
+    """
+    print("Computing document-space Laplacian L_D...")
+    L_D, W_D = compute_document_laplacian(doc_embeddings_np)
+
+    print("Computing retrieval Laplacian L_R (multi-manifold)...")
+    L_R_mm, W_R_mm = compute_retrieval_laplacian(
+        doc_embeddings_torch, query_embeddings_torch, operator, device
+    )
+
+    print("Computing retrieval Laplacian L_R (baseline)...")
+    L_R_base, W_R_base = compute_retrieval_laplacian(
+        doc_embeddings_torch, query_embeddings_torch, baseline_operator, device
+    )
+
+    print("Computing spectral discrepancy and Fiedler alignment...")
+    num_eigs = min(50, doc_embeddings_np.shape[0] - 2)
+
+    delta_mm = compute_spectral_discrepancy(L_D, L_R_mm, num_eigs)
+    delta_base = compute_spectral_discrepancy(L_D, L_R_base, num_eigs)
+    cos_theta_mm = compute_fiedler_alignment(L_D, L_R_mm)
+    cos_theta_base = compute_fiedler_alignment(L_D, L_R_base)
+
+    results = {
+        "multi_manifold": {
+            "spectral_discrepancy": delta_mm,
+            "fiedler_alignment": cos_theta_mm,
+        },
+        "baseline": {
+            "spectral_discrepancy": delta_base,
+            "fiedler_alignment": cos_theta_base,
+        },
+        "L_D": L_D,
+        "L_R_mm": L_R_mm,
+        "L_R_base": L_R_base,
+    }
+
+    print(f"\n=== Spectral Analysis Results ===")
+    print(f"Multi-Manifold:  δ = {delta_mm:.4f},  cos(θ) = {cos_theta_mm:.4f}")
+    print(f"Baseline:        δ = {delta_base:.4f},  cos(θ) = {cos_theta_base:.4f}")
+
+    return results

multi_manifold_retrieval/models/baseline.py

@@ -0,0 +1,37 @@
+"""Standard dual encoder baseline: R(q, d) = cosine_similarity(q, d)."""
+
+import torch
+import torch.nn as nn
+
+
+class BaselineOperator(nn.Module):
+    """Decomposable baseline: R(q, d) = q^T d (cosine similarity on normalized embeddings)."""
+
+    def forward(self, q: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+        """Compute cosine similarity.
+
+        Args:
+            q: Query embeddings, shape (batch_size, embedding_dim), L2-normalized.
+            d: Document embeddings, shape (batch_size, num_docs, embedding_dim) or
+               (batch_size, embedding_dim).
+
+        Returns:
+            Similarity scores.
+        """
+        if d.dim() == 2:
+            return torch.sum(q * d, dim=-1)
+        else:
+            return torch.einsum("bd,bnd->bn", q, d)
+
+    def compute_pairwise(self, q: torch.Tensor,
+                         docs: torch.Tensor) -> torch.Tensor:
+        """Compute cosine similarity for all query-document pairs.
+
+        Args:
+            q: (num_queries, embedding_dim)
+            docs: (num_docs, embedding_dim)
+
+        Returns:
+            (num_queries, num_docs)
+        """
+        return torch.mm(q, docs.t())

multi_manifold_retrieval/models/cross_manifold_operator.py

@@ -0,0 +1,142 @@
+"""Construction C: Attention-Geometric Hybrid cross-manifold operator.
+
+R(q, d) = sum_{h=1}^{H} softmax((W_Q^h q)^T (W_K^h d) / sqrt(d_h)) * v^h(q, d)
+
+where v^h(q, d) is a learned query-dependent value function parameterized as
+a small MLP taking [q; d; q * d] as input.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ValueMLP(nn.Module):
+    """Learned value function v^h(q, d) for a single attention head.
+
+    Takes concatenation [q; d; q * d] as input, outputs a scalar.
+    """
+
+    def __init__(self, embedding_dim: int, hidden_dim: int = 256,
+                 num_layers: int = 2, dropout: float = 0.1):
+        super().__init__()
+        input_dim = 3 * embedding_dim  # [q; d; q*d]
+        layers = []
+        in_dim = input_dim
+        for _ in range(num_layers):
+            layers.extend([
+                nn.Linear(in_dim, hidden_dim),
+                nn.GELU(),
+                nn.Dropout(dropout),
+            ])
+            in_dim = hidden_dim
+        layers.append(nn.Linear(hidden_dim, 1))
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, q: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+        """Compute v^h(q, d).
+
+        Args:
+            q: Query embeddings, shape (batch, embed_dim) or (batch, num_docs, embed_dim)
+            d: Document embeddings, shape (batch, embed_dim) or (batch, num_docs, embed_dim)
+
+        Returns:
+            Scalar values, shape matching the batch/doc dimensions.
+        """
+        x = torch.cat([q, d, q * d], dim=-1)
+        return self.mlp(x).squeeze(-1)
+
+
+class CrossManifoldOperator(nn.Module):
+    """Attention-Geometric Hybrid (Construction C).
+
+    Implements the cross-manifold relevance operator R(q, d) as a
+    multi-head attention mechanism with learned value functions.
+    """
+
+    def __init__(self, embedding_dim: int, num_heads: int = 4,
+                 value_hidden_dim: int = 256, value_num_layers: int = 2,
+                 dropout: float = 0.1):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.head_dim = embedding_dim // num_heads
+        assert embedding_dim % num_heads == 0, \
+            f"embedding_dim {embedding_dim} must be divisible by num_heads {num_heads}"
+
+        # Per-head query and key projections
+        self.W_Q = nn.Linear(embedding_dim, embedding_dim, bias=False)
+        self.W_K = nn.Linear(embedding_dim, embedding_dim, bias=False)
+
+        # Per-head value MLPs
+        self.value_mlps = nn.ModuleList([
+            ValueMLP(embedding_dim, value_hidden_dim, value_num_layers, dropout)
+            for _ in range(num_heads)
+        ])
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.xavier_uniform_(self.W_Q.weight)
+        nn.init.xavier_uniform_(self.W_K.weight)
+
+    def forward(self, q: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+        """Compute R(q, d) = cross-manifold relevance score.
+
+        Args:
+            q: Query embeddings, shape (batch_size, embedding_dim)
+            d: Document embeddings, shape (batch_size, num_docs, embedding_dim)
+               or (batch_size, embedding_dim) for single document.
+
+        Returns:
+            Relevance scores, shape (batch_size, num_docs) or (batch_size,).
+        """
+        single_doc = d.dim() == 2
+        if single_doc:
+            d = d.unsqueeze(1)  # (batch, 1, embed_dim)
+
+        batch_size, num_docs, _ = d.shape
+
+        # Project queries and keys: (batch, embed_dim) -> (batch, num_heads, head_dim)
+        q_proj = self.W_Q(q).view(batch_size, self.num_heads, self.head_dim)
+        # (batch, num_docs, embed_dim) -> (batch, num_docs, num_heads, head_dim)
+        d_proj = self.W_K(d).view(batch_size, num_docs, self.num_heads, self.head_dim)
+
+        # Attention scores: (batch, num_docs, num_heads)
+        scale = self.head_dim ** 0.5
+        attn = torch.einsum("bhd,bnhd->bnh", q_proj, d_proj) / scale
+
+        # Softmax over heads (not over documents) — each head contributes a
+        # weighted value, and the weighting is query-key dependent.
+        attn_weights = F.softmax(attn, dim=-1)  # (batch, num_docs, num_heads)
+
+        # Expand q for value MLPs: (batch, num_docs, embed_dim)
+        q_expanded = q.unsqueeze(1).expand(-1, num_docs, -1)
+
+        # Compute per-head values and weight them
+        total = torch.zeros(batch_size, num_docs, device=q.device)
+        for h in range(self.num_heads):
+            v_h = self.value_mlps[h](q_expanded, d)  # (batch, num_docs)
+            total = total + attn_weights[:, :, h] * v_h
+
+        if single_doc:
+            total = total.squeeze(1)
+
+        return total
+
+    def compute_pairwise(self, q: torch.Tensor,
+                         docs: torch.Tensor) -> torch.Tensor:
+        """Compute R(q, d) for all query-document pairs.
+
+        Args:
+            q: Query embeddings, shape (num_queries, embedding_dim)
+            docs: Document embeddings, shape (num_docs, embedding_dim)
+
+        Returns:
+            Relevance matrix, shape (num_queries, num_docs).
+        """
+        # Expand docs for each query
+        num_queries = q.shape[0]
+        num_docs = docs.shape[0]
+        docs_expanded = docs.unsqueeze(0).expand(num_queries, -1, -1)
+        return self.forward(q, docs_expanded)

multi_manifold_retrieval/models/encoders.py

@@ -0,0 +1,60 @@
+"""Wrapper around sentence-transformers for document and query encoders."""
+
+import torch
+import torch.nn as nn
+from sentence_transformers import SentenceTransformer
+
+
+class DualEncoder(nn.Module):
+    """Frozen pretrained encoders for query and document manifolds.
+
+    Uses the same pretrained model for both query and document encoding
+    (separate manifolds are induced by the cross-manifold operator, not
+    by separate encoder weights).
+    """
+
+    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+                 max_seq_length: int = 128, freeze: bool = True):
+        super().__init__()
+        self.model = SentenceTransformer(model_name)
+        self.model.max_seq_length = max_seq_length
+        self.embedding_dim = self.model.get_sentence_embedding_dimension()
+
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+    def encode_queries(self, texts: list[str], batch_size: int = 64,
+                       show_progress: bool = False) -> torch.Tensor:
+        """Encode query texts to embeddings on M_Q."""
+        embeddings = self.model.encode(
+            texts, batch_size=batch_size, show_progress_bar=show_progress,
+            convert_to_tensor=True, normalize_embeddings=True,
+        )
+        return embeddings
+
+    def encode_documents(self, texts: list[str], batch_size: int = 64,
+                         show_progress: bool = False) -> torch.Tensor:
+        """Encode document texts to embeddings on M_D."""
+        embeddings = self.model.encode(
+            texts, batch_size=batch_size, show_progress_bar=show_progress,
+            convert_to_tensor=True, normalize_embeddings=True,
+        )
+        return embeddings
+
+    def forward_queries(self, input_ids: torch.Tensor,
+                        attention_mask: torch.Tensor) -> torch.Tensor:
+        """Forward pass for query token IDs (for training)."""
+        features = {"input_ids": input_ids, "attention_mask": attention_mask}
+        out = self.model.forward(features)
+        embeddings = out["sentence_embedding"]
+        return nn.functional.normalize(embeddings, p=2, dim=-1)
+
+    def forward_documents(self, input_ids: torch.Tensor,
+                          attention_mask: torch.Tensor) -> torch.Tensor:
+        """Forward pass for document token IDs (for training)."""
+        return self.forward_queries(input_ids, attention_mask)
+
+    @property
+    def tokenizer(self):
+        return self.model.tokenizer

multi_manifold_retrieval/training/data.py

@@ -0,0 +1,168 @@
+"""MS MARCO data loading for training and evaluation."""
+
+import random
+from typing import Optional
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+
+
+class MSMARCOTripleDataset(Dataset):
+    """MS MARCO passage ranking dataset with hard negatives.
+
+    Each example yields (query, positive_passage, [negative_passages]).
+    """
+
+    def __init__(self, tokenizer, max_samples: int = 100_000,
+                 num_negatives: int = 7, max_seq_length: int = 128,
+                 split: str = "train", seed: int = 42):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.num_negatives = num_negatives
+
+        # Load MS MARCO dataset
+        print(f"Loading MS MARCO ({split} split, max {max_samples} samples)...")
+        dataset = load_dataset("ms_marco", "v2.1", split=split, trust_remote_code=True)
+
+        # Filter to examples with at least one selected passage
+        self.examples = []
+        for i, ex in enumerate(dataset):
+            if len(self.examples) >= max_samples:
+                break
+            passages = ex["passages"]
+            selected = [j for j, s in enumerate(passages["is_selected"]) if s == 1]
+            if selected:
+                self.examples.append({
+                    "query": ex["query"],
+                    "positive": passages["passage_text"][selected[0]],
+                    "negatives": [
+                        passages["passage_text"][j]
+                        for j in range(len(passages["passage_text"]))
+                        if j not in selected
+                    ],
+                })
+
+        print(f"Loaded {len(self.examples)} training examples.")
+        self.rng = random.Random(seed)
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, idx: int) -> dict:
+        ex = self.examples[idx]
+        # Sample negatives (from in-passage negatives, pad with random if needed)
+        available_negs = ex["negatives"]
+        if len(available_negs) >= self.num_negatives:
+            negs = self.rng.sample(available_negs, self.num_negatives)
+        else:
+            negs = available_negs[:]
+            # Pad with random negatives from other examples
+            while len(negs) < self.num_negatives:
+                rand_ex = self.examples[self.rng.randint(0, len(self.examples) - 1)]
+                if rand_ex["positive"] != ex["positive"]:
+                    negs.append(rand_ex["positive"])
+
+        return {
+            "query": ex["query"],
+            "positive": ex["positive"],
+            "negatives": negs,
+        }
+
+
+def collate_fn(batch: list[dict], tokenizer, max_seq_length: int = 128) -> dict:
+    """Collate batch into tokenized tensors."""
+    queries = [b["query"] for b in batch]
+    positives = [b["positive"] for b in batch]
+    all_negatives = []
+    for b in batch:
+        all_negatives.extend(b["negatives"])
+
+    # Tokenize
+    q_enc = tokenizer(
+        queries, padding=True, truncation=True,
+        max_length=max_seq_length, return_tensors="pt",
+    )
+    p_enc = tokenizer(
+        positives, padding=True, truncation=True,
+        max_length=max_seq_length, return_tensors="pt",
+    )
+    n_enc = tokenizer(
+        all_negatives, padding=True, truncation=True,
+        max_length=max_seq_length, return_tensors="pt",
+    )
+
+    num_negatives = len(batch[0]["negatives"])
+    return {
+        "query_input_ids": q_enc["input_ids"],
+        "query_attention_mask": q_enc["attention_mask"],
+        "pos_input_ids": p_enc["input_ids"],
+        "pos_attention_mask": p_enc["attention_mask"],
+        "neg_input_ids": n_enc["input_ids"],
+        "neg_attention_mask": n_enc["attention_mask"],
+        "num_negatives": num_negatives,
+    }
+
+
+def get_dataloader(tokenizer, max_samples: int = 100_000,
+                   num_negatives: int = 7, batch_size: int = 64,
+                   max_seq_length: int = 128, split: str = "train",
+                   seed: int = 42, num_workers: int = 0) -> DataLoader:
+    """Create a DataLoader for MS MARCO training."""
+    dataset = MSMARCOTripleDataset(
+        tokenizer=tokenizer, max_samples=max_samples,
+        num_negatives=num_negatives, max_seq_length=max_seq_length,
+        split=split, seed=seed,
+    )
+
+    def _collate(batch):
+        return collate_fn(batch, tokenizer, max_seq_length)
+
+    return DataLoader(
+        dataset, batch_size=batch_size, shuffle=True,
+        collate_fn=_collate, num_workers=num_workers,
+        drop_last=True,
+    )
+
+
+class MSMARCOEvalDataset:
+    """MS MARCO dev set for evaluation."""
+
+    def __init__(self, tokenizer, max_queries: int = 5000,
+                 max_seq_length: int = 128, seed: int = 42):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+
+        print(f"Loading MS MARCO dev set (max {max_queries} queries)...")
+        dataset = load_dataset("ms_marco", "v2.1", split="validation", trust_remote_code=True)
+
+        self.queries = []
+        self.positives = []  # list of list of positive passage texts
+        self.all_passages = []  # flat list of all passages for retrieval
+        self.passage_set = set()
+
+        rng = random.Random(seed)
+        indices = list(range(len(dataset)))
+        rng.shuffle(indices)
+
+        for i in indices:
+            if len(self.queries) >= max_queries:
+                break
+            ex = dataset[i]
+            passages = ex["passages"]
+            selected = [j for j, s in enumerate(passages["is_selected"]) if s == 1]
+            if not selected:
+                continue
+
+            self.queries.append(ex["query"])
+            pos_texts = [passages["passage_text"][j] for j in selected]
+            self.positives.append(pos_texts)
+
+            # Add all passages to the corpus
+            for text in passages["passage_text"]:
+                if text not in self.passage_set:
+                    self.passage_set.add(text)
+                    self.all_passages.append(text)
+
+        print(f"Loaded {len(self.queries)} eval queries, "
+              f"{len(self.all_passages)} unique passages.")

multi_manifold_retrieval/training/losses.py

@@ -0,0 +1,38 @@
+"""Contrastive loss for training the cross-manifold operator."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ContrastiveLoss(nn.Module):
+    """Cross-entropy contrastive loss over (query, positive, negatives) triples.
+
+    Given a query q, positive document d+, and negative documents d1-, ..., dK-,
+    the loss is:
+        -log(exp(R(q, d+)/tau) / (exp(R(q, d+)/tau) + sum_k exp(R(q, dk-)/tau)))
+    """
+
+    def __init__(self, temperature: float = 0.05):
+        super().__init__()
+        self.temperature = temperature
+
+    def forward(self, pos_scores: torch.Tensor,
+                neg_scores: torch.Tensor) -> torch.Tensor:
+        """Compute contrastive loss.
+
+        Args:
+            pos_scores: (batch_size,) — R(q, d+) for each query.
+            neg_scores: (batch_size, num_negatives) — R(q, dk-) for each query.
+
+        Returns:
+            Scalar loss.
+        """
+        # Concatenate: (batch_size, 1 + num_negatives)
+        all_scores = torch.cat([pos_scores.unsqueeze(1), neg_scores], dim=1)
+        all_scores = all_scores / self.temperature
+
+        # Target: index 0 is the positive
+        targets = torch.zeros(all_scores.shape[0], dtype=torch.long,
+                              device=all_scores.device)
+        return F.cross_entropy(all_scores, targets)

multi_manifold_retrieval/training/train.py

@@ -0,0 +1,159 @@
+"""Training loop for the cross-manifold operator."""
+
+import os
+import time
+import yaml
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import OneCycleLR
+
+from multi_manifold_retrieval.models.encoders import DualEncoder
+from multi_manifold_retrieval.models.cross_manifold_operator import CrossManifoldOperator
+from multi_manifold_retrieval.training.data import get_dataloader
+from multi_manifold_retrieval.training.losses import ContrastiveLoss
+
+
+def train(config_path: str = "configs/default.yaml",
+          device: str = "cuda" if torch.cuda.is_available() else "cpu"):
+    """Train the cross-manifold operator on MS MARCO."""
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+
+    torch.manual_seed(config["seed"])
+
+    # Initialize encoder
+    print("Initializing encoder...", flush=True)
+    encoder = DualEncoder(
+        model_name=config["encoder"]["model_name"],
+        max_seq_length=config["training"]["max_seq_length"],
+        freeze=config["encoder"]["freeze"],
+    )
+    embedding_dim = encoder.embedding_dim
+
+    # Initialize cross-manifold operator
+    cm_config = config["cross_manifold"]
+    operator = CrossManifoldOperator(
+        embedding_dim=embedding_dim,
+        num_heads=cm_config["num_heads"],
+        value_hidden_dim=cm_config["value_mlp_hidden"],
+        value_num_layers=cm_config["value_mlp_layers"],
+        dropout=cm_config["dropout"],
+    ).to(device)
+
+    print(f"Encoder initialized (dim={embedding_dim}).", flush=True)
+
+    # Loss
+    loss_fn = ContrastiveLoss(temperature=0.05)
+
+    # Data
+    print("Loading training data...", flush=True)
+    t0 = time.time()
+    train_loader = get_dataloader(
+        tokenizer=encoder.tokenizer,
+        max_samples=config["training"]["max_train_samples"],
+        num_negatives=config["training"]["num_negatives"],
+        batch_size=config["training"]["batch_size"],
+        max_seq_length=config["training"]["max_seq_length"],
+        split="train",
+        seed=config["seed"],
+    )
+
+    print(f"Data loaded in {time.time()-t0:.1f}s. Batches per epoch: {len(train_loader)}", flush=True)
+
+    # Optimizer (only train the cross-manifold operator)
+    optimizer = AdamW(
+        operator.parameters(),
+        lr=config["training"]["learning_rate"],
+        weight_decay=config["training"]["weight_decay"],
+    )
+
+    total_steps = len(train_loader) * config["training"]["epochs"]
+    scheduler = OneCycleLR(
+        optimizer,
+        max_lr=config["training"]["learning_rate"],
+        total_steps=total_steps,
+        pct_start=min(config["training"]["warmup_steps"] / total_steps, 0.1),
+    )
+
+    # Training loop
+    print(f"Moving encoder to {device}...", flush=True)
+    encoder.model.to(device)
+    encoder.model.eval()
+    operator.train()
+    print(f"Starting training: {config['training']['epochs']} epochs, {total_steps} total steps", flush=True)
+
+    save_dir = config["training"]["save_dir"]
+    os.makedirs(save_dir, exist_ok=True)
+    log_every = config["training"]["log_every"]
+
+    global_step = 0
+    best_loss = float("inf")
+
+    for epoch in range(config["training"]["epochs"]):
+        epoch_loss = 0.0
+        epoch_start = time.time()
+
+        for batch_idx, batch in enumerate(train_loader):
+            # Move to device
+            q_ids = batch["query_input_ids"].to(device)
+            q_mask = batch["query_attention_mask"].to(device)
+            p_ids = batch["pos_input_ids"].to(device)
+            p_mask = batch["pos_attention_mask"].to(device)
+            n_ids = batch["neg_input_ids"].to(device)
+            n_mask = batch["neg_attention_mask"].to(device)
+            num_neg = batch["num_negatives"]
+
+            # Encode (no grad for frozen encoder)
+            with torch.no_grad():
+                q_emb = encoder.forward_queries(q_ids, q_mask)       # (B, D)
+                p_emb = encoder.forward_documents(p_ids, p_mask)     # (B, D)
+                n_emb = encoder.forward_documents(n_ids, n_mask)     # (B*K, D)
+
+            batch_size = q_emb.shape[0]
+            n_emb = n_emb.view(batch_size, num_neg, -1)  # (B, K, D)
+
+            # Compute relevance scores via cross-manifold operator
+            pos_scores = operator(q_emb, p_emb)            # (B,)
+            neg_scores = operator(q_emb, n_emb)            # (B, K)
+
+            # Loss
+            loss = loss_fn(pos_scores, neg_scores)
+
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(operator.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+
+            epoch_loss += loss.item()
+            global_step += 1
+
+            if global_step % log_every == 0:
+                avg_loss = epoch_loss / (batch_idx + 1)
+                lr = scheduler.get_last_lr()[0]
+                print(f"  Step {global_step} | Loss: {loss.item():.4f} | "
+                      f"Avg: {avg_loss:.4f} | LR: {lr:.2e}", flush=True)
+
+        epoch_time = time.time() - epoch_start
+        avg_loss = epoch_loss / len(train_loader)
+        print(f"Epoch {epoch+1}/{config['training']['epochs']} | "
+              f"Avg Loss: {avg_loss:.4f} | Time: {epoch_time:.1f}s", flush=True)
+
+        # Save best
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save(operator.state_dict(), os.path.join(save_dir, "best_operator.pt"))
+            print(f"  Saved best model (loss={best_loss:.4f})")
+
+    # Save final
+    torch.save(operator.state_dict(), os.path.join(save_dir, "final_operator.pt"))
+    print(f"Training complete. Best loss: {best_loss:.4f}")
+
+    return encoder, operator
+
+
+if __name__ == "__main__":
+    train()

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.0
+transformers>=4.30
+sentence-transformers>=2.2
+datasets>=2.14
+faiss-cpu>=1.7
+numpy>=1.24
+scipy>=1.10
+scikit-learn>=1.3
+pyyaml>=6.0
+tqdm>=4.65

results.json

@@ -0,0 +1,136 @@
+{
+  "config": {
+    "seed": 42,
+    "encoder": {
+      "model_name": "sentence-transformers/all-MiniLM-L6-v2",
+      "embedding_dim": 384,
+      "freeze": true
+    },
+    "cross_manifold": {
+      "num_heads": 4,
+      "head_dim": 96,
+      "value_mlp_hidden": 256,
+      "value_mlp_layers": 2,
+      "dropout": 0.1
+    },
+    "training": {
+      "batch_size": 64,
+      "learning_rate": 0.0002,
+      "weight_decay": 0.01,
+      "epochs": 5,
+      "warmup_steps": 500,
+      "max_train_samples": 100000,
+      "num_negatives": 7,
+      "max_seq_length": 128,
+      "fp16": true,
+      "gradient_accumulation_steps": 1,
+      "log_every": 100,
+      "eval_every": 2000,
+      "save_dir": "checkpoints"
+    },
+    "evaluation": {
+      "max_eval_queries": 5000,
+      "metrics": [
+        "mrr@10",
+        "recall@100"
+      ]
+    },
+    "spectral": {
+      "num_documents": 1000,
+      "num_queries": 500,
+      "k_neighbors": 20
+    },
+    "attack": {
+      "target_domain": "medical",
+      "num_target_queries": 100,
+      "top_k": 10,
+      "medical_keywords": [
+        "health",
+        "medical",
+        "doctor",
+        "patient",
+        "treatment",
+        "disease",
+        "symptom",
+        "diagnosis",
+        "medicine",
+        "clinical"
+      ]
+    }
+  },
+  "device": "cuda",
+  "retrieval_multi_manifold": {
+    "mrr@10": 0.6002776984126992,
+    "recall@100": 0.9901,
+    "num_queries": 5000
+  },
+  "retrieval_baseline": {
+    "mrr@10": 0.5828701587301599,
+    "recall@100": 0.9942,
+    "num_queries": 5000
+  },
+  "mrr_ratio": 1.0298652099816936,
+  "spectral": {
+    "multi_manifold": {
+      "spectral_discrepancy": 0.05735765351097603,
+      "fiedler_alignment": 0.03973450829139222
+    },
+    "baseline": {
+      "spectral_discrepancy": 0.22395483470893326,
+      "fiedler_alignment": 0.7848795751112227
+    },
+    "num_documents": 1000,
+    "num_queries": 500
+  },
+  "attack": {
+    "bipartite_attack": {
+      "baseline_asr": 0.51,
+      "multi_manifold_asr": 0.19,
+      "baseline_margins": {
+        "mean_margin": -0.00035160839557647706,
+        "median_margin": 0.0012769699096679688,
+        "p25_margin": -0.03736262768507004,
+        "fraction_positive_margin": 0.51
+      },
+      "multi_manifold_margins": {
+        "mean_margin": -0.0419993931055069,
+        "median_margin": -0.04468509554862976,
+        "p25_margin": -0.06919527053833008,
+        "fraction_positive_margin": 0.19
+      },
+      "placement_info": {
+        "method": "bipartite_fiedler",
+        "fiedler_eigenvalue": 0.12380694040243122,
+        "weight_entropy": 4.995662314784581,
+        "max_weight": 0.016811827898423226,
+        "adv_mean_cos_to_queries": 0.2539085502225744,
+        "adv_mean_cos_to_docs": 0.2686595998305845
+      }
+    },
+    "doconly_attack": {
+      "baseline_asr": 0.03,
+      "multi_manifold_asr": 0.03,
+      "baseline_margins": {
+        "mean_margin": -0.1627955549955368,
+        "median_margin": -0.1696268543601036,
+        "p25_margin": -0.21040004305541515,
+        "fraction_positive_margin": 0.03
+      },
+      "multi_manifold_margins": {
+        "mean_margin": -0.12464183956384658,
+        "median_margin": -0.12341519445180893,
+        "p25_margin": -0.1708945743739605,
+        "fraction_positive_margin": 0.03
+      },
+      "placement_info": {
+        "method": "doconly_fiedler",
+        "fiedler_eigenvalue": 3.8886430263519287
+      }
+    },
+    "num_domain_docs": 200,
+    "num_target_queries": 100,
+    "top_k": 10,
+    "baseline_asr": 0.51,
+    "multi_manifold_asr": 0.19
+  }
+}

run_experiment.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+"""End-to-end experiment: train → evaluate → spectral analysis → attack simulation."""
+
+import os
+import json
+import random
+import time
+import yaml
+import argparse
+
+import numpy as np
+import torch
+
+from multi_manifold_retrieval.models.encoders import DualEncoder
+from multi_manifold_retrieval.models.cross_manifold_operator import CrossManifoldOperator
+from multi_manifold_retrieval.models.baseline import BaselineOperator
+from multi_manifold_retrieval.training.train import train
+from multi_manifold_retrieval.training.data import MSMARCOEvalDataset
+from multi_manifold_retrieval.evaluation.spectral_analysis import run_spectral_analysis
+from multi_manifold_retrieval.evaluation.retrieval_metrics import compute_retrieval_metrics
+from multi_manifold_retrieval.evaluation.attack_simulation import (
+    run_attack_simulation,
+    select_domain_documents,
+)
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Multi-Manifold Retrieval PoC Experiment")
+    parser.add_argument("--config", type=str, default="configs/default.yaml")
+    parser.add_argument("--skip-train", action="store_true", help="Skip training, load from checkpoint")
+    parser.add_argument("--checkpoint", type=str, default="checkpoints/best_operator.pt")
+    parser.add_argument("--output", type=str, default="results.json")
+    args = parser.parse_args()
+
+    with open(args.config) as f:
+        config = yaml.safe_load(f)
+
+    set_seed(config["seed"])
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    results = {"config": config, "device": device}
+
+    # =========================================================================
+    # Phase 1: Training
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("PHASE 1: TRAINING")
+    print("=" * 60)
+
+    if args.skip_train and os.path.exists(args.checkpoint):
+        print(f"Loading encoder and operator from checkpoint: {args.checkpoint}")
+        encoder = DualEncoder(
+            model_name=config["encoder"]["model_name"],
+            max_seq_length=config["training"]["max_seq_length"],
+            freeze=config["encoder"]["freeze"],
+        )
+        cm_config = config["cross_manifold"]
+        operator = CrossManifoldOperator(
+            embedding_dim=encoder.embedding_dim,
+            num_heads=cm_config["num_heads"],
+            value_hidden_dim=cm_config["value_mlp_hidden"],
+            value_num_layers=cm_config["value_mlp_layers"],
+            dropout=cm_config["dropout"],
+        )
+        operator.load_state_dict(torch.load(args.checkpoint, map_location=device, weights_only=True))
+        operator.to(device)
+    else:
+        encoder, operator = train(config_path=args.config, device=device)
+
+    encoder.model.to(device)
+    operator.to(device)
+    baseline_operator = BaselineOperator().to(device)
+
+    # =========================================================================
+    # Phase 2: Evaluation Data Preparation
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("PHASE 2: EVALUATION DATA PREPARATION")
+    print("=" * 60)
+
+    eval_data = MSMARCOEvalDataset(
+        tokenizer=encoder.tokenizer,
+        max_queries=config["evaluation"]["max_eval_queries"],
+        max_seq_length=config["training"]["max_seq_length"],
+        seed=config["seed"],
+    )
+
+    # Encode all evaluation passages
+    print(f"Encoding {len(eval_data.all_passages)} passages...")
+    passage_embeddings = encoder.encode_documents(
+        eval_data.all_passages, batch_size=128, show_progress=True,
+    )
+
+    # Encode evaluation queries
+    print(f"Encoding {len(eval_data.queries)} queries...")
+    query_embeddings = encoder.encode_queries(
+        eval_data.queries, batch_size=128, show_progress=True,
+    )
+
+    # =========================================================================
+    # Phase 3: Retrieval Quality Evaluation
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("PHASE 3: RETRIEVAL QUALITY")
+    print("=" * 60)
+
+    print("\n--- Multi-Manifold Model ---")
+    metrics_mm = compute_retrieval_metrics(
+        query_embeddings=query_embeddings,
+        doc_embeddings=passage_embeddings,
+        operator=operator,
+        query_texts=eval_data.queries,
+        positive_passages=eval_data.positives,
+        all_passages=eval_data.all_passages,
+        passage_embeddings=passage_embeddings,
+        device=device,
+    )
+    results["retrieval_multi_manifold"] = metrics_mm
+
+    print("\n--- Baseline (Cosine Similarity) ---")
+    metrics_base = compute_retrieval_metrics(
+        query_embeddings=query_embeddings,
+        doc_embeddings=passage_embeddings,
+        operator=baseline_operator,
+        query_texts=eval_data.queries,
+        positive_passages=eval_data.positives,
+        all_passages=eval_data.all_passages,
+        passage_embeddings=passage_embeddings,
+        device=device,
+    )
+    results["retrieval_baseline"] = metrics_base
+
+    # Check: multi-manifold within 80% of baseline
+    if metrics_base["mrr@10"] > 0:
+        ratio = metrics_mm["mrr@10"] / metrics_base["mrr@10"]
+        print(f"\nMRR@10 ratio (mm/baseline): {ratio:.4f} "
+              f"({'PASS' if ratio >= 0.8 else 'BELOW TARGET'}, target >= 0.8)")
+        results["mrr_ratio"] = ratio
+
+    # =========================================================================
+    # Phase 4: Spectral Analysis
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("PHASE 4: SPECTRAL ANALYSIS")
+    print("=" * 60)
+
+    # Sample documents for spectral analysis
+    num_spectral_docs = min(config["spectral"]["num_documents"], len(eval_data.all_passages))
+    num_spectral_queries = min(config["spectral"]["num_queries"], len(eval_data.queries))
+
+    spectral_doc_indices = np.random.choice(
+        len(eval_data.all_passages), num_spectral_docs, replace=False
+    )
+    spectral_query_indices = np.random.choice(
+        len(eval_data.queries), num_spectral_queries, replace=False
+    )
+
+    spectral_doc_emb_np = passage_embeddings[spectral_doc_indices].cpu().numpy()
+    spectral_doc_emb_torch = passage_embeddings[spectral_doc_indices]
+    spectral_query_emb_torch = query_embeddings[spectral_query_indices]
+
+    spectral_results = run_spectral_analysis(
+        doc_embeddings_np=spectral_doc_emb_np,
+        doc_embeddings_torch=spectral_doc_emb_torch,
+        query_embeddings_torch=spectral_query_emb_torch,
+        operator=operator,
+        baseline_operator=baseline_operator,
+        device=device,
+    )
+
+    results["spectral"] = {
+        "multi_manifold": spectral_results["multi_manifold"],
+        "baseline": spectral_results["baseline"],
+        "num_documents": num_spectral_docs,
+        "num_queries": num_spectral_queries,
+    }
+
+    # =========================================================================
+    # Phase 5: Attack Simulation
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("PHASE 5: ATTACK SIMULATION")
+    print("=" * 60)
+
+    attack_config = config["attack"]
+
+    # Select target queries (medical domain)
+    target_queries = []
+    for q in eval_data.queries:
+        q_lower = q.lower()
+        if any(kw in q_lower for kw in attack_config["medical_keywords"]):
+            target_queries.append(q)
+        if len(target_queries) >= attack_config["num_target_queries"]:
+            break
+
+    if len(target_queries) < 10:
+        # Fall back: use random queries if not enough medical ones
+        print(f"Only found {len(target_queries)} medical queries; "
+              f"using random queries to reach {attack_config['num_target_queries']}.")
+        remaining = attack_config["num_target_queries"] - len(target_queries)
+        other_queries = [q for q in eval_data.queries if q not in target_queries]
+        target_queries.extend(random.sample(other_queries, min(remaining, len(other_queries))))
+
+    print(f"Using {len(target_queries)} target queries for attack simulation.")
+
+    attack_results = run_attack_simulation(
+        encoder=encoder,
+        operator=operator,
+        baseline_operator=baseline_operator,
+        passages=eval_data.all_passages,
+        passage_embeddings_torch=passage_embeddings,
+        target_query_texts=target_queries,
+        medical_keywords=attack_config["medical_keywords"],
+        top_k=attack_config["top_k"],
+        device=device,
+    )
+    results["attack"] = attack_results
+
+    # =========================================================================
+    # Summary
+    # =========================================================================
+    print("\n" + "=" * 60)
+    print("EXPERIMENT SUMMARY")
+    print("=" * 60)
+
+    print(f"\n1. Retrieval Quality:")
+    print(f"   Baseline MRR@10:        {metrics_base['mrr@10']:.4f}")
+    print(f"   Multi-Manifold MRR@10:  {metrics_mm['mrr@10']:.4f}")
+    if metrics_base["mrr@10"] > 0:
+        print(f"   Ratio:                  {metrics_mm['mrr@10']/metrics_base['mrr@10']:.4f}")
+
+    print(f"\n2. Spectral Analysis:")
+    print(f"   Baseline δ:             {spectral_results['baseline']['spectral_discrepancy']:.4f}")
+    print(f"   Multi-Manifold δ:       {spectral_results['multi_manifold']['spectral_discrepancy']:.4f}")
+    print(f"   Baseline cos(θ):        {spectral_results['baseline']['fiedler_alignment']:.4f}")
+    print(f"   Multi-Manifold cos(θ):  {spectral_results['multi_manifold']['fiedler_alignment']:.4f}")
+
+    if "error" not in attack_results:
+        print(f"\n3. Attack Simulation:")
+        print(f"   Baseline ASR@{attack_config['top_k']}:        {attack_results['baseline_asr']:.4f}")
+        print(f"   Multi-Manifold ASR@{attack_config['top_k']}:  {attack_results['multi_manifold_asr']:.4f}")
+
+    # Save results (exclude numpy arrays)
+    save_results = {k: v for k, v in results.items()
+                    if k not in ("L_D", "L_R_mm", "L_R_base")}
+    with open(args.output, "w") as f:
+        json.dump(save_results, f, indent=2, default=str)
+    print(f"\nResults saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()