Training
Inspectable objectives used to train retrieval and RAG control policies.
"""Inspectable objectives used to train retrieval and RAG control policies.
This module does not train a neural network. It computes the scalar losses and
sampling diagnostics that a framework such as PyTorch would differentiate.
Keeping the mathematics in plain Python makes temperature, negative selection,
distillation, preference optimization, and policy-gradient assumptions visible
in the notebooks.
"""
import math
from dataclasses import dataclass
from typing import Iterable, List, Mapping, Optional, Sequence, Set, Tuple
@dataclass(frozen=True)
class NegativeExample:
"""A candidate negative plus provenance needed to detect false negatives."""
identifier: str
score: float
source_id: str = ""
answer_ids: Tuple[str, ...] = ()
teacher_relevance: Optional[float] = None
@dataclass(frozen=True)
class NegativeMiningReport:
"""Selected hard negatives and candidates excluded as possible positives."""
selected: Tuple[NegativeExample, ...]
excluded_false_negatives: Tuple[NegativeExample, ...]
excluded_easy: Tuple[NegativeExample, ...]
@dataclass(frozen=True)
class PolicyLoss:
"""REINFORCE loss with the returns and advantages used to obtain it."""
loss: float
returns: Tuple[float, ...]
advantages: Tuple[float, ...]
def logsumexp(values: Sequence[float]) -> float:
"""Numerically stable ``log(sum(exp(values)))``."""
if not values:
raise ValueError("logsumexp requires at least one value")
maximum = max(values)
return maximum + math.log(sum(math.exp(value - maximum) for value in values))
def softmax(logits: Sequence[float], temperature: float = 1.0) -> Tuple[float, ...]:
"""Temperature-scaled categorical probabilities."""
if temperature <= 0:
raise ValueError("temperature must be positive")
if not logits:
return ()
scaled = [value / temperature for value in logits]
normalizer = logsumexp(scaled)
return tuple(math.exp(value - normalizer) for value in scaled)
def contrastive_loss(
positive_score: float,
negative_scores: Sequence[float],
temperature: float = 1.0,
) -> float:
"""InfoNCE loss for one query, one positive, and explicit negatives."""
probabilities = softmax((positive_score, *negative_scores), temperature)
return -math.log(max(probabilities[0], 1e-300))
def in_batch_contrastive_loss(
similarities: Sequence[Sequence[float]],
positive_indices: Optional[Sequence[int]] = None,
temperature: float = 1.0,
valid_mask: Optional[Sequence[Sequence[bool]]] = None,
) -> float:
"""Mean row-wise contrastive loss with optional false-negative masking.
``similarities[i][j]`` is the query-i/document-j score. The usual paired
batch has ``positive_indices[i] == i``. ``valid_mask`` can suppress known
alternate positives from each denominator while retaining the designated
positive.
"""
if not similarities:
raise ValueError("similarity matrix must be non-empty")
width = len(similarities[0])
if width == 0 or any(len(row) != width for row in similarities):
raise ValueError("similarity matrix must be rectangular and non-empty")
positives = tuple(range(len(similarities))) if positive_indices is None else tuple(positive_indices)
if len(positives) != len(similarities):
raise ValueError("one positive index is required per query")
if valid_mask is not None and (
len(valid_mask) != len(similarities) or any(len(row) != width for row in valid_mask)
):
raise ValueError("valid mask must match the similarity matrix")
losses: List[float] = []
for row_index, (row, positive) in enumerate(zip(similarities, positives)):
if positive < 0 or positive >= width:
raise ValueError("positive index is outside the matrix")
retained = []
positive_position = -1
for column, score in enumerate(row):
valid = valid_mask is None or valid_mask[row_index][column] or column == positive
if valid:
if column == positive:
positive_position = len(retained)
retained.append(score)
probabilities = softmax(retained, temperature)
losses.append(-math.log(max(probabilities[positive_position], 1e-300)))
return sum(losses) / len(losses)
def pairwise_hinge_loss(positive_score: float, negative_score: float, margin: float = 1.0) -> float:
"""Margin ranking loss ``max(0, margin - s+ + s-)``."""
if margin < 0:
raise ValueError("margin must be non-negative")
return max(0.0, margin - positive_score + negative_score)
def pairwise_logistic_loss(positive_score: float, negative_score: float) -> float:
"""Smooth pairwise ranking loss ``softplus(s- - s+)``."""
difference = negative_score - positive_score
if difference > 0:
return difference + math.log1p(math.exp(-difference))
return math.log1p(math.exp(difference))
def listwise_cross_entropy(
student_logits: Sequence[float],
relevance: Sequence[float],
temperature: float = 1.0,
) -> float:
"""ListNet-style cross entropy between relevance and model distributions."""
if len(student_logits) != len(relevance) or not student_logits:
raise ValueError("student and relevance lists must be equally sized and non-empty")
target = softmax(relevance, temperature)
predicted = softmax(student_logits, temperature)
return -sum(expected * math.log(max(actual, 1e-300)) for expected, actual in zip(target, predicted))
def kl_distillation_loss(
student_logits: Sequence[float],
teacher_logits: Sequence[float],
temperature: float = 1.0,
) -> float:
"""Teacher-to-student KL divergence with the conventional ``T²`` scale."""
if len(student_logits) != len(teacher_logits) or not student_logits:
raise ValueError("student and teacher lists must be equally sized and non-empty")
student = softmax(student_logits, temperature)
teacher = softmax(teacher_logits, temperature)
divergence = sum(
target * (math.log(max(target, 1e-300)) - math.log(max(actual, 1e-300)))
for target, actual in zip(teacher, student)
)
return divergence * temperature * temperature
def binary_cross_entropy(logit: float, label: float) -> float:
"""Stable logistic loss for a relevance or retrieve/stop label."""
if not 0.0 <= label <= 1.0:
raise ValueError("label must be between zero and one")
return max(logit, 0.0) - logit * label + math.log1p(math.exp(-abs(logit)))
def dpo_loss(
policy_chosen_logp: float,
policy_rejected_logp: float,
reference_chosen_logp: float,
reference_rejected_logp: float,
beta: float = 0.1,
) -> float:
"""Direct Preference Optimization loss for one chosen/rejected pair."""
if beta <= 0:
raise ValueError("beta must be positive")
policy_margin = policy_chosen_logp - policy_rejected_logp
reference_margin = reference_chosen_logp - reference_rejected_logp
logit = beta * (policy_margin - reference_margin)
return binary_cross_entropy(logit, 1.0)
def discounted_returns(rewards: Sequence[float], discount: float = 1.0) -> Tuple[float, ...]:
"""Right-to-left discounted returns for a retrieval trajectory."""
if not 0.0 <= discount <= 1.0:
raise ValueError("discount must be between zero and one")
running = 0.0
output: List[float] = []
for reward in reversed(rewards):
running = reward + discount * running
output.append(running)
return tuple(reversed(output))
def reinforce_loss(
action_log_probabilities: Sequence[float],
rewards: Sequence[float],
baseline: Optional[Sequence[float]] = None,
discount: float = 1.0,
) -> PolicyLoss:
"""Monte-Carlo policy-gradient loss for search/query/stop actions."""
if len(action_log_probabilities) != len(rewards) or not rewards:
raise ValueError("one action log-probability is required per non-empty reward")
returns = discounted_returns(rewards, discount)
baselines = tuple(0.0 for _ in rewards) if baseline is None else tuple(baseline)
if len(baselines) != len(rewards):
raise ValueError("baseline must have one value per action")
advantages = tuple(value - base for value, base in zip(returns, baselines))
loss = -sum(logp * advantage for logp, advantage in zip(action_log_probabilities, advantages))
return PolicyLoss(loss=loss / len(rewards), returns=returns, advantages=advantages)
def false_negative_mask(
query_positive_sets: Sequence[Iterable[str]],
candidate_ids: Sequence[str],
) -> Tuple[Tuple[bool, ...], ...]:
"""Build a denominator mask that removes known alternate positives.
``False`` means the candidate should not act as a negative for that query.
A designated positive can still be retained by
:func:`in_batch_contrastive_loss` through its explicit positive index.
"""
masks = []
for positives in query_positive_sets:
positive_set = set(positives)
masks.append(tuple(candidate not in positive_set for candidate in candidate_ids))
return tuple(masks)
def mine_hard_negatives(
candidates: Sequence[NegativeExample],
positive_source_ids: Iterable[str] = (),
positive_answer_ids: Iterable[str] = (),
k: int = 5,
minimum_score: float = -math.inf,
teacher_positive_threshold: float = 0.5,
) -> NegativeMiningReport:
"""Choose high-scoring negatives while quarantining likely false negatives.
Source identity, answer aliases, and an optional teacher label are three
independent reasons a retrieved item may be an unlabeled positive. The
report preserves excluded items so mining quality can be audited rather
than silently trusting the training labels.
"""
if k < 0:
raise ValueError("k must be non-negative")
positive_sources: Set[str] = set(positive_source_ids)
positive_answers: Set[str] = set(positive_answer_ids)
false_negatives: List[NegativeExample] = []
eligible: List[NegativeExample] = []
easy: List[NegativeExample] = []
for candidate in candidates:
overlaps_answer = bool(set(candidate.answer_ids) & positive_answers)
teacher_positive = (
candidate.teacher_relevance is not None
and candidate.teacher_relevance >= teacher_positive_threshold
)
same_source = bool(candidate.source_id and candidate.source_id in positive_sources)
if overlaps_answer or teacher_positive or same_source:
false_negatives.append(candidate)
elif candidate.score < minimum_score:
easy.append(candidate)
else:
eligible.append(candidate)
eligible.sort(key=lambda item: (-item.score, item.identifier))
selected = eligible[:k]
easy.extend(eligible[k:])
easy.sort(key=lambda item: (-item.score, item.identifier))
false_negatives.sort(key=lambda item: (-item.score, item.identifier))
return NegativeMiningReport(tuple(selected), tuple(false_negatives), tuple(easy))