"""Independently observable multi-task objectives (PLAN sections 28 and 29). Every component is logged separately and every weight is configuration-driven. Components absent from an example (no competing paths, no edges, a non-count answer mode) report an exact zero rather than being silently folded into the total, so a flat curve for one objective is visible instead of hidden. """ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass import torch from torch import Tensor, nn from torch.nn import functional from graph.config import LossConfig COMPONENTS: tuple[str, ...] = ( "graph", "link", "next_hop", "relation", "path", "pat", "answer", "verification", "contradiction", "calibration", "count", "selection_set", "selection_rank", "node_spread", ) @dataclass(frozen=True) class LossInputs: """Tensors consumed by :class:`MultiTaskLoss`; every field is optional.""" node_embeddings: Tensor | None = None node_positive_pairs: Tensor | None = None node_negative_pairs: Tensor | None = None link_logits: Tensor | None = None link_labels: Tensor | None = None next_hop_logits: Tensor | None = None next_hop_labels: Tensor | None = None relation_logits: Tensor | None = None relation_labels: Tensor | None = None #: Path *selection* logits and their 0/1 supporting-vs-competing labels. path_scores: Tensor | None = None path_labels: Tensor | None = None #: The decision threshold the model will actually apply. Passing it lets the #: set-level objective train the threshold itself, not just the scores. selection_threshold: Tensor | None = None pat_positive: Tensor | None = None pat_negative: Tensor | None = None answer_logits: Tensor | None = None #: Multi-hot over graph nodes: a query may have several gold answers. answer_labels: Tensor | None = None verification_logits: Tensor | None = None verification_labels: Tensor | None = None contradiction_logits: Tensor | None = None contradiction_labels: Tensor | None = None count_logits: Tensor | None = None count_labels: Tensor | None = None confidence: Tensor | None = None correctness: Tensor | None = None @dataclass(frozen=True) class LossOutput: total: Tensor components: Mapping[str, Tensor] class MultiTaskLoss(nn.Module): def __init__(self, config: LossConfig) -> None: super().__init__() self.config = config self.weights: Mapping[str, float] = { "graph": config.graph, "link": config.link, "next_hop": config.next_hop, "relation": config.relation, "path": config.path, "pat": config.pat, "answer": config.answer, "verification": config.verification, "contradiction": config.contradiction, "calibration": config.calibration, "count": config.count, "selection_set": config.selection_set, "selection_rank": config.selection_rank, "node_spread": config.node_spread, } self.unsupported_weight = config.unsupported_weight missing = set(COMPONENTS) - set(self.weights) if missing: raise RuntimeError(f"loss weights missing components: {sorted(missing)}") def forward(self, inputs: LossInputs) -> LossOutput: zero = self._zero(inputs) components: dict[str, Tensor] = {} if ( inputs.node_embeddings is not None and inputs.node_positive_pairs is not None ): components["graph"] = self._contrastive( inputs.node_embeddings, inputs.node_positive_pairs, inputs.node_negative_pairs, ) if (pair := _pair(inputs.link_logits, inputs.link_labels)) is not None: logits, labels = pair components["link"] = functional.binary_cross_entropy_with_logits( logits, labels.float() ) if (pair := _pair(inputs.next_hop_logits, inputs.next_hop_labels)) is not None: logits, labels = pair components["next_hop"] = self._listwise(logits, labels) if (pair := _pair(inputs.relation_logits, inputs.relation_labels)) is not None: logits, labels = pair components["relation"] = functional.cross_entropy(logits, labels) if (pair := _pair(inputs.path_scores, inputs.path_labels)) is not None: logits, labels = pair components["path"] = functional.binary_cross_entropy_with_logits( logits, labels.float() ) if (pair := _pair(inputs.pat_positive, inputs.pat_negative)) is not None: logits, labels = pair components["pat"] = self._ranking(logits, labels) if (pair := _pair(inputs.answer_logits, inputs.answer_labels)) is not None: logits, labels = pair components["answer"] = self._multilabel(logits, labels) if ( pair := _pair(inputs.verification_logits, inputs.verification_labels) ) is not None: logits, labels = pair # Reweighted toward the rare class. Roughly 85% of curriculum examples # are answerable, and with a per-example batch the unweighted objective is # minimised well by predicting "supported" and treating abstention as # noise -- which is what the measured behaviour looked like. The weight # applies to the *negative* class, so it is expressed as a weight on # (1 - label) rather than pos_weight. weight = 1.0 + (self.unsupported_weight - 1.0) * (1.0 - labels.float()) components["verification"] = functional.binary_cross_entropy_with_logits( logits, labels.float(), weight=weight ) if ( pair := _pair(inputs.contradiction_logits, inputs.contradiction_labels) ) is not None: logits, labels = pair components["contradiction"] = functional.cross_entropy(logits, labels) if (pair := _pair(inputs.count_logits, inputs.count_labels)) is not None: logits, labels = pair components["count"] = functional.cross_entropy(logits, labels) if ( pair := _pair(inputs.path_scores, inputs.path_labels) ) is not None and inputs.selection_threshold is not None: logits, labels = pair components["selection_set"] = self._soft_set_f1_loss( logits, labels, inputs.selection_threshold ) if (pair := _pair(inputs.path_scores, inputs.path_labels)) is not None: logits, labels = pair components["selection_rank"] = self._listwise_rank_loss(logits, labels) if ( self.weights["node_spread"] and inputs.node_embeddings is not None and inputs.node_embeddings.shape[0] > 1 ): components["node_spread"] = self._node_spread(inputs.node_embeddings) if (pair := _pair(inputs.confidence, inputs.correctness)) is not None: logits, labels = pair components["calibration"] = ( (logits.sigmoid() - labels.float()) ** 2 ).mean() total = zero for name, value in components.items(): total = total + value * self.weights[name] for name in COMPONENTS: components.setdefault(name, zero.detach()) return LossOutput(total, components) # -- component implementations ------------------------------------------ @staticmethod def _contrastive( embeddings: Tensor, positives: Tensor, negatives: Tensor | None ) -> Tensor: left = functional.normalize(embeddings[positives[:, 0]], dim=-1) right = functional.normalize(embeddings[positives[:, 1]], dim=-1) positive_similarity = (left * right).sum(-1) loss = (1 - positive_similarity).mean() if negatives is not None and negatives.numel(): negative_left = functional.normalize(embeddings[negatives[:, 0]], dim=-1) negative_right = functional.normalize(embeddings[negatives[:, 1]], dim=-1) negative_similarity = (negative_left * negative_right).sum(-1) span = min(positive_similarity.shape[0], negative_similarity.shape[0]) loss = ( loss + functional.relu( 0.2 - positive_similarity[:span] + negative_similarity[:span] ).mean() ) return loss @staticmethod def _node_spread(embeddings: Tensor) -> Tensor: """Positive part of the mean pairwise cosine between node states. Zero while nodes point in unrelated directions and rising toward 1 as they collapse onto one. Computed in O(N) from the norm of the summed unit vectors: ``sum_ij u_i.u_j = |sum_i u_i|^2``. """ unit = functional.normalize(embeddings.float(), dim=-1) count = unit.shape[0] mean_cosine = (unit.sum(0).pow(2).sum() - count) / (count * (count - 1)) return functional.relu(mean_cosine) @staticmethod def _listwise(logits: Tensor, labels: Tensor) -> Tensor: """Cross-entropy over one candidate list with a 0/1 label vector.""" if labels.sum() <= 0: return logits.sum() * 0.0 target = labels.float() / labels.float().sum() return -(target * functional.log_softmax(logits, dim=-1)).sum() @staticmethod def _multilabel(logits: Tensor, labels: Tensor) -> Tensor: """Balanced BCE over all graph nodes for multi-answer supervision. Gold answers are a handful of nodes among hundreds, so unweighted BCE is dominated by the negative class and collapses to predicting "no answer". Positives are reweighted by the class ratio to counter that. """ flat_logits = logits.reshape(-1) flat_labels = labels.reshape(-1).float() positives = flat_labels.sum() if positives <= 0: return functional.binary_cross_entropy_with_logits(flat_logits, flat_labels) negatives = flat_labels.numel() - positives weight = (negatives / positives).clamp(1.0, 100.0) return functional.binary_cross_entropy_with_logits( flat_logits, flat_labels, pos_weight=weight ) @staticmethod def soft_selection( logits: Tensor, threshold: Tensor, temperature: float = 4.0 ) -> Tensor: """Differentiable stand-in for ``logits > threshold``.""" return torch.sigmoid((logits - threshold.reshape(())) * temperature) @classmethod def _soft_set_f1_loss( cls, logits: Tensor, labels: Tensor, threshold: Tensor ) -> Tensor: """``1 - soft F1`` of the thresholded selection against the gold set. Per-path BCE tells the model which paths are good but says nothing about *how many* to keep, and it never touches the threshold at all. The measured consequences were a threshold that drifted without becoming discriminative (selection precision flat near 0.20 while path recall collapsed from 0.761 to 0.368), greedy search beating a wider beam, and a regression on real KG queries where every candidate is correct. A set-level objective fixes the mismatch: it is differentiable in both the scores and the threshold, and it is minimised by committing to the right *number* of paths. """ keep = cls.soft_selection(logits, threshold) gold = labels.float() true_positive = (keep * gold).sum() false_positive = (keep * (1.0 - gold)).sum() false_negative = ((1.0 - keep) * gold).sum() denominator = 2.0 * true_positive + false_positive + false_negative if float(denominator.detach()) <= 0.0: return logits.sum() * 0.0 return 1.0 - (2.0 * true_positive) / denominator.clamp_min(1e-6) @staticmethod def _listwise_rank_loss(logits: Tensor, labels: Tensor) -> Tensor: """Multi-positive listwise cross-entropy over one query's candidate set. With ``count_guided_selection`` the committed set is the top-k candidates by score, where k comes from the count head. Nothing in the thresholded objective optimises *order*: a model can place every gold path just above the threshold alongside many non-gold ones and pay no penalty, which is consistent with the measured 0.217 path recall at a corrected set size. This term optimises exactly what top-k reads. It normalises over the candidate set, so it is scale-free in the logits and complements rather than duplicates the set-level term, which fixes the *size* of the commit. """ gold = labels.reshape(-1) > 0.5 if logits.numel() < 2 or not bool(gold.any()) or bool(gold.all()): # No contrast to learn from: a single candidate, or an all-gold set as # real KG queries produce, where every candidate is a correct answer. return logits.sum() * 0.0 log_probabilities = torch.log_softmax(logits.reshape(-1), dim=-1) return -log_probabilities[gold].mean() @staticmethod def _ranking(positive: Tensor, negative: Tensor) -> Tensor: """Pairwise softplus ranking loss over every positive/negative pair.""" if positive.ndim == 0: positive = positive.reshape(1) if negative.ndim == 0: negative = negative.reshape(1) return functional.softplus(-(positive[:, None] - negative[None, :])).mean() @staticmethod def _zero(inputs: LossInputs) -> Tensor: for value in inputs.__dict__.values(): if isinstance(value, Tensor) and value.numel(): return value.sum() * 0.0 return torch.tensor(0.0) def _pair(first: Tensor | None, second: Tensor | None) -> tuple[Tensor, Tensor] | None: """Return both tensors when both are present and non-empty. Returning the pair rather than a bool lets the type checker narrow away the ``None`` branch at each call site, which a boolean predicate cannot do. """ if first is None or second is None or first.numel() == 0 or second.numel() == 0: return None return first, second