Two-tower retrieval

Summary

A two-tower model independently encodes queries and items, then scores their vector interaction. Independent item encoding permits reuse until its inputs or representation contract change. Index approximation, candidate eligibility, negative semantics, and model-version compatibility remain separate correctness obligations.

Architecture

query  → query_tower  → q ∈ R^d
item   → item_tower   → i ∈ R^d
score  = q · i  (or cosine)

The towers can share weights; normalization and vector dimension are design choices, not definitions. For nonzero unit vectors, dot product equals cosine; otherwise magnitude changes dot-product ranking. Stored item features must not require query-specific joint computation under this architecture.

Training

With query matrix $Q\in {\mathbb{R}}^{B\times d}$ and item matrix $I\in {\mathbb{R}}^{M\times d}$, scores are $Q{I}^T/\tau$. Require positive temperature and declare acceptable positives and eligible candidates per row.

In-batch sampled softmax

The one-positive, temperature-one case uses:

For $q=(1,0)$, positive A $=(1,0)$, and valid negative B $=(0,1)$, scores are one and zero. Positive probability is $e/(e+1)\approx 0.731059$ and loss is $\log (1+e^{-1})\approx 0.313262$. If B is also acceptable, treating it as a known negative is wrong.

Deduplicate item IDs first; this implementation maximizes total probability of the acceptable-positive set, not each positive’s individual probability. It uses Boolean masks and float64 scores with representable score gaps:

import numpy as np

def retrieval_loss(query, items, positive_mask, valid_mask, temperature=1.0):
	query = np.asarray(query, dtype=np.float64)
	items = np.asarray(items, dtype=np.float64)
	positive = np.asarray(positive_mask)
	valid = np.asarray(valid_mask)
	if query.ndim != 2 or items.ndim != 2:
		raise ValueError("Expected two embedding matrices")
	if min(query.shape) == 0 or min(items.shape) == 0:
		raise ValueError("Embedding axes must be nonempty")
	if query.shape[1] != items.shape[1]:
		raise ValueError("Embedding dimensions differ")
	shape = (len(query), len(items))
	if positive.shape != shape or valid.shape != shape:
		raise ValueError("Mask shape differs from score shape")
	if positive.dtype != np.bool_ or valid.dtype != np.bool_:
		raise ValueError("Masks must be Boolean")
	if np.any(positive & ~valid) or not np.all(positive.any(axis=1)):
		raise ValueError("Each row needs eligible positive candidates")
	if not np.isfinite(temperature) or temperature <= 0:
		raise ValueError("Temperature must be finite and positive")
	if not np.isfinite(query).all() or not np.isfinite(items).all():
		raise ValueError("Embeddings must be finite")
	with np.errstate(over="raise", invalid="raise", divide="raise"):
		scores = query @ items.T / temperature
		scores = np.where(valid, scores, -np.inf)
		relative = scores - scores.max(axis=1, keepdims=True)
		denominator = np.logaddexp.reduce(relative, axis=1)
		numerator = np.logaddexp.reduce(
			np.where(positive, relative, -np.inf), axis=1
		)
		return float(np.mean(denominator - numerator))

Both candidates acceptable gives loss zero; rejecting every positive raises instead of silently creating an undefined objective.

Need Joint Token Detail

New task: Distinguish subtle query-item phrase relations lost by independent compression.

Picture and change: Evaluate a bounded cross-encoder or late-interaction reranker with its extra work and storage.

Check: A reranker cannot recover an omitted relevant item, so inspect candidate recall separately.

Use when: Added pair detail improves the measured decision under its budget.

Change the Item Model

New task: Retrain the item encoder or tokenizer.

Picture and change: Re-encode/reindex or prove a compatible migration; version the query tower, analyzer, scoring policy, and index together.

Check: Rotating only the query from $(1,0)$ to $(0,1)$ reverses A/B scores against the unchanged index. Equal vector dimension is insufficient evidence of compatibility.

Use when: Stored representation semantics change; retain an atomic rollback bundle and enforce current eligibility/deletions after retrieval.