| """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_scores: Tensor | None = None |
| path_labels: Tensor | None = None |
| |
| |
| selection_threshold: Tensor | None = None |
| pat_positive: Tensor | None = None |
| pat_negative: Tensor | None = None |
| answer_logits: Tensor | None = None |
| |
| 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 |
| |
| |
| |
| |
| |
| |
| 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) |
|
|
| |
|
|
| @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()): |
| |
| |
| 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 |
|
|
