From 2ba84ea77b7e15326771bc00ed588009874c865d Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Mon, 17 Aug 2026 13:22:39 -0400 Subject: [PATCH 1/3] Extract and document training losses --- src/langset/bridge_emission.py | 28 +- src/langset/loss.py | 648 +++++++++++++++++++++++++++++++++ src/langset/strategies.py | 156 ++++---- src/langset/trainer.py | 244 +++++-------- tests/test_replay.py | 21 +- 5 files changed, 845 insertions(+), 252 deletions(-) create mode 100644 src/langset/loss.py diff --git a/src/langset/bridge_emission.py b/src/langset/bridge_emission.py index b603d9e..541bebf 100644 --- a/src/langset/bridge_emission.py +++ b/src/langset/bridge_emission.py @@ -28,6 +28,7 @@ import torch.nn as nn import torch.nn.functional as F +from langset.loss import BridgeLossContext, bridge_loss from langset.strategies import EmissionOut, _EmissionObjective, _TargetSource if TYPE_CHECKING: @@ -196,18 +197,23 @@ def emit( vlab[r, p] = 1.0 off += mi - if matched_pred: - mp = torch.stack(matched_pred) # [K, d] - pos_t = torch.tensor(pos, device=dev) - nce = F.cross_entropy( - mp @ bank.T / self.temp, pos_t - ) # each emission retrieves its own target - else: - nce = target_lat.new_zeros(()) - vloss = F.binary_cross_entropy_with_logits( - vlog, vlab, pos_weight=torch.tensor(self.pos_weight, device=dev) + matched_predictions = torch.stack(matched_pred) if matched_pred else vecs.new_empty((0, d)) + positive_indices = torch.tensor(pos, device=dev) + losses = bridge_loss( + BridgeLossContext( + matched_predictions=matched_predictions, + target_bank=bank, + positive_indices=positive_indices, + validity_logits=vlog, + validity_labels=vlab, + temperature=self.temp, + validity_weight=self.lam_valid, + positive_weight=self.pos_weight, + ) ) - base = nce + self.lam_valid * vloss + nce = losses.info_nce.to_tensor() + vloss = losses.validity.to_tensor() + base = losses.to_tensor() zero = target_lat.new_zeros(()) return EmissionOut( recon=recon, diff --git a/src/langset/loss.py b/src/langset/loss.py new file mode 100644 index 0000000..6b84d72 --- /dev/null +++ b/src/langset/loss.py @@ -0,0 +1,648 @@ +"""Loss functions.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, TypeVar + +import torch +import torch.nn.functional as F +from torch.nn import Linear + +if TYPE_CHECKING: + from langset.modeling import LangSetModel + from langset.training_args import TrainingArguments + + +class Loss(Protocol): + """A loss result that exposes a scalar optimization objective.""" + + def to_tensor(self) -> torch.Tensor: + """Return the scalar objective to weight, combine, and backpropagate.""" + ... + + +@dataclass +class ScalarLoss: + """A leaf loss containing one scalar optimization objective.""" + + value: torch.Tensor + """Scalar objective tensor used directly for optimization.""" + + def to_tensor(self) -> torch.Tensor: + """Return this loss's scalar objective.""" + return self.value + + +ContextT = TypeVar("ContextT", contravariant=True) + +LossFn = Callable[[ContextT], Loss] + + +@dataclass +class TrainingContext: + """Dependencies shared by training losses.""" + + model: LangSetModel + """The model being trained.""" + + args: TrainingArguments + """Training configuration, including objective weights and temperatures.""" + + +@dataclass +class ReconstructionLossContext(TrainingContext): + """Inputs to :func:`recon_loss`. + + ``latent`` contains one latent vector per selected row. ``tr_ids`` and + ``tr_mask`` contain tokenized target text for every dataset row; ``rows`` + selects the targets corresponding to ``latent``. + """ + + latent: torch.Tensor + """Latents to decode, shaped ``[batch_size, latent_dim]``.""" + + rows: torch.Tensor + """Dataset-row indices selecting the target tokens for this batch.""" + + tr_ids: torch.Tensor + """Token IDs for every row's reconstruction target text.""" + + tr_mask: torch.Tensor + """Attention mask aligned with ``tr_ids``; padding tokens are ignored.""" + + connector: Linear + """Projects each latent into a sequence of synthetic prefix-token embeddings.""" + + +def recon_loss(ctx: ReconstructionLossContext) -> Loss: + """Decode each latent into its target text using next-token cross-entropy. + + The connector turns a latent into synthetic prefix tokens. The model then + predicts each following target token. Cross-entropy penalizes the model when + it assigns little probability to the actual next token, while ignoring + padding. For example, a latent for ``"The capital of France is Paris"`` is + trained to help predict the tokens in that target text. + + This auxiliary objective grounds the latent in the target's textual content. + """ + + target_ids, target_mask = ctx.tr_ids[ctx.rows], ctx.tr_mask[ctx.rows] + """training targets to predict as tokens and their attention mask""" + recon_k = ctx.connector.out_features // ctx.model.h + """number of synthetic tokens to generate before target text""" + target_embeddings = ctx.model.embed(target_ids) + """Teacher-forced embeddings of the real target tokens appended after the prefix.""" + synthetic_tokens = ( + ctx.connector(ctx.latent) + .view(ctx.latent.size(0), recon_k, ctx.model.h) + .to(target_embeddings.dtype) + ) + """Synthetic prefix embeddings generated from the latent. Prepending them lets target + token predictions depend on the latent, so reconstruction gradients update it.""" + sequence = torch.cat([synthetic_tokens, target_embeddings], dim=1) + """Combined sequence of synthetic prefix and target text embeddings""" + attention_mask = torch.cat( + [ + torch.ones( + ctx.latent.size(0), recon_k, device=target_mask.device, dtype=target_mask.dtype + ), + target_mask, + ], + dim=1, + ) + """Attention mask combining synthetic prefix and target text attention""" + backbone_output = ctx.model._run_backbone(sequence, attention_mask, target_ids, recon_k) + # Shift by one: the final prefix position predicts target_ids[:, 0]. + target_slice = slice(recon_k - 1, recon_k - 1 + target_ids.size(1)) + prediction_logits: torch.Tensor + """Vocabulary logits for target tokens, supplied by the LM head or derived from hidden states.""" + + if (logits := getattr(backbone_output, "logits", None)) is not None: + prediction_logits = logits[:, target_slice, :].float() + else: + target_hidden = ctx.model._last_hidden(backbone_output)[:, target_slice, :] + prediction_logits = F.linear(target_hidden.float(), ctx.model.embed.weight.float()) + return ScalarLoss( + F.cross_entropy( + prediction_logits.reshape(-1, ctx.model.vocab_size), + target_ids.masked_fill(target_mask == 0, -100).reshape(-1), + ignore_index=-100, + ) + ) + + +recon_loss_fn: LossFn[ReconstructionLossContext] = recon_loss +"""The reconstruction loss function.""" + + +@dataclass +class LearnLossContext(TrainingContext): + """Inputs to :func:`learn_loss` for a batch from the replay pool.""" + + pos: torch.Tensor + """Indices selecting replay examples from the tokenized replay pool.""" + + ln_doc_ids: torch.Tensor + """Left-padded document token IDs for every replay-pool row.""" + + ln_doc_mask: torch.Tensor + """Attention mask aligned with ``ln_doc_ids``.""" + + ln_tgt_ids: torch.Tensor + """Right-padded target token IDs for every replay-pool row.""" + + ln_tgt_mask: torch.Tensor + """Attention mask aligned with ``ln_tgt_ids``; padding tokens are ignored.""" + + +def learn_loss(ctx: LearnLossContext) -> Loss: + """Rehearse document-to-target knowledge with teacher-forced cross-entropy. + + This is standard next-token training: after reading a document, the model + predicts its paired target text and is penalized when it assigns low + probability to the actual next target token. For example, a replay pair of + document ``"France's capital is"`` and target ``"Paris"`` trains that + association even when the row is not in the embedding objective. + """ + + document_ids, document_mask = ctx.ln_doc_ids[ctx.pos], ctx.ln_doc_mask[ctx.pos] + """Left-padded replay documents for this batch and their attention mask.""" + target_ids, target_mask = ctx.ln_tgt_ids[ctx.pos], ctx.ln_tgt_mask[ctx.pos] + """Right-padded replay targets for this batch and their attention mask.""" + sequence = torch.cat([document_ids, target_ids], dim=1) + """Teacher-forced document followed by its target tokens.""" + attention_mask = torch.cat([document_mask, target_mask], dim=1) + """Attention mask aligned with ``sequence``; target padding is excluded.""" + hidden_states = ctx.model._last_hidden( + ctx.model._run_backbone(ctx.model.embed(sequence), attention_mask, sequence, 0) + ) # all real tokens -> real_start=0 + """Final hidden state for every document and target token in ``sequence``.""" + document_length = document_ids.size(1) + """Fixed left-padded document width, which is the target's starting offset.""" + target_hidden = hidden_states[ + :, document_length - 1 : document_length - 1 + target_ids.size(1), : + ] + """Shifted hidden states whose outputs predict the corresponding target tokens.""" + prediction_logits = F.linear(target_hidden.float(), ctx.model.embed.weight.float()) + """Vocabulary logits derived from the target-prediction hidden states.""" + return ScalarLoss( + F.cross_entropy( + prediction_logits.reshape(-1, ctx.model.vocab_size), + target_ids.masked_fill(target_mask == 0, -100).reshape(-1), + ignore_index=-100, + ) + ) + + +learn_loss_fn: LossFn[LearnLossContext] = learn_loss +"""The learn loss function.""" + + +@dataclass +class SLContext(TrainingContext): + """Inputs to the single-latent contrastive objective. + + Each row in ``pred`` should be close to the corresponding row in ``target`` + and far from the other target embeddings in the batch. This is not a + supervised-classification loss: no field here is a class label or a token + vocabulary logit vector. + """ + + pred: torch.Tensor + """Input-text embeddings being trained, shaped ``[batch_size, embedding_dim]``.""" + + target: torch.Tensor + """Matching target-text embeddings, shaped ``[batch_size, embedding_dim]``.""" + + hn: torch.Tensor | None + """Optional mined hard-negative target embeddings, one per batch row.""" + + idx: torch.Tensor + """Dataset-row indices corresponding to the batch embeddings.""" + + mask_keys: list[frozenset[str]] | None + """Per-row dataset metadata used to exclude same-key false negatives.""" + + hard_neg_text: list[str] | None + """Per-row mined hard-negative text; an empty string marks no valid negative.""" + + +def sl_loss( + ctx: SLContext, +) -> Loss: + """Align paired input and target embeddings with an InfoNCE-style objective. + + Each row of ``pred`` is an input embedding (the anchor). Its target at the + same row of ``target`` is the positive; other batch targets are negatives. + Their temperature-scaled dot products are similarity logits—not vocabulary + logits—and cross-entropy selects the same-row target as the correct + candidate. This trains each input to be more similar to its paired target + than to the permitted alternatives. + + Mined hard negatives can add one extra candidate per anchor. Matching + ``mask_keys`` exclude known false negatives from that anchor's denominator. + This function prepares those single-latent candidates, positive indices, and + exclusion mask, then delegates the primary contrastive calculation to + :func:`info_nce_loss`. + + The contrastive term has an implicit weight of ``1.0``. When + ``args.lam_uniform > 0`` and the batch contains multiple rows, the returned + loss also includes ``lam_uniform * uniformity_term``. This auxiliary term + spreads normalized input embeddings across the unit sphere. + """ + + target_matrix = ctx.target if ctx.hn is None else torch.cat([ctx.target, ctx.hn], dim=0) + """Candidate target embeddings: batch positives plus optional mined hard negatives.""" + batch_size = len(ctx.idx) + """Number of input/positive-target pairs in the current batch.""" + neg_mask = torch.zeros( + batch_size, target_matrix.size(0), dtype=torch.bool, device=ctx.pred.device + ) + """Mask marking candidate columns to exclude from each contrastive denominator.""" + if ctx.mask_keys is not None: # in-batch block: drop same-issue false negatives + batch_mask_keys = [ctx.mask_keys[j] for j in ctx.idx.tolist()] + """False-negative exclusion keys corresponding to the rows in this batch.""" + for row in range(batch_size): + if not batch_mask_keys[row]: + continue + for column in range(batch_size): + if row != column and (batch_mask_keys[row] & batch_mask_keys[column]): + neg_mask[row, column] = True + if ( + ctx.hn is not None + ): # PER-ANCHOR-ONLY hard neg: anchor i sees ONLY its own mined hard neg (col B+i) + assert ctx.hard_neg_text is not None + valid_hard_negatives = [bool(ctx.hard_neg_text[j]) for j in ctx.idx.tolist()] + """Whether each batch row has a non-empty mined hard-negative text value.""" + for row in range(batch_size): + for column in range(batch_size): + if not (row == column and valid_hard_negatives[column]): + neg_mask[row, batch_size + column] = True + positive_indices = torch.arange(batch_size, device=ctx.pred.device) + """Diagonal candidate indices: each input's paired target is its positive.""" + contrastive = info_nce_loss( + InfoNCELossContext( + anchors=ctx.pred, + candidates=target_matrix, + positive_indices=positive_indices, + temperature=ctx.args.tau, + logit_mask=neg_mask if bool(neg_mask.any()) else None, + ) + ) + """Primary contrastive loss delegated to the reusable InfoNCE kernel.""" + objective = contrastive.to_tensor() + """Scalar contrastive objective before the optional uniformity term.""" + if ctx.args.lam_uniform > 0 and batch_size > 1: # aux: uniformity + squared_distances = torch.pdist(F.normalize(ctx.pred, p=2, dim=-1), p=2).pow(2) + """Pairwise squared distances between normalized prediction embeddings.""" + objective = ( + objective + ctx.args.lam_uniform * squared_distances.mul(-2.0).exp().mean().log() + ) + return ScalarLoss(objective) + + +sl_loss_fn: LossFn[SLContext] = sl_loss +"""The single-latent contrastive loss function.""" + + +# ---- multi-latent loss kernels ------------------------------------------------------------------ +@dataclass +class InfoNCELossContext: + """Inputs to :func:`info_nce_loss`. + + The candidate at ``positive_indices[row]`` is the positive for + ``anchors[row]``. All other unmasked candidates are negatives. + """ + + anchors: torch.Tensor + """Anchor embeddings, shaped ``[anchor_count, embedding_dim]``.""" + + candidates: torch.Tensor + """Positive and negative candidate embeddings, shaped ``[candidate_count, embedding_dim]``.""" + + positive_indices: torch.Tensor + """Candidate-column index of each anchor's positive, shaped ``[anchor_count]``.""" + + temperature: float + """Positive divisor applied to similarity logits before cross-entropy.""" + + logit_mask: torch.Tensor | None = None + """Optional boolean ``[anchor_count, candidate_count]`` mask; ``True`` excludes a candidate.""" + + normalize_embeddings: bool = False + """Whether to L2-normalize anchors and candidates before calculating similarities.""" + + +def info_nce_loss(ctx: InfoNCELossContext) -> Loss: + """Compute an InfoNCE-style contrastive cross-entropy loss. + + Every anchor competes against the same candidate matrix. The positive index + selects its matching candidate, while an optional mask removes false + negatives or candidates that are not valid for that anchor. Set + ``normalize_embeddings`` for cosine-similarity InfoNCE; leave it off when + callers already provide normalized vectors or intentionally use dot product. + """ + + anchors = F.normalize(ctx.anchors, dim=-1) if ctx.normalize_embeddings else ctx.anchors + """Anchor vectors, optionally normalized to unit length.""" + candidates = F.normalize(ctx.candidates, dim=-1) if ctx.normalize_embeddings else ctx.candidates + """Candidate vectors, optionally normalized to unit length.""" + logits = anchors @ candidates.t() / ctx.temperature + """Temperature-scaled embedding similarity logits, one row per anchor.""" + if ctx.logit_mask is not None: + logits = logits.masked_fill(ctx.logit_mask, float("-inf")) + return ScalarLoss(F.cross_entropy(logits, ctx.positive_indices)) + + +info_nce_loss_fn: LossFn[InfoNCELossContext] = info_nce_loss +"""The general InfoNCE-style contrastive loss function.""" + + +@dataclass +class SoftTargetCrossEntropyContext: + """Inputs to :func:`soft_target_cross_entropy_loss`.""" + + logits: torch.Tensor + """Unnormalized class scores, shaped ``[..., class_count]``.""" + + target_probabilities: torch.Tensor + """Target probability distribution aligned with ``logits``.""" + + selection: torch.Tensor + """Boolean mask over the leading dimensions selecting supervised positions.""" + + +def soft_target_cross_entropy_loss(ctx: SoftTargetCrossEntropyContext) -> Loss: + """Compute mean cross-entropy against probability targets at selected positions. + + This is used by multi-latent code, concept, and state objectives. Unlike + ordinary classification CE, each target can spread probability mass across + several classes. It returns scalar zero when no positions are selected. + """ + + log_probabilities = F.log_softmax(ctx.logits, dim=-1) + """Log-probability for every class at every candidate position.""" + per_position_loss = -(ctx.target_probabilities * log_probabilities).sum(dim=-1) + """Soft-target cross-entropy for each position before masking and reduction.""" + objective = ( + per_position_loss[ctx.selection].mean() + if bool(ctx.selection.any()) + else log_probabilities.new_zeros(()) + ) + """Mean selected soft-target CE, or scalar zero when nothing is supervised.""" + return ScalarLoss(objective) + + +soft_target_cross_entropy_loss_fn: LossFn[SoftTargetCrossEntropyContext] = ( + soft_target_cross_entropy_loss +) +"""The masked soft-target cross-entropy loss function.""" + + +@dataclass +class StopLossContext: + """Inputs to :func:`stop_loss` for an autoregressive multi-latent emission.""" + + logits: torch.Tensor + """One stop logit per row and emission position, shaped ``[batch_size, max_items + 1]``.""" + + lengths: list[int] + """Number of real emitted items in each batch row; the next position is the stop target.""" + + +def stop_loss(ctx: StopLossContext) -> Loss: + """Train an independent sigmoid stop decision after each possible emission. + + Each row has label ``0`` through its final real item and label ``1`` at the + following position. Positions after that stop decision are excluded, so + padding never contributes to BCE. + """ + + stop_logits = ctx.logits.squeeze(-1).float() + """Stop logits normalized to ``[batch_size, max_items + 1]``.""" + labels = torch.zeros_like(stop_logits) + """Binary targets with a single stop position in each batch row.""" + selection = torch.zeros_like(stop_logits, dtype=torch.bool) + """Mask selecting real continuation decisions and the one stop decision per row.""" + for row, length in enumerate(ctx.lengths): + labels[row, length] = 1.0 + selection[row, : length + 1] = True + return ScalarLoss(F.binary_cross_entropy_with_logits(stop_logits[selection], labels[selection])) + + +stop_loss_fn: LossFn[StopLossContext] = stop_loss +"""The autoregressive emission stop loss function.""" + + +@dataclass +class SupervisedContrastiveLossContext: + """Inputs to :func:`supervised_contrastive_loss`.""" + + embeddings: torch.Tensor + """Emitted embeddings, shaped ``[item_count, embedding_dim]``.""" + + labels: list[str] + """Group labels aligned with embeddings; empty and unknown labels are excluded.""" + + temperature: float + """Positive divisor applied to pairwise cosine-similarity logits.""" + + +def supervised_contrastive_loss(ctx: SupervisedContrastiveLossContext) -> Loss: + """Pull same-label embeddings together while pushing differently labeled ones apart. + + This is supervised contrastive learning (SupCon), not classification: labels + define which other embeddings are positives rather than indexing a classifier + output. Unlabeled items are excluded. The loss is zero when fewer than two + labeled examples exist or no anchor has a same-label positive. + """ + + keep = [ + index + for index, label in enumerate(ctx.labels) + if str(label).strip().lower() not in ("", "unknown", "none", "nan") + ] + """Embedding rows with usable supervision labels.""" + if len(keep) < 2: + return ScalarLoss(ctx.embeddings.new_zeros(())) + embeddings = F.normalize(ctx.embeddings[keep], p=2, dim=-1) + """Unit-length embeddings participating in pairwise cosine similarity.""" + labels = [ctx.labels[index] for index in keep] + """Usable labels aligned with the selected embeddings.""" + item_count = len(keep) + """Number of labeled embeddings participating in the objective.""" + similarities = (embeddings @ embeddings.t() / ctx.temperature).masked_fill( + torch.eye(item_count, device=ctx.embeddings.device, dtype=torch.bool), -1e9 + ) + """Pairwise similarity logits with self-comparisons excluded from every denominator.""" + log_probabilities = similarities - torch.logsumexp(similarities, dim=1, keepdim=True) + """Log probability assigned to each other embedding by every anchor.""" + positives = torch.tensor( + [ + [ + 1.0 if row != column and labels[row] == labels[column] else 0.0 + for column in range(item_count) + ] + for row in range(item_count) + ], + device=ctx.embeddings.device, + ) + """Indicator matrix identifying same-label, non-self positive pairs.""" + positive_counts = positives.sum(dim=1) + """Number of positive embeddings available to each anchor.""" + valid_anchors = positive_counts > 0 + """Anchors with at least one same-label positive.""" + if not bool(valid_anchors.any()): + return ScalarLoss(ctx.embeddings.new_zeros(())) + objective = ( + -(positives * log_probabilities).sum(dim=1)[valid_anchors] / positive_counts[valid_anchors] + ).mean() + """Mean SupCon objective across anchors that have at least one positive.""" + return ScalarLoss(objective) + + +supervised_contrastive_loss_fn: LossFn[SupervisedContrastiveLossContext] = ( + supervised_contrastive_loss +) +"""The supervised contrastive loss function.""" + + +@dataclass +class CoTLossContext(TrainingContext): + """Inputs to :func:`cot_loss` after the strategy tokenizes seed and reasoning text.""" + + seed_ids: torch.Tensor + """Left-padded seed token IDs, shaped ``[batch_size, seed_length]``.""" + + seed_mask: torch.Tensor + """Attention mask aligned with ``seed_ids``.""" + + cot_ids: torch.Tensor + """Right-padded chain-of-thought target token IDs.""" + + cot_mask: torch.Tensor + """Attention mask aligned with ``cot_ids``; padding is ignored by cross-entropy.""" + + +def cot_loss(ctx: CoTLossContext) -> Loss: + """Generate chain-of-thought text from a teacher-forced seed using token CE. + + The left-padded seed's final real token predicts the first CoT token. CoT + tokens are right-padded, so no padding appears between seed and target text + and padded targets are ignored by cross-entropy. + """ + + sequence = torch.cat([ctx.seed_ids, ctx.cot_ids], dim=1) + """Teacher-forced seed followed by chain-of-thought target tokens.""" + attention_mask = torch.cat([ctx.seed_mask, ctx.cot_mask], dim=1) + """Attention mask aligned with the combined seed and CoT sequence.""" + hidden_states = ctx.model._last_hidden( + ctx.model._run_backbone(ctx.model.embed(sequence), attention_mask, sequence, 0) + ) + """Final backbone hidden states for every token in the combined sequence.""" + seed_length = ctx.seed_ids.size(1) + """Fixed left-padded seed width; its final position predicts the first CoT token.""" + cot_hidden = hidden_states[:, seed_length - 1 : seed_length - 1 + ctx.cot_ids.size(1), :] + """Shifted hidden states whose outputs predict the aligned CoT tokens.""" + prediction_logits = F.linear(cot_hidden, ctx.model.embed.weight) + """Vocabulary logits for each CoT target token, kept in the model's native precision.""" + return ScalarLoss( + F.cross_entropy( + prediction_logits.reshape(-1, ctx.model.vocab_size), + ctx.cot_ids.masked_fill(ctx.cot_mask == 0, -100).reshape(-1), + ignore_index=-100, + ) + ) + + +cot_loss_fn: LossFn[CoTLossContext] = cot_loss +"""The chain-of-thought generation loss function.""" + + +@dataclass +class BridgeLossContext: + """Inputs to :func:`bridge_loss` after Hungarian matching prepares bridge targets.""" + + matched_predictions: torch.Tensor + """Hungarian-matched query embeddings, shaped ``[match_count, embedding_dim]``.""" + + target_bank: torch.Tensor + """Normalized in-batch targets followed by optional hard-negative target embeddings.""" + + positive_indices: torch.Tensor + """Index into ``target_bank`` of each matched prediction's assigned target.""" + + validity_logits: torch.Tensor + """One validity logit per bridge query slot.""" + + validity_labels: torch.Tensor + """Binary labels marking the query slots selected by Hungarian matching.""" + + temperature: float + """Positive divisor applied to matched-query versus target-bank similarities.""" + + validity_weight: float + """Weight applied to the validity BCE term in the returned total.""" + + positive_weight: float + """BCE positive-class weight used to compensate for sparse valid query slots.""" + + +@dataclass +class BridgeLoss: + """A composed bridge objective with inspectable component losses.""" + + info_nce: Loss + """Matched query-to-target retrieval loss.""" + + validity: Loss + """Binary query-validity loss.""" + + validity_weight: float + """Weight applied to ``validity`` when forming the scalar optimization objective.""" + + def to_tensor(self) -> torch.Tensor: + """Return ``info_nce + validity_weight * validity`` for optimization.""" + return self.info_nce.to_tensor() + self.validity_weight * self.validity.to_tensor() + + +def bridge_loss(ctx: BridgeLossContext) -> BridgeLoss: + """Compute QueryBridge's matched-target InfoNCE and query-validity BCE. + + Hungarian matching is deliberately outside this function: it creates the + positive indices and validity labels but is non-differentiable and depends on + SciPy. This loss consumes those prepared tensors, uses all target-bank rows + as contrastive candidates, and combines retrieval with weighted validity BCE. + """ + + if ctx.matched_predictions.numel(): + info_nce = info_nce_loss( + InfoNCELossContext( + anchors=ctx.matched_predictions, + candidates=ctx.target_bank, + positive_indices=ctx.positive_indices, + temperature=ctx.temperature, + ) + ) + else: + info_nce = ScalarLoss(ctx.target_bank.new_zeros(())) + """InfoNCE over Hungarian-matched query embeddings, or zero with no matches.""" + validity = ScalarLoss( + F.binary_cross_entropy_with_logits( + ctx.validity_logits, + ctx.validity_labels, + pos_weight=torch.tensor(ctx.positive_weight, device=ctx.validity_logits.device), + ) + ) + """BCE supervising which bridge query slots should emit a latent.""" + return BridgeLoss( + info_nce=info_nce, + validity=validity, + validity_weight=ctx.validity_weight, + ) + + +bridge_loss_fn: LossFn[BridgeLossContext] = bridge_loss +"""The composed QueryBridge loss function.""" diff --git a/src/langset/strategies.py b/src/langset/strategies.py index 8254299..f1ef6e5 100644 --- a/src/langset/strategies.py +++ b/src/langset/strategies.py @@ -19,6 +19,18 @@ import torch import torch.nn.functional as F +from langset.loss import ( + CoTLossContext, + InfoNCELossContext, + SoftTargetCrossEntropyContext, + StopLossContext, + SupervisedContrastiveLossContext, + cot_loss, + info_nce_loss, + soft_target_cross_entropy_loss, + stop_loss, + supervised_contrastive_loss, +) from langset.modeling import LangSetModel from langset.sigreg import SIGReg @@ -29,35 +41,6 @@ from langset.training_args import TrainingArguments -def supcon_loss(z: torch.Tensor, labels: list[str], tau: float) -> torch.Tensor: - """Supervised-contrastive (Khosla et al.) over emitted latents: same-label items are positives (pulled together), - all others negatives (pushed apart) — a few group labels SHAPE the geometry into SEPARATE REGIONS. Being a proper - contrastive loss (each anchor has both positives and negatives) it separates without collapsing. Items whose label - is ''/'unknown'/'none'/'nan' are dropped. Returns 0 if fewer than two labelled items or no positive pair exists.""" - dev = z.device - keep = [ - k - for k, l in enumerate(labels) - if str(l).strip().lower() not in ("", "unknown", "none", "nan") - ] - if len(keep) < 2: - return z.new_zeros(()) - zz = F.normalize(z[keep], p=2, dim=-1) - lab = [labels[k] for k in keep] - b = len(keep) - sim = (zz @ zz.t() / tau).masked_fill(torch.eye(b, device=dev, dtype=torch.bool), -1e9) - logp = sim - torch.logsumexp(sim, dim=1, keepdim=True) - pos = torch.tensor( - [[1.0 if (i != j and lab[i] == lab[j]) else 0.0 for j in range(b)] for i in range(b)], - device=dev, - ) - npos = pos.sum(1) - has = npos > 0 - if not bool(has.any()): - return z.new_zeros(()) - return (-(pos * logp).sum(1)[has] / npos[has]).mean() - - # ---- per-step context handed to the aux loss terms ---------------------------------------------- @dataclass class MultiStepCtx: @@ -150,17 +133,22 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float return None if not (a.lam_multi_nce > 0 and int(c.valid.sum()) > 1): return None - rvn = F.normalize(c.recon[c.valid], dim=-1) # [N, d] emitted (gradient flows here) - tvn = F.normalize(c.target_lat[c.valid], dim=-1) # [N, d] EMA targets (already stop-grad) - nce_logits = (rvn @ tvn.t()) / a.tau # [N, N] query x key cosine / temp - n_nce = rvn.size(0) - fn_mask = torch.zeros(n_nce, n_nce, dtype=torch.bool, device=c.dev) + emitted = c.recon[c.valid] + targets = c.target_lat[c.valid] + emission_count = emitted.size(0) + fn_mask = torch.zeros(emission_count, emission_count, dtype=torch.bool, device=c.dev) for masker in self.maskers: masker(c, fn_mask) - nce_logits = nce_logits.masked_fill( - fn_mask, float("-inf") - ) # diagonal (positive) never masked - loss_nce = F.cross_entropy(nce_logits, torch.arange(n_nce, device=c.dev)) + loss_nce = info_nce_loss( + InfoNCELossContext( + anchors=emitted, + candidates=targets, + positive_indices=torch.arange(emission_count, device=c.dev), + temperature=a.tau, + logit_mask=fn_mask, + normalize_embeddings=True, + ) + ).to_tensor() return (self.key, loss_nce, a.lam_multi_nce) @@ -178,9 +166,11 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float for r, k in enumerate(c.bidx) for j in range(c.lens_l[r]) ] - loss_sup = supcon_loss( - c.recon[c.valid], sup_flat, a.sup_tau - ) # pull same-stage, push different-stage + loss_sup = supervised_contrastive_loss( + SupervisedContrastiveLossContext( + embeddings=c.recon[c.valid], labels=sup_flat, temperature=a.sup_tau + ) + ).to_tensor() # pull same-stage, push different-stage return (self.key, loss_sup, a.lam_sup) @@ -233,7 +223,7 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float self_.cot_texts[k] for k in c.bidx ): # no reasoning in this batch -> nothing to learn return None - tok, vsz = m.tokenizer, m.vocab_size + tok = m.tokenizer def _tokm(texts: list[str], mx: int, side: str) -> tuple[torch.Tensor, torch.Tensor]: e = tok( @@ -254,17 +244,16 @@ def _tokm(texts: list[str], mx: int, side: str) -> tuple[torch.Tensor, torch.Ten # silently condition short seeds' CoT on padding.) di, dm = _tokm([self_.input_text[k] for k in c.bidx], a.max_len, "left") ti, tm = _tokm([self_.cot_texts[k] or " " for k in c.bidx], a.max_len, "right") - seq = torch.cat([di, ti], dim=1) - am = torch.cat([dm, tm], dim=1) - hid = m._last_hidden(m._run_backbone(m.embed(seq), am, seq, 0)) - sd = di.size(1) - ph = hid[:, sd - 1 : sd - 1 + ti.size(1), :] # hidden that predicts each CoT token - # bf16 vocab projection (NOT .float()): the fp32 [b, T, |V|] logits were the OOM driver on the 80GB A100; - # CE reduces the softmax in fp32 internally, so bf16 logits are a numerically fine training signal. - lg = F.linear(ph, m.embed.weight) - loss_cot = F.cross_entropy( - lg.reshape(-1, vsz), ti.masked_fill(tm == 0, -100).reshape(-1), ignore_index=-100 - ) + loss_cot = cot_loss( + CoTLossContext( + model=m, + args=a, + seed_ids=di, + seed_mask=dm, + cot_ids=ti, + cot_mask=tm, + ) + ).to_tensor() return (self.key, loss_cot, a.lam_cot) @@ -408,7 +397,7 @@ def emit( ep: int, ss_mask: Optional[torch.Tensor] = None, ) -> EmissionOut: - m, a, dev = self.m, self.a, self.dev + m, a = self.m, self.a ss_prob = a.ss_prob assert ss_prob is not None # Trainer resolves the None sentinel before any emit eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) @@ -426,17 +415,15 @@ def emit( w = F.relu(target_lat.float() @ code.t()) # [b, lmax, n_codes] w = w / w.sum(-1, keepdim=True).clamp_min(1e-9) # MEMBERSHIP: soft-target CE over the codes ALONE, at real emission slots. - cell_logp = F.log_softmax(code_logits[:, :lmax, 0, :].float(), -1) # [b, lmax, n_codes] - loss_code = -(w * cell_logp).sum(-1)[valid].mean() + loss_code = soft_target_cross_entropy_loss( + SoftTargetCrossEntropyContext( + logits=code_logits[:, :lmax, 0, :].float(), + target_probabilities=w, + selection=valid, + ) + ).to_tensor() # Termination is independent of the member distribution, so set width cannot weaken its supervision. - stop_lab = torch.zeros(b, lmax + 1, device=dev) - keep = torch.zeros(b, lmax + 1, dtype=torch.bool, device=dev) - for r, nl in enumerate(lens_l): - stop_lab[r, nl] = 1.0 # continue through the tick's members, then stop after the last - keep[r, : nl + 1] = True - loss_stop = F.binary_cross_entropy_with_logits( - stop_lg.squeeze(-1).float()[keep], stop_lab[keep] - ) + loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() # `recon` is the contract every loss term reads as "the model's emission, with gradient" (MultiNCETerm # and the emitter both build their query from it). The rollout's own recon is the TARGET here -- this # path is lossless, so it carries no gradient at all, and a term querying it would compare each target @@ -561,7 +548,7 @@ def emit( ep: int, ss_mask: Optional[torch.Tensor] = None, ) -> EmissionOut: - m, a, dev = self.m, self.a, self.dev + m, a = self.m, self.a ss_prob = a.ss_prob assert ss_prob is not None eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) @@ -584,20 +571,18 @@ def emit( sel = seen[:, :, fi] & valid if not bool(sel.any()): continue - logp = F.log_softmax(flat[..., m_lo:m_hi], -1) - li = -(tgt[..., m_lo:m_hi] * logp).sum(-1)[sel].mean() + li = soft_target_cross_entropy_loss( + SoftTargetCrossEntropyContext( + logits=flat[..., m_lo:m_hi], + target_probabilities=tgt[..., m_lo:m_hi], + selection=sel, + ) + ).to_tensor() losses.append(li) per_facet[f"c_{self.facets[fi]}"] = li loss_concept = torch.stack(losses).sum() if losses else flat.new_zeros(()) - stop_lab = torch.zeros(b, lmax + 1, device=dev) - keep = torch.zeros(b, lmax + 1, dtype=torch.bool, device=dev) - for r, nl in enumerate(lens_l): - stop_lab[r, nl] = 1.0 - keep[r, : nl + 1] = True - loss_stop = F.binary_cross_entropy_with_logits( - stop_lg.squeeze(-1).float()[keep], stop_lab[keep] - ) + loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() # the emission the rest of the trainer sees — WITH gradient, or every aux term silently idles p = m.head.concept_probs(code_logits[:, :lmax]) @@ -717,23 +702,18 @@ def emit( if members: tgt[r, t, members] = 1.0 / len(members) has[r, t] = True - cell_logp = F.log_softmax(code_logits[:, :lmax, 0, :].float(), -1) - sel = has & valid - loss_state = ( - -(tgt * cell_logp).sum(-1)[sel].mean() if bool(sel.any()) else cell_logp.new_zeros(()) - ) + loss_state = soft_target_cross_entropy_loss( + SoftTargetCrossEntropyContext( + logits=code_logits[:, :lmax, 0, :].float(), + target_probabilities=tgt, + selection=has & valid, + ) + ).to_tensor() # TERMINATION: its own sigmoid, never folded into the member softmax. Folding is only fair when the # member target is one-hot; against a 1/k-diffuse target the gradient suppressing STOP scales with # P(STOP), which weakens as the set widens, so the rollout truncates exactly where sets get wide. - stop_lab = torch.zeros(b, lmax + 1, device=dev) - keep = torch.zeros(b, lmax + 1, dtype=torch.bool, device=dev) - for r, nl in enumerate(lens_l): - stop_lab[r, nl] = 1.0 - keep[r, : nl + 1] = True - loss_stop = F.binary_cross_entropy_with_logits( - stop_lg.squeeze(-1).float()[keep], stop_lab[keep] - ) + loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() # THE EMISSION the rest of the trainer sees: state mixture ++ residual, exactly what feeds back. It must # be the emission WITH gradient -- every aux term (in-batch NCE, hard negatives) builds its query from diff --git a/src/langset/trainer.py b/src/langset/trainer.py index e3f6074..1ea5bfd 100644 --- a/src/langset/trainer.py +++ b/src/langset/trainer.py @@ -19,6 +19,14 @@ from langset import selection from langset.heads import Head, RtHead, resolve_head +from langset.loss import ( + LearnLossContext, + ReconstructionLossContext, + SLContext, + learn_loss, + recon_loss, + sl_loss, +) from langset.modeling import LangSetModel from langset.strategies import ( MultiStepCtx, @@ -113,7 +121,7 @@ def _tokenize_replay( ) -> tuple[torch.Tensor, torch.Tensor]: """Tokenize replay text with an EXPLICIT padding side. The doc (the conditioning context) MUST be left-padded so every row's last real token lands in the final column — that's the position whose hidden predicts the first - target token in `_replay_ce`. The target is right-padded (its pad is masked out of the CE). Defaulting the side + target token in `learn_loss`. The target is right-padded (its pad is masked out of the CE). Defaulting the side (right-pad the doc) silently conditions the replay loss on padding for any row shorter than the batch max.""" e = tok( texts, @@ -163,33 +171,6 @@ def _restore_best(m: LangSetModel, best_state: dict[str, Any]) -> None: m.aux_heads[name].load_state_dict(sd) -def _replay_ce( - model: LangSetModel, - doc_ids: torch.Tensor, - doc_mask: torch.Tensor, - tgt_ids: torch.Tensor, - tgt_mask: torch.Tensor, - vocab_size: int, -) -> torch.Tensor: - """Teacher-forced next-token CE for a text-replay (rehearsal) row: condition on the doc, score CE on the target - tokens only, projected through the tied input embedding (no lm_head, no full-vocab OOM). REQUIRES the doc to be - left-padded (see `_tokenize_replay`) so column `sd-1` is the last REAL doc token for every row regardless of its - length; the target is right-padded and its pad positions are ignored. Shared by both learn paths.""" - seq = torch.cat([doc_ids, tgt_ids], dim=1) - am = torch.cat([doc_mask, tgt_mask], dim=1) - hid = model._last_hidden( - model._run_backbone(model.embed(seq), am, seq, 0) - ) # all real tokens -> real_start=0 - sd = doc_ids.size(1) - ph = hid[:, sd - 1 : sd - 1 + tgt_ids.size(1), :] # the hidden that predicts each target token - lg = F.linear(ph.float(), model.embed.weight.float()) # [B, St, vocab] via the tied embedding - return F.cross_entropy( - lg.reshape(-1, vocab_size), - tgt_ids.masked_fill(tgt_mask == 0, -100).reshape(-1), - ignore_index=-100, - ) - - # ---- single-latent step engines ----------------------------------------------------------------- # The real axis of the single-latent path is "where do pred/target/hard-neg features come from, and how is the # step run": the LIVE backbone (default) vs. FROZEN-POOL cached vectors (pool_mode="last" + frozen backbone). Both @@ -211,10 +192,6 @@ def featurize( """-> (pred [B,d] grad flows, target [B,d] honoring stop_grad_target, hard_neg [Hn,d] or None).""" raise NotImplementedError - def recon(self, pred: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: - """Recon aux scalar. Only defined/called when supports_recon.""" - raise NotImplementedError - def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """-> (emit_in, emit_tg) numpy for the eval MRR/retrieval block.""" raise NotImplementedError @@ -237,14 +214,12 @@ def __init__( t2_mask: torch.Tensor, hn_ids: Optional[torch.Tensor], hn_mask: Optional[torch.Tensor], - recon_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], input_text: list[str], target_text: list[str], ) -> None: self.m, self.a, self.tok = model, args, tok self.ids, self.mask, self.t2_ids, self.t2_mask = ids, mask, t2_ids, t2_mask self.hn_ids, self.hn_mask = hn_ids, hn_mask - self._recon_fn = recon_fn self.input_text, self.target_text = input_text, target_text def featurize( @@ -284,9 +259,6 @@ def featurize( hn = m(*_dyn_trim(self.hn_ids[idx], self.hn_mask[idx])) return pred, target, hn - def recon(self, pred: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: - return self._recon_fn(pred, idx) - def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: emit_in = np.asarray( self.m.encode([self.input_text[j] for j in val_idx], normalize_embeddings=True) @@ -644,7 +616,9 @@ def train(self) -> LangSetModel: rng = np.random.default_rng(a.seed) tok = m.tokenizer + # Convert all text to tokenized tensors def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: + """Tokenize text to input_ids and attention_mask tensors""" e = tok(texts, padding=True, truncation=True, max_length=mx, return_tensors="pt") return e["input_ids"].to(dev), e["attention_mask"].to(dev) @@ -659,7 +633,7 @@ def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: ): # hard-neg view (empty "" rows tokenize fine, masked below) hn_ids, hn_mask = tok_to([t or " " for t in self.hard_neg_text], a.max_len) - # knowledge-injection: learn rows go to a SEPARATE next-token-CE pool; the contrastive split is embed-only. + # texts associated with knowledge injection learn_pool = [i for i in range(len(self.input_text)) if self.is_learn[i]] ln_doc_ids = ln_doc_mask = ln_tgt_ids = ln_tgt_mask = None if learn_pool: @@ -670,54 +644,14 @@ def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: tok, [self.target_text[i] for i in learn_pool], _LEARN_TGT, "right", dev ) # target RIGHT-pad + # shuffle indices for training/validation split n = len(self.input_text) embed_all = np.array([i for i in range(n) if not self.is_learn[i]]) embed_perm = embed_all[rng.permutation(len(embed_all))] n_val = max(4, int(len(embed_perm) * a.val_frac)) val_idx, tr_idx = embed_perm[:n_val], embed_perm[n_val:] - # recon aux: latent -> K soft tokens -> backbone decodes target_text (token CE). Grounds the latent in the - # text. `connector` is TRAINING-ONLY scaffolding (not saved; inference just emits the latent). - hsz, vsz = m.h, m.vocab_size - connector = torch.nn.Linear(m.latent_dim, _RECON_K * hsz).to(dev) - - def recon_loss(latent: torch.Tensor, rows: torch.Tensor) -> torch.Tensor: - ti, tm = tr_ids[rows], tr_mask[rows] - temb = m.embed(ti) - soft = connector(latent).view(latent.size(0), _RECON_K, hsz).to(temb.dtype) - seq = torch.cat([soft, temb], dim=1) - am = torch.cat( - [torch.ones(latent.size(0), _RECON_K, device=dev, dtype=tm.dtype), tm], dim=1 - ) - out = m._run_backbone( - seq, am, ti, _RECON_K - ) # soft tokens synthetic; real target tokens at [K:] - sl = slice(_RECON_K - 1, _RECON_K - 1 + ti.size(1)) - lg = getattr(out, "logits", None) - if lg is not None: # model exposes an lm_head - pred_lg = lg[:, sl, :].float() - else: # text tower (no lm_head): project only the recon positions - hid = m._last_hidden(out)[ - :, sl, : - ] # via the tied input embedding -> avoids full-seq 262k OOM - pred_lg = F.linear(hid.float(), m.embed.weight.float()) - return F.cross_entropy( - pred_lg.reshape(-1, vsz), - ti.masked_fill(tm == 0, -100).reshape(-1), - ignore_index=-100, - ) - - def learn_loss(pos: torch.Tensor) -> torch.Tensor: - # [LEARN] rows: teacher-forced causal LM. Condition on the LEFT-padded case, CE ONLY on the substance - # tokens -> forces the backbone's hidden states to REPRESENT the substance (builds the axis the probe - # found missing). Shared _replay_ce (tied-embedding projection on the target span only, no lm_head/OOM). - assert ( - ln_doc_ids is not None and ln_doc_mask is not None - ) # learn_loss runs only when learn_pool set them - assert ln_tgt_ids is not None and ln_tgt_mask is not None - return _replay_ce( - m, ln_doc_ids[pos], ln_doc_mask[pos], ln_tgt_ids[pos], ln_tgt_mask[pos], vsz - ) + connector = torch.nn.Linear(m.latent_dim, _RECON_K * m.h).to(dev) opt = torch.optim.AdamW( [p for p in m.parameters() if p.requires_grad] + list(connector.parameters()), @@ -725,7 +659,7 @@ def learn_loss(pos: torch.Tensor) -> torch.Tensor: ) run = None if a.report_to == "wandb": - import wandb # type: ignore[import-untyped] # ty: ignore[unresolved-import] # optional dep, not installed + import wandb # type: ignore[import-untyped] # optional dep, not installed run = wandb.init(project=a.wandb_project, config=_wandb_config(a)) @@ -746,6 +680,8 @@ def learn_loss(pos: torch.Tensor) -> torch.Tensor: flush=True, ) # a DIFFERENT model/data/config can never resume us ck = None + + # code associated with resuming from a checkpoint if ck is not None: _params = dict(m.named_parameters()) for nm, t in ck["trainable"].items(): @@ -820,7 +756,6 @@ def save_resume(next_ep: int) -> None: t2_mask, hn_ids, hn_mask, - recon_loss, self.input_text, self.target_text, ) @@ -858,51 +793,6 @@ def save_resume(next_ep: int) -> None: flush=True, ) - def _sl_loss( - pred: torch.Tensor, - target: torch.Tensor, - hn: Optional[torch.Tensor], - idx: torch.Tensor, - ) -> torch.Tensor: - """Single-latent contrastive loss as a PURE FUNCTION OF THE EMBEDDINGS (pred/target/hn) — the - GradCache-compatible factoring of the inline step below. Contains every term that reads only the - pooled embeddings (in-batch-negative InfoNCE + false-neg / hard-neg masking facets + - uniformity); the recon aux is NOT here (it needs the per-token backbone graph, so it stays on the - direct path and grad_cache asserts lam_recon==0). Byte-identical math to the inline step, so - grad_cache=False is unchanged.""" - tmat = target if hn is None else torch.cat([target, hn], dim=0) - logits = (pred @ tmat.t()) / a.tau # in-batch negatives force separation (no collapse) - B = len(idx) - neg_mask = torch.zeros(B, logits.size(1), dtype=torch.bool, device=dev) - if self.mask_keys is not None: # in-batch block: drop same-issue false negatives - bkeys = [self.mask_keys[j] for j in idx.tolist()] - for r in range(B): - kr = bkeys[r] - if not kr: - continue - for c in range(B): - if r != c and (kr & bkeys[c]): - neg_mask[r, c] = True - if ( - hn_ids is not None - ): # PER-ANCHOR-ONLY hard neg: anchor i sees ONLY its own mined hard neg (col B+i) - hnt = self.hard_neg_text - assert hnt is not None - valid = [bool(hnt[j]) for j in idx.tolist()] - for r in range(B): - for c in range(B): - if not (r == c and valid[c]): - neg_mask[r, B + c] = True - if bool(neg_mask.any()): - logits = logits.masked_fill( - neg_mask, float("-inf") - ) # diagonal (positive) always kept - loss = F.cross_entropy(logits, torch.arange(B, device=dev)) # primary - if a.lam_uniform > 0 and B > 1: # aux: uniformity - sq = torch.pdist(F.normalize(pred, p=2, dim=-1), p=2).pow(2) - loss = loss + a.lam_uniform * sq.mul(-2.0).exp().mean().log() - return loss - def _grad_cache_step(idx: torch.Tensor) -> float: """GradCache: EXACT full-batch contrastive gradient with peak activation = one gc_chunk (Gao et al. 2021). Lets `batch_size` (in-batch negatives) grow far past what one graph fits. Phase 1 encodes the @@ -929,9 +819,18 @@ def _grad_cache_step(idx: torch.Tensor) -> float: if not a.stop_grad_target: tf.requires_grad_(True) hf = torch.cat(hns) if hns else None # hard negs stay no_grad (as in featurize) - loss = _sl_loss( - pf, tf, hf, idx - ) # full-batch loss -> cached rep grads (+ emb_slot head grads) + loss = sl_loss( + SLContext( + model=m, + args=a, + pred=pf, + target=tf, + hn=hf, + idx=idx, + mask_keys=self.mask_keys, + hard_neg_text=self.hard_neg_text, + ) + ).to_tensor() # full-batch loss -> cached rep grads (+ emb_slot head grads) opt.zero_grad() loss.backward() # fills pf.grad (+ tf.grad unless stop-grad, + slot-head params); backbone NOT in this graph gp = pf.grad @@ -995,7 +894,19 @@ def _grad_cache_step(idx: torch.Tensor) -> float: ), device=dev, ) - lloss = learn_loss(lp) + assert ln_doc_ids is not None and ln_doc_mask is not None + assert ln_tgt_ids is not None and ln_tgt_mask is not None + lloss = learn_loss( + LearnLossContext( + model=m, + args=a, + pos=lp, + ln_doc_ids=ln_doc_ids, + ln_doc_mask=ln_doc_mask, + ln_tgt_ids=ln_tgt_ids, + ln_tgt_mask=ln_tgt_mask, + ) + ).to_tensor() opt.zero_grad() lloss.backward() opt.step() @@ -1009,15 +920,36 @@ def _grad_cache_step(idx: torch.Tensor) -> float: pred, target, hn = engine.featurize( idx ) # engine owns WHERE features come from (backbone vs cached) - # HARD NEGATIVES flow through _sl_loss as extra negative columns; grad still reaches only `pred`. - loss = _sl_loss( - pred, target, hn, idx - ) # InfoNCE + false-neg/hard-neg masking + uniformity (factored above) + # Hard negatives are extra contrastive columns; gradients reach only ``pred``. + loss = sl_loss( + SLContext( + model=m, + args=a, + pred=pred, + target=target, + hn=hn, + idx=idx, + mask_keys=self.mask_keys, + hard_neg_text=self.hard_neg_text, + ) + ).to_tensor() if ( engine.supports_recon and a.lam_recon > 0 ): # aux: grounding. At 0 the term is zero anyway; loss = ( - loss + a.lam_recon * engine.recon(pred, idx) + loss + + a.lam_recon + * recon_loss( + ReconstructionLossContext( + model=m, + args=a, + latent=pred, + rows=idx, + tr_ids=tr_ids, + tr_mask=tr_mask, + connector=connector, + ) + ).to_tensor() ) # SKIP so recon's fp32 full-vocab ([B,S,vocab]) projection graph is NOT built every step # (that graph, not the stripped lm_head, OOM'd a 0.6B at 84GB without grad_ckpt). Also lets # frozen-pool run. Direct path only: grad_cache asserts lam_recon==0 (embedding-only caching @@ -1081,7 +1013,19 @@ def _grad_cache_step(idx: torch.Tensor) -> float: rv, tot_v = 0.0, 0 for s in range(0, len(val_idx), a.batch_size): vb = torch.tensor(val_idx[s : s + a.batch_size], device=dev) - rv += float(recon_loss(m(ids[vb], mask[vb]), vb)) * len(vb) + rv += float( + recon_loss( + ReconstructionLossContext( + model=m, + args=a, + latent=m(ids[vb], mask[vb]), + rows=vb, + tr_ids=tr_ids, + tr_mask=tr_mask, + connector=connector, + ) + ).to_tensor() + ) * len(vb) tot_v += len(vb) recon_val = rv / tot_v # recon_val is teacher-forced -> blind to collapse; hard-penalize high collapse so a collapsed epoch @@ -1250,7 +1194,7 @@ def _train_multi(self) -> LangSetModel: opt = torch.optim.AdamW(params + transient_head_params, lr=a.lr) run = None if a.report_to == "wandb": - import wandb # type: ignore[import-untyped] # ty: ignore[unresolved-import] # optional dep, not installed + import wandb # type: ignore[import-untyped] # optional dep, not installed run = wandb.init(project=a.wandb_project, config=_wandb_config(a)) @@ -1455,7 +1399,6 @@ def _rf( i for i in range(len(seeds)) if getattr(self, "is_learn", [False] * len(seeds))[i] ] ln_doc_ids = ln_doc_mask = ln_tgt_ids = ln_tgt_mask = None - vsz = m.vocab_size if learn_pool: ln_doc_ids, ln_doc_mask = _tokenize_replay( @@ -1469,15 +1412,6 @@ def _rf( dev, ) # target RIGHT-pad - def learn_loss(pos: torch.Tensor) -> torch.Tensor: # shared teacher-forced replay CE - assert ( - ln_doc_ids is not None and ln_doc_mask is not None - ) # runs only when learn_pool set them - assert ln_tgt_ids is not None and ln_tgt_mask is not None - return _replay_ce( - m, ln_doc_ids[pos], ln_doc_mask[pos], ln_tgt_ids[pos], ln_tgt_mask[pos], vsz - ) - if a.grad_cache: # multi-latent GradCache: reject the configs it cannot keep exact for the cross-batch term assert float(m.head.drop.p) == 0.0, ( "grad_cache requires dropout==0: the phase-1 (no_grad, full batch) and phase-2 (grad, per chunk) " @@ -1623,8 +1557,20 @@ def _multi_grad_cache_step( device=dev, dtype=torch.long, ) + assert ln_doc_ids is not None and ln_doc_mask is not None + assert ln_tgt_ids is not None and ln_tgt_mask is not None opt.zero_grad() - lloss = learn_loss(lp) + lloss = learn_loss( + LearnLossContext( + model=m, + args=a, + pos=lp, + ln_doc_ids=ln_doc_ids, + ln_doc_mask=ln_doc_mask, + ln_tgt_ids=ln_tgt_ids, + ln_tgt_mask=ln_tgt_mask, + ) + ).to_tensor() lloss.backward() opt.step() agg["learn_loss"] = agg.get("learn_loss", 0.0) + float(lloss.detach()) diff --git a/tests/test_replay.py b/tests/test_replay.py index 7ac3f82..9492190 100644 --- a/tests/test_replay.py +++ b/tests/test_replay.py @@ -13,8 +13,9 @@ import torch from multi_state_helpers import _args, _build_model -from langset import Trainer -from langset.trainer import _replay_ce, _tokenize_replay +from langset import Trainer, TrainingArguments +from langset.loss import LearnLossContext, learn_loss +from langset.trainer import _tokenize_replay def _row0_replay_loss(m, docs, tgts, doc_side): @@ -22,7 +23,19 @@ def _row0_replay_loss(m, docs, tgts, doc_side): di, dm = _tokenize_replay(m.tokenizer, docs, 64, doc_side, m.device) ti, tm = _tokenize_replay(m.tokenizer, tgts, 32, "right", m.device) with torch.no_grad(): - return float(_replay_ce(m, di[:1], dm[:1], ti[:1], tm[:1], m.vocab_size)) + return float( + learn_loss( + LearnLossContext( + model=m, + args=TrainingArguments(), + pos=torch.tensor([0], device=m.device), + ln_doc_ids=di, + ln_doc_mask=dm, + ln_tgt_ids=ti, + ln_tgt_mask=tm, + ) + ).to_tensor() + ) def test_replay_ce_is_padding_invariant(): @@ -50,7 +63,7 @@ def test_replay_ce_is_padding_invariant(): def test_replay_training_runs_multi_latent(): # integration: emit rows + a couple varying-length learn rows, replay firing every step -> the rewired - # multi-latent learn path must train end-to-end (exercises _tokenize_replay + _replay_ce on real batches). + # multi-latent learn path must train end-to-end (exercises _tokenize_replay + learn_loss on real batches). from multi_state_helpers import _flat_trainable m = _build_model() From 3e828f2347b8de09aeb002dfbe569979c890995d Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Mon, 17 Aug 2026 13:27:56 -0400 Subject: [PATCH 2/3] Make optional wandb import type-checkable --- src/langset/trainer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/langset/trainer.py b/src/langset/trainer.py index 1ea5bfd..0ceb1cf 100644 --- a/src/langset/trainer.py +++ b/src/langset/trainer.py @@ -10,6 +10,7 @@ import random from contextlib import AbstractContextManager +from importlib import import_module from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Optional, cast @@ -659,8 +660,7 @@ def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: ) run = None if a.report_to == "wandb": - import wandb # type: ignore[import-untyped] # optional dep, not installed - + wandb = import_module("wandb") run = wandb.init(project=a.wandb_project, config=_wandb_config(a)) best_score, best_state, no_improve = -1e9, None, 0 @@ -1194,8 +1194,7 @@ def _train_multi(self) -> LangSetModel: opt = torch.optim.AdamW(params + transient_head_params, lr=a.lr) run = None if a.report_to == "wandb": - import wandb # type: ignore[import-untyped] # optional dep, not installed - + wandb = import_module("wandb") run = wandb.init(project=a.wandb_project, config=_wandb_config(a)) @torch.no_grad() From 97db122275597511307e7ef1a9e38b4221d1f132 Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Mon, 17 Aug 2026 14:11:20 -0400 Subject: [PATCH 3/3] Test reusable InfoNCE loss --- tests/test_loss.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tests/test_loss.py diff --git a/tests/test_loss.py b/tests/test_loss.py new file mode 100644 index 0000000..10bebb2 --- /dev/null +++ b/tests/test_loss.py @@ -0,0 +1,83 @@ +"""Unit tests for reusable loss kernels.""" + +import torch +import torch.nn.functional as F + +from langset.loss import InfoNCELossContext, info_nce_loss + + +def test_info_nce_applies_logit_mask() -> None: + """Masked candidates are excluded from each anchor's contrastive denominator.""" + anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + candidates = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + positive_indices = torch.tensor([0, 1]) + logit_mask = torch.tensor([[False, True], [False, False]]) + + loss = info_nce_loss( + InfoNCELossContext( + anchors=anchors, + candidates=candidates, + positive_indices=positive_indices, + temperature=1.0, + logit_mask=logit_mask, + ) + ).to_tensor() + + expected = F.cross_entropy( + (anchors @ candidates.t()).masked_fill(logit_mask, float("-inf")), positive_indices + ) + assert torch.allclose(loss, expected) + + +def test_info_nce_normalization_toggle_changes_similarity_scale() -> None: + """Normalization switches InfoNCE from dot-product to cosine similarities.""" + anchors = torch.tensor([[2.0, 0.0]]) + candidates = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + positive_indices = torch.tensor([0]) + + dot_product_loss = info_nce_loss( + InfoNCELossContext( + anchors=anchors, + candidates=candidates, + positive_indices=positive_indices, + temperature=1.0, + ) + ).to_tensor() + cosine_loss = info_nce_loss( + InfoNCELossContext( + anchors=anchors, + candidates=candidates, + positive_indices=positive_indices, + temperature=1.0, + normalize_embeddings=True, + ) + ).to_tensor() + + assert torch.allclose( + dot_product_loss, F.cross_entropy(torch.tensor([[2.0, 0.0]]), positive_indices) + ) + assert torch.allclose( + cosine_loss, F.cross_entropy(torch.tensor([[1.0, 0.0]]), positive_indices) + ) + assert not torch.allclose(dot_product_loss, cosine_loss) + + +def test_info_nce_backpropagates_to_anchors_and_candidates() -> None: + """The scalar loss preserves autograd paths through both embedding inputs.""" + anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]], requires_grad=True) + candidates = torch.tensor([[0.8, 0.2], [0.2, 0.8]], requires_grad=True) + + loss = info_nce_loss( + InfoNCELossContext( + anchors=anchors, + candidates=candidates, + positive_indices=torch.tensor([0, 1]), + temperature=0.5, + ) + ).to_tensor() + loss.backward() + + assert anchors.grad is not None + assert candidates.grad is not None + assert bool(anchors.grad.abs().any()) + assert bool(candidates.grad.abs().any())