Two-Tower Retrieval

"""Two-tower retrieval with in-batch negatives - the objective, written out.

Every OTHER item in the batch acts as a negative for each query, so one batch
of B examples yields B*B training signals. That is the trick that makes
contrastive retrieval trainable.
"""

import math


def encode(features, weights):
    """A tower: one linear layer, then L2-normalize.

    Normalizing makes the dot product a cosine, which bounds similarity to
    [-1, 1] and is what most ANN indices assume.
    """
    out = []
    for row in weights:
        out.append(sum(f * w for f, w in zip(features, row)))

    norm = math.sqrt(sum(v * v for v in out)) or 1.0
    return [v / norm for v in out]


def in_batch_loss(query_features, item_features, q_weights, d_weights, temperature=0.05):
    """InfoNCE over the batch. The diagonal is the correct pairing."""
    batch = len(query_features)

    queries = [encode(q, q_weights) for q in query_features]
    items = [encode(d, d_weights) for d in item_features]

    total = 0.0
    for i in range(batch):
        # Similarity of query i against EVERY item in the batch.
        scores = []
        for j in range(batch):
            dot = sum(a * b for a, b in zip(queries[i], items[j]))
            scores.append(dot / temperature)

        # Softmax cross-entropy with the true label being i itself.
        largest = max(scores)
        denominator = sum(math.exp(s - largest) for s in scores)
        total -= (scores[i] - largest) - math.log(denominator)

    return total / batch

"""Two-tower retrieval - typed, with the diagnostics the loss will not give."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class TowerConfig:
    input_dim: int
    embedding_dim: int = 128
    hidden_dim: int = 512
    temperature: float = 0.05

    def __post_init__(self) -> None:
        if self.temperature <= 0:
            raise ValueError(f"temperature must be positive, got {self.temperature}")


class Tower(nn.Module):
    """One encoder. Query and item towers are SEPARATE instances - sharing
    weights only helps when data is scarce, and otherwise forces two different
    input distributions through one set of parameters."""

    def __init__(self, config: TowerConfig) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.input_dim, config.hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(config.hidden_dim, config.embedding_dim),
        )

    def forward(self, x: Tensor) -> Tensor:
        # L2-normalize so the dot product IS cosine similarity. Most ANN
        # indices assume this, and skipping it makes distances meaningless.
        return nn.functional.normalize(self.net(x), p=2, dim=-1)


@dataclass(frozen=True)
class ContrastiveStep:
    loss: Tensor
    accuracy: float
    embedding_std: float

    @property
    def has_collapsed(self) -> bool:
        """Representation collapse is invisible in the loss and fatal to
        retrieval - all embeddings drift together and the loss still looks
        fine. Embedding spread is the only signal that catches it."""
        return self.embedding_std < 0.01


def contrastive_step(
    query_tower: Tower,
    item_tower: Tower,
    query_features: Tensor,
    item_features: Tensor,
    temperature: float,
) -> ContrastiveStep:
    if query_features.size(0) != item_features.size(0):
        raise ValueError("queries and items must be paired one-to-one")
    if query_features.size(0) < 2:
        raise ValueError("in-batch negatives need a batch of at least 2")

    queries = query_tower(query_features)
    items = item_tower(item_features)

    # The whole B x B similarity matrix in one product. The diagonal is the
    # positive pair; everything off-diagonal is a negative, for free.
    similarity = (queries @ items.T) / temperature
    labels = torch.arange(similarity.size(0), device=similarity.device)

    loss = nn.functional.cross_entropy(similarity, labels)

    with torch.no_grad():
        accuracy = float((similarity.argmax(dim=1) == labels).float().mean())
        embedding_std = float(items.std(dim=0).mean())

    return ContrastiveStep(loss=loss, accuracy=accuracy, embedding_std=embedding_std)

"""Retrieval at scale - cross-device negatives, hard mining, quantized index."""

import numpy as np
import torch
import torch.distributed as dist
from numpy.typing import NDArray
from torch import Tensor


def gather_across_devices(embeddings: Tensor) -> Tensor:
    """Collect embeddings from every rank to widen the negative pool.

    Batch size IS the objective here - more in-batch negatives is a harder and
    more informative task. Gathering across N devices multiplies the effective
    negatives by N without any device holding a larger batch.
    """
    if not dist.is_initialized():
        return embeddings

    gathered = [torch.zeros_like(embeddings) for _ in range(dist.get_world_size())]
    dist.all_gather(gathered, embeddings)
    # The local slice must keep its graph connection - all_gather detaches,
    # so splicing it back in is what preserves the local gradient path.
    gathered[dist.get_rank()] = embeddings
    return torch.cat(gathered, dim=0)


def mine_hard_negatives(
    query_embeddings: Tensor,
    item_embeddings: Tensor,
    positive_indices: Tensor,
    k: int = 8,
) -> Tensor:
    """Retrieve the model's own top mistakes as negatives.

    Random negatives plateau early - separating a query from an unrelated item
    is trivial. Hard negatives are where the remaining signal is, and mining
    them from the model's CURRENT retrievals is what keeps the task hard as it
    improves.
    """
    with torch.no_grad():
        scores = query_embeddings @ item_embeddings.T
        # Mask the true positive so it cannot be mined as its own negative.
        scores.scatter_(1, positive_indices.unsqueeze(1), float("-inf"))
        return scores.topk(k, dim=1).indices


def build_quantized_index(embeddings: NDArray[np.float32], n_centroids: int = 4096):
    """IVF-PQ index: coarse partition plus product quantization.

    Flat search is O(N·d) per query and does not hold a catalogue of millions
    in memory. IVF probes only a few partitions, and PQ compresses each vector
    from d floats to a handful of bytes - together roughly 100x smaller and
    sublinear to search, at a measurable and TUNABLE recall cost.
    """
    import faiss

    dim = embeddings.shape[1]
    # Contiguous float32 is what FAISS expects; anything else is copied.
    embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
    faiss.normalize_L2(embeddings)

    quantizer = faiss.IndexFlatIP(dim)
    index = faiss.IndexIVFPQ(quantizer, dim, n_centroids, dim // 4, 8)

    index.train(embeddings)
    index.add(embeddings)
    index.nprobe = 32      # partitions probed per query - the recall/latency dial

    return index